"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 uint8_encode_1 = require("./ops/uint8-encode");
const unpack_1 = require("./ops/unpack");
const utils_1 = require("./utils");
class WebGLInferenceHandler {
    constructor(session) {
        this.session = session;
        this.packedTextureDataCache = new Map();
        this.unpackedTextureDataCache = new Map();
        this.pack2unpackMap = new Map();
        this.unpack2packMap = new Map();
    }
    run(op, inputs) {
        let artifact = this.session.programManager.getArtifact(op);
        if (!artifact) {
            const programInfo = op.createProgramInfo(this, inputs);
            artifact = this.session.programManager.build(programInfo);
            this.session.programManager.setArtifact(op, artifact);
        }
        const runData = op.createRunData(this, artifact.programInfo, inputs);
        this.runProgram(artifact, runData);
        return [runData.outputTextureData.tensor];
    }
    checkAndUpdateTextureForm(artifact, runData) {
        // pack/unpack inputs
        for (let i = 0; i < runData.inputTextureDatas.length; ++i) {
            const input = runData.inputTextureDatas[i];
            if (input.isPacked && !artifact.programInfo.expectPackedInputs) {
                runData.inputTextureDatas[i] = this.unpack(input);
            }
            else if (!input.isPacked && artifact.programInfo.expectPackedInputs) {
                runData.inputTextureDatas[i] = this.pack(input);
            }
        }
    }
    runProgram(artifact, runData) {
        this.checkAndUpdateTextureForm(artifact, runData);
        // output should match
        if (!!runData.outputTextureData.isPacked !== !!artifact.programInfo.expectPackedOutputs) {
            throw new Error('output property packed inconsistent');
        }
        this.session.programManager.run(artifact, runData);
    }
    /**
     * 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, layout, isPacked = false) {
        let td = this.getTextureData(tensor.dataId, isPacked);
        if (!td) {
            instrument_1.Logger.verbose('InferenceHandler', `Creating new TextureData for dims: [${tensor.dims}]`);
            if (!layout) {
                layout = this.createTextureLayoutFromShape(tensor.dims.slice());
            }
            // if we don't find the texture data with specific pack mode in the cache, try with the different
            // pack mode to see if the tensor is cached using that pack mode. If succeed, we can return this
            // tensor data and later apply a pack/unpack op on this texture, no need to create a new one here.
            td = this.getTextureData(tensor.dataId, !isPacked);
            if (!td) {
                if (isPacked) {
                    const unpackedTextureLayout = this.getOrCreateTextureLayout(tensor, 1, false, [], true);
                    const unpackedTextureData = this.createTextureData(unpackedTextureLayout, tensor.type, tensor.numberData, tensor, 1 /* UploadOnly */);
                    td = this.pack(unpackedTextureData);
                }
                else {
                    td = this.createTextureData(layout, tensor.type, tensor.numberData, tensor, 1 /* UploadOnly */);
                }
            }
        }
        else {
            instrument_1.Logger.verbose('InferenceHandler', `Retrieving TextureData from cache: [${tensor.dims}]`);
        }
        return td;
    }
    /**
     * Create a TextureData object from the given data type and texture layout.
     * Usage = Encoder.Usage.Default.
     * @param dataType the tensor data type
     */
    createTextureDataFromLayout(layout, dataType) {
        return this.createTextureData(layout, dataType);
    }
    /**
     * 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 /* 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);
    }
    /**
     * Create a TextureData object, using the given texture.
     * This function does not create new texture. Usually used in scenarios using texture sharing. (eg. Reshape)
     * @param dataType the tensor data type
     * @param texture the WebGLTexture object to share
     * @param tensorId the tensor ID of the shared tensor data
     */
    createSharedTextureData(layout, dataType, texture, tensorId) {
        return this.createTextureDataFromTexture(layout, dataType, texture, undefined, tensorId);
    }
    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), undefined, 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);
    }
    /**
     * Create a TextureLayout object from a tensor. If a related texture data is found, returns the cached texture layout.
     */
    getOrCreateTextureLayout(tensor, channels = 1, isPacked = false, unpackedShape, reverseWH = false) {
        const td = this.getTextureData(tensor.dataId, isPacked);
        if (td) {
            return td;
        }
        return this.createTextureLayoutFromShape(channels === 1 || isPacked ? tensor.dims : utils_1.getPackedShape(tensor.dims), channels, unpackedShape, isPacked || reverseWH ? { isPacked, reverseWH } : undefined);
    }
    /**
     * Create a TextureLayout object from shape.
     */
    createTextureLayoutFromShape(shape, channels = 1, unpackedShape, prefs) {
        const isPacked = !!(prefs && prefs.isPacked);
        const [texWidth, texHeight] = this.session.layoutStrategy.computeTextureWH(isPacked ? unpackedShape || shape : shape, prefs);
        let [width, height] = [texWidth, texHeight];
        if (prefs && prefs.reverseWH) {
            width = texHeight;
            height = texWidth;
        }
        const rank = shape.length;
        let inferredDims = shape.slice(0);
        if (rank === 0) {
            inferredDims = [1];
        }
        if (channels === 1) {
            // unpackedShape will take `shape` and not `inferredDims` so as to create a scalar Tensor if need be
            unpackedShape = shape;
        }
        else if (isPacked) {
            if (channels !== 4) {
                throw new Error('a packed texture must be 4-channel');
            }
            unpackedShape = shape;
            if (rank > 0) {
                inferredDims[rank - 1] = Math.ceil(inferredDims[rank - 1] / 2);
            }
            if (rank > 1) {
                inferredDims[rank - 2] = Math.ceil(inferredDims[rank - 2] / 2);
            }
        }
        else if (!unpackedShape) {
            throw new Error('Unpacked shape is needed when using channels > 1');
        }
        return {
            width,
            height,
            channels,
            isPacked,
            shape: inferredDims,
            strides: util_1.ShapeUtil.computeStrides(inferredDims),
            unpackedShape,
            reversedWH: (prefs && prefs.reverseWH)
        };
    }
    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) {
            const op = new uint8_encode_1.WebGLUint8Encode();
            const uint8TD = op.runInternal(this, textureData);
            return this.session.textureManager.readUint8TextureAsFloat(uint8TD);
        }
        return this.session.textureManager.readTexture(textureData, textureData.tensor.type, textureData.channels);
    }
    pack(input) {
        const cachedId = this.unpack2packMap.get(input.tensor.dataId);
        if (cachedId) {
            return this.packedTextureDataCache.get(cachedId);
        }
        const key = `${input.shape}`;
        let op = this.session.packOpCache.get(key);
        if (!op) {
            op = new pack_1.WebGLPack();
            this.session.packOpCache.set(key, op);
        }
        let artifact = this.session.programManager.getArtifact(op);
        if (!artifact) {
            const programInfo = op.createProgramInfo(this, [input.tensor]);
            artifact = this.session.programManager.build(programInfo);
            this.session.programManager.setArtifact(op, artifact);
        }
        const runData = op.createRunData(this, artifact.programInfo, [input.tensor]);
        this.runProgram(artifact, runData);
        this.unpack2packMap.set(input.tensor.dataId, runData.outputTextureData.tensor.dataId);
        return runData.outputTextureData;
    }
    unpack(input) {
        const cachedId = this.pack2unpackMap.get(input.tensor.dataId);
        if (cachedId) {
            return this.unpackedTextureDataCache.get(cachedId);
        }
        // For unpacked kernel, cache it by using input's unpackedShape as cache key.
        // Note that we need to use input.unpackedShape instead of input.shape here,
        // as the shape infers the packed texture shape. Different unpackedShape can have the
        // same packed texture shape. For example, for unpacked shape, both [2, 3] and
        // [2, 4] has the same packed shape [1, 2], but those two shapes should have different
        // unpack shaders.
        const key = `${input.unpackedShape}`;
        let op = this.session.unpackOpCache.get(key);
        if (!op) {
            op = new unpack_1.WebGLUnpack();
            this.session.unpackOpCache.set(key, op);
        }
        let artifact = this.session.programManager.getArtifact(op);
        if (!artifact) {
            const programInfo = op.createProgramInfo(this, [input.tensor]);
            artifact = this.session.programManager.build(programInfo);
            this.session.programManager.setArtifact(op, artifact);
        }
        const runData = op.createRunData(this, artifact.programInfo, [input.tensor]);
        this.runProgram(artifact, runData);
        this.pack2unpackMap.set(input.tensor.dataId, runData.outputTextureData.tensor.dataId);
        return runData.outputTextureData;
    }
}
exports.WebGLInferenceHandler = WebGLInferenceHandler;
//# sourceMappingURL=inference-handler.js.map