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 { ComputeContext, ProgramInfo } from '../types';
8

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

11
const getRepeats = (repeatsTensorView: TensorView): readonly number[] =>
12
  Array.from(repeatsTensorView.getBigInt64Array(), Number);
13

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

19
  if (
20
    inputs[0].dataType !== DataType.float &&
21
    inputs[0].dataType !== DataType.float16 &&
22
    inputs[0].dataType !== DataType.int32 &&
23
    inputs[0].dataType !== DataType.uint32
24
  ) {
25
    throw new Error('Tile only support float, float16, int32, and uint32 data types');
26
  }
27

28
  if (inputs[1].dataType !== DataType.int64) {
29
    throw new Error('Tile `repeats` input should be of int64 data type');
30
  }
31

32
  if (inputs[1].dims.length !== 1) {
33
    throw new Error('Tile `repeats` input should be 1-D');
34
  }
35

36
  const repeats: readonly number[] = getRepeats(inputs[1]);
37

38
  if (repeats.length !== inputs[0].dims.length) {
39
    throw new Error('Tile `repeats` input should have same number of elements as rank of input data tensor');
40
  }
41
};
42

43
const getOutputShape = (inputShape: readonly number[], repeats: readonly number[]): readonly number[] => {
44
  const outputShape: number[] = [];
45

46
  for (let i = 0; i < inputShape.length; ++i) {
47
    outputShape.push(inputShape[i] * repeats[i]);
48
  }
49

50
  return outputShape;
51
};
52

53
export const createTileProgramInfo = (inputs: readonly TensorView[], shape?: number[]): ProgramInfo => {
54
  const inputShape = inputs[0].dims;
55
  const repeats: readonly number[] = shape == null ? getRepeats(inputs[1]) : shape;
56
  const outputShape = getOutputShape(inputShape, repeats);
57
  const outputSize = ShapeUtil.size(outputShape);
58

59
  const dataType = inputs[0].dataType;
60
  const input = inputVariable('input', dataType, inputShape.length);
61
  const output = outputVariable('output', dataType, outputShape.length);
62

63
  const getShaderSource = (shaderHelper: ShaderHelper) => `
64
      const inputShape = ${input.indices(...inputShape)};
65
      ${shaderHelper.registerUniform('output_size', 'u32').declareVariables(input, output)}
66
      ${shaderHelper.mainStart()}
67
      ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.output_size')}
68
      let output_indices = ${output.offsetToIndices('global_idx')};
69
      var input_indices: ${input.type.indices};
70
      for (var i = 0; i < ${inputShape.length}; i++) {
71
        let input_dim_i = ${input.indicesGet('uniforms.input_shape', 'i')};
72
        let input_dim_value = ${output.indicesGet('output_indices', 'i')}  % input_dim_i;
73

74
        ${input.indicesSet('input_indices', 'i', 'input_dim_value')}
75
      }
76
      ${output.setByOffset('global_idx', input.getByIndices('input_indices'))}
77
    }`;
78

79
  return {
80
    name: 'Tile',
81
    shaderCache: { hint: `${repeats}`, inputDependencies: ['rank'] },
82
    getRunData: () => ({
83
      outputs: [{ dims: outputShape, dataType: inputs[0].dataType }],
84
      dispatchGroup: { x: Math.ceil(outputSize / 64 /* workgroup size */) },
85
      programUniforms: [
86
        { type: DataType.uint32, data: outputSize },
87
        ...createTensorShapeVariables(inputs[0].dims, outputShape),
88
      ],
89
    }),
90
    getShaderSource,
91
  };
92
};
93

94
export const tile = (context: ComputeContext): void => {
95
  validateInputs(context.inputs);
96
  context.compute(createTileProgramInfo(context.inputs), { inputs: [0] });
97
};
98