"use strict";
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
Object.defineProperty(exports, "__esModule", { value: true });
exports.WebGLSplit = void 0;
const split_1 = require("../../../ops/split");
const util_1 = require("../../../util");
class WebGLSplit extends split_1.Split {
    run(inferenceHandler, inputs) {
        if (!this.artifacts) {
            this.artifacts = [];
            const axis = util_1.ShapeUtil.normalizeAxis(this.axis, inputs[0].dims.length);
            const count = this.getProgramCount(inferenceHandler, inputs, axis);
            for (let i = 0; i < count; ++i) {
                const programInfo = this.createProgramInfo(inferenceHandler, inputs[0], axis, i);
                const artifact = inferenceHandler.session.programManager.build(programInfo);
                this.artifacts.push(artifact);
            }
        }
        const results = [];
        this.artifacts.forEach(artifact => {
            const rundata = this.createRunData(inferenceHandler, artifact.programInfo, inputs);
            inferenceHandler.session.programManager.run(artifact, rundata);
            results.push(rundata.outputTextureData.tensor);
        });
        return results;
    }
    getProgramCount(inferenceHandler, inputs, axis) {
        const [, offsets] = util_1.SplitUtil.splitShape(inputs[0].dims, axis, this.split, this.numOutputs);
        return offsets.length;
    }
    createProgramInfo(inferenceHandler, input, axis, index) {
        const [shapes, offsets] = util_1.SplitUtil.splitShape(input.dims, axis, this.split, this.numOutputs);
        const offset = offsets[index];
        const outputShape = shapes[index];
        const rank = outputShape.length;
        const shaderSource = `
      float process(int indices[${rank}]) {
        indices[${axis}] += ${offset};
        return _A(indices);
      }`;
        return {
            inputLayouts: [inferenceHandler.getOrCreateTextureLayout(input)],
            outputLayout: inferenceHandler.createTextureLayoutFromShape(outputShape),
            samplers: ['A'],
            shaderSource,
        };
    }
    createRunData(inferenceHandler, programInfo, inputs) {
        const inputTDs = [inferenceHandler.getOrCreateTextureData(inputs[0], programInfo.inputLayouts[0])];
        return {
            inputTextureDatas: inputTDs,
            outputTextureData: inferenceHandler.createTextureDataFromLayout(programInfo.outputLayout, inputTDs[0].tensor.type),
            uniformData: {}
        };
    }
}
exports.WebGLSplit = WebGLSplit;
//# sourceMappingURL=split.js.map