This is an automated email from the ASF dual-hosted git repository.
guan404ming pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 7b2ef6ad5f [Web] Avoid redundant memory byte copies (#20127)
7b2ef6ad5f is described below
commit 7b2ef6ad5fbb136c59ca780f0a491c17595a9223
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Mon Aug 17 05:15:14 2026 -0700
[Web] Avoid redundant memory byte copies (#20127)
Avoid unnecessary intermediate allocations in WASM memory operations:
- Pass a view of the cached call-stack buffer directly to
storeRawBytes(), which synchronously copies it into WASM memory
- Preserve existing ownership, length, and zero-padding behavior
The added tests are intended to target internal buffer-sharing behavior
and thus import from src instead of dist.
---
web/src/memory.ts | 8 +++---
web/tests/node/test_memory.js | 66 +++++++++++++++++++++++++++++++++++++++++++
web/tests/node/test_webgpu.js | 38 +++++++++++++++++++++++++
3 files changed, 108 insertions(+), 4 deletions(-)
diff --git a/web/src/memory.ts b/web/src/memory.ts
index 3847b8ce2a..00dc17d412 100644
--- a/web/src/memory.ts
+++ b/web/src/memory.ts
@@ -134,7 +134,7 @@ export class Memory {
this.updateViews();
}
const result = new Uint8Array(numBytes);
- result.set(this.viewU8.slice(ptr, ptr + numBytes));
+ result.set(this.viewU8.subarray(ptr, ptr + numBytes));
return result;
}
/**
@@ -217,7 +217,7 @@ export class Memory {
const length = this.loadU32(sizePtr);
const dataPtr = ffiAnyPtr + SizeOf.I32 + SizeOf.I32;
const result = new Uint8Array(length);
- result.set(this.viewU8.slice(dataPtr, dataPtr + length));
+ result.set(this.viewU8.subarray(dataPtr, dataPtr + length));
return result;
}
/**
@@ -248,7 +248,7 @@ export class Memory {
const ptr = this.loadPointer(byteArrayPtr);
const length = this.loadUSize(byteArrayPtr + this.sizeofPtr());
const result = new Uint8Array(length);
- result.set(this.viewU8.slice(ptr, ptr + length));
+ result.set(this.viewU8.subarray(ptr, ptr + length));
return result;
}
// private functions
@@ -346,7 +346,7 @@ export class CachedCallStack implements Disposable {
];
this.storePtr(targetOffset, this.ptrFromOffset(valueOffset));
}
- this.memory.storeRawBytes(this.basePtr, this.viewU8.slice(0, nbytes));
+ this.memory.storeRawBytes(this.basePtr, this.viewU8.subarray(0, nbytes));
}
/**
diff --git a/web/tests/node/test_memory.js b/web/tests/node/test_memory.js
new file mode 100644
index 0000000000..39de27835c
--- /dev/null
+++ b/web/tests/node/test_memory.js
@@ -0,0 +1,66 @@
+/*
+ * 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 { CachedCallStack, Memory } = require("../../src/memory");
+
+test("loadRawBytes returns an owned copy", () => {
+ const wasmMemory = new WebAssembly.Memory({ initial: 1 });
+ const memory = new Memory(wasmMemory);
+ const source = new Uint8Array(wasmMemory.buffer, 8, 4);
+ source.set([1, 2, 3, 4]);
+
+ const result = memory.loadRawBytes(8, 4);
+
+ expect(Array.from(result)).toEqual([1, 2, 3, 4]);
+ expect(result.buffer).not.toBe(wasmMemory.buffer);
+
+ result[0] = 10;
+ source[1] = 20;
+ expect(Array.from(result)).toEqual([10, 2, 3, 4]);
+ expect(Array.from(source)).toEqual([1, 20, 3, 4]);
+});
+
+test("loadRawBytes preserves the requested length at the end of memory", () =>
{
+ const wasmMemory = new WebAssembly.Memory({ initial: 1 });
+ const memory = new Memory(wasmMemory);
+ const source = new Uint8Array(wasmMemory.buffer);
+ source.set([5, 6], source.length - 2);
+
+ const result = memory.loadRawBytes(source.length - 2, 4);
+
+ expect(Array.from(result)).toEqual([5, 6, 0, 0]);
+});
+
+test("CachedCallStack commits a view of its cached bytes", () => {
+ const memory = {
+ wasm32: true,
+ sizeofPtr: () => 4,
+ storeRawBytes: jest.fn(),
+ };
+ const stack = new CachedCallStack(memory, () => 1024, () => {});
+ const offset = stack.allocRawBytes(4);
+ stack.storeRawBytes(offset, new Uint8Array([1, 2, 3, 4]));
+
+ stack.commitToWasmMemory(4);
+
+ expect(memory.storeRawBytes).toHaveBeenCalledTimes(1);
+ const [ptr, bytes] = memory.storeRawBytes.mock.calls[0];
+ expect(ptr).toBe(1024);
+ expect(Array.from(bytes)).toEqual([1, 2, 3, 4]);
+ expect(bytes.buffer).toBe(stack.buffer);
+});
diff --git a/web/tests/node/test_webgpu.js b/web/tests/node/test_webgpu.js
index 30c6c8c581..dcd0be2e09 100644
--- a/web/tests/node/test_webgpu.js
+++ b/web/tests/node/test_webgpu.js
@@ -110,6 +110,7 @@ function createMockDevice({
function createContext(deviceOptions) {
const gpu = createMockDevice(deviceOptions);
const memory = {
+ loadRawBytes: jest.fn(),
storeRawBytes: jest.fn(),
};
const context = new WebGPUContext(memory, gpu.device);
@@ -175,6 +176,43 @@ test("a host write flushes pending GPU copies before
writeBuffer", () => {
expect(events).toEqual(["copy", "finish", "submit", "writeBuffer"]);
});
+test("an aligned CPU to GPU copy writes the requested bytes", () => {
+ const { context, device, queue, memory, destination } = createContext();
+ const copyToGPU = context.getDeviceAPI("deviceCopyToGPU");
+ const rawBytes = new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8]);
+ memory.loadRawBytes.mockReturnValue(rawBytes);
+
+ copyToGPU(128, destination, 12, rawBytes.length);
+
+ expect(memory.loadRawBytes).toHaveBeenCalledWith(128, rawBytes.length);
+ expect(queue.writeBuffer).toHaveBeenCalledWith(
+ device.createBuffer.mock.results[1].value,
+ 12,
+ rawBytes,
+ 0,
+ rawBytes.length
+ );
+});
+
+test("an unaligned CPU to GPU copy pads the write to four bytes", () => {
+ const { context, device, queue, memory, destination } = createContext();
+ const copyToGPU = context.getDeviceAPI("deviceCopyToGPU");
+ const rawBytes = new Uint8Array([1, 2, 3]);
+ memory.loadRawBytes.mockReturnValue(rawBytes);
+
+ copyToGPU(256, destination, 4, rawBytes.length);
+
+ expect(memory.loadRawBytes).toHaveBeenCalledWith(256, rawBytes.length);
+ expect(queue.writeBuffer).toHaveBeenCalledTimes(1);
+ const [buffer, toOffset, data, dataOffset, nbytes] =
+ queue.writeBuffer.mock.calls[0];
+ expect(buffer).toBe(device.createBuffer.mock.results[1].value);
+ expect(toOffset).toBe(4);
+ expect(Array.from(data)).toEqual([1, 2, 3, 0]);
+ expect(dataOffset).toBe(0);
+ expect(nbytes).toBe(4);
+});
+
test("a GPU readback flushes pending copies before its own submission", async
() => {
const {
context,
