[Web] Improve large tensor loading in wasm runtime

web/emcc/wasm_runtime.cc

@@ -131,7 +131,8 @@ void ArrayDecodeStorage(Tensor cpu_arr, TVMFFIByteArray* bytes, const std::strin
   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);
+    uint32_t* data =
+        reinterpret_cast<uint32_t*>(static_cast<char*>(cpu_arr->data) + cpu_arr->byte_offset);
     TVM_FFI_ICHECK(cpu_arr.IsContiguous());
     size_t size = 1;
     for (int i = 0; i < cpu_arr->ndim; ++i) {

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";
 
@@ -1010,9 +1011,11 @@ export class Instance implements Disposable {
    */
   withNewScope<T>(action: () => T): T {
     this.beginScope();
-    const val = action();
-    this.endScope();
-    return val;
+    try {
+      return action();
+    } finally {
+      this.endScope();
+    }
   }
 
   /**
@@ -1323,6 +1326,10 @@ export class Instance implements Disposable {
     artifactCache: ArtifactCacheTemplate,
     signal?: AbortSignal,
   ) {
+    // CachedCallStack grows geometrically and retains its backing allocation.
+    // Cap each tensor-cache payload at 128 MiB so the retained staging
+    // allocation stays bounded independently of the final tensor.
+    const maxTensorCacheChunkBytes = 128 * 1024 * 1024;
     const perf = compact.getPerformance();
     const tstart = perf.now();
     let totalBytes = 0;
@@ -1421,9 +1428,72 @@ 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 decodeChunkPlan: TensorCacheChunkPlan | undefined;
+          let gpuCopyChunkPlan: TensorCacheChunkPlan | undefined;
+          const decodedBytes =
+            rec.format === "f32-to-bf16" && rec.dtype === "float32"
+              ? rec.nbytes * 2
+              : rec.nbytes;
+          if (
+            Number.isSafeInteger(decodedBytes) &&
+            Math.max(rec.nbytes, decodedBytes) > maxTensorCacheChunkBytes
+          ) {
+            decodeChunkPlan = planTensorCacheChunks(
+              rec.shape,
+              rec.nbytes,
+              decodedBytes,
+              maxTensorCacheChunkBytes,
+            );
+            if (device.deviceType !== DeviceStrToEnum.cpu) {
+              gpuCopyChunkPlan = planTensorCacheChunks(
+                rec.shape,
+                rec.nbytes,
+                decodedBytes,
+                maxTensorCacheChunkBytes,
+                4,
+              );
+            }
+          }
+          const decodeRecordIntoTensor = (
+            targetTensor: Tensor,
+            sourceBytes: Uint8Array,
+          ) => {
+            if (decodeChunkPlan === undefined) {
+              this.ctx.arrayDecodeStorage(targetTensor, sourceBytes, rec.format, rec.dtype);
+              return;
+            }
+            for (const chunk of decodeChunkPlan.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);
+          decodeRecordIntoTensor(cpu_arr, recSource);
           // then async stream into GPU if needed
           if (device.deviceType === DeviceStrToEnum.cpu) {
             this.tensorCacheUpdate(rec.name, cpu_arr, false);
@@ -1435,7 +1505,39 @@ export class Instance implements Disposable {
                 this.empty(rec.shape, rec.dtype, device)
               )
             });
-            gpu_arr.copyFrom(cpu_arr);
+            if (gpuCopyChunkPlan === undefined) {
+              gpu_arr.copyFrom(cpu_arr);
+            } else {
+              for (const chunk of gpuCopyChunkPlan.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();

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 {
+  outerCount: number;
+  sourceByteOffset: number;
+  sourceByteLength: number;
+  targetByteOffset: number;
+  targetByteLength: number;
+}
+
+export interface TensorCacheChunkPlan {
+  chunks: Array<TensorCacheChunk>;
+}
+
+// Tensor view offsets are currently marshalled into TVM JS's wasm32 runtime.
+const wasm32AddressSpaceBytes = 0x100000000;
+
+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<number>,
+  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 ||
+    targetBytes >= wasm32AddressSpaceBytes ||
+    !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<TensorCacheChunk>();
+  for (let outerOffset = 0; outerOffset < outerDim; outerOffset += chunkOuterCount) {
+    const outerCount = Math.min(chunkOuterCount, outerDim - outerOffset);
+    chunks.push({
+      outerCount,
+      sourceByteOffset: outerOffset * sourceStrideBytes,
+      sourceByteLength: outerCount * sourceStrideBytes,
+      targetByteOffset: outerOffset * targetStrideBytes,
+      targetByteLength: outerCount * targetStrideBytes,
+    });
+  }
+
+  return { chunks };
+}

web/tests/node/test_tensor.js

@@ -54,3 +54,27 @@ test("array copy", () => {
     testArrayCopy("float64", Float64Array);
   });
 });
+
+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]);
+  });
+});