Method getUserCode

tfjs-backend-webgpu/src/binary_op_webgpu.ts:93–145

()

  }

  getUserCode(): string {
    let userCode;
    const dType = this.outputComponent === 4 ? 'vec4<f32>' : 'f32';
    const opFnStr = `
    fn binaryOperation(a : ${dType}, b : ${dType}) -> ${dType} {
      ${getBinaryOpString(this.op, this.outputComponent === 4)}
    };
    `;

    if (this.type === 'shared') {
      const sharedIndexSnippet = this.lastDimensionSize > 1 ?
          `coords[${this.outputShape.length - 1}]` :
          '0';
      const accessDataSnippet = this.useSharedMemoryWithB ?
          `let a = getAByOutputIndex(index);
          let b = sharedBuf[${sharedIndexSnippet}];` :
          `let a = sharedBuf[${sharedIndexSnippet}];
          let b = getBByOutputIndex(index);`;
      userCode = `
        ${opFnStr}
        var<workgroup> sharedBuf : array<f32, ${this.lastDimensionSize}>;
        ${main('index')} {
          // Fill in the shared memory buffer.
          let localIndex = i32(localId.x);
          if(localIndex < ${this.lastDimensionSize}) {
            sharedBuf[localIndex] = f32(${
          this.useSharedMemoryWithB ? 'B' : 'A'}[localIndex]);
          }
          workgroupBarrier();

          if(index < uniforms.size) {
            let coords = getCoordsFromIndex(index);
            ${accessDataSnippet}
            setOutputAtIndex(index, binaryOperation(a, b));
          }
        }
        `;
    } else {
      userCode = `
       ${opFnStr}
       ${main('index')} {
         if (index < uniforms.size) {
           let coords = getCoordsFromIndex(index * ${this.outputComponent});
           let a = ${dType}(getAByOutputCoords(coords));
           let b = ${dType}(getBByOutputCoords(coords));
           setOutputAtIndex(index, binaryOperation(a, b));
         }
       }
       `;
    }

    return userCode;
  }
}