/*
 * Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "T5PluginGemm.h"
#include "src/fastertransformer/models/t5/T5Decoding.h"
#include "src/fastertransformer/models/t5/T5DecodingWeight.h"
#include "src/fastertransformer/models/t5/T5Encoder.h"
#include "src/fastertransformer/models/t5/T5EncoderWeight.h"
#include "src/fastertransformer/utils/Tensor.h"
#include "src/fastertransformer/utils/logger.h"

#include <NvInfer.h>
#include <cstdio>
#include <cstring>
#include <cublasLt.h>
#include <cublas_v2.h>
#include <cuda.h>
#include <cuda_fp16.h>
#include <vector>

// #define T5_PLUGIN_DEBUG
#ifdef T5_PLUGIN_DEBUG
#define WHERE_AM_I() printf("[%s]: this->%p\n", __func__, this);
#define PRINT_ENCODER(DATA_TYPE)                                                                                       \
    printf("[encoder::%s]Info:\n\tdatatype=%d\n", __func__, (int)DATA_TYPE);                                           \
    printf("\tmax_batch_size=%ld\n", m_.max_batch_size);                                                               \
    printf("\tmax_seq_len=%ld\n", m_.max_seq_len);                                                                     \
    printf("\tbeam_width=%ld\n", m_.beam_width);                                                                       \
    printf("\thead_num=%ld\n", m_.head_num);                                                                           \
    printf("\tsize_per_head=%ld\n", m_.size_per_head);                                                                 \
    printf("\td_model=%ld\n", m_.d_model);                                                                             \
    printf("\tinter_size=%ld\n", m_.inter_size);                                                                       \
    printf("\tnum_layer=%ld\n", m_.num_layer);                                                                         \
    printf("\tnum_bucket=%ld\n", m_.num_bucket);                                                                       \
    printf("\tmax_distance=%ld\n", m_.max_distance);                                                                   \
    printf("\tsm=%d\n", m_.sm);                                                                                        \
    printf("\tq_scaling=%f\n", m_.q_scaling);                                                                          \
    printf("\tuseFP16=%d\n", m_.useFP16);                                                                              \
    printf("\tvocab_size=%ld\n", m_.vocab_size);                                                                       \
    printf("\tis_remove_padding=%d\n", (int)m_.is_remove_padding);                                                     \
    printf("\tis_free_buffer_after_forward=%d\n", (int)m_.is_free_buffer_after_forward);                               \
    printf("\tis_sparse=%d\n", (int)m_.is_sparse);                                                                     \
    printf("\tattention_type=%d\n", (int)m_.attention_type);                                                           \
    printf("\tactivation_type=%d\n", (int)m_.activation_type);                                                         \
    printf("\tlayernorm_type=%d\n", (int)m_.layernorm_type);                                                           \
    printf("\tbatch_size=%ld\n", m_.batch_size);                                                                       \
    printf("\tseq_len=%ld\n", m_.seq_len);                                                                             \
    printf("\tckpt_path=%s\n", m_.ckpt_path);

#define PRINT_DECODING(DATA_TYPE)                                                                                      \
    printf("[decoding::%s]Info:\n\tdatatype=%d\n", __func__, (int)DATA_TYPE);                                          \
    printf("\tmax_batch_size=%ld\n", m_.max_batch_size);                                                               \
    printf("\tmax_seq_len=%ld\n", m_.max_seq_len);                                                                     \
    printf("\tmem_max_seq_len=%ld\n", m_.mem_max_seq_len);                                                             \
    printf("\tbeam_width=%ld\n", m_.beam_width);                                                                       \
    printf("\thead_num=%ld\n", m_.head_num);                                                                           \
    printf("\tsize_per_head=%ld\n", m_.size_per_head);                                                                 \
    printf("\td_model=%ld\n", m_.d_model);                                                                             \
    printf("\tinter_size=%ld\n", m_.inter_size);                                                                       \
    printf("\tnum_layer=%ld\n", m_.num_layer);                                                                         \
    printf("\tvocab_size=%ld\n", m_.vocab_size);                                                                       \
    printf("\tnum_bucket=%ld\n", m_.num_bucket);                                                                       \
    printf("\tmax_distance=%ld\n", m_.max_distance);                                                                   \
    printf("\tq_scaling=%f\n", m_.q_scaling);                                                                          \
    printf("\tstart_id=%d\n", m_.start_id);                                                                            \
    printf("\tend_id=%d\n", m_.end_id);                                                                                \
    printf("\tuseFP16=%d\n", m_.useFP16);                                                                              \
    printf("\tmem_d_model=%ld\n", m_.mem_d_model);                                                                     \
    printf("\tmem_hidden_units=%ld\n", m_.mem_hidden_units);                                                           \
    printf("\tis_free_buffer_after_forward=%d\n", m_.is_free_buffer_after_forward);                                    \
    printf("\tbatch_size=%ld\n", m_.batch_size);                                                                       \
    printf("\tseq_len=%ld\n", m_.seq_len);                                                                             \
    printf("\tckpt_path=%s\n", m_.ckpt_path);

#else
#define WHERE_AM_I()
#define PRINT_ENCODER(DATA_TYPE)
#define PRINT_DECODING(DATA_TYPE)
#endif  // DEBUG_ENABLE==1

namespace {
static const char* ENCODER_NAME{"T5EncoderPlugin"};
static const char* ENCODER_VERSION{"1"};
static const char* DECODING_NAME{"T5DecodingPlugin"};
static const char* DECODING_VERSION{"1"};
}  // namespace

using namespace fastertransformer;

namespace nvinfer1 {
// class T5EncoderPlugin ---------------------------------------------------------------------------
class T5EncoderPlugin: public IPluginV2DynamicExt {
private:
    using IPluginV2::enqueue;
    using IPluginV2::getOutputDimensions;
    using IPluginV2::getWorkspaceSize;
    using IPluginV2Ext::configurePlugin;

    const std::string name_;
    std::string       namespace_;
    bool              is_own_weight = false;
    cublasHandle_t    cublasHandle_;
    cublasLtHandle_t  cublasltHandle_;
#ifdef SPARSITY_ENABLED
    cusparseLtHandle_t cusparseltHandle_;
#endif
    cublasAlgoMap*                  pCublasAlgoMap_        = nullptr;
    std::mutex*                     pCublasWrapperMutex_   = nullptr;
    T5EncoderWeight<half>*          pT5EncoderWeightHalf_  = nullptr;
    T5EncoderWeight<float>*         pT5EncoderWeightFloat_ = nullptr;
    Allocator<AllocatorType::CUDA>* pAllocator_            = nullptr;
    cublasMMWrapper*                pCublasWrapper_        = nullptr;
    T5Encoder<half>*                pT5EncoderHalf_        = nullptr;
    T5Encoder<float>*               pT5EncoderFloat_       = nullptr;
    struct {
        // constructor parameter
        size_t max_batch_size = 128;
        size_t max_seq_len    = 384;
        size_t beam_width     = 1;
        size_t head_num       = 8;
        size_t size_per_head  = 512 / 8;
        size_t d_model        = head_num * size_per_head;
        size_t inter_size     = d_model * 4;
        size_t num_layer      = 6;
        size_t num_bucket     = 32;
        size_t max_distance   = 128;
        int    sm             = -1;  // assign later
        float  q_scaling      = 1.0f / (1.0f * sqrt(size_per_head));
        bool   useFP16        = false;
        // internal parameter
        size_t                            vocab_size                   = 32128;
        bool                              is_remove_padding            = true;
        bool                              is_free_buffer_after_forward = false;
        bool                              is_sparse                    = false;
        AttentionType                     attention_type               = AttentionType::UNFUSED_MHA;
        fastertransformer::ActivationType activation_type              = fastertransformer::ActivationType::Relu;
        LayerNormType                     layernorm_type               = LayerNormType::pre_layernorm;
        // runtime parameter
        size_t batch_size     = 0;
        size_t seq_len        = 0;
        char   ckpt_path[256] = "";
    } m_;

public:
    T5EncoderPlugin() = delete;
    T5EncoderPlugin(const std::string& name,
                    size_t             max_batch_size,
                    size_t             max_seq_len,
                    size_t             beam_width,
                    int                sm,
                    int                useFP16,
                    const std::string& ckpt_path,
                    bool               own_weight);
    T5EncoderPlugin(const std::string& name, const void* buffer, size_t length);
    ~T5EncoderPlugin();

    virtual size_t       getSerializationSize() const noexcept override;
    virtual void         serialize(void* buffer) const noexcept override;
    IPluginV2DynamicExt* clone() const noexcept override;
    int                  getNbOutputs() const noexcept override;
    DataType             getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const noexcept override;
    bool
    supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) noexcept override;
    DimsExprs    getOutputDimensions(int              index,
                                     const DimsExprs* pInputDim,
                                     int              nInputDim,
                                     IExprBuilder&    exprBuilder) noexcept override;
    virtual void configurePlugin(const DynamicPluginTensorDesc* in,
                                 int                            nbInput,
                                 const DynamicPluginTensorDesc* out,
                                 int                            nbOutput) noexcept override;
    size_t       getWorkspaceSize(const PluginTensorDesc* inputs,
                                  int32_t                 nbInputs,
                                  const PluginTensorDesc* outputs,
                                  int32_t                 nbOutputs) const noexcept override;
    int          enqueue(const PluginTensorDesc* inputDesc,
                         const PluginTensorDesc* outputDesc,
                         const void* const*      inputs,
                         void* const*            outputs,
                         void*                   workspace,
                         cudaStream_t            stream) noexcept override;
    void         setPluginNamespace(const char* szNamespace) noexcept override;
    const char*  getPluginNamespace() const noexcept override;
    const char*  getPluginType() const noexcept override;
    const char*  getPluginVersion() const noexcept override;
    int          initialize() noexcept override;
    void         terminate() noexcept override;
    void         destroy() noexcept override;
    void attachToContext(cudnnContext* /*cudnn*/, cublasContext* /*cublas*/, IGpuAllocator* /*allocator*/) noexcept;
    void detachFromContext() noexcept;
};

