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

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

import { Tensor } from '../../../tensor';
import { getGlsl } from '../glsl-source';
import { WebGLInferenceHandler } from '../inference-handler';
import { ProgramInfo, ProgramInfoLoader, ProgramMetadata, TextureType } from '../types';
import { getCoordsDataType, getGlChannels } from '../utils';

import { ConcatAttributes } from './concat';
import { getChannels, unpackFromChannel } from './packing-utils';

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

const createPackedConcatProgramInfo = (
  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 coords = getChannels('coords', rank);
  const dtype = getCoordsDataType(rank);
  const unpackChannel = unpackFromChannel();

  const shapes = inputs.map((i) => i.dims);
  const channels = getGlChannels(rank);
  const offsets: number[] = new Array(shapes.length - 1);

  offsets[0] = shapes[0][axis];
  for (let i = 1; i < offsets.length; i++) {
    offsets[i] = offsets[i - 1] + shapes[i][axis];
  }

  const channel = channels[axis];
  const lastChannels = channels.slice(-2);
  const allChannels = channels.join();

  let getValueSnippet = `if (${channel} < ${offsets[0]}) {
        return getChannel(
            getX0(${allChannels}), vec2(${lastChannels.join()}));
        }`;
  for (let i = 1; i < offsets.length; i++) {
    const shift = offsets[i - 1];
    getValueSnippet += `
            if (${channel} < ${offsets[i]}  && ${channel} >= ${offsets[i - 1]}) {
              return getChannel(
                getX${i}(${getShiftedChannelsSnippet(channels, channel, shift)}),
                vec2(${getShiftedChannelsSnippet(lastChannels, channel, shift)}));
            }`;
  }
  const lastIndex = offsets.length;
  const shift = offsets[offsets.length - 1];
  getValueSnippet += `
            return getChannel(
              getX${lastIndex}(${getShiftedChannelsSnippet(channels, channel, shift)}),
              vec2(${getShiftedChannelsSnippet(lastChannels, channel, shift)}));`;

  const glsl = getGlsl(handler.session.backend.glContext.version);

  const shaderSource = `
          ${unpackChannel}
          float getValue(${channels.map((x) => 'int ' + x)}) {
            ${getValueSnippet}
          }

          void main() {
            ${dtype} coords = getOutputCoords();
            int lastDim = coords.${channels[rank - 1]};
            coords.${channels[rank - 1]} = coords.${channels[rank - 2]};
            coords.${channels[rank - 2]} = lastDim;

            vec4 result = vec4(getValue(${coords}), 0., 0., 0.);

            ${coords[rank - 1]} = ${coords[rank - 1]} + 1;
            if (${coords[rank - 1]} < ${outputShape[rank - 1]}) {
              result.g = getValue(${coords});
            }

            ${coords[rank - 2]} = ${coords[rank - 2]} + 1;
            if (${coords[rank - 2]} < ${outputShape[rank - 2]}) {
              result.a = getValue(${coords});
            }

            ${coords[rank - 1]} = ${coords[rank - 1]} - 1;
            if (${coords[rank - 2]} < ${outputShape[rank - 2]} &&
                ${coords[rank - 1]} < ${outputShape[rank - 1]}) {
              result.b = getValue(${coords});
            }
            ${glsl.output} = result;
          }
        `;

  return {
    ...metadata,
    output: { dims: outputShape, type: inputs[0].type, textureType: TextureType.packed },
    shaderSource,
    hasMain: true,
  };
};

export const createPackedConcatProgramInfoLoader = (
  handler: WebGLInferenceHandler,
  inputs: Tensor[],
  attributes: ConcatAttributes,
): ProgramInfoLoader => {
  const metadata = createPackedConcatProgramMetadata(inputs.length, attributes.cacheKey);
  return { ...metadata, get: () => createPackedConcatProgramInfo(handler, metadata, inputs, attributes.axis) };
};

const getShiftedChannelsSnippet = (channels: string[], channel: string, shift: number): string => {
  const channelIdx = channels.indexOf(channel);
  const res = channels.map((c, idx) => {
    if (idx === channelIdx) {
      return `${c} - ${shift}`;
    } else {
      return c;
    }
  });
  return res.join();
};