Skip to content

Commit

Permalink
webgpu: Fix wgsl failure from MatMulSplitKProgram (#6838)
Browse files Browse the repository at this point in the history
Fix that Tint WGSL reader failure: index 3 out of bounds [0..2] from MatMulSplitKProgram.
  • Loading branch information
qjia7 authored Sep 16, 2022
1 parent 9317f2e commit 52d30d6
Showing 1 changed file with 7 additions and 4 deletions.
11 changes: 7 additions & 4 deletions tfjs-backend-webgpu/src/webgpu_program.ts
Original file line number Diff line number Diff line change
Expand Up @@ -629,6 +629,13 @@ function getOutputCoordsSnippet(
const {x, y = [], z = []} = dispatchLayout;

const outRank = outShape.length;
const rank = x.length + y.length + z.length;
// getOutputCoords is only meaningful when the output rank is same with
// dispatch layout rank.
if (rank !== outRank) {
return '';
}

if (x.length === outRank) {
const dtype = getCoordsDataType(outRank);
const snippet = `fn getOutputCoords() -> ${dtype}{
Expand All @@ -642,17 +649,13 @@ function getOutputCoordsSnippet(
let gatherDimensionsStr = '';
const dims = [x, y, z];

let rank = 0;

for (let i = 0; i < dims.length; i++) {
const arr = dims[i];

if (arr.length === 0) {
continue;
}

rank += arr.length;

if (arr.length === 1) {
gatherDimensionsStr += `let d${arr[0]} = i32(globalId[${i}]);`;
} else {
Expand Down

0 comments on commit 52d30d6

Please sign in to comment.