onnxruntime/js/web/lib/onnxjs/backends/webgl/ops/concat.ts

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

import { AttributeWithCacheKey, createAttributeWithCacheKey } from '../../../attribute-with-cache-key';
import { Graph } from '../../../graph';
import { OperatorImplementation, OperatorInitialization } from '../../../operators';
import { Tensor } from '../../../tensor';
import { WebGLInferenceHandler } from '../inference-handler';
import { ProgramInfo, ProgramInfoLoader, ProgramMetadata, TextureType } from '../types';

import { createPackedConcatProgramInfoLoader } from './concat-packed';

export interface ConcatAttributes extends AttributeWithCacheKey {
  readonly axis: number;
}

export const concat: OperatorImplementation<ConcatAttributes> = (
  inferenceHandler: WebGLInferenceHandler,
  inputs: Tensor[],
  attributes: ConcatAttributes,
): Tensor[] => {
  validateInputs(inputs);
  if (inferenceHandler.session.pack && inputs[0].dims.length > 1) {
    const output = inferenceHandler.run(
      createPackedConcatProgramInfoLoader(inferenceHandler, inputs, attributes),
      inputs,
    );
    return [output];
  } else {
    const output = inferenceHandler.run(
      createUnpackedConcatProgramInfoLoader(inferenceHandler, inputs, attributes),
      inputs,
    );
    return [output];
  }
};

const createUnpackedConcatProgramMetadata = (inputCount: number, cacheHint: string) => ({
  name: 'Concat',
  inputNames: Array.from({ length: inputCount }, (_v, i) => `X${i}`),
  inputTypes: Array(inputCount).fill(TextureType.unpacked),
  cacheHint,
});

const createUnpackedConcatProgramInfo = (
  _handler: WebGLInferenceHandler,
  metadata: ProgramMetadata,
  inputs: Tensor[],
  axis: number,
): ProgramInfo => {
  const inputShape = inputs[0].dims.slice();
  if (axis >= inputShape.length || axis < -1 * inputShape.length) {
    throw new Error("axis specified for concat doesn't match input dimensionality");
  }
  if (axis < 0) {
    axis = inputShape.length + axis;
  }
  // ensure all of the non-concatenated axes match each other
  // calculate the shape of the output tensor while we do that
  const outputShape = inputShape.slice(0);
  for (let i = 1; i < inputs.length; i++) {
    const dataNShape = inputs[i].dims.slice();
    for (let axisIndex = 0; axisIndex < inputShape.length; axisIndex++) {
      // add to the placeholder for computing output shape
      if (axisIndex === axis) {
        outputShape[axis] += dataNShape[axisIndex];
      }
      // ensure all non-cancatenated axes match each other
      else if (inputShape[axisIndex] !== dataNShape[axisIndex]) {
        throw new Error('non concat dimensions must match');
      }
    }
  }

  const rank = outputShape.length;

  const sizeInConcatAxis = new Array<number>(inputs.length);
  let previousSum = 0;
  for (let i = 0; i < sizeInConcatAxis.length; ++i) {
    previousSum += inputs[i].dims[axis];
    sizeInConcatAxis[i] = previousSum;
  }

  let getTextureIndexWhereDataResidesMethod = '';
  // in most cases linear search is sufficient, as in most scenarios, only 2 tensors are concatenated
  if (inputs.length < 5) {
    getTextureIndexWhereDataResidesMethod = getTextureIndexWhereDataResidesLinearSearch(sizeInConcatAxis);
  } else {
    getTextureIndexWhereDataResidesMethod = getTextureIndexWhereDataResidesBinarySearch(sizeInConcatAxis);
  }

  const fetchDataFromCorrectTextureMethod = getFetchDataFromCorrectTextureMethod(inputs.length, rank);
  const getSizeInConcatAxisValueFromIndexMethod = getGetSizeInConcatAxisValueFromIndexMethod(sizeInConcatAxis);
  const shaderSource = `
        ${fetchDataFromCorrectTextureMethod}
        ${getSizeInConcatAxisValueFromIndexMethod}
        ${getTextureIndexWhereDataResidesMethod}
        float process(int indices[${rank}]) {
          int textureIndex = getTextureWhereDataResides (indices[${axis}]);

          if(textureIndex != 0) {
            indices[${axis}] = indices[${axis}] - int(getSizeInConcatAxisValueFromIndex(textureIndex-int(1)));
          }

          return fetchDataFromCorrectTexture(textureIndex, indices);
        }`;
  return {
    ...metadata,
    output: { dims: outputShape, type: inputs[0].type, textureType: TextureType.unpacked },
    shaderSource,
  };
};

