From 56f259ee0a31ceea82b36048713a4998c1955202 Mon Sep 17 00:00:00 2001 From: MakotoUwu Date: Sun, 14 Jun 2026 22:24:35 +0200 Subject: [PATCH 1/4] [Web] Improve large tensor loading in wasm runtime --- web/emcc/wasm_runtime.cc | 39 ++++++++--- web/src/runtime.ts | 146 +++++++++++++++++++++++++++++++++++++-- 2 files changed, 169 insertions(+), 16 deletions(-) diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc index 9d3d46f18cb4..1286a514ce2f 100644 --- a/web/emcc/wasm_runtime.cc +++ b/web/emcc/wasm_runtime.cc @@ -130,32 +130,51 @@ void ArrayDecodeStorage(Tensor cpu_arr, TVMFFIByteArray* bytes, const std::strin 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]; } - TVM_FFI_ICHECK_EQ(size, byte_size / 2); - for (size_t i = 0; i < size; ++i) { - data[i] = static_cast(bf16[i]) << 16; + // The "f32-to-bf16" format encodes a float32 tensor as packed bf16 (2 + // bytes per element). When the byte_size matches that expectation, expand + // back to f32. If the byte_size matches the native float32 width + // (4 bytes per element), the payload is already raw float32; fall through + // to the generic byte copy. This makes the loader tolerant of weight + // shards produced by older / alternate quantisation pipelines that retain + // the "f32-to-bf16" tag without performing the bf16 truncation. + if (byte_size == size * sizeof(uint16_t)) { + const uint16_t* bf16 = reinterpret_cast(byte_data); + uint32_t* data = + reinterpret_cast(static_cast(cpu_arr->data) + cpu_arr->byte_offset); + for (size_t i = 0; i < size; ++i) { + data[i] = static_cast(bf16[i]) << 16; + } + return; } - } else { - cpu_arr.CopyFromBytes(byte_data, byte_size); } + cpu_arr.CopyFromBytes(byte_data, byte_size); +} + +int64_t StorageSizeBytes(int64_t num_elements, const std::string& dtype) { + TVM_FFI_ICHECK_GE(num_elements, 0); + TVMFFIByteArray dtype_bytes{dtype.data(), dtype.size()}; + DLDataType dl_dtype; + TVM_FFI_ICHECK_EQ(TVMFFIDataTypeFromString(&dtype_bytes, &dl_dtype), 0); + return static_cast( + ffi::GetDataSize(static_cast(num_elements), dl_dtype)); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def_packed( - "tvmjs.array.decode_storage", [](ffi::PackedArgs args, ffi::Any* ret) { + 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.runtime.StorageSizeBytes", StorageSizeBytes); } // Concatenate n TVMArrays diff --git a/web/src/runtime.ts b/web/src/runtime.ts index 078a0c7df21f..745f6fa7c8ea 100644 --- a/web/src/runtime.ts +++ b/web/src/runtime.ts @@ -172,6 +172,7 @@ class RuntimeContext implements Disposable { tensorCacheRemove: PackedFunc; tensorCacheClear: PackedFunc; arrayDecodeStorage: PackedFunc; + storageSizeBytes: PackedFunc; paramModuleFromCache: PackedFunc; paramModuleFromCacheByName: PackedFunc; makeShapeTuple: PackedFunc; @@ -207,6 +208,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.storageSizeBytes = getGlobalFunc("tvmjs.runtime.StorageSizeBytes"); 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"); @@ -230,6 +232,7 @@ class RuntimeContext implements Disposable { this.tensorCacheRemove.dispose(); this.tensorCacheUpdate.dispose(); this.arrayDecodeStorage.dispose(); + this.storageSizeBytes.dispose(); this.paramModuleFromCache.dispose(); this.paramModuleFromCacheByName.dispose(); this.makeShapeTuple.dispose(); @@ -1010,9 +1013,11 @@ export class Instance implements Disposable { */ withNewScope(action: () => T): T { this.beginScope(); - const val = action(); - this.endScope(); - return val; + try { + return action(); + } finally { + this.endScope(); + } } /** @@ -1323,6 +1328,19 @@ export class Instance implements Disposable { artifactCache: ArtifactCacheTemplate, signal?: AbortSignal, ) { + // Avoid a single JS-to-wasm byte-array call for multi-hundred-MiB + // tensor-cache records. The cap is a conservative per-call staging size, + // independent of the final tensor allocation size. Smaller records keep + // the existing full-record path. + const maxChunkBytes = 128 * 1024 * 1024; + const storageSizeBytes = (numElements: number, dtype: string): number | undefined => { + try { + return this.ctx.storageSizeBytes(new Scalar(numElements, "int"), dtype) as number; + } catch { + // Unknown dtypes can still use the original full-record loading path. + return undefined; + } + }; const perf = compact.getPerformance(); const tstart = perf.now(); let totalBytes = 0; @@ -1421,9 +1439,68 @@ export class Instance implements Disposable { this.empty(rec.shape, rec.dtype, this.cpu()) ) }); - const recSource = buffer.slice(rec.byteOffset, rec.byteOffset + rec.nbytes); + const shardBytes = buffer instanceof Uint8Array ? buffer : new Uint8Array(buffer); + const recSource = + rec.byteOffset === 0 && rec.nbytes === shardBytes.byteLength + ? shardBytes + : shardBytes.subarray(rec.byteOffset, rec.byteOffset + rec.nbytes); + let canChunkRecord = + rec.nbytes > maxChunkBytes && + rec.shape.length >= 1 && + Number.isInteger(rec.shape[0]) && + rec.shape[0] > 0 && + rec.nbytes % rec.shape[0] === 0; + const outerDim = canChunkRecord ? rec.shape[0] : 1; + const sourceStrideBytes = canChunkRecord ? rec.nbytes / outerDim : rec.nbytes; + let targetStrideBytes = 0; + if (canChunkRecord) { + const numElements = rec.shape.reduce((acc, value) => acc * value, 1); + const targetBytes = storageSizeBytes(numElements, rec.dtype); + canChunkRecord = + sourceStrideBytes <= maxChunkBytes && + targetBytes !== undefined && + targetBytes % outerDim === 0; + if (canChunkRecord) { + targetStrideBytes = targetBytes / outerDim; + } + } + const copyRecordToTensor = (targetTensor: Tensor, sourceBytes: Uint8Array) => { + if (!canChunkRecord) { + this.ctx.arrayDecodeStorage(targetTensor, sourceBytes, rec.format, rec.dtype); + return; + } + const chunkOuterDim = Math.max(1, Math.floor(maxChunkBytes / sourceStrideBytes)); + for (let outerOffset = 0; outerOffset < outerDim; outerOffset += chunkOuterDim) { + const outerCount = Math.min(chunkOuterDim, outerDim - outerOffset); + const sourceByteOffset = outerOffset * sourceStrideBytes; + const targetByteOffset = outerOffset * targetStrideBytes; + const chunkBytes = outerCount * sourceStrideBytes; + const chunkShape = rec.shape.slice(); + chunkShape[0] = outerCount; + const chunkView = this.withNewScope(() => { + const chunkShapeTuple = this.makeShapeTuple(chunkShape); + return this.detachFromCurrentScope( + this.ctx.tensorCreateView( + targetTensor, + chunkShapeTuple, + rec.dtype, + new Scalar(targetByteOffset, "int"), + ) + ); + }); + const chunkSource = sourceBytes.subarray( + sourceByteOffset, + sourceByteOffset + chunkBytes, + ); + try { + this.ctx.arrayDecodeStorage(chunkView, chunkSource, rec.format, rec.dtype); + } finally { + chunkView.dispose(); + } + } + }; // first sync copy to cpu. - this.ctx.arrayDecodeStorage(cpu_arr, new Uint8Array(recSource), rec.format, rec.dtype); + copyRecordToTensor(cpu_arr, recSource); // then async stream into GPU if needed if (device.deviceType === DeviceStrToEnum.cpu) { this.tensorCacheUpdate(rec.name, cpu_arr, false); @@ -1435,7 +1512,42 @@ export class Instance implements Disposable { this.empty(rec.shape, rec.dtype, device) ) }); - gpu_arr.copyFrom(cpu_arr); + if (!canChunkRecord) { + gpu_arr.copyFrom(cpu_arr); + } else { + const chunkOuterDim = Math.max(1, Math.floor(maxChunkBytes / sourceStrideBytes)); + for (let outerOffset = 0; outerOffset < outerDim; outerOffset += chunkOuterDim) { + const outerCount = Math.min(chunkOuterDim, outerDim - outerOffset); + const targetByteOffset = outerOffset * targetStrideBytes; + const chunkShape = rec.shape.slice(); + chunkShape[0] = outerCount; + const [cpuView, gpuView] = this.withNewScope(() => { + const chunkShapeTuple = this.makeShapeTuple(chunkShape); + const cView = this.ctx.tensorCreateView( + cpu_arr, + chunkShapeTuple, + rec.dtype, + new Scalar(targetByteOffset, "int"), + ); + const gView = this.ctx.tensorCreateView( + gpu_arr, + chunkShapeTuple, + rec.dtype, + new Scalar(targetByteOffset, "int"), + ); + return [ + this.detachFromCurrentScope(cView), + this.detachFromCurrentScope(gView), + ]; + }); + try { + gpuView.copyFrom(cpuView); + } finally { + cpuView.dispose(); + gpuView.dispose(); + } + } + } await device.sync(); this.tensorCacheUpdate(rec.name, gpu_arr, false); cpu_arr.dispose(); @@ -2258,6 +2370,28 @@ export class Instance implements Disposable { case TypeIndex.kTVMFFIOpaquePtr: { return this.memory.loadPointer(valuePtr); } + case TypeIndex.kTVMFFIShape: { + const shapeObjPtr = this.memory.loadPointer(valuePtr); + if (shapeObjPtr === 0) { + return null; + } + if (callbackArg) { + const shapeCellPtr = shapeObjPtr + SizeOf.ObjectHeader; + const shapeDataPtr = this.memory.loadPointer(shapeCellPtr); + const shapeLen = this.memory.loadUSize(shapeCellPtr + this.memory.sizeofPtr()); + const result = new Array(shapeLen); + for (let i = 0; i < shapeLen; ++i) { + result[i] = this.memory.loadI64(shapeDataPtr + i * SizeOf.I64); + } + this.lib.checkCall( + (this.lib.exports.TVMFFIObjectDecRef as ctypes.FTVMFFIObjectDecRef)(shapeObjPtr) + ); + return result; + } + return this.ctx.attachToCurrentScope( + new TVMObject(shapeObjPtr, this.lib, this.ctx) + ); + } case TypeIndex.kTVMFFITensor: { return this.ctx.attachToCurrentScope( new Tensor(this.memory.loadPointer(valuePtr), this.lib, this.ctx, false) From d759fc5e3e2de98d739beba13fc2e3e656d5a0e5 Mon Sep 17 00:00:00 2001 From: MakotoUwu Date: Mon, 29 Jun 2026 10:55:29 +0200 Subject: [PATCH 2/4] [Web] Apply clang-format to wasm runtime --- web/emcc/wasm_runtime.cc | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc index 1286a514ce2f..49b278ec29e4 100644 --- a/web/emcc/wasm_runtime.cc +++ b/web/emcc/wasm_runtime.cc @@ -160,20 +160,20 @@ int64_t StorageSizeBytes(int64_t num_elements, const std::string& dtype) { TVMFFIByteArray dtype_bytes{dtype.data(), dtype.size()}; DLDataType dl_dtype; TVM_FFI_ICHECK_EQ(TVMFFIDataTypeFromString(&dtype_bytes, &dl_dtype), 0); - return static_cast( - ffi::GetDataSize(static_cast(num_elements), dl_dtype)); + return static_cast(ffi::GetDataSize(static_cast(num_elements), dl_dtype)); } 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); - }) + .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.runtime.StorageSizeBytes", StorageSizeBytes); } From 2bb1a2c7bd594ae63960e39180b1bf1b7ecf183a Mon Sep 17 00:00:00 2001 From: MakotoUwu Date: Mon, 29 Jun 2026 11:44:21 +0200 Subject: [PATCH 3/4] [Web] Bound decoded tensor-cache chunk size --- web/src/runtime.ts | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/web/src/runtime.ts b/web/src/runtime.ts index 745f6fa7c8ea..26bfb7eeca31 100644 --- a/web/src/runtime.ts +++ b/web/src/runtime.ts @@ -1464,12 +1464,17 @@ export class Instance implements Disposable { targetStrideBytes = targetBytes / outerDim; } } + const chunkOuterDim = canChunkRecord + ? Math.max( + 1, + Math.floor(maxChunkBytes / Math.max(sourceStrideBytes, targetStrideBytes)) + ) + : 1; const copyRecordToTensor = (targetTensor: Tensor, sourceBytes: Uint8Array) => { if (!canChunkRecord) { this.ctx.arrayDecodeStorage(targetTensor, sourceBytes, rec.format, rec.dtype); return; } - const chunkOuterDim = Math.max(1, Math.floor(maxChunkBytes / sourceStrideBytes)); for (let outerOffset = 0; outerOffset < outerDim; outerOffset += chunkOuterDim) { const outerCount = Math.min(chunkOuterDim, outerDim - outerOffset); const sourceByteOffset = outerOffset * sourceStrideBytes; @@ -1515,7 +1520,6 @@ export class Instance implements Disposable { if (!canChunkRecord) { gpu_arr.copyFrom(cpu_arr); } else { - const chunkOuterDim = Math.max(1, Math.floor(maxChunkBytes / sourceStrideBytes)); for (let outerOffset = 0; outerOffset < outerDim; outerOffset += chunkOuterDim) { const outerCount = Math.min(chunkOuterDim, outerDim - outerOffset); const targetByteOffset = outerOffset * targetStrideBytes; From aef19199a3e9366b7c0a6028cb774119e292eacb Mon Sep 17 00:00:00 2001 From: MakotoUwu Date: Sun, 12 Jul 2026 11:25:41 +0200 Subject: [PATCH 4/4] [Web] Harden tensor-cache chunk planning --- web/emcc/wasm_runtime.cc | 18 +++- web/src/memory.ts | 20 +++-- web/src/runtime.ts | 72 +++++++-------- web/src/tensor_cache_plan.ts | 108 ++++++++++++++++++++++ web/tests/node/test_packed_func.js | 32 +++++++ web/tests/node/test_tensor.js | 50 +++++++++++ web/tests/node/test_tensor_cache_plan.js | 109 +++++++++++++++++++++++ 7 files changed, 360 insertions(+), 49 deletions(-) create mode 100644 web/src/tensor_cache_plan.ts create mode 100644 web/tests/node/test_tensor_cache_plan.js diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc index 49b278ec29e4..39861218c45c 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" @@ -156,11 +158,25 @@ void ArrayDecodeStorage(Tensor cpu_arr, TVMFFIByteArray* bytes, const std::strin } int64_t StorageSizeBytes(int64_t num_elements, const std::string& dtype) { + constexpr uint64_t kMaxSafeInteger = (uint64_t{1} << 53) - 1; TVM_FFI_ICHECK_GE(num_elements, 0); + TVM_FFI_ICHECK_LE(static_cast(num_elements), kMaxSafeInteger); TVMFFIByteArray dtype_bytes{dtype.data(), dtype.size()}; DLDataType dl_dtype; TVM_FFI_ICHECK_EQ(TVMFFIDataTypeFromString(&dtype_bytes, &dl_dtype), 0); - return static_cast(ffi::GetDataSize(static_cast(num_elements), dl_dtype)); + const uint64_t num_elements_u64 = static_cast(num_elements); + uint64_t storage_bytes; + if (dl_dtype.code == kDLUInt && dl_dtype.bits == 1 && dl_dtype.lanes == 1) { + storage_bytes = num_elements_u64; + } else { + const uint64_t bits_per_element = + static_cast(dl_dtype.bits) * static_cast(dl_dtype.lanes); + TVM_FFI_ICHECK_GT(bits_per_element, 0); + TVM_FFI_ICHECK_LE(num_elements_u64, (kMaxSafeInteger * 8) / bits_per_element); + storage_bytes = (num_elements_u64 * bits_per_element + 7) / 8; + } + TVM_FFI_ICHECK_LE(storage_bytes, kMaxSafeInteger); + return static_cast(storage_bytes); } TVM_FFI_STATIC_INIT_BLOCK() { diff --git a/web/src/memory.ts b/web/src/memory.ts index 3847b8ce2ac8..31b829ac1477 100644 --- a/web/src/memory.ts +++ b/web/src/memory.ts @@ -83,8 +83,11 @@ export class Memory { this.updateViews(); } const base = ptr >> 2; - // assumes little endian, for now truncate high. - return this.viewI32[base]; + // WebAssembly is little-endian. JavaScript numbers represent signed + // 64-bit values exactly while they remain in the safe-integer range. + const value = this.viewI32[base + 1] * 0x100000000 + this.viewU32[base]; + assert(Number.isSafeInteger(value), "64-bit integer exceeds JavaScript's safe range"); + return value; } loadF32(ptr: Pointer): number { @@ -421,13 +424,14 @@ export class CachedCallStack implements Disposable { } storeI64(offset: PtrOffset, value: number): void { - // For now, just store as 32bit - // NOTE: wasm always uses little endian. - const low = value & 0xffffffff; + assert(Number.isSafeInteger(value), "64-bit integer exceeds JavaScript's safe range"); + // WebAssembly is little-endian. floor division keeps the low word + // positive for negative values while preserving signed high bits. + const high = Math.floor(value / 0x100000000); + const low = value - high * 0x100000000; const base = offset >> 2; - this.viewI32[base] = low; - // sign extend - this.viewI32[base + 1] = value < 0 ? -1 : 0; + this.viewU32[base] = low; + this.viewI32[base + 1] = high; } storeF64(offset: PtrOffset, value: number): void { diff --git a/web/src/runtime.ts b/web/src/runtime.ts index 26bfb7eeca31..75279d08e4ef 100644 --- a/web/src/runtime.ts +++ b/web/src/runtime.ts @@ -35,6 +35,7 @@ import { TensorShardEntry, createArtifactCache, } from "./artifact_cache"; +import { TensorCacheChunkPlan, planTensorCacheChunks } from "./tensor_cache_plan"; import * as compact from "./compact"; import * as ctypes from "./ctypes"; @@ -1444,44 +1445,37 @@ export class Instance implements Disposable { rec.byteOffset === 0 && rec.nbytes === shardBytes.byteLength ? shardBytes : shardBytes.subarray(rec.byteOffset, rec.byteOffset + rec.nbytes); - let canChunkRecord = - rec.nbytes > maxChunkBytes && - rec.shape.length >= 1 && - Number.isInteger(rec.shape[0]) && - rec.shape[0] > 0 && - rec.nbytes % rec.shape[0] === 0; - const outerDim = canChunkRecord ? rec.shape[0] : 1; - const sourceStrideBytes = canChunkRecord ? rec.nbytes / outerDim : rec.nbytes; - let targetStrideBytes = 0; - if (canChunkRecord) { + let chunkPlan: TensorCacheChunkPlan | undefined; + const mayDecodeLarger = + rec.format === "f32-to-bf16" && rec.dtype === "float32"; + if (rec.nbytes > maxChunkBytes || mayDecodeLarger) { const numElements = rec.shape.reduce((acc, value) => acc * value, 1); - const targetBytes = storageSizeBytes(numElements, rec.dtype); - canChunkRecord = - sourceStrideBytes <= maxChunkBytes && - targetBytes !== undefined && - targetBytes % outerDim === 0; - if (canChunkRecord) { - targetStrideBytes = targetBytes / outerDim; + if (Number.isSafeInteger(numElements) && numElements > 0) { + const targetBytes = storageSizeBytes(numElements, rec.dtype); + if ( + targetBytes !== undefined && + Math.max(rec.nbytes, targetBytes) > maxChunkBytes + ) { + const targetAlignmentBytes = + device.deviceType === DeviceStrToEnum.cpu ? 1 : 4; + chunkPlan = planTensorCacheChunks( + rec.shape, + rec.nbytes, + targetBytes, + maxChunkBytes, + targetAlignmentBytes, + ); + } } } - const chunkOuterDim = canChunkRecord - ? Math.max( - 1, - Math.floor(maxChunkBytes / Math.max(sourceStrideBytes, targetStrideBytes)) - ) - : 1; const copyRecordToTensor = (targetTensor: Tensor, sourceBytes: Uint8Array) => { - if (!canChunkRecord) { + if (chunkPlan === undefined) { this.ctx.arrayDecodeStorage(targetTensor, sourceBytes, rec.format, rec.dtype); return; } - for (let outerOffset = 0; outerOffset < outerDim; outerOffset += chunkOuterDim) { - const outerCount = Math.min(chunkOuterDim, outerDim - outerOffset); - const sourceByteOffset = outerOffset * sourceStrideBytes; - const targetByteOffset = outerOffset * targetStrideBytes; - const chunkBytes = outerCount * sourceStrideBytes; + for (const chunk of chunkPlan.chunks) { const chunkShape = rec.shape.slice(); - chunkShape[0] = outerCount; + chunkShape[0] = chunk.outerCount; const chunkView = this.withNewScope(() => { const chunkShapeTuple = this.makeShapeTuple(chunkShape); return this.detachFromCurrentScope( @@ -1489,13 +1483,13 @@ export class Instance implements Disposable { targetTensor, chunkShapeTuple, rec.dtype, - new Scalar(targetByteOffset, "int"), + new Scalar(chunk.targetByteOffset, "int"), ) ); }); const chunkSource = sourceBytes.subarray( - sourceByteOffset, - sourceByteOffset + chunkBytes, + chunk.sourceByteOffset, + chunk.sourceByteOffset + chunk.sourceByteLength, ); try { this.ctx.arrayDecodeStorage(chunkView, chunkSource, rec.format, rec.dtype); @@ -1517,27 +1511,25 @@ export class Instance implements Disposable { this.empty(rec.shape, rec.dtype, device) ) }); - if (!canChunkRecord) { + if (chunkPlan === undefined) { gpu_arr.copyFrom(cpu_arr); } else { - for (let outerOffset = 0; outerOffset < outerDim; outerOffset += chunkOuterDim) { - const outerCount = Math.min(chunkOuterDim, outerDim - outerOffset); - const targetByteOffset = outerOffset * targetStrideBytes; + for (const chunk of chunkPlan.chunks) { const chunkShape = rec.shape.slice(); - chunkShape[0] = outerCount; + chunkShape[0] = chunk.outerCount; const [cpuView, gpuView] = this.withNewScope(() => { const chunkShapeTuple = this.makeShapeTuple(chunkShape); const cView = this.ctx.tensorCreateView( cpu_arr, chunkShapeTuple, rec.dtype, - new Scalar(targetByteOffset, "int"), + new Scalar(chunk.targetByteOffset, "int"), ); const gView = this.ctx.tensorCreateView( gpu_arr, chunkShapeTuple, rec.dtype, - new Scalar(targetByteOffset, "int"), + new Scalar(chunk.targetByteOffset, "int"), ); return [ this.detachFromCurrentScope(cView), diff --git a/web/src/tensor_cache_plan.ts b/web/src/tensor_cache_plan.ts new file mode 100644 index 000000000000..0ed619e51bab --- /dev/null +++ b/web/src/tensor_cache_plan.ts @@ -0,0 +1,108 @@ +/* + * 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. + */ + +export interface TensorCacheChunk { + outerOffset: number; + outerCount: number; + sourceByteOffset: number; + sourceByteLength: number; + targetByteOffset: number; + targetByteLength: number; +} + +export interface TensorCacheChunkPlan { + sourceStrideBytes: number; + targetStrideBytes: number; + chunks: Array; +} + +function greatestCommonDivisor(lhs: number, rhs: number): number { + while (rhs !== 0) { + const remainder = lhs % rhs; + lhs = rhs; + rhs = remainder; + } + return lhs; +} + +/** + * Plan outer-dimension chunks whose encoded and decoded sizes stay bounded. + * + * Returns undefined when outer-dimension chunking cannot satisfy the size and + * target-alignment constraints. Callers can then retain the full-record path. + */ +export function planTensorCacheChunks( + shape: Array, + sourceBytes: number, + targetBytes: number, + maxChunkBytes: number, + targetAlignmentBytes = 1, +): TensorCacheChunkPlan | undefined { + const outerDim = shape[0]; + if ( + shape.length === 0 || + !Number.isSafeInteger(outerDim) || + outerDim <= 0 || + !Number.isSafeInteger(sourceBytes) || + sourceBytes <= 0 || + !Number.isSafeInteger(targetBytes) || + targetBytes <= 0 || + !Number.isSafeInteger(maxChunkBytes) || + maxChunkBytes <= 0 || + !Number.isSafeInteger(targetAlignmentBytes) || + targetAlignmentBytes <= 0 || + sourceBytes % outerDim !== 0 || + targetBytes % outerDim !== 0 || + targetBytes % targetAlignmentBytes !== 0 + ) { + return undefined; + } + + const sourceStrideBytes = sourceBytes / outerDim; + const targetStrideBytes = targetBytes / outerDim; + const maxStrideBytes = Math.max(sourceStrideBytes, targetStrideBytes); + if (maxStrideBytes > maxChunkBytes) { + return undefined; + } + + const rowsPerAlignment = + targetAlignmentBytes / + greatestCommonDivisor(targetStrideBytes, targetAlignmentBytes); + const maxOuterCount = Math.floor(maxChunkBytes / maxStrideBytes); + const chunkOuterCount = + maxOuterCount - (maxOuterCount % rowsPerAlignment); + if (chunkOuterCount <= 0) { + return undefined; + } + + const chunks = new Array(); + for (let outerOffset = 0; outerOffset < outerDim; outerOffset += chunkOuterCount) { + const outerCount = Math.min(chunkOuterCount, outerDim - outerOffset); + chunks.push({ + outerOffset, + outerCount, + sourceByteOffset: outerOffset * sourceStrideBytes, + sourceByteLength: outerCount * sourceStrideBytes, + targetByteOffset: outerOffset * targetStrideBytes, + targetByteLength: outerCount * targetStrideBytes, + }); + } + + return { sourceStrideBytes, targetStrideBytes, chunks }; +} diff --git a/web/tests/node/test_packed_func.js b/web/tests/node/test_packed_func.js index d329b1df64ea..04a6b1e0ba17 100644 --- a/web/tests/node/test_packed_func.js +++ b/web/tests/node/test_packed_func.js @@ -125,6 +125,38 @@ test("ReturnFunc", () => { tvm.endScope(); }); +test("Shape callback argument", () => { + tvm.withNewScope(() => { + const call = tvm.getGlobalFunc("tvmjs.testing.call"); + const expectedShape = [5000000000, 3]; + const callback = tvm.toPackedFunc((shape) => { + assert.deepStrictEqual(shape, expectedShape); + }); + call(callback, tvm.makeShapeTuple(expectedShape)); + }); +}); + +test("Int64 callback arguments preserve JavaScript-safe values", () => { + tvm.withNewScope(() => { + const call = tvm.getGlobalFunc("tvmjs.testing.call"); + const values = [5000000000, -5000000000, Number.MAX_SAFE_INTEGER, Number.MIN_SAFE_INTEGER]; + const observed = []; + const callback = tvm.toPackedFunc((value) => { + observed.push(value); + }); + + for (const value of values) { + call(callback, tvm.scalar(value, "int")); + } + + assert.deepStrictEqual(observed, values); + assert.throws( + () => call(callback, tvm.scalar(Number.MAX_SAFE_INTEGER + 1, "int")), + /safe range/ + ); + }); +}); + test("RegisterGlobal", () => { tvm.beginScope(); tvm.registerFunc("xyz", function (x, y) { diff --git a/web/tests/node/test_tensor.js b/web/tests/node/test_tensor.js index 6aadc702f5fa..1e6c321f56fd 100644 --- a/web/tests/node/test_tensor.js +++ b/web/tests/node/test_tensor.js @@ -54,3 +54,53 @@ test("array copy", () => { testArrayCopy("float64", Float64Array); }); }); + +test("storage size preserves values above int32", () => { + tvm.withNewScope(() => { + const storageSizeBytes = tvm.getGlobalFunc("tvmjs.runtime.StorageSizeBytes"); + assert.strictEqual( + storageSizeBytes(tvm.scalar(536870912, "int"), "float32"), + 2147483648 + ); + assert.strictEqual( + storageSizeBytes(tvm.scalar(5000000000, "int"), "uint8"), + 5000000000 + ); + }); +}); + +test("decode storage respects tensor view byte offset", () => { + tvm.withNewScope(() => { + const decodeStorage = tvm.getGlobalFunc("tvmjs.array.decode_storage"); + const createView = tvm.getGlobalFunc("runtime.TVMTensorCreateView"); + const backing = tvm.empty([4], "float32").copyFrom([9, 9, 9, 9]); + const view = createView( + backing, + tvm.makeShapeTuple([2]), + "float32", + tvm.scalar(4, "int") + ); + + // BF16 encodings for 1.0 and -2.0, little-endian. + decodeStorage( + view, + new Uint8Array([0x80, 0x3f, 0x00, 0xc0]), + "f32-to-bf16", + "float32" + ); + + assert.deepStrictEqual(Array.from(backing.toArray()), [9, 1, -2, 9]); + }); +}); + +test("decode storage accepts raw float32 payload with bf16 tag", () => { + tvm.withNewScope(() => { + const decodeStorage = tvm.getGlobalFunc("tvmjs.array.decode_storage"); + const backing = tvm.empty([2], "float32"); + const source = new Uint8Array(new Float32Array([1.5, -2.25]).buffer); + + decodeStorage(backing, source, "f32-to-bf16", "float32"); + + assert.deepStrictEqual(Array.from(backing.toArray()), [1.5, -2.25]); + }); +}); diff --git a/web/tests/node/test_tensor_cache_plan.js b/web/tests/node/test_tensor_cache_plan.js new file mode 100644 index 000000000000..cf34bc33c88c --- /dev/null +++ b/web/tests/node/test_tensor_cache_plan.js @@ -0,0 +1,109 @@ +/* + * 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 assert = require("assert"); +const fs = require("fs"); +const path = require("path"); +const ts = require("typescript"); + +const plannerPath = path.join(__dirname, "../../src/tensor_cache_plan.ts"); +const plannerSource = fs.readFileSync(plannerPath, "utf8"); +const plannerJavaScript = ts.transpileModule(plannerSource, { + compilerOptions: { + module: ts.ModuleKind.CommonJS, + target: ts.ScriptTarget.ES2018, + }, + fileName: plannerPath, +}).outputText; +const plannerModule = { exports: {} }; +new Function("module", "exports", plannerJavaScript)( + plannerModule, + plannerModule.exports +); +const { planTensorCacheChunks } = plannerModule.exports; + +const MiB = 1024 * 1024; + +function assertValidPlan(plan, sourceBytes, targetBytes, maxChunkBytes, alignment) { + assert(plan !== undefined); + let sourceOffset = 0; + let targetOffset = 0; + for (const chunk of plan.chunks) { + assert.strictEqual(chunk.sourceByteOffset, sourceOffset); + assert.strictEqual(chunk.targetByteOffset, targetOffset); + assert(chunk.sourceByteLength <= maxChunkBytes); + assert(chunk.targetByteLength <= maxChunkBytes); + assert.strictEqual(chunk.targetByteOffset % alignment, 0); + assert.strictEqual(chunk.targetByteLength % alignment, 0); + sourceOffset += chunk.sourceByteLength; + targetOffset += chunk.targetByteLength; + } + assert.strictEqual(sourceOffset, sourceBytes); + assert.strictEqual(targetOffset, targetBytes); +} + +test("tensor-cache chunks keep WebGPU offsets aligned", () => { + const plan = planTensorCacheChunks( + [50000000, 3], + 150000000, + 150000000, + 128 * MiB, + 4 + ); + + assertValidPlan(plan, 150000000, 150000000, 128 * MiB, 4); + assert(plan.chunks.length > 1); +}); + +test("tensor-cache chunks bound Gemma 4 E2B BF16 expansion", () => { + const plan = planTensorCacheChunks( + [262144, 1120], + 587202560, + 1174405120, + 128 * MiB, + 4 + ); + + assertValidPlan(plan, 587202560, 1174405120, 128 * MiB, 4); + assert.strictEqual(plan.chunks.length, 9); +}); + +test("tensor-cache chunks preserve byte counts above int32", () => { + const totalBytes = 2147483648; + const plan = planTensorCacheChunks( + [536870912], + totalBytes, + totalBytes, + 128 * MiB, + 4 + ); + + assertValidPlan(plan, totalBytes, totalBytes, 128 * MiB, 4); + assert.strictEqual(plan.chunks.length, 16); +}); + +test("tensor-cache chunks reject an outer row above the cap", () => { + const plan = planTensorCacheChunks( + [2, 40 * MiB], + 160 * MiB, + 320 * MiB, + 128 * MiB, + 4 + ); + assert.strictEqual(plan, undefined); +});