"use strict";
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
Object.defineProperty(exports, "__esModule", { value: true });
exports.glslPRelu = exports.glslPow = exports.glslXor = exports.glslOr = exports.glslAnd = exports.glslLess = exports.glslGreater = exports.glslEqual = exports.glslSub = exports.glslMul = exports.glslDiv = exports.glslAdd = exports.WebGLBinaryOp = void 0;
const binary_op_1 = require("../../../ops/binary-op");
const util_1 = require("../../../util");
const glsl_definitions_1 = require("../glsl-definitions");
const glsl_source_1 = require("../glsl-source");
class WebGLBinaryOp extends binary_op_1.BinaryOp {
    constructor(typeConstraint, glslFunc, opType, resultType) {
        super(typeConstraint, opType, resultType);
        this.glslFunc = glslFunc;
    }
    run(inferenceHandler, inputs) {
        return inferenceHandler.run(this, inputs);
    }
    createProgramInfo(handler, inputs) {
        const inputLayouts = handler.session.pack ?
            inputs.map(t => handler.getOrCreateTextureLayout(t, 4, true, t.dims, true)) :
            inputs.map(t => handler.getOrCreateTextureLayout(t));
        const ouputLayout = handler.session.pack ?
            handler.createTextureLayoutFromShape(inputs[0].dims, 4, inputs[0].dims, { isPacked: true, reverseWH: true }) :
            handler.createTextureLayoutFromShape(inputs[0].dims);
        const isBroadcast = !util_1.ShapeUtil.areEqual(inputs[0].dims, inputs[1].dims);
        if (isBroadcast) {
            const outputShape = util_1.BroadcastUtil.calcShape(inputs[0].dims, inputs[1].dims, false);
            if (!outputShape) {
                throw new Error('Can\'t perform binary op on the given tensors');
            }
            const outputRank = outputShape.length;
            const aRank = inputs[0].dims.length !== 0 ? inputs[0].dims.length : 1;
            const bRank = inputs[1].dims.length !== 0 ? inputs[1].dims.length : 1;
            const aBcast = inputs[0].dims.length !== 0 ? 'bcastIndices_A(indices, aindices);' : 'aindices[0] = 0;';
            const bBcast = inputs[1].dims.length !== 0 ? 'bcastIndices_B(indices, bindices);' : 'bindices[0] = 0;';
            // TODO: for packed tensors, we need to implement logic to caculate textCoords for broadcast tensor
            const shaderSource = `
      ${this.glslFunc.body}
      float process(int indices[${outputRank}]) {
        int aindices[${aRank}];
        int bindices[${bRank}];
        ${aBcast}
        ${bBcast}
        return ${this.glslFunc.name}(_A(aindices), _B(bindices));
    }`;
            const outputLayout = handler.session.pack ?
                handler.createTextureLayoutFromShape(outputShape, 4, outputShape, { isPacked: true, reverseWH: true }) :
                handler.createTextureLayoutFromShape(outputShape);
            return {
                inputLayouts,
                outputLayout,
                samplers: ['A', 'B'],
                shaderSource,
                expectPackedInputs: handler.session.pack,
                expectPackedOutputs: handler.session.pack
            };
        }
        const glsl = glsl_source_1.getGlsl(handler.session.backend.glContext.version);
        const shaderSource = `
    ${this.glslFunc.body}
    void main() {
      vec4 v1 = ${glsl.texture2D}(A, TexCoords);
      vec4 v2 = ${glsl.texture2D}(B, TexCoords);
      vec4 result = ${this.glslFunc.name}(v1, v2);
      ${glsl.output} = result;
    }
    `;
        if (handler.session.pack) {
            return {
                hasMain: true,
                inputLayouts,
                outputLayout: ouputLayout,
                samplers: ['A', 'B'],
                shaderSource,
                expectPackedInputs: true,
                expectPackedOutputs: true
            };
        }
        else {
            return {
                hasMain: true,
                inputLayouts,
                outputLayout: handler.createTextureLayoutFromShape(inputs[0].dims),
                samplers: ['A', 'B'],
                shaderSource,
            };
        }
    }
    createRunData(handler, programInfo, inputs) {
        const inputTDs = handler.session.pack ?
            inputs.map((t) => handler.getOrCreateTextureData(t, handler.getOrCreateTextureLayout(t, 1, false, [], true))) :
            inputs.map((t, i) => handler.getOrCreateTextureData(t, programInfo.inputLayouts[i]));
        return {
            inputTextureDatas: inputTDs,
            outputTextureData: handler.createTextureDataFromLayout(programInfo.outputLayout, this.resultType ? this.resultType : inputs[0].type),
            uniformData: {}
        };
    }
}
exports.WebGLBinaryOp = WebGLBinaryOp;
function glslAdd() {
    const name = 'add_';
    const body = `
  float ${name}(float a, float b) {
    return a + b;
  }
  vec4 ${name}(vec4 v1, vec4 v2) {
    return v1 + v2;
  }
  `;
    return { body, name, type: glsl_definitions_1.FunctionType.ValueBased };
}
exports.glslAdd = glslAdd;
function glslDiv() {
    const name = 'div_';
    const body = `
  float ${name}(float a, float b) {
    return a / b;
  }
  vec4 ${name}(vec4 v1, vec4 v2) {
    return v1 / v2;
  }
  `;
    return { body, name, type: glsl_definitions_1.FunctionType.ValueBased };
}
exports.glslDiv = glslDiv;
function glslMul() {
    const name = 'mul_';
    const body = `
  float ${name}(float a, float b) {
    return a * b;
  }
  vec4 ${name}(vec4 v1, vec4 v2) {
    return v1 * v2;
  }
  `;
    return { body, name, type: glsl_definitions_1.FunctionType.ValueBased };
}
exports.glslMul = glslMul;
function glslSub() {
    const name = 'sub_';
    const body = `
  float ${name}(float a, float b) {
    return a - b;
  }
  vec4 ${name}(vec4 v1, vec4 v2) {
    return v1 - v2;
  }
  `;
    return { body, name, type: glsl_definitions_1.FunctionType.ValueBased };
}
exports.glslSub = glslSub;
function glslEqual() {
    const name = 'equal_';
    const body = `
  float ${name}(float a, float b) {
    return float(a == b);
  }
  vec4 ${name}(vec4 v1, vec4 v2) {
    return vec4( v1 == v2 );
  }
  `;
    return { body, name, type: glsl_definitions_1.FunctionType.ValueBased };
}
exports.glslEqual = glslEqual;
function glslGreater() {
    const name = 'greater_';
    const body = `
  float ${name}(float a, float b) {
    return float(a > b);
  }
  vec4 ${name}(vec4 v1, vec4 v2) {
    return vec4( v1.r > v2.r ,
      v1.g > v2.g,
      v1.b > v2.b,
      v1.a > v2.a );
  }
  `;
    return { body, name, type: glsl_definitions_1.FunctionType.ValueBased };
}
exports.glslGreater = glslGreater;
function glslLess() {
    const name = 'less_';
    const body = `
  float ${name}(float a, float b) {
    return float(a < b);
  }
  vec4 ${name}(vec4 v1, vec4 v2) {
    return vec4( v1.r < v2.r ,
                v1.g < v2.g,
                v1.b < v2.b,
                v1.a < v2.a );
  }
  `;
    return { body, name, type: glsl_definitions_1.FunctionType.ValueBased };
}
exports.glslLess = glslLess;
function glslAnd() {
    const name = 'and_';
    const body = `
  float ${name}(float a, float b) {
    return float( bool(a) && bool(b) );
  }
  vec4 ${name}(vec4 v1, vec4 v2) {
    bvec4 b1 = bvec4(v1);
    bvec4 b2 = bvec4(v2);
    return vec4( b1.r && b2.r ,
                b1.g && b2.g,
                b1.b && b2.b,
                b1.a && b2.a );
  }
  `;
    return { body, name, type: glsl_definitions_1.FunctionType.ValueBased };
}
exports.glslAnd = glslAnd;
function glslOr() {
    const name = 'or_';
    const body = `
  float ${name}(float a, float b) {
    return float( bool(a) || bool(b) );
  }
  vec4 ${name}(vec4 v1, vec4 v2) {
    bvec4 b1 = bvec4(v1);
    bvec4 b2 = bvec4(v2);
    return vec4( b1.r || b2.r ,
                b1.g || b2.g,
                b1.b || b2.b,
                b1.a || b2.a );
  }
  `;
    return { body, name, type: glsl_definitions_1.FunctionType.ValueBased };
}
exports.glslOr = glslOr;
function glslXor() {
    const name = 'xor_';
    const body = `
  float ${name}(float a, float b) {
    return float( bool(a) ^^ bool(b) );
  }
  vec4 ${name}(vec4 v1, vec4 v2) {
    bvec4 b1 = bvec4(v1);
    bvec4 b2 = bvec4(v2);
    return vec4( b1.r ^^ b2.r ,
                b1.g ^^ b2.g,
                b1.b ^^ b2.b,
                b1.a ^^ b2.a );
  }
  `;
    return { body, name, type: glsl_definitions_1.FunctionType.ValueBased };
}
exports.glslXor = glslXor;
function glslPow() {
    return glslBuiltinBinary('pow');
}
exports.glslPow = glslPow;
function glslPRelu() {
    const name = 'prelu_';
    const body = `
  float ${name}(float a, float b) {
    return a < 0.0 ? a * b: a;
  }
  vec4 ${name}(vec4 v1, vec4 v2) {
    return vec4(
      v1.r < 0.0 ? v1.r * v2.r: v1.r,
      v1.g < 0.0 ? v1.g * v2.g: v1.g,
      v1.b < 0.0 ? v1.b * v2.b: v1.b,
      v1.a < 0.0 ? v1.a * v2.a: v1.a
      );
  }
  `;
    return { body, name, type: glsl_definitions_1.FunctionType.ValueBased };
}
exports.glslPRelu = glslPRelu;
function glslBuiltinBinary(fname) {
    const name = `${fname}_`;
    const body = `
  float ${name}(float a, float b) {
    return ${fname}(a, b);
  }
  vec4 ${name}(vec4 v1, vec4 v2) {
    return ${fname}(v1, v2);
  }
  `;
    return { body, name, type: glsl_definitions_1.FunctionType.ValueBased };
}
//# sourceMappingURL=binary-op.js.map