[Web] Avoid tensor-cache record copies

web/src/artifact_cache.ts

@@ -37,6 +37,35 @@ export interface TensorShardEntry {
   records: Array<TensorCacheEntry>;
 }
 
+/**
+ * Return a borrowed view of one tensor record within a shard.
+ *
+ * The returned view aliases the shard and is only valid for as long as the
+ * shard data remains alive. Tensor-cache decoding consumes it synchronously.
+ */
+export function getTensorCacheRecordBytes(
+  shardData: ArrayBuffer | Uint8Array,
+  record: Pick<TensorCacheEntry, "byteOffset" | "nbytes">,
+): Uint8Array {
+  const shardBytes =
+    shardData instanceof Uint8Array ? shardData : new Uint8Array(shardData);
+  const { byteOffset, nbytes } = record;
+  if (!Number.isSafeInteger(byteOffset) || byteOffset < 0) {
+    throw new Error(`Invalid tensor-cache byteOffset: ${byteOffset}`);
+  }
+  if (!Number.isSafeInteger(nbytes) || nbytes < 0) {
+    throw new Error(`Invalid tensor-cache nbytes: ${nbytes}`);
+  }
+  const endOffset = byteOffset + nbytes;
+  if (!Number.isSafeInteger(endOffset) || endOffset > shardBytes.byteLength) {
+    throw new Error(
+      `Tensor-cache record range [${byteOffset}, ${endOffset}) exceeds ` +
+      `shard size ${shardBytes.byteLength}`,
+    );
+  }
+  return shardBytes.subarray(byteOffset, endOffset);
+}
+
 /**
  *   Common Interface for the artifact cache
  */

web/src/runtime.ts

@@ -39,6 +39,7 @@ import {
   TensorCacheAccessOptions,
   TensorShardEntry,
   createArtifactCache,
+  getTensorCacheRecordBytes,
 } from "./artifact_cache";
 import * as compact from "./compact";
 import * as ctypes from "./ctypes";
@@ -1417,18 +1418,20 @@ export class Instance implements Disposable {
         this.env.logger("Error: Cannot fetch " + dataUrl + " err= " + err);
         throw err;
       }
