[webgpu] Use i32 as coords type in WGSL

tfjs-backend-webgpu/src/kernels/BatchMatMul_impl.ts

@@ -144,8 +144,8 @@ export function batchMatMulImpl({
     const dimInner = transposeA === true ? a3d.shape[1] : a3d.shape[2];
     const dimBOuter = transposeB === true ? b3d.shape[1] : b3d.shape[2];
     dimensions = [
-      {type: 'uint32', data: [dimAOuter]}, {type: 'uint32', data: [dimBOuter]},
-      {type: 'uint32', data: [dimInner]}
+      {type: 'int32', data: [dimAOuter]}, {type: 'int32', data: [dimBOuter]},
+      {type: 'int32', data: [dimInner]}
     ];
   }
   const out = backend.runWebGPUProgram(program, inputs, a.dtype, dimensions);

tfjs-backend-webgpu/src/kernels/Conv2D.ts

@@ -76,9 +76,9 @@ export function conv2d(
     const dimInner =
         convInfo.filterHeight * convInfo.filterWidth * convInfo.inShape[3];
     dimensions.push(
-        {type: 'uint32', data: [dimAOuter]},
-        {type: 'uint32', data: [dimBOuter]},
-        {type: 'uint32', data: [dimInner]});
+        {type: 'int32', data: [dimAOuter]},
+        {type: 'int32', data: [dimBOuter]},
+        {type: 'int32', data: [dimInner]});
   }
 
   return backend.runWebGPUProgram(program, [x, filter], x.dtype, dimensions);

tfjs-backend-webgpu/src/kernels/Conv2D_impl.ts

