This is an automated email from the ASF dual-hosted git repository.
akaashrp 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 f8e5f6aa29 [Web] Decode packed BF16 tensor records in place (#20167)
f8e5f6aa29 is described below
commit f8e5f6aa2916261b07c4fd4c3ac94bcf022de396
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Thu Aug 27 13:41:58 2026 -0700
[Web] Decode packed BF16 tensor records in place (#20167)
Reduce peak WASM memory during BF16 tensor-cache loading by decoding
directly in the final `float32` CPU tensor allocation. The packed BF16
input is copied into the lower half of the destination tensor and
expanded backward, preventing output writes from overwriting unread
input. This removes the separate FFI byte-array allocation previously
required during decoding. Leave the existing decoder as a fallback when
the new WASM entry point is unavailable.
---
web/emcc/wasm_runtime.cc | 88 +++++--
web/src/runtime.ts | 89 ++++++-
web/tests/node/test_bf16_decode.js | 406 +++++++++++++++++++++++++++++
web/tests/node/test_tensor_cache_webgpu.js | 11 +-
4 files changed, 565 insertions(+), 29 deletions(-)
diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc
index 15282e04a6..537e80d748 100644
--- a/web/emcc/wasm_runtime.cc
+++ b/web/emcc/wasm_runtime.cc
@@ -32,6 +32,9 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
+#include <cstring>
+#include <limits>
+
#include "src/runtime/cpu_device_api.cc"
#include "src/runtime/device_api.cc"
#include "src/runtime/extra/contrib/sort/sort.cc"
@@ -126,38 +129,91 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
+size_t GetCheckedTensorElementCount(const Tensor& tensor) {
+ size_t size = 1;
+ for (int i = 0; i < tensor->ndim; ++i) {
+ TVM_FFI_ICHECK_GE(tensor->shape[i], 0);
+ const uint64_t dim_u64 = static_cast<uint64_t>(tensor->shape[i]);
+ TVM_FFI_ICHECK_LE(dim_u64, std::numeric_limits<size_t>::max());
+ const size_t dim = static_cast<size_t>(dim_u64);
+ if (size != 0) {
+ TVM_FFI_ICHECK_LE(dim, std::numeric_limits<size_t>::max() / size);
+ }
+ size *= dim;
+ }
+ return size;
+}
+
+void CheckF32CPUTensor(const Tensor& tensor) {
+ TVM_FFI_ICHECK_EQ(tensor->device.device_type, kDLCPU);
+ TVM_FFI_ICHECK(tensor.IsContiguous());
+ TVM_FFI_ICHECK_EQ(tensor->dtype.code, kDLFloat);
+ TVM_FFI_ICHECK_EQ(tensor->dtype.bits, 32);
+ TVM_FFI_ICHECK_EQ(tensor->dtype.lanes, 1);
+}
+
void ArrayDecodeStorage(Tensor cpu_arr, TVMFFIByteArray* bytes, const
std::string& format,
const std::string& dtype) {
TVM_FFI_ICHECK_NE(bytes, nullptr);
const char* byte_data = bytes->data;
const size_t byte_size = bytes->size;
if (format == "f32-to-bf16" && dtype == "float32") {
- const uint16_t* bf16 = reinterpret_cast<const uint16_t*>(byte_data);
- uint32_t* data = static_cast<uint32_t*>(cpu_arr->data);
- TVM_FFI_ICHECK(cpu_arr.IsContiguous());
- size_t size = 1;
- for (int i = 0; i < cpu_arr->ndim; ++i) {
- size *= cpu_arr->shape[i];
+ CheckF32CPUTensor(cpu_arr);
+ const size_t size = GetCheckedTensorElementCount(cpu_arr);
+ TVM_FFI_ICHECK_LE(size, std::numeric_limits<size_t>::max() / 4);
+ TVM_FFI_ICHECK_EQ(byte_size, size * 2);
+ if (size == 0) {
+ return;
}
- TVM_FFI_ICHECK_EQ(size, byte_size / 2);
+ TVM_FFI_ICHECK_NE(cpu_arr->data, nullptr);
+ TVM_FFI_ICHECK_NE(byte_data, nullptr);
+ const uint8_t* bf16 = reinterpret_cast<const uint8_t*>(byte_data);
+ uint8_t* data = static_cast<uint8_t*>(cpu_arr->data) +
cpu_arr->byte_offset;
for (size_t i = 0; i < size; ++i) {
- data[i] = static_cast<uint32_t>(bf16[i]) << 16;
+ uint16_t bf16_bits;
+ std::memcpy(&bf16_bits, bf16 + 2 * i, sizeof(bf16_bits));
+ const uint32_t f32_bits = static_cast<uint32_t>(bf16_bits) << 16;
+ std::memcpy(data + 4 * i, &f32_bits, sizeof(f32_bits));
}
} else {
cpu_arr.CopyFromBytes(byte_data, byte_size);
}
}
+void ArrayDecodeBF16ToF32Inplace(Tensor cpu_arr, int64_t encoded_nbytes) {
+ CheckF32CPUTensor(cpu_arr);
+ TVM_FFI_ICHECK_GE(encoded_nbytes, 0);
+
+ const size_t size = GetCheckedTensorElementCount(cpu_arr);
+ TVM_FFI_ICHECK_LE(size, std::numeric_limits<size_t>::max() / 4);
+ TVM_FFI_ICHECK_EQ(static_cast<uint64_t>(encoded_nbytes),
static_cast<uint64_t>(size) * 2);
+ if (size == 0) {
+ return;
+ }
+ TVM_FFI_ICHECK_NE(cpu_arr->data, nullptr);
+
+ uint8_t* data = static_cast<uint8_t*>(cpu_arr->data) + cpu_arr->byte_offset;
+ for (size_t i = size; i != 0; --i) {
+ const size_t j = i - 1;
+ uint16_t bf16_bits;
+ std::memcpy(&bf16_bits, data + 2 * j, sizeof(bf16_bits));
+ const uint32_t f32_bits = static_cast<uint32_t>(bf16_bits) << 16;
+ std::memcpy(data + 4 * j, &f32_bits, sizeof(f32_bits));
+ }
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def_packed(
- "tvmjs.array.decode_storage", [](ffi::PackedArgs args, ffi::Any* ret) {
- Tensor cpu_arr = args[0].cast<Tensor>();
- TVMFFIByteArray* bytes = args[1].cast<TVMFFIByteArray*>();
- std::string format = args[2].cast<ffi::String>().operator
std::string();
- std::string dtype = args[3].cast<ffi::String>().operator std::string();
- ArrayDecodeStorage(cpu_arr, bytes, format, dtype);
- });
+ refl::GlobalDef()
+ .def_packed("tvmjs.array.decode_storage",
+ [](ffi::PackedArgs args, ffi::Any* ret) {
+ Tensor cpu_arr = args[0].cast<Tensor>();
+ TVMFFIByteArray* bytes = args[1].cast<TVMFFIByteArray*>();
+ std::string format = args[2].cast<ffi::String>().operator
std::string();
+ std::string dtype = args[3].cast<ffi::String>().operator
std::string();
+ ArrayDecodeStorage(cpu_arr, bytes, format, dtype);
+ })
+ .def("tvmjs.array.decode_bf16_to_f32_inplace",
ArrayDecodeBF16ToF32Inplace);
}
// Concatenate n TVMArrays
diff --git a/web/src/runtime.ts b/web/src/runtime.ts
index 1ce05b7dde..abc49be594 100644
--- a/web/src/runtime.ts
+++ b/web/src/runtime.ts
@@ -179,6 +179,7 @@ class RuntimeContext implements Disposable {
tensorCacheRemove: PackedFunc;
tensorCacheClear: PackedFunc;
arrayDecodeStorage: PackedFunc;
+ arrayDecodeBF16ToF32Inplace: PackedFunc | undefined;
paramModuleFromCache: PackedFunc;
paramModuleFromCacheByName: PackedFunc;
makeShapeTuple: PackedFunc;
@@ -214,6 +215,7 @@ class RuntimeContext implements Disposable {
this.tensorCacheUpdate = getGlobalFunc("vm.builtin.tensor_cache.update");
this.tensorCacheClear = getGlobalFunc("vm.builtin.tensor_cache.clear");
this.arrayDecodeStorage = getGlobalFunc("tvmjs.array.decode_storage");
+ this.arrayDecodeBF16ToF32Inplace = undefined;
this.paramModuleFromCache =
getGlobalFunc("vm.builtin.param_module_from_cache");
this.paramModuleFromCacheByName =
getGlobalFunc("vm.builtin.param_module_from_cache_by_name");
this.makeShapeTuple = getGlobalFunc("ffi.Shape");
@@ -237,6 +239,7 @@ class RuntimeContext implements Disposable {
this.tensorCacheRemove.dispose();
this.tensorCacheUpdate.dispose();
this.arrayDecodeStorage.dispose();
+ this.arrayDecodeBF16ToF32Inplace?.dispose();
this.paramModuleFromCache.dispose();
this.paramModuleFromCacheByName.dispose();
this.makeShapeTuple.dispose();
@@ -606,6 +609,22 @@ export class Tensor extends TVMObject {
return this.dataPtr;
}
+ /**
+ * Return the effective address of a CPU Tensor's storage.
+ * @returns The address in Wasm linear memory.
+ * @internal
+ */
+ getCPUDataAddress(): Pointer {
+ if (this.device.deviceType !== DeviceStrToEnum.cpu) {
+ throw new Error("Can only obtain a linear-memory address for a CPU
Tensor");
+ }
+ const address = this.getDataPtr() + this.byteOffset;
+ if (!Number.isSafeInteger(address) || address < 0) {
+ throw new Error("Invalid CPU Tensor storage address");
+ }
+ return address;
+ }
+
/**
* Copy data from another Tensor or javascript array.
* The number of elements must match.
@@ -953,6 +972,10 @@ export class Instance implements Disposable {
return this.getGlobalFuncInternal(name, autoAttachToScope);
}
);
+ this.ctx.arrayDecodeBF16ToF32Inplace = this.getGlobalFuncInternalOptional(
+ "tvmjs.array.decode_bf16_to_f32_inplace",
+ /*autoAttachToScope=*/ false,
+ );
this.registerEnvGlobalPackedFuncs();
this.registerObjectFactoryFuncs();
this.rng = new LinearCongruentialGenerator();
@@ -1156,6 +1179,17 @@ export class Instance implements Disposable {
}
private getGlobalFuncInternal(name: string, autoAttachToScope = true):
PackedFunc {
+ const ret = this.getGlobalFuncInternalOptional(name, autoAttachToScope);
+ if (ret === undefined) {
+ throw Error("Cannot find global function " + name);
+ }
+ return ret;
+ }
+
+ private getGlobalFuncInternalOptional(
+ name: string,
+ autoAttachToScope = true,
+ ): PackedFunc | undefined {
const stack = this.lib.getOrAllocCallStack();
const nameOffset = stack.allocByteArrayForString(name);
const outOffset = stack.allocPtrArray(1);
@@ -1172,7 +1206,7 @@ export class Instance implements Disposable {
const handle = this.memory.loadPointer(outPtr);
this.lib.recycleCallStack(stack);
if (handle === 0) {
- throw Error("Cannot find global function " + name);
+ return undefined;
}
const ret = this.makePackedFunc(handle);
if (autoAttachToScope) this.ctx.attachToCurrentScope(ret);
@@ -1357,12 +1391,22 @@ export class Instance implements Disposable {
): void {
const recBytes = getTensorCacheRecordBytes(shardBytes, rec);
if (cpuArray !== undefined) {
- this.ctx.arrayDecodeStorage(
- cpuArray,
- recBytes,
- rec.format,
- rec.dtype,
- );
+ const isPackedBF16 =
+ rec.format === "f32-to-bf16" && rec.dtype === "float32";
+ if (isPackedBF16 && this.ctx.arrayDecodeBF16ToF32Inplace !== undefined) {
+ this.memory.storeRawBytes(cpuArray.getCPUDataAddress(), recBytes);
+ this.ctx.arrayDecodeBF16ToF32Inplace(
+ cpuArray,
+ new Scalar(recBytes.byteLength, "int64"),
+ );
+ } else {
+ this.ctx.arrayDecodeStorage(
+ cpuArray,
+ recBytes,
+ rec.format,
+ rec.dtype,
+ );
+ }
}
if (gpuArray !== undefined) {
if (cpuArray === undefined) {
@@ -1373,6 +1417,24 @@ export class Instance implements Disposable {
}
}
+ /** Return the exact byte size of a packed BF16 tensor. */
+ private getPackedBF16Bytes(shape: Array<number>): number {
+ let numElements = 1;
+ for (const dim of shape) {
+ if (!Number.isSafeInteger(dim) || dim < 0) {
+ throw new Error(`Invalid tensor dimension: ${dim}`);
+ }
+ numElements *= dim;
+ if (!Number.isSafeInteger(numElements)) {
+ throw new Error("Tensor element count exceeds JavaScript's safe
integer range");
+ }
+ }
+ const nbytes = numElements * 2;
+ if (!Number.isSafeInteger(nbytes)) {
+ throw new Error("Packed BF16 size exceeds JavaScript's safe integer
range");
+ }
+ return nbytes;
+ }
/**
* Fetch list of Tensor into the TensorCache.
@@ -1487,11 +1549,20 @@ export class Instance implements Disposable {
let gpu_arr: Tensor | undefined;
try {
const rec = shardRecords[j];
- const requiresDecode =
+ const isPackedBF16 =
rec.format === "f32-to-bf16" && rec.dtype === "float32";
+ if (isPackedBF16) {
+ const expectedBytes = this.getPackedBF16Bytes(rec.shape);
+ if (rec.nbytes !== expectedBytes) {
+ throw new Error(
+ `Packed BF16 record has ${rec.nbytes} bytes, ` +
+ `but shape requires ${expectedBytes}`,
+ );
+ }
+ }
const directToWebGPU =
device.deviceType === DeviceStrToEnum.webgpu &&
- !requiresDecode &&
+ !isPackedBF16 &&
rec.nbytes % 4 === 0;
if (!directToWebGPU) {
diff --git a/web/tests/node/test_bf16_decode.js
b/web/tests/node/test_bf16_decode.js
new file mode 100644
index 0000000000..cad335c453
--- /dev/null
+++ b/web/tests/node/test_bf16_decode.js
@@ -0,0 +1,406 @@
+/*
+ * 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");
+
+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 encodeBF16(bits) {
+ const bytes = new Uint8Array(bits.length * 2);
+ for (let i = 0; i < bits.length; ++i) {
+ bytes[2 * i] = bits[i] & 0xff;
+ bytes[2 * i + 1] = bits[i] >>> 8;
+ }
+ return bytes;
+}
+
+function decodedF32Bytes(bits) {
+ const bytes = new Uint8Array(bits.length * 4);
+ for (let i = 0; i < bits.length; ++i) {
+ bytes[4 * i + 2] = bits[i] & 0xff;
+ bytes[4 * i + 3] = bits[i] >>> 8;
+ }
+ return bytes;
+}
+
+function createArtifactCache(manifest, shard) {
+ return {
+ hasAllKeys: async () => true,
+ addToCache: async () => {},
+ deleteInCache: async () => {},
+ fetchWithCache: async (_url, storeType) => {
+ return storeType === "json" ? manifest : shard;
+ },
+ };
+}
+
+let tvm;
+
+beforeAll(() => {
+ tvm = createInstance();
+});
+
+afterAll(() => {
+ tvm.dispose();
+});
+
+test("in-place BF16 expansion preserves exact floating-point bits", () => {
+ const decode = tvm.ctx.arrayDecodeBF16ToF32Inplace;
+ const bits = [
+ 0x0000,
+ 0x8000,
+ 0x3f80,
+ 0xc020,
+ 0x0001,
+ 0x7f80,
+ 0xff80,
+ 0x7fc1,
+ ];
+ const encoded = encodeBF16(bits);
+
+ expect(decode).toBeDefined();
+ tvm.withNewScope(() => {
+ const tensor = tvm.empty([bits.length], "float32");
+ tvm.memory.storeRawBytes(tensor.getCPUDataAddress(), encoded);
+ decode(tensor, new tvmjs.Scalar(encoded.byteLength, "int64"));
+ expect(Array.from(tensor.toRawBytes())).toEqual(
+ Array.from(decodedF32Bytes(bits)),
+ );
+ });
+});
+
+test.each([
+ [[], [0x3f80]],
+ [[0], []],
+ [[1], [0x8000]],
+ [[3], [0x3f80, 0xc000, 0x7fc1]],
+ [[2, 2], [0x0001, 0x3f00, 0x7f80, 0xff80]],
+])("in-place BF16 expansion supports shape %j", (shape, bits) => {
+ const encoded = encodeBF16(bits);
+
+ tvm.withNewScope(() => {
+ const tensor = tvm.empty(shape, "float32");
+ tvm.memory.storeRawBytes(tensor.getCPUDataAddress(), encoded);
+ tvm.ctx.arrayDecodeBF16ToF32Inplace(
+ tensor,
+ new tvmjs.Scalar(encoded.byteLength, "int64"),
+ );
+ expect(Array.from(tensor.toRawBytes())).toEqual(
+ Array.from(decodedF32Bytes(bits)),
+ );
+ });
+});
+
+test("both BF16 decoders honor an unaligned tensor byte offset", () => {
+ const bits = [0x3f80, 0xc000, 0x7fc1];
+ const encoded = encodeBF16(bits);
+ const expected = decodedF32Bytes(bits);
+
+ tvm.withNewScope(() => {
+ for (const inplace of [true, false]) {
+ const base = tvm.empty([expected.byteLength + 2], "uint8");
+ base.copyFromRawBytes(
+ new Uint8Array(expected.byteLength + 2).fill(0xa5),
+ );
+ const view = tvm.ctx.tensorCreateView(
+ base,
+ tvm.ctx.makeShapeTuple(new tvmjs.Scalar(bits.length, "int")),
+ "float32",
+ new tvmjs.Scalar(1, "int64"),
+ );
+ if (inplace) {
+ tvm.memory.storeRawBytes(view.getCPUDataAddress(), encoded);
+ tvm.ctx.arrayDecodeBF16ToF32Inplace(
+ view,
+ new tvmjs.Scalar(encoded.byteLength, "int64"),
+ );
+ } else {
+ tvm.ctx.arrayDecodeStorage(
+ view,
+ encoded,
+ "f32-to-bf16",
+ "float32",
+ );
+ }
+
+ const result = base.toRawBytes();
+ expect(result[0]).toBe(0xa5);
+ expect(Array.from(result.subarray(1, 1 + expected.byteLength))).toEqual(
+ Array.from(expected),
+ );
+ expect(result[result.length - 1]).toBe(0xa5);
+ }
+ });
+});
+
+test("in-place BF16 expansion rejects invalid tensor contracts", () => {
+ const error = jest.spyOn(console, "error").mockImplementation(() => {});
+ const originalExitCode = process.exitCode;
+ const tvm = createInstance();
+
+ try {
+ tvm.withNewScope(() => {
+ const tensor = tvm.empty([2], "float32");
+ const encoded = encodeBF16([0x3f80, 0x4000]);
+ tvm.memory.storeRawBytes(tensor.getCPUDataAddress(), encoded);
+ expect(() => tvm.ctx.arrayDecodeBF16ToF32Inplace(
+ tensor,
+ new tvmjs.Scalar(encoded.byteLength - 1, "int64"),
+ )).toThrow();
+ expect(() => tvm.ctx.arrayDecodeBF16ToF32Inplace(
+ tensor,
+ new tvmjs.Scalar(encoded.byteLength + 1, "int64"),
+ )).toThrow();
+ expect(() => tvm.ctx.arrayDecodeBF16ToF32Inplace(
+ tensor,
+ new tvmjs.Scalar(-1, "int64"),
+ )).toThrow();
+
+ const wrongDtype = tvm.empty([2], "int32");
+ expect(() => tvm.ctx.arrayDecodeBF16ToF32Inplace(
+ wrongDtype,
+ new tvmjs.Scalar(encoded.byteLength, "int64"),
+ )).toThrow();
+
+ expect(() => tvm.ctx.arrayDecodeStorage(
+ tensor,
+ encoded.subarray(0, encoded.byteLength - 1),
+ "f32-to-bf16",
+ "float32",
+ )).toThrow();
+ expect(() => tvm.ctx.arrayDecodeStorage(
+ tensor,
+ new Uint8Array(encoded.byteLength * 2),
+ "f32-to-bf16",
+ "float32",
+ )).toThrow();
+ expect(() => tvm.ctx.arrayDecodeStorage(
+ wrongDtype,
+ encoded,
+ "f32-to-bf16",
+ "float32",
+ )).toThrow();
+ });
+ } finally {
+ tvm.dispose();
+ error.mockRestore();
+ // Emscripten marks expected native contract failures as process failures.
+ process.exitCode = originalExitCode;
+ }
+});
+
+test("tensor cache uses in-place expansion only for packed BF16 records",
async () => {
+ const packedBits = [0x3f80, 0xc000];
+ const packed = encodeBF16(packedBits);
+ const shard = Uint8Array.from([1, 2, 3, 4, ...packed]).buffer;
+ const manifest = {
+ metadata: {},
+ records: [{
+ dataPath: "params.bin",
+ format: "raw-shard",
+ nbytes: shard.byteLength,
+ records: [
+ {
+ name: "test.decode.raw",
+ shape: [4],
+ dtype: "uint8",
+ format: "raw",
+ byteOffset: 0,
+ nbytes: 4,
+ },
+ {
+ name: "test.decode.packed_bf16",
+ shape: [2],
+ dtype: "float32",
+ format: "f32-to-bf16",
+ byteOffset: 4,
+ nbytes: packed.byteLength,
+ },
+ ],
+ }],
+ };
+ const originalInplace = tvm.ctx.arrayDecodeBF16ToF32Inplace;
+ const originalStorage = tvm.ctx.arrayDecodeStorage;
+ const inplace = jest.fn((...args) => originalInplace(...args));
+ const storage = jest.fn((...args) => originalStorage(...args));
+ tvm.ctx.arrayDecodeBF16ToF32Inplace = inplace;
+ tvm.ctx.arrayDecodeStorage = storage;
+
+ try {
+ await tvm.fetchTensorCache(
+ "https://example.test/model/";,
+ tvm.cpu(),
+ { artifactCache: createArtifactCache(manifest, shard) },
+ );
+
+ expect(inplace).toHaveBeenCalledTimes(1);
+ expect(storage).toHaveBeenCalledTimes(1);
+ tvm.withNewScope(() => {
+ expect(Array.from(tvm.tensorCacheGet("test.decode.raw").toRawBytes()))
+ .toEqual([1, 2, 3, 4]);
+ expect(Array.from(
+ tvm.tensorCacheGet("test.decode.packed_bf16").toRawBytes(),
+ )).toEqual(Array.from(decodedF32Bytes(packedBits)));
+ });
+ } finally {
+ tvm.ctx.arrayDecodeBF16ToF32Inplace = originalInplace;
+ tvm.ctx.arrayDecodeStorage = originalStorage;
+ tvm.tensorCacheClear();
+ }
+});
+
+test("tensor cache falls back when in-place expansion is unavailable", async
() => {
+ const bits = [0x3f80, 0xc000];
+ const packed = encodeBF16(bits);
+ const manifest = {
+ metadata: {},
+ records: [{
+ dataPath: "params.bin",
+ format: "raw-shard",
+ nbytes: packed.byteLength,
+ records: [{
+ name: "test.decode.fallback",
+ shape: [2],
+ dtype: "float32",
+ format: "f32-to-bf16",
+ byteOffset: 0,
+ nbytes: packed.byteLength,
+ }],
+ }],
+ };
+ const originalInplace = tvm.ctx.arrayDecodeBF16ToF32Inplace;
+ const originalStorage = tvm.ctx.arrayDecodeStorage;
+ const storage = jest.fn((...args) => originalStorage(...args));
+ tvm.ctx.arrayDecodeBF16ToF32Inplace = undefined;
+ tvm.ctx.arrayDecodeStorage = storage;
+
+ try {
+ await tvm.fetchTensorCache(
+ "https://example.test/model/";,
+ tvm.cpu(),
+ { artifactCache: createArtifactCache(manifest, packed.buffer) },
+ );
+ expect(storage).toHaveBeenCalledTimes(1);
+ tvm.withNewScope(() => {
+ expect(Array.from(
+ tvm.tensorCacheGet("test.decode.fallback").toRawBytes(),
+ )).toEqual(Array.from(decodedF32Bytes(bits)));
+ });
+ } finally {
+ tvm.ctx.arrayDecodeBF16ToF32Inplace = originalInplace;
+ tvm.ctx.arrayDecodeStorage = originalStorage;
+ tvm.tensorCacheClear();
+ }
+});
+
+test("tensor cache disposes its CPU tensor when in-place expansion fails",
async () => {
+ const packed = encodeBF16([0x3f80, 0xc000]);
+ const manifest = {
+ metadata: {},
+ records: [{
+ dataPath: "params.bin",
+ format: "raw-shard",
+ nbytes: packed.byteLength,
+ records: [{
+ name: "test.decode.failure",
+ shape: [2],
+ dtype: "float32",
+ format: "f32-to-bf16",
+ byteOffset: 0,
+ nbytes: packed.byteLength,
+ }],
+ }],
+ };
+ const originalLogger = tvm.env.logger;
+ const originalInplace = tvm.ctx.arrayDecodeBF16ToF32Inplace;
+ const empty = jest.spyOn(tvm, "empty");
+ tvm.env.logger = () => {};
+ tvm.ctx.arrayDecodeBF16ToF32Inplace = jest.fn(() => {
+ throw new Error("in-place decode failed");
+ });
+
+ try {
+ await expect(tvm.fetchTensorCache(
+ "https://example.test/model/";,
+ tvm.cpu(),
+ { artifactCache: createArtifactCache(manifest, packed.buffer) },
+ )).rejects.toThrow("in-place decode failed");
+ expect(empty).toHaveBeenCalledTimes(1);
+ expect(empty.mock.results[0].value.getHandle(false)).toBe(0);
+ } finally {
+ empty.mockRestore();
+ tvm.ctx.arrayDecodeBF16ToF32Inplace = originalInplace;
+ tvm.env.logger = originalLogger;
+ }
+});
+
+test.each([
+ [[2], 8, "shape requires 4"],
+ [[1.5], 3, "Invalid tensor dimension"],
+ [[Number.MAX_SAFE_INTEGER, 2], 0, "safe integer range"],
+])(
+ "tensor cache rejects malformed packed BF16 metadata for shape %j",
+ async (shape, nbytes, message) => {
+ const originalLogger = tvm.env.logger;
+ tvm.env.logger = () => {};
+ const empty = jest.spyOn(tvm, "empty");
+ const shard = new ArrayBuffer(nbytes);
+ const slice = jest.spyOn(shard, "slice");
+ const manifest = {
+ metadata: {},
+ records: [{
+ dataPath: "params.bin",
+ format: "raw-shard",
+ nbytes,
+ records: [{
+ name: "test.decode.invalid",
+ shape,
+ dtype: "float32",
+ format: "f32-to-bf16",
+ byteOffset: 0,
+ nbytes,
+ }],
+ }],
+ };
+
+ try {
+ await expect(tvm.fetchTensorCache(
+ "https://example.test/model/";,
+ tvm.cpu(),
+ { artifactCache: createArtifactCache(manifest, shard) },
+ )).rejects.toThrow(message);
+ expect(empty).not.toHaveBeenCalled();
+ expect(slice).not.toHaveBeenCalled();
+ } finally {
+ slice.mockRestore();
+ empty.mockRestore();
+ tvm.env.logger = originalLogger;
+ }
+ },
+);
diff --git a/web/tests/node/test_tensor_cache_webgpu.js
b/web/tests/node/test_tensor_cache_webgpu.js
index 9df60f4230..3383a2e41e 100644
--- a/web/tests/node/test_tensor_cache_webgpu.js
+++ b/web/tests/node/test_tensor_cache_webgpu.js
@@ -94,7 +94,7 @@ function createArtifactCache(manifest, shard) {
};
}
-test("WebGPU tensor cache uploads pass-through records directly", async () => {
+test("WebGPU tensor cache uploads pass-through records and decodes BF16 in
place", async () => {
const tvm = createInstance();
const gpu = createMockGPUDevice();
tvm.initWebGPU(gpu.device);
@@ -150,8 +150,11 @@ test("WebGPU tensor cache uploads pass-through records
directly", async () => {
};
const originalDecode = tvm.ctx.arrayDecodeStorage;
+ const originalInplace = tvm.ctx.arrayDecodeBF16ToF32Inplace;
const decode = jest.fn((...args) => originalDecode(...args));
+ const inplace = jest.fn((...args) => originalInplace(...args));
tvm.ctx.arrayDecodeStorage = decode;
+ tvm.ctx.arrayDecodeBF16ToF32Inplace = inplace;
try {
await tvm.fetchTensorCache(
"https://example.test/model/";,
@@ -159,9 +162,8 @@ test("WebGPU tensor cache uploads pass-through records
directly", async () => {
{ 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(decode).not.toHaveBeenCalled();
+ expect(inplace).toHaveBeenCalledTimes(1);
expect(gpu.writes.map((write) => Array.from(write.snapshot))).toEqual([
[1, 2, 3, 4, 5, 6, 7, 8],
[9, 10, 11, 12],
@@ -170,6 +172,7 @@ test("WebGPU tensor cache uploads pass-through records
directly", async () => {
]);
} finally {
tvm.ctx.arrayDecodeStorage = originalDecode;
+ tvm.ctx.arrayDecodeBF16ToF32Inplace = originalInplace;
tvm.tensorCacheClear();
tvm.dispose();
}
