Class ArgMinMaxProgram

tfjs-backend-webgpu/src/argminmax_webgpu.ts:22–167

import {computeDispatch, flatDispatchLayout} from './webgpu_util';

export class ArgMinMaxProgram implements WebGPUProgram {
  outputShape: number[];
  shaderKey: string;
  dispatchLayout: {x: number[]};
  dispatch: [number, number, number];
  workgroupSize: [number, number, number] = [64, 1, 1];
  variableNames = ['x'];
  uniforms = 'infinityValue : f32,';
  inputShape: number[];
  reductionFactor: number;
  op: string;
  size = true;
  private type: string;

  constructor(inputShape: number[], axis: number, reduceType: 'min'|'max') {
    const axes = [axis];

    this.op = reduceType === 'min' ? '<' : '>';

    // |outShape| is the shape with the removed axis
    const [outputShape, reduceShape] =
        backend_util.computeOutAndReduceShapes(inputShape, axes);

    this.outputShape = outputShape.length === 0 ? [1] : outputShape;
    this.dispatchLayout = flatDispatchLayout(this.outputShape);
    // The shared algorithm is mainly used for large reduce size. It fully
    // utilizes the threads in one workgroup to do the reduction. However,
    // when the reduce size is very small, it's better to use the plain
    // algorithm to reduce the number of workgroups to speedup. The threthold
    // can be further tuned.
    if (util.sizeFromShape(reduceShape) < 32) {
      this.type = 'plain';
      this.dispatch = computeDispatch(
          this.dispatchLayout, this.outputShape, this.workgroupSize);
    } else {
      this.type = 'shared';
      // A work group only outputs a data, so we transfer [1, 1, 1] to compute
      // dispatch size.
      this.dispatch =
          computeDispatch(this.dispatchLayout, this.outputShape, [1, 1, 1]);
    }

    this.inputShape = inputShape;
    this.shaderKey = `argMinMax_${this.op}_${this.type}`;
  }

  getUserCode(): string {
    const workgroupSizeX = this.workgroupSize[0];
    const getInputShapeLastDim = () => {
      if (this.inputShape.length === 1) {
        return 'uniforms.xShape';
      } else {
        return `uniforms.xShape.${getCoordsXYZ(this.inputShape.length - 1)}`;
      }
    };

    const splitOutputCoords = () => {
      let snippet = '';