4
import { DataType } from '../../../wasm-common';
5
import { TensorView } from '../../tensor-view';
6
import { ShapeUtil } from '../../util';
7
import { AttributeWithCacheKey, createAttributeWithCacheKey } from '../attribute-with-cache-key';
8
import { ComputeContext, ProgramInfo, ProgramUniform } from '../types';
10
import { createTensorShapeVariables, inputVariable, outputVariable, ShaderHelper } from './common';
12
export interface GatherAttributes extends AttributeWithCacheKey {
16
const validateInputs = (inputs: readonly TensorView[]): void => {
17
if (!inputs || inputs.length !== 2) {
18
throw new Error('Gather requires 2 inputs.');
22
const createGatherProgramInfo = (inputs: readonly TensorView[], attributes: GatherAttributes): ProgramInfo => {
23
const inputShape = inputs[0].dims;
24
const indicesShape = inputs[1].dims;
26
const inputRank = inputShape.length;
27
const axis = ShapeUtil.normalizeAxis(attributes.axis, inputRank);
29
const outputShape = inputShape.slice(0);
30
outputShape.splice(axis, 1, ...indicesShape);
32
const axisDimLimit = inputShape[axis];
33
const components = inputs[0].dataType === DataType.bool ? 4 : 1;
34
const outputSize = Math.ceil(ShapeUtil.size(outputShape) / components);
36
const programUniforms: ProgramUniform[] = [
37
{ type: DataType.uint32, data: outputSize },
38
{ type: DataType.int32, data: axisDimLimit },
39
{ type: DataType.uint32, data: axis },
40
...createTensorShapeVariables(inputs[0].dims, inputs[1].dims, outputShape),
43
const getShaderSource = (shaderHelper: ShaderHelper) => {
44
const data = inputVariable('data', inputs[0].dataType, inputs[0].dims.length, components);
45
const indices = inputVariable('inputIndices', inputs[1].dataType, inputs[1].dims.length);
46
const output = outputVariable('output', inputs[0].dataType, outputShape.length, components);
48
const calcDataIndices = (x: number | string): string => {
49
const indicesRank = indicesShape.length;
50
let calcStr = `var indicesIndices${x} = ${indices.type.indices}(0);`;
51
for (let i = 0; i < indicesRank; i++) {
52
calcStr += `${indicesRank > 1 ? `indicesIndices${x}[${i}]` : `indicesIndices${x}`} = ${
53
outputShape.length > 1 ? `outputIndices${x}[uniforms.axis + ${i}]` : `outputIndices${x}`
57
var idx${x} = ${indices.getByIndices(`indicesIndices${x}`)};
59
idx${x} = idx${x} + uniforms.axisDimLimit;
61
var dataIndices${x} : ${data.type.indices};
63
for (let i = 0, j = 0; i < inputRank; i++) {
65
calcStr += `${inputRank > 1 ? `dataIndices${x}[${i}]` : `dataIndices${x}`} = u32(idx${x});`;
68
calcStr += `${inputRank > 1 ? `dataIndices${x}[${i}]` : `dataIndices${x}`} = ${
69
outputShape.length > 1 ? `outputIndices${x}[${j}]` : `outputIndices${x}`
76
let assignment: string;
77
if (inputs[0].dataType === DataType.bool) {
78
const singleAssignment = (resStr: string, x: number, typeCast = '') => `
79
let outputIndices${x} = ${output.offsetToIndices(`outputOffset + ${x}u`)};
80
${calcDataIndices(x)};
81
let offset${x} = ${data.indicesToOffset(`dataIndices${x}`)};
82
let index${x} = offset${x} / 4u;
83
let component${x} = offset${x} % 4u;
84
${resStr}[${x}] = ${typeCast}(${data.getByOffset(`index${x}`)}[component${x}]);
87
let outputOffset = global_idx * ${components};
88
var value = vec4<u32>(0);
89
${singleAssignment('value', 0, 'u32')}
90
${singleAssignment('value', 1, 'u32')}
91
${singleAssignment('value', 2, 'u32')}
92
${singleAssignment('value', 3, 'u32')}
93
${output.setByOffset('global_idx', 'value')}
97
let outputIndices = ${output.offsetToIndices('global_idx')};
98
${calcDataIndices('')};
99
let value = ${data.getByIndices('dataIndices')};
100
${output.setByOffset('global_idx', 'value')};
105
.registerUniform('outputSize', 'u32')
106
.registerUniform('axisDimLimit', 'i32')
107
.registerUniform('axis', 'u32')
108
.declareVariables(data, indices, output)}
109
${shaderHelper.mainStart()}
110
${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.outputSize')}
116
shaderCache: { hint: attributes.cacheKey, inputDependencies: ['rank', 'rank'] },
118
outputs: [{ dims: outputShape, dataType: inputs[0].dataType }],
119
dispatchGroup: { x: Math.ceil(outputSize / 64 ) },
126
export const parseGatherAttributes = (attributes: Record<string, unknown>): GatherAttributes =>
127
createAttributeWithCacheKey({ axis: attributes.axis as number });
129
export const gather = (context: ComputeContext, attributes: GatherAttributes): void => {
130
const inputs = context.inputs;
131
validateInputs(inputs);
132
context.compute(createGatherProgramInfo(context.inputs, attributes));