[JS/WebGPU] support Concat.int32 operator

microsoft · Aug 15, 2023 · 412cb26 · 412cb26
1 parent 24e0bd3
commit 412cb26
Show file tree

Hide file tree

Showing 5 changed files with 434 additions and 22 deletions.
diff --git a/js/web/lib/wasm/jsep/backend-webgpu.ts b/js/web/lib/wasm/jsep/backend-webgpu.ts
@@ -18,15 +18,16 @@ import {ComputeContext, GpuData, GpuDataType, ProgramInfo, ProgramInfoLoader} fr
  *
  */
 const getProgramInfoUniqueKey =
-    (programInfo: ProgramInfo|ProgramInfoLoader, inputTensorShapes: ReadonlyArray<TensorView['dims']>,
+    (programInfo: ProgramInfo|ProgramInfoLoader, inputs: readonly TensorView[],
      inputGpuDataTypes: readonly GpuDataType[]): string => {
-      const inputTensorShapesToString = inputTensorShapes.map(d => `${d.join(',')}`).join('_');
+      const inputDataTypesToString = inputs.map(i => `${i.dataType}`).join('_');
+      const inputTensorShapesToString = inputs.map(i => `${i.dims.join(',')}`).join('_');
       const inputGpuDataTypesToString = inputGpuDataTypes.join('_');
       let key = programInfo.name;
       if (programInfo.cacheHint) {
         key += '[' + programInfo.cacheHint + ']';
       }
-      key += ':' + inputTensorShapesToString + ';' + inputGpuDataTypesToString;
+      key += ':' + inputDataTypesToString + ';' + inputTensorShapesToString + ';' + inputGpuDataTypesToString;
       return key;
     };
 
@@ -221,7 +222,7 @@ export class WebGpuBackend {
       inputDatas[i] = gpuData;
     }
 
-    const key = getProgramInfoUniqueKey(program, inputs.map(i => i.dims), inputDatas.map(i => i.type));
+    const key = getProgramInfoUniqueKey(program, inputs, inputDatas.map(i => i.type));
     let artifact = this.programManager.getArtifact(key);
     const programInfo = artifact ?
         artifact.programInfo :

diff --git a/js/web/lib/wasm/jsep/webgpu/ops/concat.ts b/js/web/lib/wasm/jsep/webgpu/ops/concat.ts
@@ -118,7 +118,7 @@ const createConcatProgramInfo =
   ${shaderHelper.mainStart()}
     ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(outputSize)}
 
-    let indices = ${output.offsetToIndices('global_idx')};
+    var indices = ${output.offsetToIndices('global_idx')};
 
     let inputIndex = calculateInputIndex(${indicesAxis});
     if (inputIndex != 0u) {