"use strict";
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
Object.defineProperty(exports, "__esModule", { value: true });
exports.WebGLInferenceHandler = void 0;
const instrument_1 = require("../../instrument");
const tensor_1 = require("../../tensor");
const util_1 = require("../../util");
const pack_1 = require("./ops/pack");
const reshape_packed_1 = require("./ops/reshape-packed");
const uint8_encode_1 = require("./ops/uint8-encode");
const unpack_1 = require("./ops/unpack");
const texture_layout_1 = require("./texture-layout");
const types_1 = require("./types");
const getProgramInfoUniqueKey = (programInfo, inputTextureDatas) => {
    const inputs = inputTextureDatas.map(texture => `${texture.unpackedShape.join(',')};${texture.width}x${texture.height}`)
        .join('_');
    let key = programInfo.name;
    if (programInfo.cacheHint) {
        key += '[' + programInfo.cacheHint + ']';
    }
    key += ':' + inputs;
    return key;
};
class WebGLInferenceHandler {
    constructor(session) {
        this.session = session;
        this.packedTextureDataCache = new Map();
        this.unpackedTextureDataCache = new Map();
    }
    /**
     * @returns [width, height]
     */
    calculateTextureWidthAndHeight(shape, textureType) {
        return (0, texture_layout_1.calculateTextureWidthAndHeight)(this.session.layoutStrategy, shape, textureType);
    }
    executeProgram(program, inputs) {
        if (inputs.length < program.inputNames.length) {
            throw new Error(`Input size mustn't be less than ${program.inputNames.length}.`);
        }
        if (program.inputNames.length !== program.inputTypes.length) {
            throw new Error('input names size does not match input types');
        }
        // create texture info for input
        const inputTextureDatas = [];
        for (let i = 0; i < program.inputNames.length; ++i) {
            inputTextureDatas[i] = this.getOrCreateTextureData(inputs[i], program.inputTypes[i]);
        }
        const key = getProgramInfoUniqueKey(program, inputTextureDatas);
        let artifact = this.session.programManager.getArtifact(key);
        const programInfo = artifact ?
            artifact.programInfo :
            (typeof program.get === 'function' ? program.get() :
                program);
        // create texture info for output
        const outputTextureLayout = (0, texture_layout_1.createTextureLayoutFromTextureType)(this.session.layoutStrategy, programInfo.output.dims, programInfo.output.textureType);
        const outputTextureData = this.createTextureData(outputTextureLayout, programInfo.output.type);
        if (!artifact) {
            artifact = this.session.programManager.build(programInfo, inputTextureDatas, outputTextureData);
            this.session.programManager.setArtifact(key, artifact);
        }
        this.runProgram(artifact, inputTextureDatas, outputTextureData);
        return outputTextureData;
    }
    run(program, inputs) {
        const outputTextureData = this.executeProgram(program, inputs);
        return outputTextureData.tensor;
    }
    runProgram(artifact, inputs, output) {
        // input should match
        for (let i = 0; i < inputs.length; ++i) {
            if (!!inputs[i].isPacked !== (artifact.programInfo.inputTypes[i] === types_1.TextureType.packed)) {
                throw new Error(`input[${i}] property packed inconsistent`);
            }
        }
        // output should match
        if (!!output.isPacked !== (artifact.programInfo.output.textureType === types_1.TextureType.packed)) {
            throw new Error('output property packed inconsistent');
        }
        this.session.programManager.run(artifact, inputs, output);
    }
    /**
     * Create a TextureData object from a tensor.
     * Usage = Encoder.Usage.UploadOnly.
     * If a related texture data is found in cache, returns it;
     * Otherwise:
     *   Creates a new texture layout if not provided;
     *   Creates WebGLTexture with the layout;
     *   Upload tensor data to the texture;
     *   Creates a texture data object associated with the given tensor.
     * @param tensor the tensor with data to upload
     */
    getOrCreateTextureData(tensor, textureType) {
        let td = this.getTextureData(tensor.dataId, textureType === types_1.TextureType.packed);
        if (!td) {
            // check if we have texture data in different type
            td = this.getTextureData(tensor.dataId, textureType !== types_1.TextureType.packed);
            if (td) {
                if (textureType === types_1.TextureType.packed) {
                    return this.pack(td);
                }
                else {
                    return this.unpack(td);
                }
            }
        }
        if (!td) {
            const layout = (0, texture_layout_1.createTextureLayoutFromTextureType)(this.session.layoutStrategy, tensor.dims, textureType);
            if (textureType === types_1.TextureType.packedLastDimension) {
                const group = 1;
                const channels = 4;
                const shape = tensor.dims;
                if (shape.length === 4) {
                    // pre-processing for kernel data of Conv.
                    //
                    // TODO: currently this is a hacking to overwrite Conv's weight. The correct way to do this should be:
                    // 1. implement texture based const-folding
                    // 2. create a WebGL program "preprocessConvWeight" to do the same work as below
                    // 3. run the program before dotProduct.
                    //
                    const adjustedKernelShape = [shape[0], Math.ceil((shape[1] * shape[2] * shape[3]) / channels)];
                    const adjustedLayout = (0, texture_layout_1.createTextureLayoutFromTextureType)(this.session.layoutStrategy, adjustedKernelShape, textureType);
                    let buffer = tensor.numberData;
                    if (shape[1] * shape[2] * shape[3] % channels !== 0) {
                        const numFeatureMaps = shape[0];
                        const oldRowSize = shape[1] * shape[2] * shape[3];
                        const newRowSize = Math.ceil(oldRowSize * group / channels) * channels;
                        const newSize = numFeatureMaps * newRowSize;
                        buffer = new Float32Array(newSize);
                        for (let f = 0; f < numFeatureMaps; ++f) {
                            const oldOffset = f * oldRowSize;
                            const newOffset = f * newRowSize + f % group * oldRowSize;
                            buffer.set(tensor.numberData.subarray(oldOffset, oldOffset + oldRowSize), newOffset);
                        }
                    }
                    return this.createTextureData(adjustedLayout, tensor.type, buffer, tensor, 1 /* Encoder.Usage.UploadOnly */);
                }
            }
            if (textureType === types_1.TextureType.packed) {
                const unpackedTextureLayout = (0, texture_layout_1.createTextureLayoutFromShape)(this.session.layoutStrategy, tensor.dims, 1, [], { reverseWH: true });
                const unpackedTextureData = this.createTextureData(unpackedTextureLayout, tensor.type, tensor.numberData, tensor, 1 /* Encoder.Usage.UploadOnly */);
                td = this.pack(unpackedTextureData);
            }
            else {
                td = this.createTextureData(layout, tensor.type, tensor.numberData, tensor, 1 /* Encoder.Usage.UploadOnly */);
            }
        }
        return td;
    }
    /**
     * Create a TextureData object using the given data and bind to the given tensor.
     * Usage = Encoder.Usage.UploadOnly.
     * NOTE: this function is a hack for Conv implementation. should remove this function, after rewriting Conv
     * implementation by Graph.Transformer
     * @param dataType the tensor data type
     * @param data the actual data to upload
     * @param tensor the tensor to bind. tensor's data is ignored.
     */
    createTextureDataFromLayoutBindTensor(layout, dataType, data, tensor) {
        return this.createTextureData(layout, dataType, data, tensor, 1 /* Encoder.Usage.UploadOnly */);
    }
    createTextureData(layout, dataType, data, tensor, usage) {
        instrument_1.Logger.verbose('InferenceHandler', `Creating TextureData: layout:[${JSON.stringify(layout)}]`);
        const texture = this.session.textureManager.createTextureFromLayout(dataType, layout, data, usage);
        return this.createTextureDataFromTexture(layout, dataType, texture, tensor);
    }
    reshapeUnpacked(input, reshapedDims) {
        const inputTD = this.getOrCreateTextureData(input, types_1.TextureType.unpacked);
        const newTextureLayout = {
            channels: inputTD.channels,
            height: inputTD.height,
            width: inputTD.width,
            // handle reshaping into scalar Tensors
            shape: reshapedDims.length !== 0 ? reshapedDims : [1],
            strides: util_1.ShapeUtil.computeStrides(reshapedDims),
            unpackedShape: reshapedDims,
        };
        const newTextureData = this.createTextureDataFromTexture(newTextureLayout, input.type, inputTD.texture);
        return newTextureData.tensor;
    }
    reshapePacked(input, reshapedDims) {
        const inputTD = this.getOrCreateTextureData(input, types_1.TextureType.packed);
        // check if the reshape is 'cheap'
        if ((0, reshape_packed_1.isReshapeCheap)(input.dims, reshapedDims)) {
            const newTextureLayout = {
                channels: inputTD.channels,
                height: inputTD.height,
                width: inputTD.width,
                // handle reshaping into scalar Tensors
                shape: reshapedDims.length !== 0 ? reshapedDims : [1],
                strides: util_1.ShapeUtil.computeStrides(reshapedDims),
                unpackedShape: reshapedDims,
                isPacked: true
            };
            const newTextureData = this.createTextureDataFromTexture(newTextureLayout, input.type, inputTD.texture);
            return newTextureData.tensor;
        }
        const squeezedInputShape = (0, reshape_packed_1.processDims3D)(input.dims);
        const squeezedOutputShape = (0, reshape_packed_1.processDims3D)(reshapedDims);
        const squeezedInputTensor = this.reshapePacked(input, squeezedInputShape);
        const squeezedOutputTensor = this.run((0, reshape_packed_1.createPackedReshape3DProgramInfoLoader)(this, squeezedInputTensor, squeezedOutputShape), [squeezedInputTensor]);
        const outputTensor = this.reshapePacked(squeezedOutputTensor, reshapedDims);
        return outputTensor;
    }
    cast(input, type) {
        const inputTD = this.getOrCreateTextureData(input, types_1.TextureType.unpacked);
        const newTextureData = this.createTextureDataFromTexture(inputTD, type, inputTD.texture);
        return newTextureData.tensor;
    }
    createTextureDataFromTexture(layout, dataType, texture, tensor, tensorId) {
        const textureData = Object.assign(Object.assign({}, layout), { tensor: tensor ||
                new tensor_1.Tensor(layout.unpackedShape, dataType, (_id) => this.readTexture(textureData), async (_id) => this.readTextureAsync(textureData), undefined, tensorId), texture });
        this.setTextureData(textureData.tensor.dataId, textureData, layout.isPacked);
        return textureData;
    }
    getTextureData(tensorId, isPacked = false) {
        return this.session.isInitializer(tensorId) ? this.session.getTextureData(tensorId, isPacked) :
            isPacked ? this.packedTextureDataCache.get(tensorId) :
                this.unpackedTextureDataCache.get(tensorId);
    }
    setTextureData(tensorId, td, isPacked = false) {
        if (this.session.isInitializer(tensorId)) {
            this.session.setTextureData(tensorId, td, isPacked);
        }
        else {
            (isPacked ? this.packedTextureDataCache : this.unpackedTextureDataCache).set(tensorId, td);
        }
    }
    isTextureLayoutCached(tensor, isPacked = false) {
        return !!this.getTextureData(tensor.dataId, isPacked);
    }
    dispose() {
        this.session.textureManager.clearActiveTextures();
        this.packedTextureDataCache.forEach(td => this.session.textureManager.releaseTexture(td));
        this.packedTextureDataCache = new Map();
        this.unpackedTextureDataCache.forEach(td => this.session.textureManager.releaseTexture(td));
        this.unpackedTextureDataCache = new Map();
    }
    readTexture(textureData) {
        if (textureData.isPacked) {
            return this.readTexture(this.unpack(textureData));
        }
        if (!this.session.backend.glContext.isFloat32DownloadSupported) {
            return this.session.textureManager.readUint8TextureAsFloat((0, uint8_encode_1.encodeAsUint8)(this, textureData));
        }
        return this.session.textureManager.readTexture(textureData, textureData.tensor.type, textureData.channels);
    }
    async readTextureAsync(textureData) {
        if (textureData.isPacked) {
            return this.readTextureAsync(this.unpack(textureData));
        }
        if (!this.session.backend.glContext.isFloat32DownloadSupported) {
            return this.session.textureManager.readUint8TextureAsFloat((0, uint8_encode_1.encodeAsUint8)(this, textureData));
        }
        return this.session.textureManager.readTextureAsync(textureData, textureData.tensor.type, textureData.channels);
    }
    pack(input) {
        const outputTextureData = this.executeProgram((0, pack_1.createPackProgramInfoLoader)(this, input.tensor), [input.tensor]);
        return outputTextureData;
    }
    unpack(input) {
        const outputTextureData = this.executeProgram((0, unpack_1.createUnpackProgramInfoLoader)(this, input.tensor), [input.tensor]);
        return outputTextureData;
    }
}
exports.WebGLInferenceHandler = WebGLInferenceHandler;
//# sourceMappingURL=inference-handler.js.map