// class T5EncoderPluginCreator --------------------------------------------------------------------
class T5EncoderPluginCreator: public IPluginCreator {
private:
    static PluginFieldCollection    fc_;
    static std::vector<PluginField> attr_;
    std::string                     namespace_;

public:
    T5EncoderPluginCreator();
    ~T5EncoderPluginCreator();
    IPluginV2*  createPlugin(const char* name, const PluginFieldCollection* fc) noexcept override;
    IPluginV2*  deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept override;
    void        setPluginNamespace(const char* szNamespace) noexcept override;
    const char* getPluginNamespace() const noexcept override;
    const char* getPluginName() const noexcept override;
    const char* getPluginVersion() const noexcept override;
    const PluginFieldCollection* getFieldNames() noexcept override;
};
// class T5DecodingPlugin --------------------------------------------------------------------------
class T5DecodingPlugin: public IPluginV2DynamicExt {
private:
    using IPluginV2::enqueue;
    using IPluginV2::getOutputDimensions;
    using IPluginV2::getWorkspaceSize;
    using IPluginV2Ext::configurePlugin;

    const std::string               name_;
    std::string                     namespace_;
    bool                            is_own_weight = false;
    cublasHandle_t                  cublasHandle_;
    cublasLtHandle_t                cublasltHandle_;
    cudaDeviceProp                  cuda_device_prop_;
    cublasAlgoMap*                  pCublasAlgoMap_         = nullptr;
    Allocator<AllocatorType::CUDA>* pAllocator_             = nullptr;
    std::mutex*                     pCublasWrapperMutex_    = nullptr;
    cublasMMWrapper*                pCublasWrapper_         = nullptr;
    T5DecodingWeight<half>*         pT5DecodingWeightHalf_  = nullptr;
    T5DecodingWeight<float>*        pT5DecodingWeightFloat_ = nullptr;
    T5Decoding<half>*               pT5DecodingHalf_        = nullptr;
    T5Decoding<float>*              pT5DecodingFloat_       = nullptr;
    struct {
        // constructor parameter
        size_t max_batch_size             = 128;
        size_t max_seq_len                = 384;
        size_t mem_max_seq_len            = max_seq_len;
        size_t beam_width                 = 4;
        size_t head_num                   = 8;
        size_t size_per_head              = 512 / 8;
        size_t d_model                    = head_num * size_per_head;
        size_t inter_size                 = d_model * 4;
        size_t num_layer                  = 6;
        size_t vocab_size                 = 32128;
        size_t num_bucket                 = 32;
        size_t max_distance               = 128;
        float  q_scaling                  = 1.0f / (1.0f * sqrt(size_per_head));
        int    start_id                   = 0;
        int    end_id                     = 1;
        float  beam_search_diversity_rate = 0.0f;
        int    top_k                      = 0;
        float  top_p                      = 0.0f;
        float  temperature                = 1.0f;
        float  len_penalty                = 2.0f;
        float  repetition_penalty         = 1.0f;
        bool   useFP16                    = false;
        // internal parameter
        size_t mem_d_model                  = d_model;
        size_t mem_hidden_units             = d_model;
        bool   is_free_buffer_after_forward = false;
        // runtime parameter
        size_t batch_size     = 128;
        size_t seq_len        = 384;
        char   ckpt_path[256] = "";
    } m_;

public:
    T5DecodingPlugin() = delete;
    T5DecodingPlugin(const std::string& name,
                     size_t             max_batch_size,
                     size_t             max_seq_len,
                     size_t             mem_max_seq_len,
                     size_t             beam_width,
                     int                useFP16,
                     const std::string& ckpt_path,
                     bool               own_weight);
    T5DecodingPlugin(const std::string& name, const void* buffer, size_t length);
    ~T5DecodingPlugin();

