Class ReduceProgram

tfjs-backend-webgpu/src/reduce_webgpu.ts:22–145

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

export class ReduceProgram implements WebGPUProgram {
  outputShape: number[];
  shaderKey: string;
  dispatchLayout: {x: number[]};
  dispatch: [number, number, number];
  workgroupSize: [number, number, number];
  variableNames = ['x'];
  uniforms = 'reduceSize : i32,';
  reduceType: 'all'|'any'|'max'|'mean'|'min'|'prod'|'sum';
  inputShape: number[];
  size = true;

  constructor(
      reduceInfo: backend_util.ReduceInfo,
      reduceType: 'all'|'any'|'max'|'mean'|'min'|'prod'|'sum',
      maxComputeWorkgroupSizeX: number) {
    this.inputShape = [reduceInfo.batchSize, reduceInfo.inSize];
    const [outputShape, ] =
        backend_util.computeOutAndReduceShapes(this.inputShape, [1]);
    this.outputShape = outputShape.length === 0 ? [1] : outputShape;
    // If reduceSize |reduceInfo.inSize| is very large, the I/O accessing will
    // become the bottleneck. Increasing workgroupSize can reduce the times of
    // accessing global memory. The threshold value is just to make sure the
    // reduceSize is large enough for a bigger workgroupSize.
    if (reduceInfo.inSize >= 32768 && maxComputeWorkgroupSizeX >= 512) {
      this.workgroupSize = [512, 1, 1];
    } else if (reduceInfo.inSize >= 4096) {
      this.workgroupSize = [256, 1, 1];
    } else {
      this.workgroupSize = [64, 1, 1];
    }
    this.dispatchLayout = flatDispatchLayout(this.outputShape);
    // 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.reduceType = reduceType;
    this.shaderKey = `reduce_${reduceType}`;
  }

  getUserCode(): string {
    let reduceOp = ``;
    let initValue = '0.0';
    const workgroupSizeX = this.workgroupSize[0];
    if (this.reduceType === 'min' || this.reduceType === 'max') {
      reduceOp = `
         if (isnan(candidate)) {
          bestValue = uniforms.NAN;
         } else if (!isnan(bestValue) && candidate ${
          this.reduceType === 'min' ? '<' : '>'} bestValue)
           {  bestValue = candidate; }`;
      initValue = 'f32(x[offset])';
    } else if (this.reduceType === 'sum' || this.reduceType === 'mean') {
      reduceOp = ' bestValue = bestValue + candidate; ';
    } else if (this.reduceType === 'prod') {
      reduceOp = ' bestValue = bestValue * candidate; ';
      initValue = '1.0';