[webgpu] Fix alignment missing for 5D and 6D

tfjs-backend-webgpu/src/setup_test.ts

@@ -44,14 +44,12 @@ const TEST_FILTERS: TestFilter[] = [
     startsWith: 'avgPool ',
     excludes: [
       'gradient',  // Not yet implemented.
-      //'avgPool3d',  // Not yet implemented.
     ]
   },
   {
     startsWith: 'batchToSpaceND ',
     excludes: [
-      'tensor3d', 'tensor4d', 'gradient',
-      'accepts a tensor-like object',  // tensor6d not yet implemented
+      'gradient',  // Not yet implemented.
     ]
   },
   {
@@ -84,13 +82,6 @@ const TEST_FILTERS: TestFilter[] = [
       'gradient',  // gradient function not found.
     ]
   },
-  {
-    startsWith: 'einsum ',
-    excludes: [
-      '4d tensors',               // rank 5 is not yet supported.
-      '4d tensor and 3d tensor',  // rank 5 is not yet supported.
-    ]
-  },
   {
     startsWith: 'elu ',
     excludes: [
@@ -210,14 +201,6 @@ const TEST_FILTERS: TestFilter[] = [
       'gradient'            // gradient function not found.
     ]
   },
-  {
-    startsWith: 'slice ',
-    excludes: [
-      'slice5d',             // Rank 5 is not yet implemented.
-      'slice6d',             // Rank 6 is not yet implemented.
-      'strided slice with',  // Rank 6 is not yet implemented.
-    ]
-  },
   {
     startsWith: 'softmax ',
     excludes: [
@@ -251,13 +234,6 @@ const TEST_FILTERS: TestFilter[] = [
       'dilation2d',  // 'dilation2d' not yet implemented.
     ]
   },
-  {
-    startsWith: 'stridedSlice ',
-    excludes: [
-      'strided slice with several new axes',  // Rank 6 is not yet implemented.
-      'strided slice with new axes and',      // Rank 6 is not yet implemented.
-    ]
-  },
   {
     startsWith: 'tensor ',
     excludes: [

tfjs-backend-webgpu/src/slice_webgpu.ts

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {getCoordsDataType, getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
+import {getCoordsDataType, getCoordsXYZ, getMainHeaderAndGlobalIndexString, WebGPUProgram} from './webgpu_program';
 import {computeDispatch, flatDispatchLayout} from './webgpu_util';
 
 export class SliceProgram implements WebGPUProgram {
@@ -54,8 +54,8 @@ export class SliceProgram implements WebGPUProgram {
       });
     } else {
       coordSum = this.outputShape.map((_, i) => {
-        return `sourceLoc.${coords[i]} = uniforms.start[${i}] + coords.${
-            coords[i]};`;
+        return `sourceLoc.${coords[i]} = uniforms.start.${
+            getCoordsXYZ(i)} + coords.${coords[i]};`;
       });
     }
 

tfjs-backend-webgpu/src/webgpu_program.ts

@@ -179,55 +179,28 @@ function makeShader(
     ].join('\n');
   }
 
-  let preMemberIsStruct = false;
-  let currentMemberIsStruct = false;
   let uniformDeclaration = 'struct Uniforms { NAN : f32, ';
   program.variableNames.forEach((x, i) => {
     const perDataType = getCoordsDataType(inputInfo[i].shape.length);
-    if (perDataType === 'vec5' || perDataType === 'vec6') {
-      currentMemberIsStruct = true;
-    }
-    if (preMemberIsStruct || currentMemberIsStruct) {
-      uniformDeclaration += `@align(16) `;
-    }
-    preMemberIsStruct = currentMemberIsStruct;
     uniformDeclaration +=
         `${x.charAt(0).toLowerCase() + x.slice(1)}Shape : ${perDataType}, `;
   });
   const outputDataType = getCoordsDataType(outputData.shape.length);
-  currentMemberIsStruct =
-      outputDataType === 'vec5' || outputDataType === 'vec6';
-  if (preMemberIsStruct || currentMemberIsStruct) {
-    uniformDeclaration += `@align(16) `;
-  }
-  preMemberIsStruct = currentMemberIsStruct;
   uniformDeclaration += `outShape : ${outputDataType}, `;
   const stridesLength = outputData.shape.length - 1;
   const stridesDataType = getCoordsDataType(stridesLength);
-  currentMemberIsStruct =
-      stridesDataType === 'vec5' || stridesDataType === 'vec6';
-  if (preMemberIsStruct || currentMemberIsStruct) {
-    uniformDeclaration += `@align(16) `;
-  }
-  preMemberIsStruct = currentMemberIsStruct;
   uniformDeclaration += `
          outShapeStrides: ${stridesDataType}, `;
 
   if (program.size) {
-    if (preMemberIsStruct) {
-      uniformDeclaration += `@align(16) `;
-    }
-    preMemberIsStruct = false;
     uniformDeclaration += 'size : i32, ';
   }
 
   if (program.uniforms) {
-    if (preMemberIsStruct) {
-      uniformDeclaration += `@align(16) `;
-    }
     uniformDeclaration += program.uniforms;
   }
   uniformDeclaration += '};';
+  uniformDeclaration = insertAlignment(uniformDeclaration);
 
   prefixSnippets.push(uniformDeclaration);
 
@@ -832,3 +805,18 @@ function setOutputSnippet(
 
   return snippet;
 }
+
+function insertAlignment(uniformShader: string) {
+  // insert alignment when current pattern is vec5 or vec6
+  const curInsertRe = /(\w+)\s*:\s*vec(5|6)/g;
+  uniformShader = uniformShader.replace(curInsertRe, (match) => {
+    return '@align(16) ' + match;
+  });
+
+  // insert alignment when previous pattern is vec5 or vec6
+  const preInsertRe = /vec(5|6)\s*,\s*(\w+)/g;
+  uniformShader = uniformShader.replace(preInsertRe, (_, p1, p2) => {
+    return `vec${p1}, @align(16) ${p2}`;
+  });
+  return uniformShader;
+}