diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc index 9d3d46f18cb4..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" @@ -130,32 +132,65 @@ 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; } + } + cpu_arr.CopyFromBytes(byte_data, byte_size); +} + +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); + 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 { - cpu_arr.CopyFromBytes(byte_data, byte_size); + 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() { 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.runtime.StorageSizeBytes", StorageSizeBytes); } // Concatenate n TVMArrays 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 078a0c7df21f..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"; @@ -172,6 +173,7 @@ class RuntimeContext implements Disposable { tensorCacheRemove: PackedFunc; tensorCacheClear: PackedFunc; arrayDecodeStorage: PackedFunc; + storageSizeBytes: PackedFunc; paramModuleFromCache: PackedFunc; paramModuleFromCacheByName: PackedFunc; makeShapeTuple: PackedFunc; @@ -207,6 +209,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 +233,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 +1014,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 +1329,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 +1440,66 @@ 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 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); + 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 copyRecordToTensor = (targetTensor: Tensor, sourceBytes: Uint8Array) => { + if (chunkPlan === undefined) { + this.ctx.arrayDecodeStorage(targetTensor, sourceBytes, rec.format, rec.dtype); + return; + } + for (const chunk of chunkPlan.chunks) { + const chunkShape = rec.shape.slice(); + chunkShape[0] = chunk.outerCount; + const chunkView = this.withNewScope(() => { + const chunkShapeTuple = this.makeShapeTuple(chunkShape); + return this.detachFromCurrentScope( + this.ctx.tensorCreateView( + targetTensor, + chunkShapeTuple, + rec.dtype, + new Scalar(chunk.targetByteOffset, "int"), + ) + ); + }); + const chunkSource = sourceBytes.subarray( + chunk.sourceByteOffset, + chunk.sourceByteOffset + chunk.sourceByteLength, + ); + 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 +1511,39 @@ export class Instance implements Disposable { this.empty(rec.shape, rec.dtype, device) ) }); - gpu_arr.copyFrom(cpu_arr); + if (chunkPlan === undefined) { + gpu_arr.copyFrom(cpu_arr); + } else { + for (const chunk of chunkPlan.chunks) { + const chunkShape = rec.shape.slice(); + 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(chunk.targetByteOffset, "int"), + ); + const gView = this.ctx.tensorCreateView( + gpu_arr, + chunkShapeTuple, + rec.dtype, + new Scalar(chunk.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 +2366,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) 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); +});