"use strict";
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
Object.defineProperty(exports, "__esModule", { value: true });
exports.WebGLTranspose = void 0;
const transpose_1 = require("../../../ops/transpose");
const util_1 = require("../../../util");
const glsl_definitions_1 = require("../glsl-definitions");
class WebGLTranspose extends transpose_1.Transpose {
    run(inferenceHandler, inputs) {
        return inferenceHandler.run(this, inputs);
    }
    getOutputShape(inputShapes) {
        const perm = this.getAdjustedPerm(inputShapes[0]);
        return util_1.ShapeUtil.sortBasedOnPerm(inputShapes[0], perm);
    }
    createProgramInfo(handler, inputs) {
        const inputShapes = inputs.map(t => t.dims.slice());
        const perm = this.getAdjustedPerm(inputShapes[0]);
        const unpackedOutputShape = this.getOutputShape(inputShapes);
        const rank = inputs[0].dims.length;
        // A dims=[${inputs[0].dims.toString()}]
        // out Dims=[${unpackedOutputShape.toString()}]
        // based on perm=[${perm.toString()}]
        const shaderSource = `
      ${this.getPermFunctionBody('perm', perm, rank)}
      float process(int indices[${rank}]) {
        int a[${rank}];
        perm(a, indices);
        return _A(a);
      }`;
        const outputLayout = handler.createTextureLayoutFromShape(unpackedOutputShape, 1, unpackedOutputShape);
        return { inputLayouts: [handler.getOrCreateTextureLayout(inputs[0])], outputLayout, samplers: ['A'], shaderSource };
    }
    createRunData(handler, programInfo, inputs) {
        const inputTDs = [handler.getOrCreateTextureData(inputs[0], programInfo.inputLayouts[0])];
        return {
            inputTextureDatas: inputTDs,
            outputTextureData: handler.createTextureDataFromLayout(programInfo.outputLayout, inputTDs[0].tensor.type),
            uniformData: {}
        };
    }
    getPositionalFunction(handler, inputShape, name) {
        const outputShape = this.getOutputShape([inputShape]);
        if (!name) {
            name = 'perm';
        }
        return {
            name,
            body: this.getPermFunctionBody(name, this.getAdjustedPerm(inputShape), outputShape.length),
            type: glsl_definitions_1.FunctionType.Positional,
            inputShape,
            outputShape
        };
    }
    getAdjustedPerm(inputShape) {
        let perm = this.perm;
        if (perm && perm.length !== inputShape.length) {
            perm = [...(inputShape.keys())].reverse();
        }
        return perm;
    }
    getPermFunctionBody(name, perm, rank) {
        const reverseFunc = [];
        reverseFunc.push(`void ${name}(out int a[${rank}], int src[${rank}]) {`);
        for (let i = 0; i < rank; ++i) {
            reverseFunc.push(`\ta[${perm[i]}]=src[${i}];`);
        }
        reverseFunc.push('\t}');
        return reverseFunc.join('\n');
    }
}
exports.WebGLTranspose = WebGLTranspose;
//# sourceMappingURL=transpose.js.map