[js/webgpu] add bool type for Expand

BUG #18584
microsoft · Nov 29, 2023 · a8694a1 · a8694a1
1 parent 14a3434
commit a8694a1
Show file tree

Hide file tree

Showing 4 changed files with 126 additions and 27 deletions.
diff --git a/js/web/lib/wasm/jsep/webgpu/ops/expand.ts b/js/web/lib/wasm/jsep/webgpu/ops/expand.ts
@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+import {DataType} from '../../../wasm-common';
 import {TensorView} from '../../tensor-view';
 import {ShapeUtil} from '../../util';
 import {ComputeContext, ProgramInfo, ProgramUniform} from '../types';
@@ -44,34 +45,51 @@ const createExpandProgramInfo = (inputs: readonly TensorView[]): ProgramInfo =>
   const inputShape = inputs[0].dims;
   const shape = Array.from(inputs[1].getBigInt64Array(), Number);
   const outputShape: number[] = calculateOutputShape(inputShape, shape);
-  const outputSize = ShapeUtil.size(outputShape);
-
   const dataType = inputs[0].dataType;
+  const components = dataType === DataType.bool ? 4 : 1;
+  const outputSize = ShapeUtil.size(outputShape) / components;
+
   const enableInputShapeUniform = enableShapesUniforms(inputShape.length);
-  const inputShapeOrRank = enableInputShapeUniform ? inputShape.length : inputShape;
-  const input = inputVariable('input', dataType, inputShapeOrRank);
   const enableOutputShapeUniform = enableShapesUniforms(outputShape.length);
-  const outputShapeOrRank = enableOutputShapeUniform ? outputShape.length : outputShape;
-  const output = outputVariable('output', dataType, outputShapeOrRank);
 
