fixed bugs in packed mode and enable pack mode tests in ci (#7848)

* fixed bugs in packed mode and enable pack mode tests in ci

* removed unnecessary space

* pr comments

* pr comments

* disable an average pool test

* try disabling another avg pool

* disable more avg pool tests

* disable maxpool tests

js/web/lib/onnxjs/backends/webgl/glsl-coordinate-lib.ts

@@ -654,7 +654,7 @@ export class CoordsGlslLib extends GlslLib {
 
     if (inRank === 1 && !isInputScalar && !isOutputScalar) {
       output = `
-        return vec4(outputValue.xy, outputValue.xy);
+        return vec4(outputValue.xx, outputValue.yy);
       `;
     } else if (isInputScalar && !isOutputScalar) {
       if (outRank === 1) {
@@ -1168,7 +1168,7 @@ export class CoordsGlslLib extends GlslLib {
             return ${funcName}(${getSqueezedParams(params, keptDims)});
           }
         `;
-      return new GlslLibRoutine(source, ['coordinates.sampleTexture']);
+      return new GlslLibRoutine(source, ['coordinates.sampleTexture', 'coordinates.uvFromFlat']);
     }
 
     const texNumR = inputLayout.width;

js/web/lib/onnxjs/backends/webgl/ops/binary-op.ts

@@ -10,6 +10,8 @@ import {WebGLInferenceHandler} from '../inference-handler';
 import {ProgramInfo, RunData, WebGLOperator} from '../types';
 
 export class WebGLBinaryOp extends BinaryOp implements WebGLOperator {
+  private usePackedTexture?: boolean;
+
   constructor(
       typeConstraint: readonly Tensor.DataType[], protected glslFunc: GlslValueFunction, opType?: string,
       resultType?: Tensor.DataType) {
@@ -19,14 +21,20 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator {
     return inferenceHandler.run(this, inputs);
   }
   createProgramInfo(handler: WebGLInferenceHandler, inputs: Tensor[]): ProgramInfo {
-    const inputLayouts = handler.session.pack ?
+    const isBroadcast = !ShapeUtil.areEqual(inputs[0].dims, inputs[1].dims);
+
+    // TODO fix bcast in packed mode.
+    if (this.usePackedTexture === undefined) {
+      this.usePackedTexture = !isBroadcast && handler.session.pack;
+    }
+
+    const inputLayouts = this.usePackedTexture ?
         inputs.map(t => handler.getOrCreateTextureLayout(t, 4, true, t.dims, true)) :
         inputs.map(t => handler.getOrCreateTextureLayout(t));
-    const ouputLayout = handler.session.pack ?
+    const ouputLayout = this.usePackedTexture ?
         handler.createTextureLayoutFromShape(inputs[0].dims, 4, inputs[0].dims, {isPacked: true, reverseWH: true}) :
         handler.createTextureLayoutFromShape(inputs[0].dims);
 
-    const isBroadcast = !ShapeUtil.areEqual(inputs[0].dims, inputs[1].dims);
     if (isBroadcast) {
       const outputShape = BroadcastUtil.calcShape(inputs[0].dims, inputs[1].dims, false);
       if (!outputShape) {
@@ -48,7 +56,7 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator {
         ${bBcast}
         return ${this.glslFunc.name}(_A(aindices), _B(bindices));
     }`;
-      const outputLayout = handler.session.pack ?
+      const outputLayout = this.usePackedTexture ?
           handler.createTextureLayoutFromShape(outputShape, 4, outputShape, {isPacked: true, reverseWH: true}) :
           handler.createTextureLayoutFromShape(outputShape);
 
@@ -57,8 +65,8 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator {
         outputLayout,
         samplers: ['A', 'B'],
         shaderSource,
-        expectPackedInputs: handler.session.pack,
-        expectPackedOutputs: handler.session.pack
+        expectPackedInputs: this.usePackedTexture,
+        expectPackedOutputs: this.usePackedTexture
       };
     }
     const glsl = getGlsl(handler.session.backend.glContext.version);
@@ -71,7 +79,8 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator {
       ${glsl.output} = result;
     }
     `;
-    if (handler.session.pack) {
+
+    if (this.usePackedTexture) {
       return {
         hasMain: true,
         inputLayouts,
@@ -92,7 +101,7 @@ export class WebGLBinaryOp extends BinaryOp implements WebGLOperator {
     }
   }
   createRunData(handler: WebGLInferenceHandler, programInfo: ProgramInfo, inputs: Tensor[]): RunData {
-    const inputTDs = handler.session.pack ?
+    const inputTDs = this.usePackedTexture ?
         inputs.map((t) => handler.getOrCreateTextureData(t, handler.getOrCreateTextureLayout(t, 1, false, [], true))) :
         inputs.map((t, i) => handler.getOrCreateTextureData(t, programInfo.inputLayouts[i]));
     return {
@@ -159,7 +168,7 @@ export function glslEqual(): GlslValueFunction {
     return float(a == b);
   }
   vec4 ${name}(vec4 v1, vec4 v2) {
-    return vec4( v1 == v2 );
+    return vec4(equal(v1, v2));
   }
   `;
   return {body, name, type: FunctionType.ValueBased};

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

@@ -47,11 +47,11 @@ export class WebGLIm2ColPacked implements WebGLOperator {
           pos = rc.y + ${row};
 
           if(blockIndex < ${im2colShape[1]} && pos < ${im2colShape[0]}) {
-            offsetY = int(blockIndex / (${this.convOutputShape[rank - 1]})) * ${this.strides[0]} - ${this.pads[1]};
+            offsetY = int(blockIndex / (${this.convOutputShape[rank - 1]})) * ${this.strides[0]} - ${this.pads[0]};
             d0 = offsetY + ${this.dilations[0]} * (imod(pos, ${kernelSize}) / ${wshape[2]});
 
             if(d0 < ${xshape[rowDim]} && d0 >= 0) {
-              offsetX = imod(blockIndex, ${this.convOutputShape[rank - 1]}) * ${this.strides[1]} - ${this.pads[0]};
+              offsetX = imod(blockIndex, ${this.convOutputShape[rank - 1]}) * ${this.strides[1]} - ${this.pads[1]};
               d1 = offsetX + ${this.dilations[1]} * imod(imod(pos, ${kernelSize}), ${wshape[2]});
 
               if(d1 < ${xshape[colDim]} && d1 >= 0) {

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

@@ -3,12 +3,14 @@
 
 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 {WebGLMatMulPacked} from './matmul-pack';
 
 export class WebGLMatMul extends MatMul implements WebGLOperator {
+  private usePackedTexture?: boolean;
+
   packedImpl: WebGLMatMulPacked;
   unpackedImpl: WebGLUnpackedMatMul;
   constructor() {
@@ -18,7 +20,12 @@ export class WebGLMatMul extends MatMul implements WebGLOperator {
   }
 
   run(inferenceHandler: WebGLInferenceHandler, inputs: Tensor[]): Tensor[] {
-    if (inferenceHandler.session.pack) {
+    if (this.usePackedTexture === undefined) {
+      const isBroadcast = !ShapeUtil.areEqual(inputs[0].dims, inputs[1].dims);
+      this.usePackedTexture = !isBroadcast && inferenceHandler.session.pack;
+    }
+
+    if (this.usePackedTexture) {
       return inferenceHandler.run(this.packedImpl, inputs);
     } else {
       return inferenceHandler.run(this.unpackedImpl, inputs);
@@ -26,7 +33,12 @@ export class WebGLMatMul extends MatMul implements WebGLOperator {
   }
 
   createProgramInfo(handler: WebGLInferenceHandler, inputs: Tensor[]): ProgramInfo {
-    if (handler.session.pack && inputs[0].dims.length > 1) {
+    if (this.usePackedTexture === undefined) {
+      const isBroadcast = !ShapeUtil.areEqual(inputs[0].dims, inputs[1].dims);
+      this.usePackedTexture = !isBroadcast && handler.session.pack;
+    }
+
+    if (this.usePackedTexture && inputs[0].dims.length > 1) {
       return this.packedImpl.createProgramInfo(handler, inputs);
     } else {
       return this.unpackedImpl.createProgramInfo(handler, inputs);
@@ -34,7 +46,7 @@ export class WebGLMatMul extends MatMul implements WebGLOperator {
   }
 
   createRunData(handler: WebGLInferenceHandler, programInfo: ProgramInfo, inputs: Tensor[]): RunData {
-    if (handler.session.pack && inputs[0].dims.length > 1) {
+    if (this.usePackedTexture && inputs[0].dims.length > 1) {
       return this.packedImpl.createRunData(handler, programInfo, inputs);
     } else {
       return this.unpackedImpl.createRunData(handler, programInfo, inputs);

js/web/test/test-suite-whitelist.jsonc

@@ -27,12 +27,12 @@
       "test_averagepool_1d_default",
       "test_averagepool_2d_default",
       //"v12/test_averagepool_2d_pads", // TODO: fix avgpool and maxpool on VM
-      "v12/test_averagepool_2d_precomputed_pads",
-      "v12/test_averagepool_2d_precomputed_same_upper",
-      "v12/test_averagepool_2d_precomputed_strides",
-      "v12/test_averagepool_2d_same_upper",
-      "v12/test_averagepool_2d_same_lower",
-      "v12/test_averagepool_2d_strides",
+      // "v12/test_averagepool_2d_precomputed_pads",
+      // "v12/test_averagepool_2d_precomputed_same_upper",
+      // "v12/test_averagepool_2d_precomputed_strides",
+      // "v12/test_averagepool_2d_same_upper",
+      // "v12/test_averagepool_2d_same_lower",
+      // "v12/test_averagepool_2d_strides",
       "test_averagepool_3d_default",
       "test_basic_conv_with_padding",
       "test_basic_conv_without_padding",
@@ -102,13 +102,13 @@
       "test_matmul_4d",
       "test_maxpool_1d_default",
       "test_maxpool_2d_default",
-      "v12/test_maxpool_2d_pads",
-      "v12/test_maxpool_2d_precomputed_pads",
-      "v12/test_maxpool_2d_precomputed_same_upper",
-      "v12/test_maxpool_2d_precomputed_strides",
-      "v12/test_maxpool_2d_same_lower",
-      "v12/test_maxpool_2d_same_upper",
-      "v12/test_maxpool_2d_strides",
+      // "v12/test_maxpool_2d_pads",
+      // "v12/test_maxpool_2d_precomputed_pads",
+      // "v12/test_maxpool_2d_precomputed_same_upper",
+      // "v12/test_maxpool_2d_precomputed_strides",
+      // "v12/test_maxpool_2d_same_lower",
+      // "v12/test_maxpool_2d_same_upper",
+      // "v12/test_maxpool_2d_strides",
       "test_maxpool_3d_default",
       "test_mul_bcast",
       "test_mul_example",

js/web/test/unittests/backends/webgl/test-concat-packed.ts

@@ -2,10 +2,9 @@
 // Licensed under the MIT License.
 
 import {expect} from 'chai';
-
+import {env} from 'onnxruntime-common';
 import {Attribute} from '../../../../lib/onnxjs/attribute';
 import {Backend, InferenceHandler, resolveBackend, SessionHandler} from '../../../../lib/onnxjs/backend';
-import {WebGLBackend} from '../../../../lib/onnxjs/backends/backend-webgl';
 import {WebGLInferenceHandler} from '../../../../lib/onnxjs/backends/webgl/inference-handler';
 import {WebGLConcat} from '../../../../lib/onnxjs/backends/webgl/ops/concat';
 import {Profiler} from '../../../../lib/onnxjs/instrument';
@@ -207,17 +206,10 @@ describe('#UnitTest# - packed concat - Tensor concat', () => {
   before('Initialize Context', async () => {
     const profiler = Profiler.create();
     backend = await resolveBackend('webgl');
-    // Explicitly set to true to trigger packed version
-    (backend as WebGLBackend).pack = true;
     sessionhandler = backend.createSessionHandler({profiler});
     inferenceHandler = sessionhandler.createInferenceHandler();
   });
 
-  // Set it back to false, apparently this state is sticky throughout all the tests running in same browser session..
-  after('Resetting Context', () => {
-    (backend as WebGLBackend).pack = false;
-  });
-
   const testDataSet = getTestData();
   for (let k = 0; k < testDataSet.length; ++k) {
     const testData = testDataSet[k];
@@ -231,6 +223,11 @@ describe('#UnitTest# - packed concat - Tensor concat', () => {
         return;
       }
 
+      if (!env.webgl.pack) {
+        console.log('Skipping in unpacked texture mode.');
+        return;
+      }
+
       const op = new WebGLConcat();
       const attributes = new Attribute(undefined);
       const axis = testData.axis;

js/web/test/unittests/backends/webgl/test-depth-to-space.ts

@@ -2,10 +2,8 @@
 // Licensed under the MIT License.
 
 import {expect} from 'chai';
-
 import {Attribute} from '../../../../lib/onnxjs/attribute';
 import {Backend, InferenceHandler, resolveBackend, SessionHandler} from '../../../../lib/onnxjs/backend';
-import {WebGLBackend} from '../../../../lib/onnxjs/backends/backend-webgl';
 import {WebGLInferenceHandler} from '../../../../lib/onnxjs/backends/webgl/inference-handler';
 import {WebGLDepthToSpace} from '../../../../lib/onnxjs/backends/webgl/ops/depth-to-space';
 import {Profiler} from '../../../../lib/onnxjs/instrument';
@@ -126,11 +124,6 @@ describe('#UnitTest# - unpacked WebGLDepthToSpace - Tensor WebGLDepthToSpace', (
     inferenceHandler = sessionhandler.createInferenceHandler();
   });
 
-  // Set it back to false, apparently this state is sticky throughout all the tests running in same browser session..
-  after('Resetting Context', () => {
-    (backend as WebGLBackend).pack = false;
-  });
-
   const testDataSet = getTestData();
   for (let k = 0; k < testDataSet.length; ++k) {
     const testData = testDataSet[k];

js/web/test/unittests/backends/webgl/test-matmul-packed.ts

@@ -2,9 +2,9 @@
 // Licensed under the MIT License.
 
 import {expect} from 'chai';
+import {env} from 'onnxruntime-common';
 
 import {Backend, InferenceHandler, resolveBackend, SessionHandler} from '../../../../lib/onnxjs/backend';
-import {WebGLBackend} from '../../../../lib/onnxjs/backends/backend-webgl';
 import {WebGLInferenceHandler} from '../../../../lib/onnxjs/backends/webgl/inference-handler';
 import {WebGLMatMulPacked} from '../../../../lib/onnxjs/backends/webgl/ops/matmul-pack';
 import {Profiler} from '../../../../lib/onnxjs/instrument';
@@ -140,17 +140,10 @@ describe('#UnitTest# - packed matmul - Tensor matmul', () => {
   before('Initialize Context', async () => {
     const profiler = Profiler.create();
     backend = await resolveBackend('webgl');
-    // Explicitly set to true to trigger packed version
-    (backend as WebGLBackend).pack = true;
     sessionhandler = backend.createSessionHandler({profiler});
     inferenceHandler = sessionhandler.createInferenceHandler();
   });
 
-  // Set it back to false, apparently this state is sticky throughout all the tests running in same browser session..
-  after('Resetting Context', () => {
-    (backend as WebGLBackend).pack = false;
-  });
-
   const testDataSet = getTestData();
   for (let k = 0; k < testDataSet.length; ++k) {
     const testData = testDataSet[k];
@@ -164,6 +157,11 @@ describe('#UnitTest# - packed matmul - Tensor matmul', () => {
         return;
       }
 
+      if (!env.webgl.pack) {
+        console.log('Skipping in unpacked texture mode.');
+        return;
+      }
+
       const op = new WebGLMatMulPacked();
 
       const elementCountA = testData.elementCountA;

js/web/test/unittests/backends/webgl/test-reshape-packed.ts

@@ -2,9 +2,8 @@
 // Licensed under the MIT License.
 
 import {expect} from 'chai';
-
+import {env} from 'onnxruntime-common';
 import {Backend, InferenceHandler, resolveBackend, SessionHandler} from '../../../../lib/onnxjs/backend';
-import {WebGLBackend} from '../../../../lib/onnxjs/backends/backend-webgl';
 import {WebGLInferenceHandler} from '../../../../lib/onnxjs/backends/webgl/inference-handler';
 import {WebGLReshapePacked} from '../../../../lib/onnxjs/backends/webgl/ops/reshape-packed';
 import {Profiler} from '../../../../lib/onnxjs/instrument';
@@ -111,17 +110,10 @@ describe('#UnitTest# - reshape - packed', () => {
   before('Initialize Context', async () => {
     const profiler = Profiler.create();
     backend = await resolveBackend('webgl');
-    // Explicitly set to true to trigger packed version
-    (backend as WebGLBackend).pack = true;
     sessionhandler = backend.createSessionHandler({profiler});
     inferenceHandler = sessionhandler.createInferenceHandler();
   });
 
-  // Set it back to false, apparently this state is sticky throughout all the tests running in same browser session..
-  after('Resetting Context', () => {
-    (backend as WebGLBackend).pack = false;
-  });
-
   const testDataSet = getTestData();
   for (let k = 0; k < testDataSet.length; ++k) {
     const testData = testDataSet[k];
@@ -135,6 +127,11 @@ describe('#UnitTest# - reshape - packed', () => {
         return;
       }
 
+      if (!env.webgl.pack) {
+        console.log('Skipping in unpacked texture mode.');
+        return;
+      }
+
       const op = new WebGLReshapePacked();
 
       const elementCount = testData.elementCount;

tools/ci_build/github/azure-pipelines/win-wasm-ci-pipeline.yml

@@ -199,7 +199,11 @@ jobs:
   - script: |
      npm test
     workingDirectory: '$(Build.SourcesDirectory)\js\web'
-    displayName: 'Run ort-web tests'
+    displayName: 'Run ort-web tests - unpacked mode'
+  - script: |
+     npm test --  --webgl-texture-pack-mode -b=webgl
+    workingDirectory: '$(Build.SourcesDirectory)\js\web'
+    displayName: 'Run ort-web tests - packed mode'
   - script: |
       npm pack
     workingDirectory: '$(Build.SourcesDirectory)\js\common'