+      const shardBytes =
+        buffer instanceof Uint8Array ? buffer : new Uint8Array(buffer);
       const shardRecords = shard.records;
       for (let j = 0; j < shardRecords.length; ++j) {
         try {
           const rec = shardRecords[j];
+          const recSource = getTensorCacheRecordBytes(shardBytes, rec);
           const cpu_arr = this.withNewScope(() => {
             return this.detachFromCurrentScope(
               this.empty(rec.shape, rec.dtype, this.cpu())
             )
           });
-          const recSource = buffer.slice(rec.byteOffset, rec.byteOffset + rec.nbytes);
           // first sync copy to cpu.
-          this.ctx.arrayDecodeStorage(cpu_arr, new Uint8Array(recSource), rec.format, rec.dtype);
+          this.ctx.arrayDecodeStorage(cpu_arr, recSource, rec.format, rec.dtype);
           // then async stream into GPU if needed
           if (device.deviceType === DeviceStrToEnum.cpu) {
             this.tensorCacheUpdate(rec.name, cpu_arr, false);

web/tests/node/test_tensor.js

@@ -54,3 +54,60 @@ test("array copy", () => {
     testArrayCopy("float64", Float64Array);
   });
 });
+
+test("tensor cache loads adjacent records from a Uint8Array shard", async () => {
+  const backing = new Uint8Array([90, 91, 1, 2, 3, 4, 5, 6, 7, 8, 92]);
+  const shard = backing.subarray(2, 10);
+  const manifest = {
+    metadata: {},
+    records: [{
+      dataPath: "params.bin",
+      format: "raw-shard",
+      nbytes: shard.byteLength,
+      records: [
+        {
+          name: "test.record_view.first",
+          shape: [4],
+          dtype: "uint8",
+          format: "raw",
+          byteOffset: 0,
+          nbytes: 4,
+        },
+        {
+          name: "test.record_view.second",
+          shape: [4],
+          dtype: "uint8",
+          format: "raw",
+          byteOffset: 4,
+          nbytes: 4,
+        },
+      ],
+    }],
+  };
+  const artifactCache = {
+    hasAllKeys: async () => true,
+    addToCache: async () => {},
+    deleteInCache: async () => {},
+    fetchWithCache: async (_url, storeType) => {
+      return storeType === "json" ? manifest : shard;
+    },
+  };
+
+  await tvm.fetchTensorCache(
+    "https://example.test/model/",
+    tvm.cpu(),
+    { artifactCache },
+  );
+
+  tvm.withNewScope(() => {
+    assert.deepStrictEqual(
+      Array.from(tvm.tensorCacheGet("test.record_view.first").toArray()),
+      [1, 2, 3, 4],
+    );
+    assert.deepStrictEqual(
+      Array.from(tvm.tensorCacheGet("test.record_view.second").toArray()),
+      [5, 6, 7, 8],
+    );
+  });
+  tvm.tensorCacheClear();
+});

web/tests/node/test_tensor_cache.js

@@ -0,0 +1,86 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+const { getTensorCacheRecordBytes } = require("../../src/artifact_cache");
+
+test("tensor-cache record is a borrowed ArrayBuffer view", () => {
+  const shard = new Uint8Array([1, 2, 3, 4, 5, 6]).buffer;
+
+  const record = getTensorCacheRecordBytes(shard, {
+    byteOffset: 2,
+    nbytes: 3,
+  });
+
+  expect(Array.from(record)).toEqual([3, 4, 5]);
+  expect(record.buffer).toBe(shard);
+  expect(record.byteOffset).toBe(2);
+});
+
+test("tensor-cache record respects a Uint8Array shard offset", () => {
+  const backing = new Uint8Array([90, 91, 1, 2, 3, 4, 92]);
+  const shard = backing.subarray(2, 6);
+
+  const record = getTensorCacheRecordBytes(shard, {
+    byteOffset: 1,
+    nbytes: 2,
+  });
+
+  expect(Array.from(record)).toEqual([2, 3]);
+  expect(record.buffer).toBe(backing.buffer);
+  expect(record.byteOffset).toBe(shard.byteOffset + 1);
+});
+
+test("tensor-cache record may cover the full shard", () => {
+  const shard = new Uint8Array([1, 2, 3, 4]);
+
+  const record = getTensorCacheRecordBytes(shard, {
+    byteOffset: 0,
+    nbytes: shard.byteLength,
+  });
+
+  expect(record).toEqual(shard);
+  expect(record.buffer).toBe(shard.buffer);
+});
+
+test("tensor-cache record may be empty at the end of the shard", () => {
+  const shard = new Uint8Array([1, 2, 3, 4]);
+
+  const record = getTensorCacheRecordBytes(shard, {
+    byteOffset: shard.byteLength,
+    nbytes: 0,
+  });
+
+  expect(record.byteLength).toBe(0);
+  expect(record.byteOffset).toBe(shard.byteOffset + shard.byteLength);
+});
+
+test.each([
+  [{ byteOffset: -1, nbytes: 1 }, "byteOffset"],
+  [{ byteOffset: 0.5, nbytes: 1 }, "byteOffset"],
+  [{ byteOffset: Number.MAX_SAFE_INTEGER + 1, nbytes: 1 }, "byteOffset"],
+  [{ byteOffset: 0, nbytes: -1 }, "nbytes"],
+  [{ byteOffset: 0, nbytes: 0.5 }, "nbytes"],
+  [{ byteOffset: 0, nbytes: Number.MAX_SAFE_INTEGER + 1 }, "nbytes"],
+  [{ byteOffset: 5, nbytes: 0 }, "exceeds shard size"],
+  [{ byteOffset: 3, nbytes: 2 }, "exceeds shard size"],
+  [{ byteOffset: Number.MAX_SAFE_INTEGER, nbytes: 1 }, "exceeds shard size"],
+])("tensor-cache record rejects invalid range %j", (range, message) => {
+  expect(() => getTensorCacheRecordBytes(new ArrayBuffer(4), range)).toThrow(
+    message,
+  );
+});