Function getBinaryOpString

tfjs-backend-webgpu/src/binary_op_util.ts:165–282

(
    type: BinaryOpType, useVec4?: boolean)

const SUB = 'let resultTemp = a - b;';

export function getBinaryOpString(
    type: BinaryOpType, useVec4?: boolean): string {
  let doOpSnippet: string;

  // Ops with NaN check
  do {
    switch (type) {
      case BinaryOpType.ATAN2:
        doOpSnippet = ATAN2;
        break;
      case BinaryOpType.MAX:
        doOpSnippet = MAX;
        break;
      case BinaryOpType.MIN:
        doOpSnippet = MIN;
        break;
      case BinaryOpType.MOD:
        doOpSnippet = useVec4 ? MOD_VEC4 : MOD;
        break;
      case BinaryOpType.NOT_EQUAL:
        doOpSnippet = useVec4 ? NOT_EQUAL_VEC4 : NOT_EQUAL;
        break;
      case BinaryOpType.POW:
        doOpSnippet = useVec4 ? POW_VEC4 : POW;
        break;
      default:
        continue;
    }

    let isNaN: string;
    let dTypeN: string;
    let boolN: string;
    if (useVec4) {
      isNaN = 'isnanVec4';
      dTypeN = 'vec4<f32>';
      boolN = 'vec4<bool>';
    } else {
      isNaN = 'isnan';
      dTypeN = 'f32';
      boolN = 'bool';
    }

    return `
      let aIsNaN = ${isNaN}(a);
      let aPostLegalization = select(a, ${dTypeN}(42), aIsNaN);
      let bIsNaN = ${isNaN}(b);
      let bPostLegalization = select(b, ${dTypeN}(42), bIsNaN);
      let isNaN = false;
      let valueForNaN = uniforms.NAN;
      {
        let a = aPostLegalization;
        let b = bPostLegalization;
        ${doOpSnippet}
        return select(
            resultTemp, ${dTypeN}(valueForNaN),
            ${boolN}(isNaN) | aIsNaN | bIsNaN);
      }
    `;