diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc index 9d3d46f18cb4..bd032b6eed10 100644 --- a/web/emcc/wasm_runtime.cc +++ b/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(byte_data); - uint32_t* data = static_cast(cpu_arr->data); + uint32_t* data = + reinterpret_cast(static_cast(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) { diff --git a/web/src/runtime.ts b/web/src/runtime.ts index 078a0c7df21f..20d32e8ce6ce 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"; @@ -1010,9 +1011,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 +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(); diff --git a/web/src/tensor_cache_plan.ts b/web/src/tensor_cache_plan.ts new file mode 100644 index 000000000000..06859ce2c677 --- /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 { + outerCount: number; + sourceByteOffset: number; + sourceByteLength: number; + targetByteOffset: number; + targetByteLength: number; +} + +export interface TensorCacheChunkPlan { + chunks: Array; +} + +// 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, + 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(); + 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 }; +} diff --git a/web/tests/node/test_tensor.js b/web/tests/node/test_tensor.js index 6aadc702f5fa..9b99f2495346 100644 --- a/web/tests/node/test_tensor.js +++ b/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]); + }); +}); 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..f7902b763f29 --- /dev/null +++ b/web/tests/node/test_tensor_cache_plan.js @@ -0,0 +1,158 @@ +/* + * 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 keep the final WebGPU copy aligned", () => { + const plan = planTensorCacheChunks([6, 1], 12, 12, 8, 4); + + assertValidPlan(plan, 12, 12, 8, 4); + assert.deepStrictEqual( + plan.chunks.map((chunk) => chunk.targetByteLength), + [8, 4] + ); +}); + +test("tensor-cache chunks reject an unaligned WebGPU total", () => { + const plan = planTensorCacheChunks([5, 1], 10, 10, 8, 4); + + assert.strictEqual(plan, undefined); +}); + +test("tensor-cache CPU decode does not depend on WebGPU copy alignment", () => { + const recordBytes = 128 * MiB + 1; + const decodePlan = planTensorCacheChunks( + [recordBytes, 1], + recordBytes, + recordBytes, + 128 * MiB + ); + const gpuCopyPlan = planTensorCacheChunks( + [recordBytes, 1], + recordBytes, + recordBytes, + 128 * MiB, + 4 + ); + + assertValidPlan(decodePlan, recordBytes, recordBytes, 128 * MiB, 1); + assert.strictEqual(decodePlan.chunks.length, 2); + assert.strictEqual(gpuCopyPlan, undefined); +}); + +test("tensor-cache chunks bound a large raw record", () => { + const plan = planTensorCacheChunks( + [262144, 1120], + 1174405120, + 1174405120, + 128 * MiB, + 4 + ); + + assertValidPlan(plan, 1174405120, 1174405120, 128 * MiB, 4); + assert.strictEqual(plan.chunks.length, 9); +}); + +test("tensor-cache chunks bound encoded and decoded BF16 sizes", () => { + 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 reject an outer row above the cap", () => { + const plan = planTensorCacheChunks( + [2, 40 * MiB], + 160 * MiB, + 320 * MiB, + 128 * MiB, + 4 + ); + assert.strictEqual(plan, undefined); +}); + +test("tensor-cache chunks reject targets at the wasm32 address-space limit", () => { + for (const targetBytes of [0x100000000, 0x100000000 + 1]) { + const plan = planTensorCacheChunks( + [targetBytes], + targetBytes, + targetBytes, + 128 * MiB + ); + + assert.strictEqual(plan, undefined); + } +});