Skip to content

Commit

Permalink
[JS/WebGPU] Improve MatMulNBits perf (#19974)
Browse files Browse the repository at this point in the history
### Description
<!-- Describe your changes. -->
Improve performance using shared memory


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
  • Loading branch information
satyajandhyala authored Apr 12, 2024
1 parent 794d39a commit b33216b
Show file tree
Hide file tree
Showing 4 changed files with 1,097 additions and 128 deletions.
11 changes: 11 additions & 0 deletions js/web/lib/wasm/jsep/init.ts
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,17 @@ class ComputeContextImpl implements ComputeContext {
this.inputs = inputs;
}

getMaxComputeWorkgroupSizes(): [number, number, number] {
return [
this.backend.device.limits.maxComputeWorkgroupSizeX, this.backend.device.limits.maxComputeWorkgroupSizeY,
this.backend.device.limits.maxComputeWorkgroupSizeZ
];
}

getMaxComputeWorkgroupStoragesize(): number {
return this.backend.device.limits.maxComputeWorkgroupStorageSize;
}

compute(program: ProgramInfo, inputsOutputsMapping?: ComputeContextInputsOutputsMapping): TensorView[] {
// prepare inputs. inputs should always be valid data.
const mappedInputs =
Expand Down
Loading

0 comments on commit b33216b

Please sign in to comment.