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

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

12
export interface TransposeAttributes extends AttributeWithCacheKey {
13
  readonly perm: number[];
14
}
15

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

22
const getAdjustedPerm = (inputRank: number, perm: number[]): number[] =>
23
  perm && perm.length !== inputRank ? [...new Array(inputRank).keys()].reverse() : perm;
24

25
const getOutputShape = (inputShape: readonly number[], perm: number[]): readonly number[] =>
26
  ShapeUtil.sortBasedOnPerm(inputShape, getAdjustedPerm(inputShape.length, perm));
27

28
const permFunctionBody = (perm: number[], rank: number, input: IndicesHelper, output: IndicesHelper): string => {
29
  const reverseFunc = [];
30
  reverseFunc.push(`fn perm(i: ${output.type.indices}) -> ${input.type.indices} {
31
    var a: ${input.type.indices};`);
32
  for (let i = 0; i < rank; ++i) {
33
    reverseFunc.push(input.indicesSet('a', perm[i], `i[${i}]`));
34
  }
35
  reverseFunc.push('return a;}');
36
  return reverseFunc.join('\n');
37
};
38

39
export const createTransposeProgramInfo = (inputTensor: TensorView, permAttr: number[]): ProgramInfo => {
40
  const inputDataType = inputTensor.dataType;
41
  const inputRank = inputTensor.dims.length;
42
  const perm = getAdjustedPerm(inputRank, permAttr);
43
  const outputShape = getOutputShape(inputTensor.dims, perm);
44
  const output = outputVariable('output', inputDataType, outputShape.length);
45
  const input = inputVariable('a', inputDataType, inputRank);
46
  let getShaderSource;
47
  if (perm.length === 2 && perm[0] === 1 && perm[1] === 0) {
48
    const wgslType = output.type.value;
49
    const workgroupSize: [number, number, number] = [16, 16, 1];
50
    getShaderSource = (shaderHelper: ShaderHelper) => `
51
  ${shaderHelper.registerUniform('output_size', 'u32').declareVariables(input, output)}
52
  var<workgroup> tile : array<array<${wgslType}, ${workgroupSize[0] + 1}>, ${workgroupSize[0]}>;
53
  ${shaderHelper.mainStart(workgroupSize)}
54
    var x = workgroup_id.x * ${workgroupSize[0]}u + local_id.x;
55
    var y = workgroup_id.y * ${workgroupSize[0]}u + local_id.y;
56
    let width = uniforms.output_shape[0];
57
    let height = uniforms.output_shape[1];
58
    if (x < width && y < height) {
59
      tile[local_id.y][local_id.x] = ${input.getByOffset('y * width + x')};
60
    }
61
    workgroupBarrier();
62
    x = workgroup_id.y * ${workgroupSize[0]}u + local_id.x;
63
    y = workgroup_id.x * ${workgroupSize[0]}u + local_id.y;
64
    if (x < height && y < width) {
65
      ${output.setByOffset('y * height + x', 'tile[local_id.x][local_id.y]')}
66
    }
67
  }`;
68
  } else {
69
    getShaderSource = (shaderHelper: ShaderHelper) => `
70
  ${shaderHelper.registerUniform('output_size', 'u32').declareVariables(input, output)}
71

72
  ${permFunctionBody(perm, inputRank, input, output)}
73

74
  ${shaderHelper.mainStart()}
75
    ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.output_size')}
76

77
    let indices = ${output.offsetToIndices('global_idx')};
78
    let aIndices = perm(indices);
79

80
    ${output.setByOffset('global_idx', input.getByIndices('aIndices'))}
81
  }`;
82
  }
83
  return {
84
    name: 'Transpose',
85
    shaderCache: { hint: `${permAttr}`, inputDependencies: ['rank'] },
86
    getRunData: () => {
87
      const outputSize = ShapeUtil.size(outputShape);
88
      return {
89
        outputs: [{ dims: outputShape, dataType: inputTensor.dataType }],
90
        dispatchGroup: { x: Math.ceil(outputSize / 64 /* workgroup size */) },
91
        programUniforms: [
92
          { type: DataType.uint32, data: outputSize },
93
          ...createTensorShapeVariables(inputTensor.dims, outputShape),
94
        ],
95
      };
96
    },
97
    getShaderSource,
98
  };
99
};
100

101
export const transpose = (context: ComputeContext, attributes: TransposeAttributes): void => {
102
  validateInputs(context.inputs);
103
  context.compute(createTransposeProgramInfo(context.inputs[0], attributes.perm));
104
};
105

106
export const parseTransposeAttributes = (attributes: Record<string, unknown>): TransposeAttributes =>
107
  createAttributeWithCacheKey({ perm: attributes.perm as number[] });
108