gemini-code-assist[bot] commented on code in PR #19771:
URL: https://github.com/apache/tvm/pull/19771#discussion_r3414229424
##########
web/src/runtime.ts:
##########
@@ -2258,6 +2359,25 @@ export class Instance implements Disposable {
case TypeIndex.kTVMFFIOpaquePtr: {
return this.memory.loadPointer(valuePtr);
}
+ case TypeIndex.kTVMFFIShape: {
+ const shapeObjPtr = this.memory.loadPointer(valuePtr);
+ if (callbackArg) {
Review Comment:
Add a defensive null check for `shapeObjPtr` before performing pointer
arithmetic and loading memory. If `shapeObjPtr` is null (`0`), accessing it
will cause an out-of-bounds memory access or return garbage.
```suggestion
case TypeIndex.kTVMFFIShape: {
const shapeObjPtr = this.memory.loadPointer(valuePtr);
if (shapeObjPtr === 0) {
return null;
}
if (callbackArg) {
```
##########
web/emcc/wasm_runtime.cc:
##########
@@ -130,20 +130,28 @@ 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<const uint16_t*>(byte_data);
- uint32_t* data = static_cast<uint32_t*>(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<uint32_t>(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 (size == byte_size / 2) {
+ const uint16_t* bf16 = reinterpret_cast<const uint16_t*>(byte_data);
+ uint32_t* data = static_cast<uint32_t*>(cpu_arr->data);
Review Comment:
The custom `f32-to-bf16` decoding path directly uses `cpu_arr->data` as the
destination pointer, completely ignoring `cpu_arr->byte_offset`. When loading
chunked records, `chunkView` is a tensor view with a non-zero `byte_offset`.
Ignoring this offset causes all decoded chunks to be written to the very
beginning of the underlying tensor, leading to silent data corruption. To fix
this, the destination pointer must be offset by `cpu_arr->byte_offset`.
```suggestion
uint32_t* data =
reinterpret_cast<uint32_t*>(static_cast<char*>(cpu_arr->data) +
cpu_arr->byte_offset);
```
##########
web/src/runtime.ts:
##########
@@ -1435,7 +1499,44 @@ 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;
+ const chunkShapeTuple = this.makeShapeTuple(chunkShape);
+ const [cpuView, gpuView] = this.withNewScope(() => {
+ return [
+ this.detachFromCurrentScope(
+ this.ctx.tensorCreateView(
+ cpu_arr,
+ chunkShapeTuple,
+ rec.dtype,
+ new Scalar(targetByteOffset, "int"),
+ )
+ ),
+ this.detachFromCurrentScope(
+ this.ctx.tensorCreateView(
+ gpu_arr,
+ chunkShapeTuple,
+ rec.dtype,
+ new Scalar(targetByteOffset, "int"),
+ )
+ ),
+ ];
+ });
Review Comment:
In the GPU copy path, `cpuView` and `gpuView` are created and immediately
detached from the current scope inside the `withNewScope` block. If the second
`tensorCreateView` (for `gpuView`) throws an exception, the first view
(`cpuView`) will have already been detached from the scope's auto-disposal
list, but the assignment to `[cpuView, gpuView]` will never complete. This
results in a resource leak of `cpuView` because it is neither auto-disposed by
the scope nor disposed in the `finally` block. To prevent this, detach both
views only after both have been successfully created.
```typescript
const [cpuView, gpuView] = this.withNewScope(() => {
const cView = this.ctx.tensorCreateView(
cpu_arr,
chunkShapeTuple,
rec.dtype,
new Scalar(targetByteOffset, "int"),
);
const gView = this.ctx.tensorCreateView(
gpu_arr,
chunkShapeTuple,
rec.dtype,
new Scalar(targetByteOffset, "int"),
);
return [
this.detachFromCurrentScope(cView),
this.detachFromCurrentScope(gView),
];
});
```
##########
web/src/runtime.ts:
##########
@@ -1323,6 +1323,20 @@ 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 storageBytes = (dtype: string) => {
+ const match = dtype.match(/(\d+)(?:x(\d+))?$/);
+ if (match === null) {
+ throw new Error("Cannot determine storage width of dtype " + dtype);
+ }
+ const bits = Number(match[1]);
+ const lanes = match[2] === undefined ? 1 : Number(match[2]);
+ return (bits * lanes + 7) >> 3;
+ };
Review Comment:
The `storageBytes` helper uses a regular expression to extract the bit width
from the `dtype` string. However, if the `dtype` is `"bool"`, the regex will
fail to match any digits and throw an error. In TVM, boolean tensors are
commonly represented with the `"bool"` dtype (which is 1 byte / 8 bits in
DLPack). Adding an explicit fallback for `"bool"` makes the helper more robust.
```suggestion
const storageBytes = (dtype: string) => {
if (dtype === "bool") {
return 1;
}
const match = dtype.match(/(\d+)(?:x(\d+))?$/);
if (match === null) {
throw new Error("Cannot determine storage width of dtype " + dtype);
}
const bits = Number(match[1]);
const lanes = match[2] === undefined ? 1 : Number(match[2]);
return (bits * lanes + 7) >> 3;
};
```
--
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]
