"use strict";
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
Object.defineProperty(exports, "__esModule", { value: true });
exports.WebGLUnpackedMatMul = exports.WebGLMatMul = void 0;
const matmul_1 = require("../../../ops/matmul");
const util_1 = require("../../../util");
const matmul_pack_1 = require("./matmul-pack");
class WebGLMatMul extends matmul_1.MatMul {
    constructor() {
        super();
        this.packedImpl = new matmul_pack_1.WebGLMatMulPacked();
        this.unpackedImpl = new WebGLUnpackedMatMul();
    }
    run(inferenceHandler, inputs) {
        if (inferenceHandler.session.pack) {
            return inferenceHandler.run(this.packedImpl, inputs);
        }
        else {
            return inferenceHandler.run(this.unpackedImpl, inputs);
        }
    }
    createProgramInfo(handler, inputs) {
        if (handler.session.pack && inputs[0].dims.length > 1) {
            return this.packedImpl.createProgramInfo(handler, inputs);
        }
        else {
            return this.unpackedImpl.createProgramInfo(handler, inputs);
        }
    }
    createRunData(handler, programInfo, inputs) {
        if (handler.session.pack && inputs[0].dims.length > 1) {
            return this.packedImpl.createRunData(handler, programInfo, inputs);
        }
        else {
            return this.unpackedImpl.createRunData(handler, programInfo, inputs);
        }
    }
}
exports.WebGLMatMul = WebGLMatMul;
class WebGLUnpackedMatMul extends matmul_1.MatMul {
    run(inferenceHandler, inputs) {
        return inferenceHandler.run(this, inputs);
    }
    createProgramInfo(handler, inputs) {
        const aShape = inputs[0].dims;
        const bShape = inputs[1].dims;
        const outputShape = util_1.BroadcastUtil.calcShape(aShape, bShape, true);
        if (!outputShape) {
            throw new Error('Can\'t use matmul on the given tensors');
        }
        const rank = outputShape.length;
        const arank = aShape.length;
        const brank = bShape.length;
        const sharedDim = aShape[aShape.length - 1];
        const shaderSource = `
      float process(int indices[${rank}]) {
          int a[${arank}];
          int b[${brank}];
          bcastMatmulIndices_A(indices, a);
          bcastMatmulIndices_B(indices, b);

          float value;
          for (int k=0; k<${sharedDim}; ++k) {
              a[${arank - 1}] = k;
              b[${brank - 2}] = k;
              value += _A(a) * _B(b);
          }
          return value;
      }`;
        return {
            inputLayouts: inputs.map(t => handler.getOrCreateTextureLayout(t)),
            outputLayout: handler.createTextureLayoutFromShape(outputShape),
            samplers: ['A', 'B'],
            shaderSource,
        };
    }
    createRunData(handler, programInfo, inputs) {
        const inputTDs = inputs.map((t, i) => handler.getOrCreateTextureData(t, programInfo.inputLayouts[i]));
        return {
            inputTextureDatas: inputTDs,
            outputTextureData: handler.createTextureDataFromLayout(programInfo.outputLayout, inputTDs[0].tensor.type),
            uniformData: {}
        };
    }
}
exports.WebGLUnpackedMatMul = WebGLUnpackedMatMul;
//# sourceMappingURL=matmul.js.map