gemini-code-assist[bot] commented on code in PR #18871:
URL: https://github.com/apache/tvm/pull/18871#discussion_r2884999162
##########
web/src/webgpu.ts:
##########
@@ -727,17 +812,37 @@ export class WebGPUContext {
}
});
- compute.setBindGroup(0, this.device.createBindGroup({
- layout: bindGroupLayout,
- entries: bindGroupEntries
- }));
+ // Bind group cache: build key from pipeline identity + buffer ptrs +
uniform key
+ let bgCacheKey = pipelineId;
+ for (let i = 0; i < bufferArgIndices.length; ++i) {
+ bgCacheKey += ":" + args[bufferArgIndices[i]];
+ }
+ bgCacheKey += ":" + uniformKey;
Review Comment:
The `bgCacheKey` used for caching WebGPU bind groups is constructed by
concatenating the `pipelineId`, buffer pointers, and the uniform key using a
colon (`:`) as a separator. However, the `pipelineId` itself is constructed
from `finfo.name` (the shader name) and the number of buffer arguments, also
using a colon. If a shader name contains a colon, it is possible to craft two
different shader/argument combinations that result in the same `bgCacheKey`.
For example, a collision can occur between:
1. Shader A: name=`"a:2"`, 1 buffer argument (pointer `10`), uniform key
`"100"` -> `bgCacheKey = "a:2:1:10:100"`.
2. Shader B: name=`"a"`, 2 buffer arguments (pointer `1`, pointer `10`),
uniform key `"100"` -> `bgCacheKey = "a:2:1:10:100"`.
When a collision occurs, a shader will use the bind group (and thus the GPU
buffers) of a different shader. This can lead to unauthorized reading or
writing of GPU memory, potentially leaking sensitive data or causing memory
corruption on the GPU.
##########
web/src/runtime.ts:
##########
@@ -1674,8 +1679,17 @@ export class Instance implements Disposable {
* @returns The created shape tuple.
*/
makeShapeTuple(shape: Array<number>): TVMObject {
+ const key = shape.toString();
+ const cached = this.shapeTupleCache.get(key);
+ if (cached !== undefined) {
+ return cached;
+ }
const shapeArray = shape.map((value) => new Scalar(value, "int"));
- return this.ctx.makeShapeTuple(...shapeArray);
+ const tuple = this.ctx.makeShapeTuple(...shapeArray);
+ // Detach from scope so the cached object survives across scopes.
+ this.detachFromCurrentScope(tuple);
+ this.shapeTupleCache.set(key, tuple);
+ return tuple;
Review Comment:
The `shapeTupleCache` caches `TVMObject` (`ShapeTuple`) instances without an
eviction policy or size limit. This can lead to unbounded memory usage and
memory exhaustion, potentially causing a Denial of Service by crashing the
browser tab. It is recommended to implement a simple FIFO eviction policy with
a reasonable limit (e.g., 512) to prevent this.
```suggestion
if (this.shapeTupleCache.size >= 512) { // Using 512 as a reasonable
limit.
const oldestKey = this.shapeTupleCache.keys().next().value;
if (oldestKey !== undefined) {
this.shapeTupleCache.get(oldestKey)?.dispose();
this.shapeTupleCache.delete(oldestKey);
}
}
this.shapeTupleCache.set(key, tuple);
return tuple;
```
--
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]
