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 FormatAttributes {
13
  readonly format: 'NHWC' | 'NCHW';
14
}
15

16
export interface DepthToSpaceAttributes extends FormatAttributes, AttributeWithCacheKey {
17
  readonly blocksize: number;
18
  readonly mode: string;
19
}
20

21
const validateInputs = (inputs: readonly TensorView[]): void => {
22
  if (!inputs || inputs.length !== 1) {
23
    throw new Error('DepthToSpace requires 1 input.');
24
  }
25
  if (inputs[0].dims.length !== 4) {
26
    throw new Error('DepthToSpace requires 4D input.');
27
  }
28
};
29

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

41
const createDepthToSpaceProgramInfo = (inputTensor: TensorView, attributes: DepthToSpaceAttributes): ProgramInfo => {
42
  let n: number, h: number, w: number, c: number;
43
  let shape: number[];
44
  let perm: number[];
45
  const isChannelLast = attributes.format === 'NHWC';
46
  const blocksize = attributes.blocksize;
47
  const isDCRmode = attributes.mode === 'DCR';
48
  if (isChannelLast) {
49
    [n, h, w, c] = inputTensor.dims;
50
    shape = isDCRmode
51
      ? [n, h, w, blocksize, blocksize, c / blocksize ** 2]
52
      : [n, h, w, c / blocksize ** 2, blocksize, blocksize];
53
    perm = isDCRmode ? [0, 1, 3, 2, 4, 5] : [0, 1, 4, 2, 5, 3];
54
  } else {
55
    [n, h, w, c] = [inputTensor.dims[0], inputTensor.dims[2], inputTensor.dims[3], inputTensor.dims[1]];
56
    shape = isDCRmode
57
      ? [n, blocksize, blocksize, c / blocksize ** 2, h, w]
58
      : [n, c / blocksize ** 2, blocksize, blocksize, h, w];
59
    perm = isDCRmode ? [0, 3, 4, 1, 5, 2] : [0, 1, 4, 2, 5, 3];
60
  }
61
  const reshapedInputTensor = inputTensor.reshape(shape);
62
  const reshapedInputRank = reshapedInputTensor.dims.length;
63
  const inputDataType = inputTensor.dataType;
64

65
  const reshapedInput = inputVariable('a', inputDataType, reshapedInputRank);
66
  const permedOutput = outputVariable('output', inputDataType, reshapedInputRank);
67

68
  const getShaderSource = (shaderHelper: ShaderHelper) => `
69
  ${shaderHelper.registerUniform('output_size', 'u32').declareVariables(reshapedInput, permedOutput)}
70

71
  ${permFunctionBody(perm, reshapedInputRank, reshapedInput, permedOutput)}
72

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

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

79
    ${permedOutput.setByOffset('global_idx', reshapedInput.getByIndices('aIndices'))}
80
  }`;
81

82
  return {
83
    name: 'DepthToSpace',
84
    shaderCache: {
85
      hint: `${inputTensor.dims};${attributes.blocksize};${attributes.mode}`,
86
      inputDependencies: ['rank'],
87
    },
88
    getRunData: (inputs) => {
89
      const outputShape = isChannelLast
90
        ? [n, h * blocksize, w * blocksize, c / blocksize ** 2]
91
        : [n, c / blocksize ** 2, h * blocksize, w * blocksize];
92
      const outputSize = ShapeUtil.size(outputShape);
93
      const shapeBeforePerm = reshapedInputTensor.dims;
94
      const shapeAfterPerm = ShapeUtil.sortBasedOnPerm(shapeBeforePerm, perm);
95
      return {
96
        outputs: [{ dims: outputShape, dataType: inputs[0].dataType }],
97
        dispatchGroup: { x: Math.ceil(outputSize / 64 /* workgroup size */) },
98
        programUniforms: [
99
          { type: DataType.uint32, data: outputSize },
100
          ...createTensorShapeVariables(shapeBeforePerm, shapeAfterPerm),
101
        ],
102
      };
103
    },
104
    getShaderSource,
105
  };
106
};
107

108
export const depthToSpace = (context: ComputeContext, attributes: DepthToSpaceAttributes): void => {
109
  validateInputs(context.inputs);
110
  context.compute(createDepthToSpaceProgramInfo(context.inputs[0], attributes));
111
};
112

113
export const parseDepthToSpaceAttributes = (attributes: Record<string, unknown>): DepthToSpaceAttributes =>
114
  createAttributeWithCacheKey({
115
    blocksize: attributes.blocksize as number,
116
    mode: attributes.mode as string,
117
    format: attributes.format as 'NHWC' | 'NCHW',
118
  });
119