[Web] Avoid tensor-cache record copies (#20156)
Tensor-cache shards can contain many records. Using `ArrayBuffer.slice`
for each entry duplicates the record before it is decoded or uploaded.
Validate each record range and use a `Uint8Array` view into the shard
instead. Preserve the original byte offset when the cache already
returns a view.
diff --git a/web/src/artifact_cache.ts b/web/src/artifact_cache.ts
index da655da..4c29327 100644
--- a/web/src/artifact_cache.ts
+++ b/web/src/artifact_cache.ts
@@ -38,6 +38,35 @@
}
/**
+ * 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
*/
export interface ArtifactCacheTemplate {
diff --git a/web/src/runtime.ts b/web/src/runtime.ts
index 9c83a4d..2cd5ea1 100644
--- a/web/src/runtime.ts
+++ b/web/src/runtime.ts
@@ -39,6 +39,7 @@
TensorCacheAccessOptions,
TensorShardEntry,
createArtifactCache,
+ getTensorCacheRecordBytes,
} from "./artifact_cache";
import * as compact from "./compact";
import * as ctypes from "./ctypes";
@@ -1417,18 +1418,20 @@
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);
diff --git a/web/tests/node/test_tensor.js b/web/tests/node/test_tensor.js
index 6aadc70..2790bd1 100644
--- a/web/tests/node/test_tensor.js
+++ b/web/tests/node/test_tensor.js
@@ -54,3 +54,60 @@
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();
+});
diff --git a/web/tests/node/test_tensor_cache.js b/web/tests/node/test_tensor_cache.js
new file mode 100644
index 0000000..5ab92e3
--- /dev/null
+++ b/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,
+ );
+});