    virtual size_t       getSerializationSize() const noexcept override;
    virtual void         serialize(void* buffer) const noexcept override;
    IPluginV2DynamicExt* clone() const noexcept override;
    int                  getNbOutputs() const noexcept override;
    DataType             getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const noexcept override;
    bool
    supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) noexcept override;
    DimsExprs    getOutputDimensions(int              index,
                                     const DimsExprs* pInputDim,
                                     int              nInputDim,
                                     IExprBuilder&    exprBuilder) noexcept override;
    virtual void configurePlugin(const DynamicPluginTensorDesc* in,
                                 int                            nbInput,
                                 const DynamicPluginTensorDesc* out,
                                 int                            nbOutput) noexcept override;
    size_t       getWorkspaceSize(const PluginTensorDesc* inputs,
                                  int32_t                 nbInputs,
                                  const PluginTensorDesc* outputs,
                                  int32_t                 nbOutputs) const noexcept override;
    int          enqueue(const PluginTensorDesc* inputDesc,
                         const PluginTensorDesc* outputDesc,
                         const void* const*      inputs,
                         void* const*            outputs,
                         void*                   workspace,
                         cudaStream_t            stream) noexcept override;
    void         setPluginNamespace(const char* szNamespace) noexcept override;
    const char*  getPluginNamespace() const noexcept override;
    const char*  getPluginType() const noexcept override;
    const char*  getPluginVersion() const noexcept override;
    int          initialize() noexcept override;
    void         terminate() noexcept override;
    void         destroy() noexcept override;
    void attachToContext(cudnnContext* /*cudnn*/, cublasContext* /*cublas*/, IGpuAllocator* /*allocator*/) noexcept;
    void detachFromContext() noexcept;
};

// class T5DecodingPluginCreator -------------------------------------------------------------------
class T5DecodingPluginCreator: public IPluginCreator {
private:
    static PluginFieldCollection    fc_;
    static std::vector<PluginField> attr_;
    std::string                     namespace_;

public:
    T5DecodingPluginCreator();
    ~T5DecodingPluginCreator();
    IPluginV2*  createPlugin(const char* name, const PluginFieldCollection* fc) noexcept override;
    IPluginV2*  deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept override;
    void        setPluginNamespace(const char* szNamespace) noexcept override;
    const char* getPluginNamespace() const noexcept override;
    const char* getPluginName() const noexcept override;
    const char* getPluginVersion() const noexcept override;
    const PluginFieldCollection* getFieldNames() noexcept override;
};

}  // namespace nvinfer1
