Method getUserCode

tfjs-backend-webgpu/src/cum_webgpu.ts:53–88

()

  }

  getUserCode(): string {
    const rank = this.outputShape.length;
    const initVal = this.op === CumOpType.Prod ? '1.0' : '0.0';
    const val = this.exclusive ? initVal :
                                 `getX(${getCoords(rank, 'coords', this.op)})`;
    const length = this.outputShape[this.outputShape.length - 1];
    let condition = '';
    let idxString = '';
    // When exclusive is set, the cum op becomes roll op that copies the
    // value from the previous index based on the direction specified by the
    // reverse flag.
    if (this.exclusive) {
      condition = this.reverse ? `end != ${length - 1}` : 'end != 0';
      idxString = this.reverse ? 'end + 1' : 'end - 1';
    } else {
      condition = this.reverse ? `end + pow2 < ${length}` : 'end >= pow2';
      idxString = (this.reverse ? 'end + pow2' : 'end - pow2');
    }
    return `
      ${main('index')} {
       if (index < uniforms.size) {
         var coords = getCoordsFromIndex(index);

         let end = ${getFinalCoord(rank, 'coords', this.op)};
         var val = ${val};
         let pow2 = i32(pow(2.0, uniforms.index));
         if (${condition}) {
           let idx = ${idxString};
           ${getFinalCoord(rank, 'coords', this.op)} = idx;
           val ${this.op}= getX(${getCoords(rank, 'coords', this.op)});
         }
         setOutputAtIndex(index, val);
       }
      }
    `;
  }
}

function getCoords(rank: number, name: string, op: CumOpType): string {