onnxruntime

1
// Copyright (c) Microsoft Corporation. All rights reserved.
2
// Licensed under the MIT License.
3

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, ProgramInputTensorInfoDependency, ProgramUniform } from '../types';
9

10
import { createTensorShapeVariables, inputVariable, outputVariable, ShaderHelper } from './common';
11

12
export interface GatherElementsAttributes extends AttributeWithCacheKey {
13
  axis: number;
14
}
15

16
const validateInputs = (inputs: readonly TensorView[]): void => {
17
  if (!inputs || inputs.length !== 2) {
18
    throw new Error('GatherElements requires 2 inputs.');
19
  }
20

21
  if (inputs[0].dims.length < 1) {
22
    throw new Error('GatherElements requires that the data input be rank >= 1.');
23
  }
24

25
  if (inputs[0].dims.length !== inputs[1].dims.length) {
26
    throw new Error(`GatherElements requires that the data input and
27
                     indices input tensors be of same rank.`);
28
  }
29
};
30

31
const createGatherElementsProgramInfo = (
32
  inputs: readonly TensorView[],
33
  attributes: GatherElementsAttributes,
34
): ProgramInfo => {
35
  const inputShape = inputs[0].dims;
36
  const inputOutputDataType = inputs[0].dataType;
37
  const inputRank = inputShape.length;
38

39
  const indicesShape = inputs[1].dims;
40
  const indicesDataType = inputs[1].dataType;
41
  const axis = ShapeUtil.normalizeAxis(attributes.axis, inputRank);
42
  const axisDimLimit = inputShape[axis];
43

44
  const outputShape = indicesShape.slice(0);
45
  const outputSize = ShapeUtil.size(outputShape);
46

47
  const input = inputVariable('input', inputOutputDataType, inputRank);
48
  const indices = inputVariable('indicesInput', indicesDataType, indicesShape.length);
49
  const output = outputVariable('output', inputOutputDataType, outputShape.length);
50

51
  const programUniforms: ProgramUniform[] = [
52
    { type: DataType.uint32, data: outputSize },
53
    { type: DataType.int32, data: axisDimLimit },
54
    { type: DataType.uint32, data: axis },
55
  ];
56
  programUniforms.push(...createTensorShapeVariables(inputShape, indicesShape, outputShape));
57
  const inputDependencies: ProgramInputTensorInfoDependency[] = ['rank', 'rank'];
58

59
  // int64 indices would be treated as little endian i32 with assumption they fall in i32 limits
60
  // That assumption is safe as it's not possible to allocate >2gb buffer for input tensor
61
  // Input data will be treated as u32 or two u32 for 8-byte tensors
62
  const getShaderSource = (shaderHelper: ShaderHelper) => `
63
      ${shaderHelper
64
        .registerUniform('outputSize', 'u32')
65
        .registerUniform('axisDimLimit', 'i32')
66
        .registerUniform('axis', 'u32')
67
        .declareVariables(input, indices, output)}
68
      ${shaderHelper.mainStart()}
69
      ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.outputSize')}
70

71
      let outputIndices = ${output.offsetToIndices('global_idx')};
72

73
      var idx = ${indices.getByOffset('global_idx')};
74
      if (idx < 0) {
75
        idx = idx + uniforms.axisDimLimit;
76
      }
77
      var inputIndices = ${input.type.indices}(outputIndices);
78
      ${input.indicesSet('inputIndices', 'uniforms.axis', 'u32(idx)')};
79
      let value = ${input.getByIndices('inputIndices')};
80

81
      ${output.setByOffset('global_idx', 'value')};
82
  }`;
83

84
  return {
85
    name: 'GatherElements',
86
    shaderCache: { inputDependencies },
87
    getRunData: () => ({
88
      outputs: [{ dims: outputShape, dataType: inputs[0].dataType }],
89
      dispatchGroup: { x: Math.ceil(outputSize / 64 /* workgroup size */) },
90
      programUniforms,
91
    }),
92
    getShaderSource,
93
  };
94
};
95

96
export const parseGatherElementsAttributes = (attributes: Record<string, unknown>): GatherElementsAttributes =>
97
  createAttributeWithCacheKey({ axis: attributes.axis as number });
98

99
export const gatherElements = (context: ComputeContext, attributes: GatherElementsAttributes): void => {
100
  const inputs = context.inputs;
101
  validateInputs(inputs);
102
  context.compute(createGatherElementsProgramInfo(context.inputs, attributes));
103
};
104