-  const getShaderSource = (shaderHelper: ShaderHelper) => `
-  const inputShape = ${input.indices(...inputShape)};
-  ${shaderHelper.registerUniform('vec_size', 'u32').declareVariables(input, output)}
-  ${shaderHelper.mainStart()}
-  ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.vec_size')}
-    let outputIndices = ${output.offsetToIndices('global_idx')};
-    var inputIndices: ${input.type.indices};
-    for (var i = 0; i < ${inputShape.length}; i++) {
-      if (${input.indicesGet('inputShape', 'i')} == 1) {
-        ${input.indicesSet('inputIndices', 'i', 0)}
-      } else {
-        ${
-      input.indicesSet(
-          'inputIndices', 'i', output.indicesGet('outputIndices', `i + ${outputShape.length - inputShape.length}`))}
-      }
+
+  const getShaderSource = (shaderHelper: ShaderHelper) => {
+    const inputShapeOrRank = enableInputShapeUniform ? inputShape.length : inputShape;
+    const outputShapeOrRank = enableOutputShapeUniform ? outputShape.length : outputShape;
+    const input = inputVariable('input', dataType, inputShapeOrRank, components);
+    const output = outputVariable('output', dataType, outputShapeOrRank, components);
+    let assignment: string;
+    if (dataType === DataType.bool) {
+      const singleAssignment = (resStr: string, x: number, typeCast = '') => `
+          let outputIndices${x} = ${output.offsetToIndices(`outputOffset + ${x}u`)};
+          let offset${x} = ${input.broadcastedIndicesToOffset(`outputIndices${x}`, output)};
+          let index${x} = offset${x} / 4u;
+          let component${x} = offset${x} % 4u;
+          ${resStr}[${x}] = ${typeCast}(${input.getByOffset(`index${x}`)}[component${x}]);
+        `;
+      assignment = `
+        let outputOffset = global_idx * ${components};
+        var data = vec4<u32>(0);
+        ${singleAssignment('data', 0, 'u32')}
+        ${singleAssignment('data', 1, 'u32')}
+        ${singleAssignment('data', 2, 'u32')}
+        ${singleAssignment('data', 3, 'u32')}
+        ${output.setByOffset('global_idx', 'data')}
+      }`;
+    } else {
+      assignment = `
+        let outputIndices = ${output.offsetToIndices('global_idx')};
+        let inputOffset = ${input.broadcastedIndicesToOffset('outputIndices', output)};
+        ${output.setByOffset('global_idx', input.getByOffset('inputOffset'))}
+      }`;
     }
-    ${output.setByOffset('global_idx', input.getByIndices('inputIndices'))}
-  }`;
+    return `
+    ${shaderHelper.registerUniform('vec_size', 'u32').declareVariables(input, output)}
+    ${shaderHelper.mainStart()}
+    ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.vec_size')}
+    ${assignment}`;
+  };
+
   const programUniforms: ProgramUniform[] = [{type: 'uint32', data: outputSize}];
   if (enableInputShapeUniform) {
     programUniforms.push(...createTensorShapeVariables(inputShape));
@@ -81,7 +99,7 @@ const createExpandProgramInfo = (inputs: readonly TensorView[]): ProgramInfo =>
   }
   return {
     name: 'Expand',
-    shaderCache: {hint: `${outputShape}`, inputDependencies: [enableInputShapeUniform ? 'rank' : 'dims']},
+    shaderCache: {hint: `${outputShape.length}`, inputDependencies: [enableInputShapeUniform ? 'rank' : 'dims']},
     getShaderSource,
     getRunData: () => ({
       outputs: [{dims: outputShape, dataType: inputs[0].dataType}],

diff --git a/js/web/test/data/ops/expand.jsonc b/js/web/test/data/ops/expand.jsonc
@@ -112,6 +112,79 @@
             "type": "float32"
           }
         ]
+      },
+      {
+        "name": "Expand 5 - shape < input.size()",
+        "inputs": [
+          {
+            "data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
+            "dims": [1, 1, 1, 2, 6],
+            "type": "float32"
+          },
+          {
+            "data": [2, 1, 6],
+            "dims": [3],
+            "type": "int64"
+          }
+        ],
+        "outputs": [
+          {
+            "data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
+            "dims": [1, 1, 2, 2, 6],
+            "type": "float32"
+          }
+        ]
+      }
+    ]
+  },
+  {
+    "name": "Expand - bool",
+    "operator": "Expand",
+    "attributes": [],
+    "cases": [
+      {
+        "name": "Expand - last dim is divisible by 4",
+        "inputs": [
+          {
+            "data": [true, false, false, true],
+            "dims": [4],
+            "type": "bool"
+          },
+          {
+            "data": [2, 4],
+            "dims": [2],
+            "type": "int64"
+          }
+        ],
+        "outputs": [
+          {
+            "data": [true, false, false, true, true, false, false, true],
+            "dims": [2, 4],
+            "type": "bool"
+          }
+        ]
+      },
+      {
+        "name": "Expand - last dim is not divisible by 4",
+        "inputs": [
+          {
+            "data": [true, false, false, true, true, true, false, false, false, true, true, true],
+            "dims": [2, 6],
+            "type": "bool"
+          },
+          {
+            "data": [2, 1],
+            "dims": [2],
+            "type": "int64"
+          }
+        ],
+        "outputs": [
+          {
+            "data": [true, false, false, true, true, true, false, false, false, true, true, true],
+            "dims": [2, 6],
+            "type": "bool"
+          }
+        ]
       }
     ]
   }

diff --git a/onnxruntime/core/providers/js/js_data_types.cc b/onnxruntime/core/providers/js/js_data_types.cc
@@ -29,4 +29,4 @@ const std::vector<MLDataType>& JsepSupportedFloatTypes() {
 }
 
 }  // namespace js
-}  // namespace onnxruntime
+}  // namespace onnxruntime
diff --git a/onnxruntime/core/providers/js/operators/expand.cc b/onnxruntime/core/providers/js/operators/expand.cc
@@ -13,7 +13,11 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
     12,
     kJsExecutionProvider,
     KernelDefBuilder()
-        .TypeConstraint("T", DataTypeImpl::GetTensorType<float>())
+        .TypeConstraint("T", BuildKernelDefConstraintsFromTypeList<TypeList<float,
+                                                                            MLFloat16,
+                                                                            int32_t,
+                                                                            uint32_t,
+                                                                            bool>>())
         .InputMemoryType(OrtMemTypeCPU, 1),
     Expand);
 
@@ -23,7 +27,11 @@ ONNX_OPERATOR_KERNEL_EX(
     13,
     kJsExecutionProvider,
     KernelDefBuilder()
-        .TypeConstraint("T", DataTypeImpl::GetTensorType<float>())
+        .TypeConstraint("T", BuildKernelDefConstraintsFromTypeList<TypeList<float,
+                                                                            MLFloat16,
+                                                                            int32_t,
+                                                                            uint32_t,
+                                                                            bool>>())
         .InputMemoryType(OrtMemTypeCPU, 1),
     Expand);
 }  // namespace js