[js/web] adding webgl pointwise conv kernel (#8418)

js/web/lib/onnxjs/backends/webgl/ops/conv.ts

@@ -5,7 +5,7 @@ import {Attribute} from '../../../attribute';
 import {Logger} from '../../../instrument';
 import {Conv} from '../../../ops/conv';
 import {Tensor} from '../../../tensor';
-import {PoolConvUtil} from '../../../util';
+import {assert, PoolConvUtil} from '../../../util';
 import {getGlsl} from '../glsl-source';
 import {WebGLInferenceHandler} from '../inference-handler';
 import {Artifact, ProgramInfo, RunData, TextureLayout, WebGLOperator} from '../types';
@@ -13,17 +13,21 @@ import {WebGLContext} from '../webgl-context';
 
 import {WebGLConvPacked} from './conv-pack';
 import {getActicationSnippet} from './fuse-utils';
+import {WebGLUnpackedMatMul} from './matmul';
+import {reshape} from './reshape';
 
 export class WebGLConv extends Conv {
   unpackedGroupedConvImpl: WebGLUnpackedGroupedConv;
   unpackedConvImpl: WebGLUnpackedConv;
   packedConvImpl: WebGLConvPacked;
+  unpackedPointwiseConvImpl: WebGLUnpackedPointwiseConv;
 
   constructor() {
     super();
     this.unpackedGroupedConvImpl = new WebGLUnpackedGroupedConv();
     this.unpackedConvImpl = new WebGLUnpackedConv();
     this.packedConvImpl = new WebGLConvPacked();
+    this.unpackedPointwiseConvImpl = new WebGLUnpackedPointwiseConv();
   }
 
   initialize(attributes: Attribute): void {
@@ -31,6 +35,7 @@ export class WebGLConv extends Conv {
     this.unpackedGroupedConvImpl.initialize(attributes);
     this.unpackedConvImpl.initialize(attributes);
     this.packedConvImpl.initialize(attributes);
+    this.unpackedPointwiseConvImpl.initialize(attributes);
   }
 
   run(inferenceHandler: WebGLInferenceHandler, inputs: Tensor[]): Tensor[] {
@@ -38,7 +43,9 @@ export class WebGLConv extends Conv {
     const isPointwise = this.kernelShape[0] === 1 && this.kernelShape[1] === 1;
     if (this.group > 1) {
       return this.unpackedGroupedConvImpl.run(inferenceHandler, inputs);
-    } else if (packMode && inputs[0].dims.length === 4 && inputs[0].dims[0] === 1 && !isPointwise) {
+    } else if (isPointwise && packMode) {
+      return this.unpackedPointwiseConvImpl.run(inferenceHandler, inputs);
+    } else if (packMode && inputs[0].dims.length === 4 && inputs[0].dims[0] === 1) {
       return this.packedConvImpl.run(inferenceHandler, inputs);
     } else {
       return this.unpackedConvImpl.run(inferenceHandler, inputs);
@@ -62,6 +69,64 @@ export class WebGLConv extends Conv {
   }
 }
 
+export class WebGLUnpackedPointwiseConv extends Conv {
+  private programInfo: ProgramInfo[];
+  private artifacts: Artifact[];
+  private matmul = new WebGLUnpackedMatMul();
+  private outputShape: number[];
+
+  run(inferenceHandler: WebGLInferenceHandler, inputs: Tensor[]): Tensor[] {
+    const programManager = inferenceHandler.session.programManager;
+    const xShape = inputs[0].dims.slice();
+    const wShape = inputs[1].dims.slice();
+    PoolConvUtil.adjustPadsBasedOnAutoPad(
+        inputs[0].dims, this.strides, this.dilations, this.kernelShape, this.pads, this.autoPad);
+    Logger.verbose(
+        'Conv',
+        `autpPad:${this.autoPad}, dilations:${this.dilations}, group:${this.group}, kernelShape:${
+            this.kernelShape}, pads:${this.pads}, strides:${this.strides}`);
+
+    if (!this.outputShape) {
+      this.outputShape = WebGLConv.calcOutputShape(xShape, wShape, this.dilations, this.pads, this.strides);
+    }
+
+    if (this.activation) {
+      const attributes = new Attribute(undefined);
+      attributes.set('__internal_activation', 'string', (this.activation));
+      if (this.activation === 'Clip') {
+        attributes.set('__clip_max', 'float', this.clipMax);
+        attributes.set('__clip_min', 'float', this.clipMin);
+      }
+      this.matmul.initialize(attributes);
+    }
+    const reshapedX = reshape(inferenceHandler, inputs[0], [xShape[1], xShape[2] * xShape[3]]);
+    const reshapedW = reshape(inferenceHandler, inputs[1], [wShape[0], wShape[1]]);
+    const hasBias = (inputs.length === 3);
+
+    if (!this.artifacts) {
+      this.artifacts = [];
+      this.programInfo = [];
+      this.programInfo[0] = hasBias ?
+          this.matmul.createProgramInfo(inferenceHandler, [reshapedW, reshapedX, inputs[2]]) :
+          this.matmul.createProgramInfo(inferenceHandler, [reshapedW, reshapedX]);
+      this.artifacts[0] = programManager.build(this.programInfo[0]);
+    }
+
+    assert(this.artifacts.length === 1, () => 'expect matmul artifacts was created');
+    const runDataMatmul = this.matmul.createRunData(
+        inferenceHandler, this.programInfo[0], hasBias ? [reshapedW, reshapedX, inputs[2]] : [reshapedW, reshapedX]);
+    programManager.run(this.artifacts[0], runDataMatmul);
+    const matmulOutput = runDataMatmul.outputTextureData.tensor;
+
+    const res = [reshape(inferenceHandler, matmulOutput, this.outputShape)];
+    // eslint-disable-next-line no-console
+    // console.log(
+    //     'res[0]=', res[0].data[0], ' res[1]=', res[0].data[1], ' res[1]=', res[0].data[2], ' res[3]=',
+    //     res[0].data[3]);
+    return res;
+  }
+}
+
 export class WebGLUnpackedGroupedConv extends Conv implements WebGLOperator {
   run(inferenceHandler: WebGLInferenceHandler, inputs: Tensor[]): Tensor[] {
     return inferenceHandler.run(this, inputs);

js/web/lib/onnxjs/backends/webgl/ops/matmul-pack.ts

@@ -11,6 +11,7 @@ import {ProgramInfo, RunData, WebGLOperator} from '../types';
 import {getCoordsDataType} from '../utils';
 
 import {getActicationSnippet} from './fuse-utils';
+import {getBiasForMatmul} from './matmul';
 
 export class WebGLMatMulPacked extends MatMul implements WebGLOperator {
   run(inferenceHandler: WebGLInferenceHandler, inputs: Tensor[]): Tensor[] {
@@ -98,38 +99,6 @@ export class WebGLMatMulPacked extends MatMul implements WebGLOperator {
   }
 }
 
-function getBiasForMatmul(
-    coordsDataType: string, allGlChannels: readonly string[], inShape: readonly number[],
-    outShape: readonly number[]): string {
-  let unpackedCoordsSnippet = '';
-  const inRank = inShape.length;
-  const outRank = outShape.length;
-  const rankDiff = outRank - inRank;
-  if (outRank < 2 && inRank > 0) {
-    unpackedCoordsSnippet = 'coords';
-  } else {
-    unpackedCoordsSnippet = inShape.map((s, i) => `coords.${allGlChannels[i + rankDiff]}`).join(', ');
-  }
-  const broadcastDims = BroadcastUtil.getBroadcastDims(inShape, outShape);
-  const coordsSnippet = broadcastDims.map(d => `coords.${allGlChannels[d + rankDiff]} = 0;`).join('\n');
-  const inSize = ShapeUtil.size(inShape);
-  const isInputScalar = inSize === 1;
-  let output = 'vec4(outputValue.xx, outputValue.yy)';
-  if (isInputScalar) {
-    output = 'vec4(outputValue.x)';
-  }
-  const getBiasForMatmulSource = `
-  vec4 getBiasForMatmul() {
-    ${coordsDataType} coords = getOutputCoords();
-    ${coordsSnippet}
-    vec4 outputValue = getBias(${unpackedCoordsSnippet});
-    return ${output};
-
-  }`;
-
-  return getBiasForMatmulSource;
-}
-
 function getBcastSamplerForMatmul(
     coordsDataType: string, allGlChannels: readonly string[], inputs: Tensor[], outShape: readonly number[]): string {
   let unpackedACoordsSnippet = [];

js/web/lib/onnxjs/backends/webgl/ops/matmul.ts

@@ -3,9 +3,11 @@
 
 import {MatMul} from '../../../ops/matmul';
 import {Tensor} from '../../../tensor';
-import {BroadcastUtil} from '../../../util';
+import {BroadcastUtil, ShapeUtil} from '../../../util';
 import {WebGLInferenceHandler} from '../inference-handler';
 import {ProgramInfo, RunData, WebGLOperator} from '../types';
+import {getCoordsDataType} from '../utils';
+import {getActicationSnippet} from './fuse-utils';
 import {WebGLMatMulPacked} from './matmul-pack';
 
 export class WebGLMatMul extends MatMul implements WebGLOperator {
@@ -57,17 +59,34 @@ export class WebGLUnpackedMatMul extends MatMul implements WebGLOperator {
     return inferenceHandler.run(this, inputs);
   }
   createProgramInfo(handler: WebGLInferenceHandler, inputs: Tensor[]): ProgramInfo {
+    const hasBias = inputs.length > 2;
+    const processBias = hasBias ? 'value += getBiasForMatmul();' : '';
     const aShape = inputs[0].dims;
     const bShape = inputs[1].dims;
     const outputShape = BroadcastUtil.calcShape(aShape, bShape, true);
     if (!outputShape) {
       throw new Error('Can\'t use matmul on the given tensors');
     }
+    const coordsDataType = getCoordsDataType(outputShape.length);
+    const allGlChannels = ['x', 'y', 'z', 'w', 'u', 'v'];
+
+    const {activationFunction, applyActivation} = getActicationSnippet(this.activation);
+    const additionalVars = this.activation === 'Clip' ? `
+    float min = float(${this.clipMin});
+    float max = float(${this.clipMax});` :
+                                                        '';
+    const getBiasForMatmulSnippet =
+        hasBias ? `${getBiasForMatmul(coordsDataType, allGlChannels, inputs[2].dims, outputShape, false)}` : '';
+
     const rank = outputShape.length;
     const arank = aShape.length;
     const brank = bShape.length;
     const sharedDim = aShape[aShape.length - 1];
     const shaderSource = `
+    ${additionalVars}
+    ${activationFunction}
+    ${getBiasForMatmulSnippet}
+
       float process(int indices[${rank}]) {
           int a[${arank}];
           int b[${brank}];
@@ -80,12 +99,15 @@ export class WebGLUnpackedMatMul extends MatMul implements WebGLOperator {
               b[${brank - 2}] = k;
               value += _A(a) * _B(b);
           }
+          ${processBias}
+          ${applyActivation}
+
           return value;
       }`;
     return {
       inputLayouts: inputs.map(t => handler.getOrCreateTextureLayout(t)),
       outputLayout: handler.createTextureLayoutFromShape(outputShape),
-      samplers: ['A', 'B'],
+      samplers: hasBias ? ['A', 'B', 'Bias'] : ['A', 'B'],
       shaderSource,
     };
   }
@@ -98,3 +120,41 @@ export class WebGLUnpackedMatMul extends MatMul implements WebGLOperator {
     };
   }
 }
+
+export function getBiasForMatmul(
+    coordsDataType: string, allGlChannels: readonly string[], inShape: readonly number[], outShape: readonly number[],
+    isPacked = true): string {
+  let unpackedCoordsSnippet = '';
+  const inRank = inShape.length;
+  const outRank = outShape.length;
+  const rankDiff = outRank - inRank;
+  if (outRank < 2 && inRank > 0) {
+    unpackedCoordsSnippet = 'coords';
+  } else {
+    unpackedCoordsSnippet = inShape.map((s, i) => `coords.${allGlChannels[i + rankDiff]}`).join(', ');
+  }
+  const broadcastDims = BroadcastUtil.getBroadcastDims(inShape, outShape);
+  const coordsSnippet = broadcastDims.map(d => `coords.${allGlChannels[d + rankDiff]} = 0;`).join('\n');
+  const inSize = ShapeUtil.size(inShape);
+  const isInputScalar = inSize === 1;
+  let output = 'vec4(outputValue.xx, outputValue.yy)';
+  if (isInputScalar) {
+    output = 'vec4(outputValue.x)';
+  }
+  const getBiasForMatmulSource = isPacked ? `
+vec4 getBiasForMatmul() {
+  ${coordsDataType} coords = getOutputCoords();
+  ${coordsSnippet}
+  vec4 outputValue = getBias(${unpackedCoordsSnippet});
+  return ${output};
+}` :
+                                            `
+float getBiasForMatmul() {
+  ${coordsDataType} coords = getOutputCoords();
+  ${coordsSnippet}
+  return getBias(coords.x);
+}
+`;
+
+  return getBiasForMatmulSource;
+}

js/web/lib/onnxjs/backends/webgl/ops/reshape.ts

@@ -7,8 +7,10 @@ import {ShapeUtil} from '../../../util';
 import {WebGLInferenceHandler} from '../inference-handler';
 import {TextureLayout} from '../types';
 import {getPackedShape} from '../utils';
+
 import {WebGLReshapePacked} from './reshape-packed';
 
+
 export class WebGLReshape extends Reshape {
   packedImpl: WebGLReshapePacked;
   constructor() {
@@ -28,10 +30,15 @@ export class WebGLReshape extends Reshape {
 
 export function reshape(
     inferenceHandler: WebGLInferenceHandler, input: Tensor, reshapedDims: readonly number[]): Tensor {
-  const inputTD = inferenceHandler.getOrCreateTextureData(input);
+  let inputTD = inferenceHandler.getOrCreateTextureData(input);
   let packedShape = reshapedDims;
+
   if (inputTD.channels === 4) {
-    packedShape = getPackedShape(reshapedDims);
+    if (!inputTD.isPacked) {
+      packedShape = getPackedShape(reshapedDims);
+    } else {
+      inputTD = inferenceHandler.unpack(inputTD);
+    }
   }
   const newTextureLayout: TextureLayout = {
     channels: inputTD.channels,

js/web/test/data/ops/conv.jsonc

@@ -416,5 +416,92 @@
         ]
       }
     ]
+  },
+
+  {
+    "name": "conv - pointwise",
+    "operator": "Conv",
+    "attributes": [{ "name": "kernel_shape", "data": [1, 1], "type": "ints" }],
+    "cases": [
+      {
+        "name": "T[0]",
+        "inputs": [
+          {
+            "data": [
+              0.0,
+              1.0,
+              2.0,
+              3.0,
+              4.0,
+              5.0,
+              6.0,
+              7.0,
+              8.0,
+              9.0,
+              10.0,
+              11.0,
+              12.0,
+              13.0,
+              14.0,
+              15.0,
+              16.0,
+              17.0,
+              18.0,
+              19.0,
+              20.0,
+              21.0,
+              22.0,
+              23.0,
+              24.0,
+              25.0,
+              26.0,
+              27.0,
+              28.0,
+              29.0,
+              30.0,
+              31.0
+            ],
+            "dims": [1, 8, 2, 2],
+            "type": "float32"
+          },
+          {
+            "data":[
+              0.0,
+              1.0,
+              2.0,
+              3.0,
+              4.0,
+              5.0,
+              6.0,
+              7.0,
+              8.0,
+              9.0,
+              10.0,
+              11.0,
+              12.0,
+              13.0,
+              14.0,
+              15.0
+            ],
+            "dims":[2, 8, 1, 1],
+            "type":"float32"
+          },
+          {
+            "data": [0.5, 0.4],
+            "dims": [2],
+            "type": "float32"
+          }
+        ],
+        "outputs":[
+          {
+            "data": [
+              560.5,588.5,616.5,644.5,1456.4,1548.4,1640.4,1732.4
+            ],
+            "dims": [1, 2, 2, 2],
+            "type": "float32"
+          }
+        ]
+      }
+    ]
   }
 ]