gemini-code-assist[bot] commented on code in PR #19771:
URL: https://github.com/apache/tvm/pull/19771#discussion_r3410190319
##########
web/src/runtime.ts:
##########
@@ -1435,7 +1498,48 @@ 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;
+ // Use withNewScope so the shape tuple is auto-disposed,
+ // and detach the views we need for manual lifetime control.
+ const [cpuView, gpuView] = this.withNewScope(() => {
+ const chunkShapeTuple = this.ctx.makeShapeTuple(
+ ...chunkShape.map((value) => new Scalar(value, "int")),
+ );
Review Comment:
We can leverage the cached `makeShapeTuple` method on the `Instance` class
instead of directly calling the FFI `this.ctx.makeShapeTuple` on every chunk.
This avoids redundant FFI round-trips to create the same shape tuple multiple
times across chunks and records, improving performance.
```typescript
const chunkShapeTuple = this.makeShapeTuple(chunkShape);
```
##########
web/src/runtime.ts:
##########
@@ -1421,9 +1431,62 @@ 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);
+ const 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;
+ const targetBytes = rec.shape.reduce((acc, value) => acc * value, 1)
*
+ storageBytes(rec.dtype);
+ const targetStrideBytes = canChunkRecord ? targetBytes / outerDim :
targetBytes;
+ 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;
+ // Wrap in withNewScope so TVM intermediate objects (shape tuple)
+ // are disposed after each chunk, but detach the view we need.
+ const chunkView = this.withNewScope(() => {
+ return this.detachFromCurrentScope(
+ this.ctx.tensorCreateView(
+ targetTensor,
+ this.ctx.makeShapeTuple(
+ ...chunkShape.map((value) => new Scalar(value, "int")),
+ ),
Review Comment:
We can leverage the cached `makeShapeTuple` method on the `Instance` class
instead of directly calling the FFI `this.ctx.makeShapeTuple` on every chunk.
This avoids redundant FFI round-trips to create the same shape tuple multiple
times across chunks and records, improving performance.
```typescript
this.makeShapeTuple(chunkShape),
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