@@ -166,8 +166,8 @@ export function conv2dWithIm2Col({
     const dimInner = a3dShape[2];
     const dimBOuter = convInfo.outChannels;
     matmulDimensions = [
-      {type: 'uint32', data: [dimAOuter]}, {type: 'uint32', data: [dimBOuter]},
-      {type: 'uint32', data: [dimInner]}
+      {type: 'int32', data: [dimAOuter]}, {type: 'int32', data: [dimBOuter]},
+      {type: 'int32', data: [dimInner]}
     ];
   }
 

tfjs-backend-webgpu/src/kernels/FromPixels_utils/from_pixels_webgpu.ts

@@ -74,9 +74,9 @@ export class FromPixelsProgram implements WebGPUProgram {
       ${getMainHeaderStringWgsl()} {
         ${getGlobalIndexStringWgsl()}
         let flatIndexBase = index * uniforms.numChannels;
-        let coords: vec3<u32> = getCoordsFromFlatIndex(flatIndexBase);
+        let coords = getCoordsFromFlatIndex(flatIndexBase);
         let values = ${textureLoad};
-        for (var i: u32 = 0u; i < uniforms.numChannels; i = i + 1u) {
+        for (var i = 0; i < uniforms.numChannels; i = i + 1) {
           let flatIndex = flatIndexBase + i;
           if (flatIndex < uniforms.size) {
             result.numbers[flatIndex] = i32(floor(255.0 * values[i]));

tfjs-backend-webgpu/src/kernels/FusedConv2D.ts

@@ -98,9 +98,9 @@ export function fusedConv2d(args: {
         convInfo.filterHeight * convInfo.filterWidth * convInfo.inShape[3];
     if (program.useWgsl) {
       dimensions.push(
-          {type: 'uint32', data: [dimAOuter]},
-          {type: 'uint32', data: [dimBOuter]},
-          {type: 'uint32', data: [dimInner]});
+          {type: 'int32', data: [dimAOuter]},
+          {type: 'int32', data: [dimBOuter]},
+          {type: 'int32', data: [dimInner]});
     }
   }
 

tfjs-backend-webgpu/src/kernels/addn_packed_webgpu.ts

@@ -93,8 +93,8 @@ export class AddNPackedProgram implements WebGPUProgram {
     const userCode = `
       ${getMainHeaderStringWgsl()} {
         ${getGlobalIndexStringWgsl()}
-        for (var i = 0u; i < ${this.workPerThread}u; i = i + 1u) {
-          let flatIndex = index * ${this.workPerThread}u + i;
+        for (var i = 0; i < ${this.workPerThread}; i = i + 1) {
+          let flatIndex = index * ${this.workPerThread} + i;
           if (flatIndex < uniforms.size) {
             let coords = getCoordsFromFlatIndex(flatIndex);
             ${snippets.join('\n        ')}

tfjs-backend-webgpu/src/kernels/argminmax_webgpu.ts

@@ -31,7 +31,7 @@ export class ArgMinMaxProgram implements WebGPUProgram {
   workGroupSize: [number, number, number];
   variableNames = ['x'];
   uniforms = 'int axis;';
-  uniformsWgsl = 'axis : u32;';
+  uniformsWgsl = 'axis : i32;';
   inputShape: number[];
   reductionFactor: number;
   op: string;
@@ -203,20 +203,20 @@ export class ArgMinMaxProgram implements WebGPUProgram {
     // and iteratively reduced.
     const reduceInSharedMemory = this.workGroupSize[0] > 1;
     const sharedMemorySnippet = `
-      var<workgroup> xBestIndices : array<u32, ${this.workGroupSize[0]}>;
+      var<workgroup> xBestIndices : array<i32, ${this.workGroupSize[0]}>;
       var<workgroup> xBestValues : array<f32, ${this.workGroupSize[0]}>;
     `;
 
     const sharedMemoryReduceSnippet = `
       xBestIndices[localId.x] = bestIndex;
       xBestValues[localId.x] = bestValue;
 
-      for(var currentSize = WorkGroupSize; currentSize > 1u; currentSize = DIV_CEIL(currentSize, ${
-        this.reductionFactor}u)) {
+      for(var currentSize = WorkGroupSize; currentSize > 1; currentSize = DIV_CEIL(currentSize, ${
+        this.reductionFactor})) {
         workgroupBarrier();
 
-        for (var w = 0u; w < ${this.reductionFactor}u; w = w + 1u) {
-          let i = localId.x * ${this.reductionFactor}u + w;
+        for (var w = 0; w < ${this.reductionFactor}; w = w + 1) {
+          let i = i32(localId.x) * ${this.reductionFactor} + w;
           if (i < currentSize) {
             let candidateIndex = xBestIndices[i];
             let candidate = xBestValues[i];
@@ -255,58 +255,58 @@ export class ArgMinMaxProgram implements WebGPUProgram {
     };
 
     const userCode = `
-      fn DIV_CEIL(a : u32, b : u32) -> u32 {
-        return ((a - 1u) / b + 1u);
+      fn DIV_CEIL(a : i32, b : i32) -> i32 {
+        return ((a - 1) / b + 1);
       }
 
-      let WorkGroupSize = ${this.workGroupSize[0]}u;
+      let WorkGroupSize = ${this.workGroupSize[0]};
 
       ${reduceInSharedMemory ? sharedMemorySnippet : ''}
 
       // In order to get a flattened index into the input tensor, we need to
       // add back the index along the reduced dimension to |outputCoords|.
       // This function outputs the offset to the first value along
       // |axis| and the stride to get the next value of the input along |axis|.
-      fn getInputCoordInfo(globalId : vec3<u32>, globalIndex : u32) -> vec2<u32>{
+      fn getInputCoordInfo(globalId : vec3<u32>, globalIndex : i32) -> vec2<i32>{
         let outputCoords : ${
         outputCoordsType} = getOutputCoords(globalId, globalIndex);
-        var i = ${this.outputShape.length - 1}u;
+        var i = ${this.outputShape.length - 1};
 
-        var stride = 1u;
-        var inputStride = 1u;
-        var offset = 0u;
+        var stride = 1;
+        var inputStride = 1;
+        var offset = 0;
 
-        for (var r = 1u; r <= ${this.inputShape.length}u; r = r + 1u) {
-          let length = ${indexInputShape(`${this.inputShape.length}u - r`)};
-          if (${this.inputShape.length}u - r == uniforms.axis) {
+        for (var r = 1; r <= ${this.inputShape.length}; r = r + 1) {
+          let length = ${indexInputShape(`${this.inputShape.length} - r`)};
+          if (${this.inputShape.length} - r == uniforms.axis) {
             inputStride = stride;
           } else {
             offset = offset + ${
         indexOutputCoords('outputCoords', 'i')} * stride;
-            i = i - 1u;
+            i = i - 1;
           }
           stride = stride * length;
         }
 
-        return vec2<u32>(offset, inputStride);
+        return vec2<i32>(offset, inputStride);
       }
 
-      fn getInputIndex(coordInfo : vec2<u32>, index : u32) -> u32{
+      fn getInputIndex(coordInfo : vec2<i32>, index : i32) -> i32{
         return coordInfo[0] + coordInfo[1] * index;
       }
 
       ${getMainHeaderStringWgsl()} {
         ${getGlobalIndexStringWgsl()}
         let coordInfo = getInputCoordInfo(globalId, index);
 
-        var bestIndex = 0u;
+        var bestIndex = 0;
         var bestValue = x.numbers[getInputIndex(coordInfo, bestIndex)];
 
         let Length = ${indexInputShape('uniforms.axis')};
         let WorkPerThread = DIV_CEIL(Length, WorkGroupSize);
 
-        for (var w = 0u; w < WorkPerThread; w = w + 1u) {
-          let i = globalId.x * WorkPerThread + w;
+        for (var w = 0; w < WorkPerThread; w = w + 1) {
+          let i = i32(globalId.x) * WorkPerThread + w;
           if (i < Length) {
             let candidate = x.numbers[getInputIndex(coordInfo, i)];
             if (candidate ${
@@ -317,11 +317,11 @@ export class ArgMinMaxProgram implements WebGPUProgram {
           }
         }
 
-        let flatOutputIndex = globalId.y;
+        let flatOutputIndex = i32(globalId.y);
         ${
         reduceInSharedMemory ?
             sharedMemoryReduceSnippet :
-            'setOutputFlatI32(flatOutputIndex, i32(bestIndex));'}
+            'setOutputFlatI32(flatOutputIndex, bestIndex);'}
       }
     `;
     return userCode;

tfjs-backend-webgpu/src/kernels/binary_op_shared_webgpu.ts

@@ -159,16 +159,16 @@ export class BinaryOpSharedProgram implements WebGPUProgram {
           // Fill in the shared memory buffer. Here we need a loop to make sure
           // that all data in A|B are uploaded when |sharedMemorySize| is larger
           // than work group size.
-          for(var localIndex = localId.x; localIndex < ${
-        this.lastDimensionSize}u; localIndex = localIndex + ${
-        this.workGroupSize[0]}u) {
+          for(var localIndex = i32(localId.x); localIndex < ${
+        this.lastDimensionSize}; localIndex = localIndex + ${
+        this.workGroupSize[0]}) {
             sharedBuf[localIndex] = f32(${
         this.useSharedMemoryWithB ? 'B' : 'A'}.numbers[localIndex]);
           }
           workgroupBarrier();
 
-          for(var i = 0u; i < ${this.workPerThread}u; i = i + 1u) {
-            let flatIndex = index * ${this.workPerThread}u + i;
+          for(var i = 0; i < ${this.workPerThread}; i = i + 1) {
+            let flatIndex = index * ${this.workPerThread} + i;
 
             ${writeDataSnippet}
           }

tfjs-backend-webgpu/src/kernels/binary_op_webgpu.ts

@@ -152,8 +152,8 @@ export class BinaryOpProgram implements WebGPUProgram {
       ${miscStr}
       ${getMainHeaderStringWgsl()} {
         ${getGlobalIndexStringWgsl()}
-        for (var i = 0u; i < ${this.workPerThread}u; i = i + 1u ) {
-          let flatIndex = index * ${this.workPerThread}u + i;
+        for (var i = 0; i < ${this.workPerThread}; i = i + 1 ) {
+          let flatIndex = index * ${this.workPerThread} + i;
 
           if(flatIndex < uniforms.size) {
             let coords = getCoordsFromFlatIndex(flatIndex);

tfjs-backend-webgpu/src/kernels/clip_vec4_webgpu.ts

@@ -76,7 +76,7 @@ export class ClipVec4Program implements WebGPUProgram {
         if(index < uniforms.size) {
           let value = getAAtOutCoordsByGlobalId(globalId, index);
           var clampedValue : vec4<f32>;
-          for (var i = 0u; i < 4u; i = i + 1u) {
+          for (var i = 0; i < 4; i = i + 1) {
             if (isNanCustom(value[i])) {
               clampedValue[i] = value[i];
             } else {

tfjs-backend-webgpu/src/kernels/concat_webgpu.ts

@@ -104,26 +104,26 @@ export class ConcatProgram implements WebGPUProgram {
       }
 
       snippets.push(`if (yC < ${
-          offsets[0]}u){ setOutput(coords.x, coords.y, getT0(yR, yC)); }`);
+          offsets[0]}){ setOutput(coords.x, coords.y, getT0(yR, yC)); }`);
       for (let i = 1; i < offsets.length; i++) {
         const shift = offsets[i - 1];
         snippets.push(
-            `elseif (yC < ${offsets[i]}u){ ` +
-            `setOutput(coords.x, coords.y, getT${i}(yR, yC - ${shift}u)); }`);
+            `elseif (yC < ${offsets[i]}){ ` +
+            `setOutput(coords.x, coords.y, getT${i}(yR, yC - ${shift})); }`);
       }
       const lastIndex = offsets.length;
       const lastShift = offsets[offsets.length - 1];
       snippets.push(`else { setOutput(coords.x, coords.y, getT${
-          lastIndex}(yR, yC - ${lastShift}u)); }`);
+          lastIndex}(yR, yC - ${lastShift})); }`);
     } else {
       snippets.push(`setOutput(coords.x, coords.y, getT0(yR, yC));`);
     }
 
     const userCode = `
       ${getMainHeaderStringWgsl()} {
         ${getGlobalIndexStringWgsl()}
-        for(var i = 0u; i < ${this.workPerThread}u; i = i + 1u) {
-          let flatIndex = index * ${this.workPerThread}u + i;
+        for(var i = 0; i < ${this.workPerThread}; i = i + 1) {
+          let flatIndex = index * ${this.workPerThread} + i;
           if(flatIndex < uniforms.size) {
             let coords = getCoordsFromFlatIndex(flatIndex);
             let yR = coords.x;