1
// Copyright (c) Microsoft Corporation. All rights reserved.
2
// Licensed under the MIT License.
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';
10
import { createTensorShapeVariables, IndicesHelper, inputVariable, outputVariable, ShaderHelper } from './common';
12
export interface ConcatAttributes extends AttributeWithCacheKey {
13
readonly axis: number;
16
const validateInputs = (inputs: readonly TensorView[], axis: number): void => {
17
if (!inputs || inputs.length < 1) {
18
throw new Error('too few inputs');
20
const referenceIndex = 0;
21
const referenceInput = inputs[referenceIndex];
22
const inputType = referenceInput.dataType;
23
const inputRank = referenceInput.dims.length;
24
inputs.forEach((input, i) => {
25
if (i === referenceIndex) {
28
// make sure types of all inputs match
29
if (input.dataType !== inputType) {
30
throw new Error('input tensors should be one type');
32
// make sure the dimensionality of all inputs are the same
33
if (input.dims.length !== inputRank) {
34
throw new Error('input tensors should have the same shape');
36
input.dims.forEach((dim, i) => {
37
if (i !== axis && dim !== referenceInput.dims[i]) {
38
throw new Error('non concat dimensions must match');
44
const calculateInputIndexImpl = (numberOfTensors: number, sizeInConcatAxisStr: string): string => `
45
fn calculateInputIndex(index: u32) -> u32 {
46
let sizeInConcatAxis = array<u32, ${numberOfTensors}u>(${sizeInConcatAxisStr});
47
for (var i: u32 = 0u; i < ${numberOfTensors}; i += 1u ) {
48
if (index < sizeInConcatAxis[i]) {
52
return ${numberOfTensors}u;
55
const assignOutputData = (inputs: readonly IndicesHelper[], output: IndicesHelper) => {
56
const numberOfTensors = inputs.length;
58
const codeLines: string[] = [];
59
for (let i = 0; i < numberOfTensors; ++i) {
60
const returnSnippet = output.setByOffset('global_idx', inputs[i].getByIndices('indices'));
61
if (numberOfTensors === 1) {
62
codeLines.push(returnSnippet);
64
codeLines.push(`if (inputIndex == ${i}u) { ${returnSnippet} }`);
65
} else if (i === numberOfTensors - 1) {
66
codeLines.push(`else { ${returnSnippet} }`);
68
codeLines.push(`else if (inputIndex == ${i}) { ${returnSnippet} }`);
71
return codeLines.join('\n');
74
const createConcatProgramInfo = (
75
inputs: readonly TensorView[],
77
outputShape: number[],
80
const outputSize = ShapeUtil.size(outputShape);
82
const sizeInConcatAxis = new Array<number>(inputs.length);
83
const inputVars = new Array<IndicesHelper>(inputs.length);
86
const inputDependencies: ProgramInputTensorInfoDependency[] = [];
87
const inputRanks = [];
88
const programUniforms: ProgramUniform[] = [{ type: DataType.uint32, data: outputSize }];
89
for (let i = 0; i < inputs.length; ++i) {
90
previousSum += inputs[i].dims[adjustedAxis];
91
sizeInConcatAxis[i] = previousSum;
92
inputRanks.push(inputs[i].dims.length);
93
inputVars[i] = inputVariable(`input${i}`, dataType, inputRanks[i]);
94
inputDependencies.push('rank');
95
programUniforms.push({ type: DataType.uint32, data: sizeInConcatAxis[i] });
97
for (let i = 0; i < inputs.length; ++i) {
98
programUniforms.push(...createTensorShapeVariables(inputs[i].dims));
100
programUniforms.push(...createTensorShapeVariables(outputShape));
102
const output = outputVariable('output', dataType, outputShape.length);
103
const indicesAxis = output.indicesGet('indices', adjustedAxis);
104
const sizeInConcatAxisStr = Array.from(Array(sizeInConcatAxis.length).keys())
105
.map((i) => `uniforms.sizeInConcatAxis${i}`)
107
const getShaderSource = (shaderHelper: ShaderHelper) => `
110
shaderHelper.registerUniform('outputSize', 'u32');
111
for (let i = 0; i < inputs.length; i++) {
112
shaderHelper.registerUniform(`sizeInConcatAxis${i}`, 'u32');
114
return shaderHelper.declareVariables(...inputVars, output);
117
${calculateInputIndexImpl(sizeInConcatAxis.length, sizeInConcatAxisStr)}
119
${shaderHelper.mainStart()}
120
${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes('uniforms.outputSize')}
122
var indices = ${output.offsetToIndices('global_idx')};
124
let inputIndex = calculateInputIndex(${indicesAxis});
125
if (inputIndex != 0u) {
126
let sizeInConcatAxis = array<u32, ${sizeInConcatAxis.length}u>(${sizeInConcatAxisStr});
127
${indicesAxis} -= sizeInConcatAxis[inputIndex - 1u];
130
${assignOutputData(inputVars, output)}
135
shaderCache: { hint: `${adjustedAxis}`, inputDependencies },
137
outputs: [{ dims: outputShape, dataType }],
138
dispatchGroup: { x: Math.ceil(outputSize / 64 /* workgroup size */) },
145
export const concat = (context: ComputeContext, attributes: ConcatAttributes): void => {
146
const inputs = context.inputs;
147
const inputShape = inputs[0].dims;
148
const adjustedAxis = ShapeUtil.normalizeAxis(attributes.axis, inputShape.length);
149
validateInputs(inputs, adjustedAxis);
150
const outputShape = inputShape.slice();
151
outputShape[adjustedAxis] = inputs.reduce(
152
(sum, input) => sum + (input.dims.length > adjustedAxis ? input.dims[adjustedAxis] : 0),
155
// 0 length tensors are valid for concat, remove them
156
const nonEmptyInputs = inputs.filter((input) => ShapeUtil.size(input.dims) > 0);
157
context.compute(createConcatProgramInfo(nonEmptyInputs, adjustedAxis, outputShape, inputs[0].dataType), {
158
inputs: nonEmptyInputs,
162
export const parseConcatAttributes = (attributes: Record<string, unknown>): ConcatAttributes =>
163
createAttributeWithCacheKey({ axis: attributes.axis as number });