Skip to content

Commit

Permalink
[js/webgpu] add bool type for Gather
Browse files Browse the repository at this point in the history
  • Loading branch information
@qjia7
qjia7 committed Nov 29, 2023
1 parent a8694a1 commit fd23e94
Show file tree
Hide file tree
Showing 3 changed files with 109 additions and 41 deletions.
103 changes: 65 additions & 38 deletions js/web/lib/wasm/jsep/webgpu/ops/gather.ts
View file
Original file line number Diff line number Diff line change
@@ -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 {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key';
Expand Down Expand Up @@ -29,7 +30,8 @@ const createGatherProgramInfo = (inputs: readonly TensorView[], attributes: Gath
outputShape.splice(axis, 1, ...indicesShape);

const axisDimLimit = inputShape[axis];
const outputSize = ShapeUtil.size(outputShape);
const components = inputs[0].dataType === DataType.bool ? 4 : 1;
const outputSize = ShapeUtil.size(outputShape) / components;

const enableInputShapesUniforms = enableShapesUniforms(inputs[0].dims.length);
const inputShapeOrRank = enableInputShapesUniforms ? inputs[0].dims.length : inputs[0].dims;
Expand All @@ -38,10 +40,6 @@ const createGatherProgramInfo = (inputs: readonly TensorView[], attributes: Gath
const enableOutputShapesUniforms = enableShapesUniforms(outputShape.length);
const outputShapeOrRank = enableOutputShapesUniforms ? outputShape.length : outputShape;

const data = inputVariable('data', inputs[0].dataType, inputShapeOrRank);
const indices = inputVariable('inputIndices', inputs[1].dataType, indicesShapeOrRank);
const output = outputVariable('output', inputs[0].dataType, outputShapeOrRank);

const programUniforms: ProgramUniform[] =
[{type: 'uint32', data: outputSize}, {type: 'int32', data: axisDimLimit}, {type: 'uint32', data: axis}];
if (enableInputShapesUniforms) {
Expand All @@ -58,46 +56,75 @@ const createGatherProgramInfo = (inputs: readonly TensorView[], attributes: Gath
inputDependencies.push(enableInputShapesUniforms ? 'rank' : 'dims');
inputDependencies.push(enableIndicesShapesUniforms ? 'rank' : 'dims');

const calcDataIndices = (): string => {
const indicesRank = indicesShape.length;
let calcStr = `var indicesIndices = ${indices.type.indices}(0);`;
for (let i = 0; i < indicesRank; i++) {
calcStr += `${indicesRank > 1 ? `indicesIndices[${i}]` : 'indicesIndices'} = ${
outputShape.length > 1 ? `outputIndices[uniforms.axis + ${i}]` : 'outputIndices'};`;
}
calcStr += `
var idx = ${indices.getByIndices('indicesIndices')};
if (idx < 0) {
idx = idx + uniforms.axisDimLimit;
const getShaderSource = (shaderHelper: ShaderHelper) => {
const data = inputVariable('data', inputs[0].dataType, inputShapeOrRank, components);
const indices = inputVariable('inputIndices', inputs[1].dataType, indicesShapeOrRank);
const output = outputVariable('output', inputs[0].dataType, outputShapeOrRank, components);

const calcDataIndices = (x: number|string): string => {
const indicesRank = indicesShape.length;
let calcStr = `var indicesIndices${x} = ${indices.type.indices}(0);`;
for (let i = 0; i < indicesRank; i++) {
calcStr += `${indicesRank > 1 ? `indicesIndices${x}[${i}]` : `indicesIndices${x}`} = ${
outputShape.length > 1 ? `outputIndices${x}[uniforms.axis + ${i}]` : `outputIndices${x}`};`;
}
calcStr += `
var idx${x} = ${indices.getByIndices(`indicesIndices${x}`)};
if (idx${x} < 0) {
idx${x} = idx${x} + uniforms.axisDimLimit;
}
var dataIndices${x} = ${data.type.indices}(0);
`;
for (let i = 0, j = 0; i < inputRank; i++) {
if (i === axis) {
calcStr += `${inputRank > 1 ? `dataIndices${x}[${i}]` : `dataIndices${x}`} = u32(idx${x});`;
j += indicesRank;
} else {
calcStr += `${inputRank > 1 ? `dataIndices${x}[${i}]` : `dataIndices${x}`} = ${
outputShape.length > 1 ? `outputIndices${x}[${j}]` : `outputIndices${x}`};`;
j++;
}
var dataIndices = ${data.type.indices}(0);
`;
for (let i = 0, j = 0; i < inputRank; i++) {
if (i === axis) {
calcStr += `${inputRank > 1 ? `dataIndices[${i}]` : 'dataIndices'} = u32(idx);`;
j += indicesRank;
} else {
calcStr += `${inputRank > 1 ? `dataIndices[${i}]` : 'dataIndices'} = ${
outputShape.length > 1 ? `outputIndices[${j}]` : 'outputIndices'};`;
j++;
}
return calcStr;
};
let assignment: string;
if (inputs[0].dataType === DataType.bool) {
const singleAssignment = (resStr: string, x: number, typeCast = '') => `
let outputIndices${x} = ${output.offsetToIndices(`outputOffset + ${x}u`)};
${calcDataIndices(x)};
let offset${x} = ${data.indicesToOffset(`dataIndices${x}`)};
let index${x} = offset${x} / 4u;
let component${x} = offset${x} % 4u;
${resStr}[${x}] = ${typeCast}(${data.getByOffset(`index${x}`)}[component${x}]);
`;
assignment = `
let outputOffset = global_idx * ${components};
var value = vec4<u32>(0);
${singleAssignment('value', 0, 'u32')}
${singleAssignment('value', 1, 'u32')}
${singleAssignment('value', 2, 'u32')}
${singleAssignment('value', 3, 'u32')}
${output.setByOffset('global_idx', 'value')}
`;
} else {
assignment = `
let outputIndices = ${output.offsetToIndices('global_idx')};
${calcDataIndices('')};
let value = ${data.getByIndices('dataIndices')};
${output.setByOffset('global_idx', 'value')};
`;
}
return calcStr;
};

const getShaderSource = (shaderHelper: ShaderHelper) => `
return `
${
shaderHelper.registerUniform('outputSize', 'u32')
.registerUniform('axisDimLimit', 'i32')
.registerUniform('axis', 'u32')
.declareVariables(data, indices, output)}
shaderHelper.registerUniform('outputSize', 'u32')
.registerUniform('axisDimLimit', 'i32')
.registerUniform('axis', 'u32')
.declareVariables(data, indices, output)}
${shaderHelper.mainStart()}
${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.outputSize')}
let outputIndices = ${output.offsetToIndices('global_idx')};
${calcDataIndices()};
let value = ${data.getByIndices('dataIndices')};
${output.setByOffset('global_idx', 'value')};
${assignment}
}`;
};
return {
name: 'Gather',
shaderCache: {hint: attributes.cacheKey, inputDependencies},
Expand Down
29 changes: 29 additions & 0 deletions js/web/test/data/ops/gather.jsonc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,5 +93,34 @@
]
}
]
},
{
"name": "Gather - bool",
"operator": "Gather",
"attributes": [],
"cases": [
{
"name": "data[2,4] indices[1]",
"inputs": [
{
"data": [true, false, false, true, false, false, true, true],
"dims": [2, 4],
"type": "bool"
},
{
"data": [1],
"dims": [1],
"type": "int32"
}
],
"outputs": [
{
"data": [false, false, true, true],
"dims": [1, 4],
"type": "bool"
}
]
}
]
}
]
18 changes: 15 additions & 3 deletions onnxruntime/core/providers/js/operators/gather.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,11 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
10,
kJsExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", JsepSupportedDataTypes())
.TypeConstraint("T", BuildKernelDefConstraintsFromTypeList<TypeList<float,
MLFloat16,
int32_t,
uint32_t,
bool>>())
.TypeConstraint("Tind", BuildKernelDefConstraintsFromTypeList<TypeList<int32_t, int64_t>>()),
Gather);

Expand All @@ -26,7 +30,11 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
12,
kJsExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", JsepSupportedDataTypes())
.TypeConstraint("T", BuildKernelDefConstraintsFromTypeList<TypeList<float,
MLFloat16,
int32_t,
uint32_t,
bool>>())
.TypeConstraint("Tind", BuildKernelDefConstraintsFromTypeList<TypeList<int32_t, int64_t>>()),
Gather);

Expand All @@ -36,7 +44,11 @@ ONNX_OPERATOR_KERNEL_EX(
13,
kJsExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", JsepSupportedDataTypes())
.TypeConstraint("T", BuildKernelDefConstraintsFromTypeList<TypeList<float,
MLFloat16,
int32_t,
uint32_t,
bool>>())
.TypeConstraint("Tind", BuildKernelDefConstraintsFromTypeList<TypeList<int32_t, int64_t>>()),
Gather);

Expand Down

0 comments on commit fd23e94

Please sign in to comment.