//
//  ConvolutionCommon.hpp
//  MNN
//
//  Created by MNN on 2020/03/02.
//  Copyright © 2018, Alibaba Group Holding Limited
//

#ifndef ConvolutionCommon_hpp
#define ConvolutionCommon_hpp
#include "AutoStorage.h"
#include "Execution.hpp"
#include "MNN_generated.h"
namespace MNN {
class MNN_PUBLIC ConvolutionCommon : public Execution {
public:
    struct Int8Common {
        AutoStorage<int8_t> weight;
        // alpha (fp32) and alphaHalf (fp16 raw bits) are alternative views of the same scale/zero
        // data. alphaIsFp16 records the canonical (disk-side) form. Use getAlphaFloat() /
        // getAlphaHalf() to access; they lazily fill the alternate view on first request.
        AutoStorage<float> alpha;
        AutoStorage<int16_t> alphaHalf;
        AutoStorage<float> weightFloat;
        const IDSTQuan* quan;
        bool asymmetric;
        bool alphaIsFp16 = false;
        std::vector<int8_t> weightMap;
        bool canUseInt4 = false;
        bool canUseInt3 = false;
        bool canUseInt2 = false;
        Backend* backend = nullptr;
        int originBits = 8;
        int alphaSize;
        // Lazy accessors. Element count is alphaSize.
        const float* getAlphaFloat();
        const int16_t* getAlphaHalf();
    };
    static std::shared_ptr<Int8Common> load(const Op* op, Backend* backend = nullptr, bool forceFloat = false, bool forceInt8 = false, void* weightPtr = nullptr);
    // if can not get quant bits, return 0
    static int getQuantBitFromExternalFile(const Op* op);
    static void getConvParameters(std::shared_ptr<ConvolutionCommon::Int8Common> *quanCommon, Backend* backend, const MNN::Op *op, const float** originWeight, int* originWeightSize);
    static bool getConvInt8Parameters(const MNN::Op* op, std::shared_ptr<Int8Common>& quanCommon, Backend* backend,
                                      const int8_t*& weight, int& weightSize, float* scale, int32_t* bias, int ocUpHp);

    // Return padX, padY
    static std::pair<int, int> convolutionPad(const Tensor* input, const Tensor* output,
                                              const Convolution2DCommon* common);
    // Return padLeft, padTop, padRight, padBottom
    static std::tuple<int, int, int, int> convolutionPadFull(const Tensor* input, const Tensor* output,
                                              const Convolution2DCommon* common);
    static std::pair<int, int> convolutionTransposePad(const Tensor* input, const Tensor* output,
                                                       const Convolution2DCommon* common);
    struct Im2ColParameter {
        int32_t padX;
        int32_t padY;
        int32_t dilateX;
        int32_t dilateY;
        int32_t strideX;
        int32_t strideY;
        int32_t kernelX;
        int32_t kernelY;
        int32_t icDiv4;
        int32_t kernelCountUnit;
        int32_t iw;
        int32_t ih;
        int32_t ow;
        int32_t oh;
        int32_t srcZStep;
        int32_t srcYStep;
        int32_t packCUnit;
        int32_t destICStride;
        int32_t ic; // ic packed by LP, used by im2col dst data.
        int32_t icup4; // ic packed by LP, used by im2col src data.
    };
};
} // namespace MNN
#endif