"use strict";
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
Object.defineProperty(exports, "__esModule", { value: true });
exports.WebGLConvPacked = void 0;
const instrument_1 = require("../../../instrument");
const conv_1 = require("../../../ops/conv");
const tensor_1 = require("../../../tensor");
const util_1 = require("../../../util");
const conv_2 = require("./conv");
const im2col_pack_1 = require("./im2col-pack");
const matmul_pack_1 = require("./matmul-pack");
const reshape_packed_1 = require("./reshape-packed");
class WebGLConvPacked extends conv_1.Conv {
    run(inferenceHandler, inputs) {
        const programManager = inferenceHandler.session.programManager;
        const xshape = inputs[0].dims.slice();
        const kshape = inputs[1].dims.slice();
        // if kernelShape is not specified in the attributes of this op, infer it from the weight tensor dims
        if (this.kernelShape.length === 0) {
            for (let i = 2; i < kshape.length; ++i) {
                this.kernelShape.push(kshape[i]);
            }
        }
        util_1.PoolConvUtil.adjustPadsBasedOnAutoPad(inputs[0].dims, this.strides, this.dilations, this.kernelShape, this.pads, this.autoPad);
        instrument_1.Logger.verbose('Conv', `autpPad:${this.autoPad}, dilations:${this.dilations}, group:${this.group}, kernelShape:${this.kernelShape}, pads:${this.pads}, strides:${this.strides}`);
        const outputShape = conv_2.WebGLConv.calcOutputShape(xshape, kshape, this.dilations, this.pads, this.strides);
        const im2col = new im2col_pack_1.WebGLIm2ColPacked(outputShape, kshape, this.dilations, this.pads, this.strides);
        const matmul = new matmul_pack_1.WebGLMatMulPacked();
        const reshape = new reshape_packed_1.WebGLReshapePacked();
        // shape for kernel reshape
        const shape = new tensor_1.Tensor([2], 'int32', undefined, undefined, new Int32Array([kshape[0], kshape[1] * kshape[2] * kshape[3]]));
        if (!this.artifacts) {
            this.artifacts = [];
            this.programInfo = [];
            this.programInfo[0] = im2col.createProgramInfo(inferenceHandler, [inputs[0], inputs[1]]);
            this.artifacts[0] = programManager.build(this.programInfo[0]);
            this.programInfo[1] = reshape.createProgramInfo(inferenceHandler, [inputs[1], shape]);
            this.artifacts[1] = programManager.build(this.programInfo[1]);
        }
        // run im2col
        const runDataIm2col = im2col.createRunData(inferenceHandler, this.programInfo[0], [inputs[0], inputs[1]]);
        inferenceHandler.checkAndUpdateTextureForm(this.artifacts[0], runDataIm2col);
        programManager.run(this.artifacts[0], runDataIm2col);
        const im2colOutput = runDataIm2col.outputTextureData.tensor;
        // reshape kernel
        const runDataKernelReshape = reshape.createRunData(inferenceHandler, this.programInfo[1], [inputs[1], shape]);
        inferenceHandler.checkAndUpdateTextureForm(this.artifacts[1], runDataKernelReshape);
        programManager.run(this.artifacts[1], runDataKernelReshape);
        const kernelReshaped = runDataKernelReshape.outputTextureData.tensor;
        // run matmul
        const hasBias = (inputs.length === 3);
        util_1.assert(this.artifacts.length > 1, () => 'expect at least 2 artifacts created');
        if (this.artifacts.length === 2) {
            this.programInfo[2] = matmul.createProgramInfo(inferenceHandler, hasBias ? [kernelReshaped, im2colOutput, inputs[2]] : [kernelReshaped, im2colOutput]);
            this.artifacts[2] = programManager.build(this.programInfo[2]);
        }
        const runDataMatmul = matmul.createRunData(inferenceHandler, this.programInfo[2], hasBias ? [kernelReshaped, im2colOutput, inputs[2]] : [kernelReshaped, im2colOutput]);
        inferenceHandler.checkAndUpdateTextureForm(this.artifacts[2], runDataMatmul);
        programManager.run(this.artifacts[2], runDataMatmul);
        const matmulOutput = runDataMatmul.outputTextureData.tensor;
        // reshape output
        const outputShapeTensor = new tensor_1.Tensor([outputShape.length], 'int32', undefined, undefined, new Int32Array([outputShape[0], outputShape[1], outputShape[2], outputShape[3]]));
        util_1.assert(this.artifacts.length > 2, () => 'expect at least 3 artifacts created');
        if (this.artifacts.length === 3) {
            this.programInfo[3] = reshape.createProgramInfo(inferenceHandler, [matmulOutput, outputShapeTensor]);
            this.artifacts[3] = programManager.build(this.programInfo[3]);
        }
        const runDataOutputReshape = reshape.createRunData(inferenceHandler, this.programInfo[3], [matmulOutput, outputShapeTensor]);
        inferenceHandler.checkAndUpdateTextureForm(this.artifacts[3], runDataOutputReshape);
        programManager.run(this.artifacts[3], runDataOutputReshape);
        return [runDataOutputReshape.outputTextureData.tensor];
    }
}
exports.WebGLConvPacked = WebGLConvPacked;
//# sourceMappingURL=conv-pack.js.map