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

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

11
const validateInputs = (inputs: readonly TensorView[]): void => {
12
  if (!inputs || inputs.length !== 2) {
13
    throw new Error('Expand requires 2 input.');
14
  }
15
  const inputShape = inputs[0].dims;
16
  const shape = Array.from(inputs[1].getBigInt64Array(), Number);
17

18
  let shapeIndex = shape.length < inputShape.length ? 0 : shape.length - inputShape.length;
19
  let inputShapeIndex = inputShape.length < shape.length ? 0 : inputShape.length - shape.length;
20
  for (; shapeIndex < shape.length && inputShapeIndex < inputShape.length; ++shapeIndex, ++inputShapeIndex) {
21
    if (
22
      shape[shapeIndex] !== inputShape[inputShapeIndex] &&
23
      shape[shapeIndex] !== 1 &&
24
      inputShape[inputShapeIndex] !== 1
25
    ) {
26
      throw new Error('Expand requires shape to be broadcastable to input');
27
    }
28
  }
29
};
30

31
const getAdjustedShape = (shape1: readonly number[], shape2: readonly number[]): number[] => {
32
  const diff = shape1.length - shape2.length;
33
  const shape: number[] = [];
34
  for (let i = 0; i < diff; ++i) {
35
    shape.push(shape1[i]);
36
  }
37
  for (let i = 0; i < shape2.length; ++i) {
38
    shape.push(shape2[i] === 1 ? shape1[i + diff] : shape2[i]);
39
  }
40
  return shape;
41
};
42

43
const calculateOutputShape = (inputShape: readonly number[], shape: readonly number[]): number[] =>
44
  inputShape.length > shape.length ? getAdjustedShape(inputShape, shape) : getAdjustedShape(shape, inputShape);
45

46
const createExpandProgramInfo = (inputs: readonly TensorView[]): ProgramInfo => {
47
  const inputShape = inputs[0].dims;
48
  const shape = Array.from(inputs[1].getBigInt64Array(), Number);
49
  const outputShape: number[] = calculateOutputShape(inputShape, shape);
50
  const dataType = inputs[0].dataType;
51
  const components = dataType === DataType.bool ? 4 : 1;
52
  const outputSize = Math.ceil(ShapeUtil.size(outputShape) / components);
53

54
  const getShaderSource = (shaderHelper: ShaderHelper) => {
55
    const input = inputVariable('input', dataType, inputShape.length, components);
56
    const output = outputVariable('output', dataType, outputShape.length, components);
57
    let assignment: string;
58
    if (dataType === DataType.bool) {
59
      const singleAssignment = (resStr: string, x: number, typeCast = '') => `
60
          let outputIndices${x} = ${output.offsetToIndices(`outputOffset + ${x}u`)};
61
          let offset${x} = ${input.broadcastedIndicesToOffset(`outputIndices${x}`, output)};
62
          let index${x} = offset${x} / 4u;
63
          let component${x} = offset${x} % 4u;
64
          ${resStr}[${x}] = ${typeCast}(${input.getByOffset(`index${x}`)}[component${x}]);
65
        `;
66
      assignment = `
67
        let outputOffset = global_idx * ${components};
68
        var data = vec4<u32>(0);
69
        ${singleAssignment('data', 0, 'u32')}
70
        ${singleAssignment('data', 1, 'u32')}
71
        ${singleAssignment('data', 2, 'u32')}
72
        ${singleAssignment('data', 3, 'u32')}
73
        ${output.setByOffset('global_idx', 'data')}
74
      }`;
75
    } else {
76
      assignment = `
77
        let outputIndices = ${output.offsetToIndices('global_idx')};
78
        let inputOffset = ${input.broadcastedIndicesToOffset('outputIndices', output)};
79
        ${output.setByOffset('global_idx', input.getByOffset('inputOffset'))}
80
      }`;
81
    }
82
    return `
83
    ${shaderHelper.registerUniform('vec_size', 'u32').declareVariables(input, output)}
84
    ${shaderHelper.mainStart()}
85
    ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.vec_size')}
86
    ${assignment}`;
87
  };
88

89
  const programUniforms: ProgramUniform[] = [
90
    { type: DataType.uint32, data: outputSize },
91
    ...createTensorShapeVariables(inputShape, outputShape),
92
  ];
93
  return {
94
    name: 'Expand',
95
    shaderCache: { hint: `${outputShape.length}`, inputDependencies: ['rank'] },
96
    getShaderSource,
97
    getRunData: () => ({
98
      outputs: [{ dims: outputShape, dataType: inputs[0].dataType }],
99
      dispatchGroup: { x: Math.ceil(outputSize / 64 /* workgroup size */) },
100
      programUniforms,
101
    }),
102
  };
103
};
104

105
export const expand = (context: ComputeContext): void => {
106
  validateInputs(context.inputs);
107
  context.compute(createExpandProgramInfo(context.inputs), { inputs: [0] });
108
};
109