const createUnpackedConcatProgramInfoLoader = (
  handler: WebGLInferenceHandler,
  inputs: Tensor[],
  attributes: ConcatAttributes,
): ProgramInfoLoader => {
  const metadata = createUnpackedConcatProgramMetadata(inputs.length, attributes.cacheKey);
  return { ...metadata, get: () => createUnpackedConcatProgramInfo(handler, metadata, inputs, attributes.axis) };
};

const getTextureIndexWhereDataResidesLinearSearch = (sizeInConcatAxis: number[]): string => {
  const searchAxis = sizeInConcatAxis.map(
    (size, i) => `if(index<${size}) {return ${i};}
`,
  );
  return `int getTextureWhereDataResides(int index) {
      ${searchAxis.join('')}
    }`;
};

// TODO: Implement BinarySearch in GLSL
const getTextureIndexWhereDataResidesBinarySearch = (sizeInConcatAxis: number[]): string =>
  getTextureIndexWhereDataResidesLinearSearch(sizeInConcatAxis);

const getFetchDataFromCorrectTextureMethod = (numberOfTensors: number, tensorRank: number) => {
  const codeLines: string[] = [`float fetchDataFromCorrectTexture(int textureIndex, int indices[${tensorRank}]) {`];
  for (let i = 0; i < numberOfTensors; ++i) {
    if (i === 0) {
      codeLines.push('\t' + `if (textureIndex == ${i}) { return _X${i}(indices); }`);
    } else if (i === numberOfTensors - 1) {
      codeLines.push('\t' + `else { return _X${i}(indices); }`);
    } else {
      codeLines.push('\t' + `else if (textureIndex == ${i}) { return _X${i}(indices); }`);
    }
  }
  codeLines.push('\t' + '}');
  return codeLines.join('\n');
};

const getGetSizeInConcatAxisValueFromIndexMethod = (sizeInConcatAxis: number[]): string => {
  const codeLines: string[] = ['int getSizeInConcatAxisValueFromIndex(int index) {'];
  for (let i = 0; i < sizeInConcatAxis.length; ++i) {
    if (i === 0) {
      codeLines.push('\t' + `if (index == ${i}) { return ${sizeInConcatAxis[i]}; }`);
    } else if (i === sizeInConcatAxis.length - 1) {
      codeLines.push('\t' + `else { return ${sizeInConcatAxis[i]}; }`);
    } else {
      codeLines.push('\t' + `else if (index == ${i}) { return ${sizeInConcatAxis[i]}; }`);
    }
  }
  codeLines.push('\t' + '}');

  return codeLines.join('\n');
};

export const parseConcatAttributes: OperatorInitialization<ConcatAttributes> = (node: Graph.Node): ConcatAttributes =>
  createAttributeWithCacheKey({ axis: node.attributes.getInt('axis') });

const validateInputs = (inputs: Tensor[]): void => {
  if (!inputs || inputs.length < 1) {
    throw new Error('too few inputs');
  }

  const inputType = inputs[0].type;
  const inputDimensionality = inputs[0].dims.length;

  // TODO: Support string concat
  if (inputType === 'string') {
    throw new Error('string tensor is not supported yet');
  }

  for (const input of inputs) {
    // make sure types of all inputs match
    if (input.type !== inputType) {
      throw new Error('input tensors should be one type');
    }

    // make sure the dimensionality of all inputs are the same
    if (input.dims.length !== inputDimensionality) {
      throw new Error('input tensors should have the same shape');
    }
  }
};