diff --git a/web/src/artifact_cache.ts b/web/src/artifact_cache.ts index da655da4e317..4c2932739e9f 100644 --- a/web/src/artifact_cache.ts +++ b/web/src/artifact_cache.ts @@ -37,6 +37,35 @@ export interface TensorShardEntry { records: Array; } +/** + * 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, +): 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 */ diff --git a/web/src/runtime.ts b/web/src/runtime.ts index 9c83a4d9acb6..2cd5ea1827f8 100644 --- a/web/src/runtime.ts +++ b/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); diff --git a/web/tests/node/test_tensor.js b/web/tests/node/test_tensor.js index 6aadc702f5fa..2790bd17db5e 100644 --- a/web/tests/node/test_tensor.js +++ b/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(); +}); diff --git a/web/tests/node/test_tensor_cache.js b/web/tests/node/test_tensor_cache.js new file mode 100644 index 000000000000..5ab92e333f7c --- /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, + ); +});