Class StridedSliceProgram

tfjs-backend-webgpu/src/strided_slice_webgpu.ts:21–74

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

export class StridedSliceProgram implements WebGPUProgram {
  variableNames = ['x'];
  uniforms: string;
  outputShape: number[];
  shaderKey: string;
  dispatchLayout: {x: number[]};
  dispatch: [number, number, number];
  // TODO(xing.xu): Increase the workPerThread.
  workPerThread = 1;
  workgroupSize: [number, number, number] = [64, 1, 1];
  size = true;

  constructor(destSize: number[]) {
    this.outputShape = destSize;
    this.dispatchLayout = flatDispatchLayout(this.outputShape);
    this.dispatch = computeDispatch(
        this.dispatchLayout, this.outputShape, this.workgroupSize,
        [this.workPerThread, 1, 1]);

    const dtype = getCoordsDataType(this.outputShape.length);
    this.uniforms = `begin : ${dtype},  strides : ${dtype}, `;
    this.shaderKey = 'stridedSlice';
  }

  getUserCode(): string {
    const rank = this.outputShape.length;
    let newCoords = '';
    if (rank === 1) {
      newCoords = 'coords * uniforms.strides + uniforms.begin';
    } else {
      let outputAxis = 0;
      newCoords =
          this.outputShape
              .map((_, i) => {
                outputAxis++;
                return this.outputShape.length === 1 ?
                    `coords * uniforms.strides[${i}] + uniforms.begin[${i}]` :
                    `coords[${outputAxis - 1}] * uniforms.strides[${
                        i}] + uniforms.begin[${i}]`;
              })
              .join(',');
    }

    const userCode = `
       ${main('index')} {
         if (index < uniforms.size) {
           let coords = getCoordsFromIndex(index);
           setOutputAtIndex(index, getX(${newCoords}));
         }
       }
     `;
    return userCode;
  }
}