webgpu: add a naive implementation of conv3d

tfjs-backend-webgpu/src/conv3d_naive_webgpu.ts

@@ -0,0 +1,129 @@
+/**
+ * @license
+ * Copyright 2023 Google LLC.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {backend_util} from '@tensorflow/tfjs-core';
+
+import {getMainHeaderString as main, WebGPUProgram} from './webgpu_program';
+import {computeDispatch, flatDispatchLayout} from './webgpu_util';
+
+export class Conv3DNaiveProgram implements WebGPUProgram {
+  outputShape: number[];
+  shaderKey: string;
+  dispatchLayout: {x: number[]};
+  dispatch: [number, number, number];
+  variableNames = ['x', 'W'];
+  uniforms =
+      'filterDims: vec3<i32>, pads: vec3<i32>, strides: vec3<i32>, dilations: vec3<i32>,';
+  workgroupSize: [number, number, number] = [64, 1, 1];
+  size = true;
+
+  constructor(convInfo: backend_util.Conv3DInfo) {
+    this.outputShape = convInfo.outShape;
+    this.dispatchLayout = flatDispatchLayout(this.outputShape);
+    this.dispatch = computeDispatch(
+        this.dispatchLayout, this.outputShape, this.workgroupSize);
+
+    this.shaderKey = `conv3dnaive`;
+  }
+
+  getUserCode(): string {
+    const userCode = `
+    ${main('index')} {
+      if (index < uniforms.size) {
+        let coords = getOutputCoords();
+        let batch = coords.x;
+        let d2 = coords.u;
+
+        let xFRCCorner = vec3<i32>(coords.y, coords.z, coords.w) * uniforms.strides - uniforms.pads;
+        let xFCorner = xFRCCorner.x;
+        let xRCorner = xFRCCorner.y;
+        let xCCorner = xFRCCorner.z;
+
+        let inputDepthNearestVec4 = (uniforms.xShape.u / 4) * 4;
+        let inputDepthVec4Remainder = uniforms.xShape.u % 4;
+
+        var dotProd = 0.0;
+        for (var wF = 0; wF < uniforms.filterDims[0]; wF++) {
+          let xF = xFCorner + wF * uniforms.dilations[0];
+          if (xF < 0 || xF >= uniforms.xShape.y) {
+            continue;
+          }
+
+          for (var wR = 0; wR < uniforms.filterDims[1]; wR++) {
+            let xR = xRCorner + wR * uniforms.dilations[1];
+            if (xR < 0 || xR >= uniforms.xShape.z) {
+              continue;
+            }
+
+            for (var wC = 0; wC < uniforms.filterDims[2]; wC++) {
+              let xC = xCCorner + wC * uniforms.dilations[2];
+              if (xC < 0 || xC >= uniforms.xShape.w) {
+                continue;
+              }
+
+              for (var d1 = 0; d1 < inputDepthNearestVec4; d1 += 4) {
+                let xValues = vec4<f32>(
+                  getX(batch, xF, xR, xC, d1),
+                  getX(batch, xF, xR, xC, d1 + 1),
+                  getX(batch, xF, xR, xC, d1 + 2),
+                  getX(batch, xF, xR, xC, d1 + 3)
+                );
+                let wValues = vec4<f32>(
+                  getW(wF, wR, wC, d1, d2),
+                  getW(wF, wR, wC, d1 + 1, d2),
+                  getW(wF, wR, wC, d1 + 2, d2),
+                  getW(wF, wR, wC, d1 + 3, d2)
+                );
+
+                dotProd += dot(xValues, wValues);
+              }
+
+              if (inputDepthVec4Remainder == 1) {
+                dotProd += getX(batch, xF, xR, xC, inputDepthNearestVec4) *
+                  getW(wF, wR, wC, inputDepthNearestVec4, d2);
+              } else if (inputDepthVec4Remainder == 2) {
+                let xValues = vec2<f32>(
+                  getX(batch, xF, xR, xC, inputDepthNearestVec4),
+                  getX(batch, xF, xR, xC, inputDepthNearestVec4 + 1)
+                );
+                let wValues = vec2<f32>(
+                  getW(wF, wR, wC, inputDepthNearestVec4, d2),
+                  getW(wF, wR, wC, inputDepthNearestVec4 + 1, d2)
+                );
+                dotProd += dot(xValues, wValues);
+              } else if (inputDepthVec4Remainder == 3) {
+                let xValues = vec3<f32>(
+                  getX(batch, xF, xR, xC, inputDepthNearestVec4),
+                  getX(batch, xF, xR, xC, inputDepthNearestVec4 + 1),
+                  getX(batch, xF, xR, xC, inputDepthNearestVec4 + 2)
+                );
+                let wValues = vec3<f32>(
+                  getW(wF, wR, wC, inputDepthNearestVec4, d2),
+                  getW(wF, wR, wC, inputDepthNearestVec4 + 1, d2),
+                  getW(wF, wR, wC, inputDepthNearestVec4 + 2, d2)
+                );
+                dotProd += dot(xValues, wValues);
+              }
+            }
+          }
+        }
+        setOutputAtIndex(index, dotProd);
+      }
+    }`;
+    return userCode;
+  }
+}

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

@@ -0,0 +1,61 @@
+/**
+ * @license
+ * Copyright 2023 Google LLC.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {backend_util, Conv3D, Conv3DAttrs, Conv3DInputs, KernelConfig, KernelFunc, upcastType} from '@tensorflow/tfjs-core';
+
+import {WebGPUBackend} from '../backend_webgpu';
+import {Conv3DNaiveProgram} from '../conv3d_naive_webgpu';
+
+export function conv3D(
+    args: {inputs: Conv3DInputs, attrs: Conv3DAttrs, backend: WebGPUBackend}) {
+  const {inputs, backend, attrs} = args;
+  const {x, filter} = inputs;
+  const {strides, pad, dilations} = attrs;
+
+  const convInfo = backend_util.computeConv3DInfo(
+      x.shape as [number, number, number, number, number],
+      filter.shape as [number, number, number, number, number], strides,
+      dilations, pad);
+
+  const padInfo =
+      [convInfo.padInfo.front, convInfo.padInfo.top, convInfo.padInfo.left];
+  const dimensions = [
+    {
+      type: 'int32',
+      data: [convInfo.filterDepth, convInfo.filterHeight, convInfo.filterWidth]
+    },
+    {type: 'int32', data: [...padInfo]}, {
+      type: 'int32',
+      data: [convInfo.strideDepth, convInfo.strideHeight, convInfo.strideWidth]
+    },
+    {
+      type: 'int32',
+      data: [
+        convInfo.dilationDepth, convInfo.dilationHeight, convInfo.dilationWidth
+      ]
+    }
+  ];
+  const program = new Conv3DNaiveProgram(convInfo);
+  const dtype = upcastType(x.dtype, filter.dtype);
+  return backend.runWebGPUProgram(program, [x, filter], dtype, dimensions);
+}
+
+export const conv3DConfig: KernelConfig = {
+  kernelName: Conv3D,
+  backendName: 'webgpu',
+  kernelFunc: conv3D as {} as KernelFunc,
+};

tfjs-backend-webgpu/src/register_all_kernels.ts

@@ -44,6 +44,7 @@ import {concatConfig} from './kernels/Concat';
 import {conv2DConfig} from './kernels/Conv2D';
 import {conv2DBackpropFilterConfig} from './kernels/Conv2DBackpropFilter';
 import {conv2DBackpropInputConfig} from './kernels/Conv2DBackpropInput';
+import {conv3DConfig} from './kernels/Conv3D';
 import {cosConfig} from './kernels/Cos';
 import {coshConfig} from './kernels/Cosh';
 import {cropAndResizeConfig} from './kernels/CropAndResize';
@@ -189,6 +190,7 @@ const kernelConfigs: KernelConfig[] = [
   conv2DConfig,
   conv2DBackpropFilterConfig,
   conv2DBackpropInputConfig,
+  conv3DConfig,
   cosConfig,
   coshConfig,
   cropAndResizeConfig,

tfjs-backend-webgpu/src/setup_test.ts

@@ -57,6 +57,12 @@ const TEST_FILTERS: TestFilter[] = [
       'gradient',  // gradient function not found.
     ]
   },
+  {
+    startsWith: 'conv3d ',
+    excludes: [
+      'gradient',  // Not yet implemented.
+    ]
+  },
   {
     startsWith: 'cumprod ',
     excludes: [
@@ -238,7 +244,6 @@ const TEST_FILTERS: TestFilter[] = [
       'conv2DBackpropFilter ',
       'gradient with clones, input=2x2x1,d2=1,f=1,s=1,d=1,p=same',  // Conv2DBackpropFilter
       'conv1d gradients',  // Conv2DBackpropFilter
-      'conv3d ',
       'conv3dTranspose ',
       'maxPool3d ',
       'maxPool3dBackprop ',

tfjs-backend-webgpu/src/shader_util.ts

@@ -18,12 +18,13 @@
 // Generates WGSL that computes strides.
 export function symbolicallyComputeStrides(
     indicesArr: number[], variableName: string): string[] {
-  if (Math.max(...indicesArr) > 3) {
-    throw new Error('Cannot symbolically compute strides for rank > 4 tensor.');
+  if (Math.max(...indicesArr) > 5) {
+    throw new Error('Cannot symbolically compute strides for rank > 6 tensor.');
   }
 
   const numCoords = indicesArr.length;
-  const shape = indicesArr.map(d => `${variableName}[${d}]`);
+  const indicesStr = 'xyzwuv';
+  const shape = indicesArr.map(d => `${variableName}.${indicesStr[d]}`);
   const strides = new Array(numCoords - 1);
   strides[numCoords - 2] = shape[numCoords - 1];
   for (let i = numCoords - 3; i >= 0; --i) {