This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 4e9a099d15 [Web] Upload pass-through tensor-cache records directly to
WebGPU (#20166)
4e9a099d15 is described below
commit 4e9a099d154d7c4644a40a1a9c00b8873226468e
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Sat Aug 22 21:29:02 2026 -0700
[Web] Upload pass-through tensor-cache records directly to WebGPU (#20166)
Building on #20156's borrowed tensor-cache record views, upload aligned
records that require no decoding directly into their final WebGPU
tensors. This avoids allocating a temporary CPU tensor, passing the
record through the FFI decoder, and copying the CPU tensor to WebGPU.
BF16, unaligned, and CPU records retain the existing decoder path.
Also:
- Validate exact tensor storage sizes, including packed sub-byte dtypes
- Honor tensor byte offsets during WebGPU writes
- Read `DLTensor::byte_offset` as an aligned unsigned 64-bit value and
reject values outside JavaScript’s safe integer range
- Dispose partially constructed tensors when loading fails
---
web/src/memory.ts | 13 +
web/src/runtime.ts | 132 +++++++---
web/tests/node/test_memory.js | 12 +
web/tests/node/test_tensor_cache_webgpu.js | 372 +++++++++++++++++++++++++++++
4 files changed, 499 insertions(+), 30 deletions(-)
diff --git a/web/src/memory.ts b/web/src/memory.ts
index 5da8b59640..290826d6a9 100644
--- a/web/src/memory.ts
+++ b/web/src/memory.ts
@@ -87,6 +87,19 @@ export class Memory {
return this.viewI32[base];
}
+ loadU64(ptr: Pointer): number {
+ if (this.buffer != this.memory.buffer) {
+ this.updateViews();
+ }
+ // WebAssembly is little-endian. Reject values JavaScript cannot represent
exactly.
+ const base = ptr >> 2;
+ const value = this.viewU32[base] + this.viewU32[base + 1] * 0x100000000;
+ if (!Number.isSafeInteger(value)) {
+ throw new Error("Cannot represent uint64 value as a JavaScript number");
+ }
+ return value;
+ }
+
loadF32(ptr: Pointer): number {
if (this.buffer != this.memory.buffer) {
this.updateViews();
diff --git a/web/src/runtime.ts b/web/src/runtime.ts
index 2cd5ea1827..1ce05b7dde 100644
--- a/web/src/runtime.ts
+++ b/web/src/runtime.ts
@@ -37,6 +37,7 @@ import {
ArtifactCacheTemplate,
ArtifactCacheType,
TensorCacheAccessOptions,
+ TensorCacheEntry,
TensorShardEntry,
createArtifactCache,
getTensorCacheRecordBytes,
@@ -546,7 +547,9 @@ export class Tensor extends TVMObject {
const arrayOffsetDtypeLanes = arrayOffsetDtypeBits + SizeOf.U8;
const arrayOffsetShape = arrayOffsetDtype + SizeOf.DLDataType;
const arrayOffsetStrides = arrayOffsetShape + this.lib.sizeofPtr();
- const arrayOffsetByteOffset = arrayOffsetStrides + this.lib.sizeofPtr();
+ const byteOffsetUnaligned = arrayOffsetStrides + this.lib.sizeofPtr();
+ const arrayOffsetByteOffset =
+ Math.ceil(byteOffsetUnaligned / SizeOf.I64) * SizeOf.I64;
// dataPtr
this.dataPtr = lib.memory.loadPointer(this.dltensor);
// ndim
@@ -570,7 +573,7 @@ export class Tensor extends TVMObject {
this.device = new DLDevice(deviceType, deviceId, lib);
// byte_offset
- this.byteOffset = lib.memory.loadI64(this.dltensor +
arrayOffsetByteOffset);
+ this.byteOffset = lib.memory.loadU64(this.dltensor +
arrayOffsetByteOffset);
}
/**
@@ -654,22 +657,50 @@ export class Tensor extends TVMObject {
* @returns this
*/
copyFromRawBytes(data: Uint8Array): this {
+ const nbytes = this.numStorageBytes();
+ if (nbytes != data.length) {
+ throw new Error("Expect the data's length equals nbytes=" + nbytes);
+ }
// short cut for gpu copy
if (this.device.deviceType === DeviceStrToEnum.webgpu) {
- this.lib.webGPUContext?.copyRawBytesToBuffer(data, this.getDataPtr(), 0,
data.length);
+ if (this.byteOffset % 4 != 0 || data.length % 4 != 0) {
+ throw new Error(
+ "WebGPU raw byte copies require four-byte-aligned offsets and sizes",
+ );
+ }
+ const webGPUContext = this.lib.webGPUContext;
+ if (webGPUContext === undefined) {
+ throw new Error("WebGPU context is not initialized");
+ }
+ webGPUContext.copyRawBytesToBuffer(
+ data,
+ this.getDataPtr(),
+ this.byteOffset,
+ data.length,
+ );
return this;
}
// CPU copy
- const size = this.shape.reduce((a, b) => {
- return a * b;
- }, 1);
- const nbytes = this.dlDataType.numStorageBytes() * size;
- if (nbytes != data.length) {
- throw new Error("Expect the data's length equals nbytes=" + nbytes);
- }
this.ctx.tensorCopyFromJSBytes(this, data);
return this;
}
+
+ private numStorageBytes(): number {
+ let totalBits = this.dlDataType.bits * this.dlDataType.lanes;
+ if (!Number.isSafeInteger(totalBits) || totalBits < 0) {
+ throw new Error(`Invalid tensor dtype bit width: ${totalBits}`);
+ }
+ for (const dim of this.shape) {
+ if (!Number.isSafeInteger(dim) || dim < 0) {
+ throw new Error(`Invalid tensor dimension: ${dim}`);
+ }
+ totalBits *= dim;
+ if (!Number.isSafeInteger(totalBits)) {
+ throw new Error("Tensor storage size exceeds JavaScript's safe integer
range");
+ }
+ }
+ return Math.ceil(totalBits / 8);
+ }
/**
* Return a copied Uint8Array of the raw bytes in the Tensor.
* @returns The result array.
@@ -1312,6 +1343,36 @@ export class Instance implements Disposable {
this.cacheMetadata = { ...this.cacheMetadata, ...(list["metadata"] as
Record<string, any>) };
}
+ /**
+ * Consume a tensor-cache record synchronously.
+ *
+ * Keeping the borrowed record view local to this non-async helper avoids
+ * capturing it across the caller's WebGPU synchronization point.
+ */
+ private loadTensorCacheRecordData(
+ shardBytes: Uint8Array,
+ rec: TensorCacheEntry,
+ cpuArray?: Tensor,
+ gpuArray?: Tensor,
+ ): void {
+ const recBytes = getTensorCacheRecordBytes(shardBytes, rec);
+ if (cpuArray !== undefined) {
+ this.ctx.arrayDecodeStorage(
+ cpuArray,
+ recBytes,
+ rec.format,
+ rec.dtype,
+ );
+ }
+ if (gpuArray !== undefined) {
+ if (cpuArray === undefined) {
+ gpuArray.copyFromRawBytes(recBytes);
+ } else {
+ gpuArray.copyFrom(cpuArray);
+ }
+ }
+ }
+
/**
* Fetch list of Tensor into the TensorCache.
@@ -1422,32 +1483,40 @@ export class Instance implements Disposable {
buffer instanceof Uint8Array ? buffer : new Uint8Array(buffer);
const shardRecords = shard.records;
for (let j = 0; j < shardRecords.length; ++j) {
+ let cpu_arr: Tensor | undefined;
+ let gpu_arr: Tensor | undefined;
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())
- )
- });
- // first sync copy to cpu.
- 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);
- cpu_arr.dispose();
- } else {
- // allocate a gpu arr and async copy to it.
- const gpu_arr = this.withNewScope(() => {
+ const requiresDecode =
+ rec.format === "f32-to-bf16" && rec.dtype === "float32";
+ const directToWebGPU =
+ device.deviceType === DeviceStrToEnum.webgpu &&
+ !requiresDecode &&
+ rec.nbytes % 4 === 0;
+
+ if (!directToWebGPU) {
+ cpu_arr = this.withNewScope(() => {
+ return this.detachFromCurrentScope(
+ this.empty(rec.shape, rec.dtype, this.cpu()),
+ );
+ });
+ }
+
+ if (device.deviceType !== DeviceStrToEnum.cpu) {
+ gpu_arr = this.withNewScope(() => {
return this.detachFromCurrentScope(
- this.empty(rec.shape, rec.dtype, device)
- )
+ this.empty(rec.shape, rec.dtype, device),
+ );
});
- gpu_arr.copyFrom(cpu_arr);
+ }
+
+ this.loadTensorCacheRecordData(shardBytes, rec, cpu_arr, gpu_arr);
+
+ if (device.deviceType === DeviceStrToEnum.cpu) {
+ this.tensorCacheUpdate(rec.name, cpu_arr!, false);
+ } else {
await device.sync();
this.tensorCacheUpdate(rec.name, gpu_arr, false);
- cpu_arr.dispose();
- gpu_arr.dispose();
}
} catch (err) {
this.env.logger(
@@ -1455,6 +1524,9 @@ export class Instance implements Disposable {
"Error: " + err
);
throw err;
+ } finally {
+ cpu_arr?.dispose();
+ gpu_arr?.dispose();
}
}
fetchedBytes += shard.nbytes;
diff --git a/web/tests/node/test_memory.js b/web/tests/node/test_memory.js
index 2daa9ff9b4..1d39179d8e 100644
--- a/web/tests/node/test_memory.js
+++ b/web/tests/node/test_memory.js
@@ -18,6 +18,18 @@
*/
const { CachedCallStack, Memory } = require("../../src/memory");
+test("loadU64 reads unsigned values and rejects unsafe integers", () => {
+ const wasmMemory = new WebAssembly.Memory({ initial: 1 });
+ const memory = new Memory(wasmMemory);
+ const words = new Uint32Array(wasmMemory.buffer);
+
+ words.set([0x80000001, 1], 2);
+ expect(memory.loadU64(8)).toBe(0x180000001);
+
+ words.set([0, 0x200000], 2);
+ expect(() => memory.loadU64(8)).toThrow("Cannot represent uint64 value");
+});
+
test("loadRawBytes returns an owned copy", () => {
const wasmMemory = new WebAssembly.Memory({ initial: 1 });
const memory = new Memory(wasmMemory);
diff --git a/web/tests/node/test_tensor_cache_webgpu.js
b/web/tests/node/test_tensor_cache_webgpu.js
new file mode 100644
index 0000000000..9df60f4230
--- /dev/null
+++ b/web/tests/node/test_tensor_cache_webgpu.js
@@ -0,0 +1,372 @@
+/*
+ * 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 fs = require("fs");
+const path = require("path");
+const tvmjs = require("../../dist/tvmjs.bundle");
+
+global.GPUBufferUsage = {
+ MAP_READ: 1 << 0,
+ COPY_DST: 1 << 1,
+ COPY_SRC: 1 << 2,
+ STORAGE: 1 << 3,
+ UNIFORM: 1 << 4,
+};
+
+const wasmSource = fs.readFileSync(
+ path.join(__dirname, "../../dist/wasm/tvmjs_runtime.wasm"),
+);
+
+function createInstance() {
+ return new tvmjs.Instance(
+ new WebAssembly.Module(wasmSource),
+ tvmjs.createPolyfillWASI(),
+ );
+}
+
+function createMockGPUDevice({ detachWriteSources = false } = {}) {
+ const buffers = [];
+ const writes = [];
+ const queue = {
+ submit: jest.fn(),
+ onSubmittedWorkDone: jest.fn(() => Promise.resolve()),
+ writeBuffer: jest.fn(
+ (buffer, bufferOffset, data, dataOffset = 0, size = data.byteLength) => {
+ const source = new Uint8Array(data.buffer, data.byteOffset,
data.byteLength);
+ const snapshot = source.slice(dataOffset, dataOffset + size);
+ if (bufferOffset + snapshot.byteLength > buffer.contents.byteLength) {
+ throw new RangeError(
+ `write of ${snapshot.byteLength} bytes at ${bufferOffset} exceeds
` +
+ `${buffer.contents.byteLength}-byte mock buffer`,
+ );
+ }
+ new Uint8Array(buffer.contents).set(snapshot, bufferOffset);
+ writes.push({ buffer, bufferOffset, data, dataOffset, size, snapshot
});
+ if (detachWriteSources) {
+ structuredClone(data.buffer, { transfer: [data.buffer] });
+ }
+ },
+ ),
+ };
+ const device = {
+ queue,
+ lost: new Promise(() => {}),
+ addEventListener: jest.fn(),
+ pushErrorScope: jest.fn(),
+ popErrorScope: jest.fn(() => Promise.resolve(null)),
+ createBuffer: jest.fn((descriptor) => {
+ const buffer = {
+ size: descriptor.size,
+ contents: new ArrayBuffer(descriptor.size),
+ destroy: jest.fn(),
+ };
+ buffers.push(buffer);
+ return buffer;
+ }),
+ destroy: jest.fn(),
+ };
+ return { device, buffers, writes };
+}
+
+function createArtifactCache(manifest, shard) {
+ return {
+ hasAllKeys: async () => true,
+ addToCache: async () => {},
+ deleteInCache: async () => {},
+ fetchWithCache: async (_url, storeType) => {
+ return storeType === "json" ? manifest : shard;
+ },
+ };
+}
+
+test("WebGPU tensor cache uploads pass-through records directly", async () => {
+ const tvm = createInstance();
+ const gpu = createMockGPUDevice();
+ tvm.initWebGPU(gpu.device);
+
+ const records = [
+ {
+ name: "test.direct.raw_uint32",
+ shape: [2],
+ dtype: "uint32",
+ format: "raw",
+ byteOffset: 0,
+ nbytes: 8,
+ },
+ {
+ name: "test.direct.passthrough_float16",
+ shape: [2],
+ dtype: "float16",
+ format: "f32-to-bf16",
+ byteOffset: 8,
+ nbytes: 4,
+ },
+ {
+ name: "test.direct.raw_float32",
+ shape: [1],
+ dtype: "float32",
+ format: "raw",
+ byteOffset: 12,
+ nbytes: 4,
+ },
+ {
+ name: "test.decode.packed_bf16",
+ shape: [2],
+ dtype: "float32",
+ format: "f32-to-bf16",
+ byteOffset: 16,
+ nbytes: 4,
+ },
+ ];
+ const shard = Uint8Array.from([
+ 1, 2, 3, 4, 5, 6, 7, 8,
+ 9, 10, 11, 12,
+ 13, 14, 15, 16,
+ 0x80, 0x3f, 0x00, 0xc0,
+ ]).buffer;
+ const manifest = {
+ metadata: {},
+ records: [{
+ dataPath: "params.bin",
+ format: "raw-shard",
+ nbytes: shard.byteLength,
+ records,
+ }],
+ };
+
+ const originalDecode = tvm.ctx.arrayDecodeStorage;
+ const decode = jest.fn((...args) => originalDecode(...args));
+ tvm.ctx.arrayDecodeStorage = decode;
+ try {
+ await tvm.fetchTensorCache(
+ "https://example.test/model/";,
+ tvm.webgpu(),
+ { artifactCache: createArtifactCache(manifest, shard) },
+ );
+
+ expect(decode).toHaveBeenCalledTimes(1);
+ expect(decode.mock.calls[0][2]).toBe("f32-to-bf16");
+ expect(decode.mock.calls[0][3]).toBe("float32");
+ expect(gpu.writes.map((write) => Array.from(write.snapshot))).toEqual([
+ [1, 2, 3, 4, 5, 6, 7, 8],
+ [9, 10, 11, 12],
+ [13, 14, 15, 16],
+ [0, 0, 128, 63, 0, 0, 0, 192],
+ ]);
+ } finally {
+ tvm.ctx.arrayDecodeStorage = originalDecode;
+ tvm.tensorCacheClear();
+ tvm.dispose();
+ }
+});
+
+test("direct upload snapshots a borrowed shard view before synchronization",
async () => {
+ const tvm = createInstance();
+ const gpu = createMockGPUDevice({ detachWriteSources: true });
+ tvm.initWebGPU(gpu.device);
+
+ const shard = new Uint8Array([1, 2, 3, 4]).buffer;
+ const manifest = {
+ metadata: {},
+ records: [{
+ dataPath: "params.bin",
+ format: "raw-shard",
+ nbytes: shard.byteLength,
+ records: [{
+ name: "test.direct.release_before_sync",
+ shape: [1],
+ dtype: "uint32",
+ format: "raw",
+ byteOffset: 0,
+ nbytes: 4,
+ }],
+ }],
+ };
+
+ try {
+ await tvm.fetchTensorCache(
+ "https://example.test/model/";,
+ tvm.webgpu(),
+ { artifactCache: createArtifactCache(manifest, shard) },
+ );
+
+ expect(gpu.writes).toHaveLength(1);
+ expect(gpu.writes[0].data.buffer).toBe(shard);
+ expect(gpu.writes[0].data.byteLength).toBe(0);
+ expect(Array.from(gpu.writes[0].snapshot)).toEqual([1, 2, 3, 4]);
+ } finally {
+ tvm.tensorCacheClear();
+ tvm.dispose();
+ }
+});
+
+test("direct tensor-cache upload rejects a record with the wrong size", async
() => {
+ const log = jest.spyOn(console, "log").mockImplementation(() => {});
+ const tvm = createInstance();
+ const gpu = createMockGPUDevice();
+ tvm.initWebGPU(gpu.device);
+ const shard = new Uint8Array([1, 2, 3, 4]).buffer;
+ const manifest = {
+ metadata: {},
+ records: [{
+ dataPath: "params.bin",
+ format: "raw-shard",
+ nbytes: shard.byteLength,
+ records: [{
+ name: "test.direct.invalid_size",
+ shape: [2],
+ dtype: "uint32",
+ format: "raw",
+ byteOffset: 0,
+ nbytes: 4,
+ }],
+ }],
+ };
+
+ try {
+ await expect(tvm.fetchTensorCache(
+ "https://example.test/model/";,
+ tvm.webgpu(),
+ { artifactCache: createArtifactCache(manifest, shard) },
+ )).rejects.toThrow("nbytes=8");
+ expect(gpu.writes).toHaveLength(0);
+ expect(gpu.buffers).toHaveLength(1);
+ expect(gpu.buffers[0].destroy).toHaveBeenCalledTimes(1);
+ } finally {
+ tvm.tensorCacheClear();
+ tvm.dispose();
+ log.mockRestore();
+ }
+});
+
+test("CPU tensor-cache loading keeps the decoder path", async () => {
+ const tvm = createInstance();
+ const shard = new Uint8Array([
+ 1, 2, 3, 4,
+ 5, 6, 7, 8,
+ ]).buffer;
+ const manifest = {
+ metadata: {},
+ records: [{
+ dataPath: "params.bin",
+ format: "raw-shard",
+ nbytes: shard.byteLength,
+ records: [
+ {
+ name: "test.cpu.raw",
+ shape: [4],
+ dtype: "uint8",
+ format: "raw",
+ byteOffset: 0,
+ nbytes: 4,
+ },
+ {
+ name: "test.cpu.passthrough",
+ shape: [4],
+ dtype: "uint8",
+ format: "f32-to-bf16",
+ byteOffset: 4,
+ nbytes: 4,
+ },
+ ],
+ }],
+ };
+ const originalDecode = tvm.ctx.arrayDecodeStorage;
+ const decode = jest.fn((...args) => originalDecode(...args));
+ tvm.ctx.arrayDecodeStorage = decode;
+
+ try {
+ await tvm.fetchTensorCache(
+ "https://example.test/model/";,
+ tvm.cpu(),
+ { artifactCache: createArtifactCache(manifest, shard) },
+ );
+
+ expect(decode).toHaveBeenCalledTimes(2);
+ tvm.withNewScope(() => {
+ expect(Array.from(tvm.tensorCacheGet("test.cpu.raw").toRawBytes()))
+ .toEqual([1, 2, 3, 4]);
+
expect(Array.from(tvm.tensorCacheGet("test.cpu.passthrough").toRawBytes()))
+ .toEqual([5, 6, 7, 8]);
+ });
+ } finally {
+ tvm.ctx.arrayDecodeStorage = originalDecode;
+ tvm.tensorCacheClear();
+ tvm.dispose();
+ }
+});
+
+test("copyFromRawBytes uses packed tensor size and tensor byte offset", () => {
+ const tvm = createInstance();
+ const gpu = createMockGPUDevice();
+ tvm.initWebGPU(gpu.device);
+
+ tvm.withNewScope(() => {
+ const packed = tvm.empty([3], "uint4", tvm.webgpu());
+ expect(() => packed.copyFromRawBytes(new
Uint8Array(3))).toThrow("nbytes=2");
+
+ const base = tvm.empty([2], "uint32", tvm.webgpu());
+ const view = tvm.ctx.tensorCreateView(
+ base,
+ tvm.ctx.makeShapeTuple(new tvmjs.Scalar(1, "int")),
+ "uint32",
+ new tvmjs.Scalar(4, "int"),
+ );
+ view.copyFromRawBytes(new Uint8Array([21, 22, 23, 24]));
+
+ expect(gpu.writes).toHaveLength(1);
+ expect(gpu.writes[0].bufferOffset).toBe(4);
+ expect(Array.from(gpu.writes[0].snapshot)).toEqual([21, 22, 23, 24]);
+ });
+ tvm.dispose();
+});
+
+test("Tensor parses byte offsets as checked unsigned 64-bit values", () => {
+ const tvm = createInstance();
+
+ tvm.withNewScope(() => {
+ const tensor = tvm.empty([1], "uint32", tvm.cpu());
+ // DLTensor::byte_offset starts at byte 32 in the wasm32 ABI.
+ const byteOffsetPtr = tensor.dltensor + 32;
+ const words = new Uint32Array(tvm.memory.memory.buffer);
+ const wordOffset = byteOffsetPtr >>> 2;
+ const original = [words[wordOffset], words[wordOffset + 1]];
+
+ try {
+ words.set([0x80000001, 1], wordOffset);
+ const parsed = new tvmjs.Tensor(
+ tensor.dltensor,
+ tensor.lib,
+ tensor.ctx,
+ true,
+ );
+ expect(parsed.byteOffset).toBe(0x180000001);
+
+ words.set([0, 0x200000], wordOffset);
+ expect(() => new tvmjs.Tensor(
+ tensor.dltensor,
+ tensor.lib,
+ tensor.ctx,
+ true,
+ )).toThrow("Cannot represent uint64 value");
+ } finally {
+ words.set(original, wordOffset);
+ }
+ });
+ tvm.dispose();
+});
