From d6cbf31f373f754dbe448c7e4e61075f1e455a0c Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Sun, 16 Aug 2026 08:26:18 -0400 Subject: [PATCH 1/2] [Web] Decode packed BF16 tensor records in place Packed BF16 decoding currently uses separate input and output storage. Add a Wasm entry point that expands values backward in the destination buffer after the packed bytes have been copied into its lower half. Use the in-place path when the entry point is available and keep the existing decoder otherwise. Both paths account for DLTensor byte_offset. --- web/emcc/wasm_runtime.cc | 89 ++++- 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, 566 insertions(+), 29 deletions(-) create mode 100644 web/tests/node/test_bf16_decode.js diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc index 15282e04a6ed..5f9cc25146af 100644 --- a/web/emcc/wasm_runtime.cc +++ b/web/emcc/wasm_runtime.cc @@ -32,6 +32,8 @@ #include #include +#include + #include "src/runtime/cpu_device_api.cc" #include "src/runtime/device_api.cc" #include "src/runtime/extra/contrib/sort/sort.cc" @@ -126,38 +128,93 @@ 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(tensor->shape[i]); + TVM_FFI_ICHECK_LE(dim_u64, std::numeric_limits::max()); + const size_t dim = static_cast(dim_u64); + if (size != 0) { + TVM_FFI_ICHECK_LE(dim, std::numeric_limits::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(byte_data); - uint32_t* data = static_cast(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::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(byte_data); + uint8_t* data = static_cast(cpu_arr->data) + cpu_arr->byte_offset; for (size_t i = 0; i < size; ++i) { - data[i] = static_cast(bf16[i]) << 16; + data[4 * i] = 0; + data[4 * i + 1] = 0; + data[4 * i + 2] = bf16[2 * i]; + data[4 * i + 3] = bf16[2 * i + 1]; } } 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::max() / 4); + TVM_FFI_ICHECK_EQ(static_cast(encoded_nbytes), static_cast(size) * 2); + if (size == 0) { + return; + } + TVM_FFI_ICHECK_NE(cpu_arr->data, nullptr); + + uint8_t* data = static_cast(cpu_arr->data) + cpu_arr->byte_offset; + for (size_t i = size; i != 0; --i) { + const size_t j = i - 1; + const uint8_t low = data[2 * j]; + const uint8_t high = data[2 * j + 1]; + data[4 * j] = 0; + data[4 * j + 1] = 0; + data[4 * j + 2] = low; + data[4 * j + 3] = high; + } +} + 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(); - TVMFFIByteArray* bytes = args[1].cast(); - std::string format = args[2].cast().operator std::string(); - std::string dtype = args[3].cast().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(); + TVMFFIByteArray* bytes = args[1].cast(); + std::string format = args[2].cast().operator std::string(); + std::string dtype = args[3].cast().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 1ce05b7dde40..abc49be59490 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 { + 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 000000000000..cad335c453fc --- /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 9df60f42308a..3383a2e41eeb 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(); } From 98b4c67238799b42d5d5f65702bce26203e00cba Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Wed, 26 Aug 2026 00:26:41 -0400 Subject: [PATCH 2/2] [WEB] Use 32-bit stores for in-place BF16 decode --- web/emcc/wasm_runtime.cc | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc index 5f9cc25146af..159c0780294f 100644 --- a/web/emcc/wasm_runtime.cc +++ b/web/emcc/wasm_runtime.cc @@ -32,6 +32,7 @@ #include #include +#include #include #include "src/runtime/cpu_device_api.cc" @@ -194,12 +195,10 @@ void ArrayDecodeBF16ToF32Inplace(Tensor cpu_arr, int64_t encoded_nbytes) { uint8_t* data = static_cast(cpu_arr->data) + cpu_arr->byte_offset; for (size_t i = size; i != 0; --i) { const size_t j = i - 1; - const uint8_t low = data[2 * j]; - const uint8_t high = data[2 * j + 1]; - data[4 * j] = 0; - data[4 * j + 1] = 0; - data[4 * j + 2] = low; - data[4 * j + 3] = high; + uint16_t bf16_bits; + std::memcpy(&bf16_bits, data + 2 * j, sizeof(bf16_bits)); + const uint32_t f32_bits = static_cast(bf16_bits) << 16; + std::memcpy(data + 4 * j, &f32_bits, sizeof(f32_bits)); } }