/*
 * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * 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.
 */

// cspell:ignore accu CUBLASLT splitk

// cublasLT was introduced in CUDA 10.1
#include <cuda.h>
#if CUDA_VERSION >= 10010

#include "NvInfer.h"
#include "common/serialize.hpp"
#include "fcPlugin.h"

#include <algorithm>
#include <cstdio>
#include <cstring>
#include <cuda_runtime.h>
#include <memory>
#include <string_view>
#include <vector>

using namespace nvinfer1;
using namespace nvinfer1::plugin;
using namespace nvinfer1::plugin::bert;
using namespace nvinfer1::pluginInternal;

// plugin specific constants
namespace
{
using namespace std::string_view_literals;
char const* const kFC_VERSION{"1"};
char const* const kFC_NAME{"CustomFCPluginDynamic"};
constexpr size_t kMAX_WORKSPACE_BYTES = 4 * 1024 * 1024; // 4MiB
} // namespace

REGISTER_TENSORRT_PLUGIN(FCPluginDynamicCreator);

// Utility function to print customMatmulPerf_t structure
static void printPerfStructure(customMatmulPerf_t const& perf, int32_t const m, int32_t const n, int32_t const k)
{
    AlgoProps p;
    p.populate(perf.algo);
    // Calculate GFLOPS
    double timeAvg
        = perf.time * 1e-3; // Convert to seconds. It has been divided by kNB_KERNEL_REPEATS in customMatmulRun().
    double gflop = (2 * static_cast<uint64_t>(m * n) * k) * 1e-9; // Real

    gLogVerbose << "Algo=" << p.algoId << " Tile=" << p.tile << " (" << matmulTileName[p.tile] << ") K=" << p.numSplitsK
                << " Red.Sch.=" << p.reductionScheme << " Swiz=" << p.swizzle << " Cust=" << p.customOption
                << " Stat=" << perf.status << " Time=" << perf.time << " WSbytes=" << perf.workspaceSize
                << " math=" << p.numericImpl << " waves=" << perf.wavesCount << "GFlops=" << (gflop / timeAvg)
                << std::endl;
}

static bool timeCompare(customMatmulPerf_t const& perf_a, customMatmulPerf_t const& perf_b)
{
    return ((perf_a.status == CUBLAS_STATUS_SUCCESS) && (perf_a.time < perf_b.time));
}

static cublasStatus_t customMatmulRun(cublasLtHandle_t ltHandle, // to get the capabilities (required a GPU)
    cublasLtMatmulDesc_t operationDesc, void const* alpha,       // host or device pointer
    void const* A, cublasLtMatrixLayout_t Adesc, void const* B, cublasLtMatrixLayout_t Bdesc,
    void const* beta, // host or device pointer
    void const* C, cublasLtMatrixLayout_t Cdesc, void* D, cublasLtMatrixLayout_t Ddesc,
    cublasLtMatmulAlgo_t const& algo, void* workSpace, size_t workSpaceSizeInBytes, customMatmulPerf_t& perfResults,
    cudaStream_t stream, cudaEvent_t& startEvent, cudaEvent_t& stopEvent)
{

    cublasLtMatmulHeuristicResult_t heurResult;

    CublasLtWrapper& cublasLtWrapper = getCublasLtWrapper();
    // Looping over the Algo
    cublasStatus_t algoStatus = cublasLtWrapper.cublasLtMatmulAlgoCheck(
        ltHandle, operationDesc, Adesc, Bdesc, Cdesc, Ddesc, &algo, &heurResult);

    if (algoStatus == CUBLAS_STATUS_SUCCESS)
    {
        if (heurResult.workspaceSize <= workSpaceSizeInBytes)
        {
            if (cudaEventRecord(startEvent, stream) != cudaSuccess)
            {
                return CUBLAS_STATUS_INTERNAL_ERROR;
            }
            for (int32_t loop = 0; loop < kNB_KERNEL_REPEATS; loop++)
            {
                cublasStatus_t oneRunStatus
                    = cublasLtWrapper.cublasLtMatmul(ltHandle, operationDesc, alpha, // host or device pointer
                        A, Adesc, B, Bdesc, beta,                                    // host or device pointer
                        C, Cdesc, D, Ddesc, &algo, workSpace, workSpaceSizeInBytes, stream);
                if (oneRunStatus != CUBLAS_STATUS_SUCCESS)
                {
                    algoStatus = oneRunStatus;
                    break;
                }
            }
            if (cudaEventRecord(stopEvent, stream) != cudaSuccess)
            {
                return CUBLAS_STATUS_INTERNAL_ERROR;
            }
            if (cudaEventSynchronize(stopEvent) != cudaSuccess)
            {
                return CUBLAS_STATUS_INTERNAL_ERROR;
            }
            float time;
            if (cudaEventElapsedTime(&time, startEvent, stopEvent) != cudaSuccess)
            {
                return CUBLAS_STATUS_INTERNAL_ERROR;
            }
            // For the moment only add successful findings
            perfResults.algo = algo;
            perfResults.time = time / kNB_KERNEL_REPEATS; // Average time
            perfResults.workspaceSize = heurResult.workspaceSize;
            perfResults.wavesCount = heurResult.wavesCount;
        }
        else
        {
            algoStatus = CUBLAS_STATUS_NOT_SUPPORTED; // Not enough workspace
        }
    }
    return algoStatus;
}

// Sample wrapper running through multiple algo and config attributes
// combination for single precision gemm using cublasLt low-level API
// NOLINTNEXTLINE(readability-function-cognitive-complexity)
void nvinfer1::plugin::bert::LtGemmSearch(cublasLtHandle_t ltHandle, cublasOperation_t transa, cublasOperation_t transb,
    int32_t const& m, int32_t const& n, int32_t const& k, void const* alpha,                // host pointer
    void const* A, int32_t const& lda, void const* B, int32_t const& ldb, void const* beta, // host pointer
    void* C, int32_t const& ldc, void* workSpace, size_t workSpaceSize, cublasComputeType_t computeType,
    cudaDataType_t scaleType, cudaDataType_t Atype, cudaDataType_t Btype, cudaDataType_t Ctype,
    std::vector<customMatmulPerf_t>& perfResults, cudaStream_t stream)
{

    cublasStatus_t status = CUBLAS_STATUS_SUCCESS;

    cublasLtMatmulDesc_t operationDesc = nullptr;
    cublasLtMatrixLayout_t Adesc = nullptr;
    cublasLtMatrixLayout_t Bdesc = nullptr;
    cublasLtMatrixLayout_t Cdesc = nullptr;
    cublasLtMatmulPreference_t preference = nullptr;

    cudaEvent_t startEvent = nullptr;
    cudaEvent_t stopEvent = nullptr;

    CublasLtWrapper& cublasLtWrapper = getCublasLtWrapper();

    // SplitK value that we are going to try when SplitK is supported for a given algo.
    int32_t const splitKSequenceA[] = {2, 3, 4, 5, 6, 8, 12, 16, 32};

    // Let try a fixed number of combinations
    int32_t algoCount = 0;
    int32_t nbAlgoIds = 0;
    int32_t algoIdA[kNB_ALGO_IDS];

    PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulPreferenceCreate(&preference));
    PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulPreferenceSetAttribute(
        preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workSpaceSize, sizeof(workSpaceSize)));

    uint64_t const numericImplPrefer
        = Ctype == CUDA_R_16F ? CUBLASLT_NUMERICAL_IMPL_FLAGS_HMMA : CUBLASLT_NUMERICAL_IMPL_FLAGS_FMA;
    PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulPreferenceSetAttribute(
        preference, CUBLASLT_MATMUL_PREF_IMPL_MASK, &numericImplPrefer, sizeof(numericImplPrefer)));

    // Create operation descriptor; see cublasLtMatmulDescAttributes_t for details
    // about defaults; here we just need to set the transforms for A and B
    PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulDescCreate(&operationDesc, computeType, scaleType));
    PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulDescSetAttribute(
        operationDesc, nvinfer1::pluginInternal::CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa)));
    PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulDescSetAttribute(
        operationDesc, nvinfer1::pluginInternal::CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa)));

    // Create matrix descriptors. We are good with the details here so no need to
    // set any extra attributes
    PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutCreate(
        &Adesc, Atype, transa == CUBLAS_OP_N ? m : k, transa == CUBLAS_OP_N ? k : m, lda));
    PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutCreate(
        &Bdesc, Btype, transb == CUBLAS_OP_N ? k : n, transb == CUBLAS_OP_N ? n : k, ldb));
    PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutCreate(&Cdesc, Ctype, m, n, ldc));

    // Request the 4 first AlgoId available for SGEMM ( computeType = scaleType =
    // Atype = Btype = Ctype = Dtype = CUDA_R_32F)
    PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoGetIds(
        ltHandle, computeType, scaleType, Atype, Btype, Ctype, Ctype, kNB_ALGO_IDS, algoIdA, &nbAlgoIds));

    gLogVerbose << "Number of algos" << nbAlgoIds << std::endl;

    // Create CUDA event to time the execution time of each algo
    PLUGIN_CUASSERT(cudaEventCreate(&startEvent, cudaEventBlockingSync));
    PLUGIN_CUASSERT(cudaEventCreate(&stopEvent, cudaEventBlockingSync));

    // Loop over the Algo IDs
    for (int32_t idx = 0; (idx < nbAlgoIds) && (algoCount < kNB_ALGO_COMBINATIONS); idx++)
    {
        cublasLtMatmulAlgo_t algo;
        size_t nbBytesWritten = 0;
        // Initialize algo structure with given Algp ID.
        status = cublasLtWrapper.cublasLtMatmulAlgoInit(
            ltHandle, computeType, scaleType, Atype, Btype, Ctype, Ctype, algoIdA[idx], &algo);
        if (status != CUBLAS_STATUS_SUCCESS)
        {
            continue;
        }

        uint64_t numericImpl = -1;
        PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoCapGetAttribute(
            &algo, CUBLASLT_ALGO_CAP_NUMERICAL_IMPL_FLAGS, &numericImpl, sizeof(numericImpl), nullptr));
        if (Ctype == CUDA_R_32F && numericImpl == CUBLASLT_NUMERICAL_IMPL_FLAGS_HMMA)
        {
            // skip HMMA-fp32accu kernels
            continue;
        }

        // Query the tiles enums supported by that algo
        PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoCapGetAttribute(
            &algo, CUBLASLT_ALGO_CAP_TILE_IDS, nullptr, 0, &nbBytesWritten));
        auto tileA = std::vector<int32_t>(nbBytesWritten / sizeof(int32_t));
        if (tileA.empty())
        {
            tileA = {CUBLASLT_MATMUL_TILE_UNDEFINED};
        }

        // Retrieve Algo Capabilities attributes to be able to setup loop over the
        // different combinations
        PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoCapGetAttribute(&algo, CUBLASLT_ALGO_CAP_TILE_IDS,
            static_cast<void*>(tileA.data()), sizeof(int32_t) * tileA.size(), &nbBytesWritten));

        //! Read an int32_t attribute from `cublasLtWrapper` and return the value and update nbBytesWritten.
        auto getInt32Attr = [&](cublasLtMatmulAlgoCapAttributes attr, size_t& nbBytesWritten) -> int32_t {
            int32_t value;
            PLUGIN_CUBLASASSERT(
                cublasLtWrapper.cublasLtMatmulAlgoCapGetAttribute(&algo, attr, &value, sizeof(value), &nbBytesWritten));
            return value;
        };
        auto const splitkSupport = getInt32Attr(CUBLASLT_ALGO_CAP_SPLITK_SUPPORT, nbBytesWritten);
        auto const redMask = getInt32Attr(CUBLASLT_ALGO_CAP_REDUCTION_SCHEME_MASK, nbBytesWritten);
        auto const swizzlingMax = getInt32Attr(CUBLASLT_ALGO_CAP_CTA_SWIZZLING_SUPPORT, nbBytesWritten);
        auto const customOptionMax = getInt32Attr(CUBLASLT_ALGO_CAP_CUSTOM_OPTION_MAX, nbBytesWritten);
        std::ignore = getInt32Attr(CUBLASLT_ALGO_CAP_EPILOGUE_MASK, nbBytesWritten);

        // Loop over the different tiles
        for (uint32_t tileIdx = 0; tileIdx < tileA.size(); tileIdx++)
        {
            // Loop over the different custom option if any
            for (int32_t customOption = 0; customOption <= customOptionMax; customOption++)
            {
                PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoConfigSetAttribute(
                    &algo, CUBLASLT_ALGO_CONFIG_CUSTOM_OPTION, &customOption, sizeof(customOption)));
                // Loop over the CTAs swizzling support
                for (int32_t k = 0; k <= swizzlingMax; k++)
                {
                    int32_t const splitkTrial
                        = splitkSupport ? (sizeof(splitKSequenceA) / sizeof(splitKSequenceA[0])) : 0;
                    // Loop over the splitK value over a fixed sequence splitKSequenceA in
                    // addition to the case where splitK is not enabled
                    for (int32_t l = 0; (l < (1 + splitkTrial)) && (algoCount < kNB_ALGO_COMBINATIONS); l++)
                    {
                        //! Set an attribute of the algo to a given value.
                        auto setAttr = [&](cublasLtMatmulAlgoConfigAttributes attr, auto value) {
                            static_assert(std::is_same_v<decltype(value), int32_t>, "value must be an int32_t");
                            PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoConfigSetAttribute(
                                &algo, attr, &value, sizeof(value)));
                        };
                        // Setup attribute of the algo to run
                        setAttr(CUBLASLT_ALGO_CONFIG_TILE_ID, tileA[tileIdx]);
                        int32_t splitK_val = 0;
                        setAttr(CUBLASLT_ALGO_CONFIG_SPLITK_NUM, splitK_val);
                        setAttr(CUBLASLT_ALGO_CONFIG_CTA_SWIZZLING, k);
                        int32_t redScheme = CUBLASLT_REDUCTION_SCHEME_NONE;
                        setAttr(CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME, redScheme);

                        if (l > 0)
                        { // Split-K case
                            splitK_val = splitKSequenceA[l - 1];
                            setAttr(CUBLASLT_ALGO_CONFIG_SPLITK_NUM, splitKSequenceA[l - 1]);
                            // Going over all the reduction scheme
                            for (redScheme = 1; redScheme < static_cast<int32_t>(CUBLASLT_REDUCTION_SCHEME_MASK)
                                 && (algoCount < kNB_ALGO_COMBINATIONS);
                                 redScheme <<= 1)
                            {
                                if (redScheme & redMask)
                                {
                                    setAttr(CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME, redScheme);

                                    status = customMatmulRun(ltHandle, operationDesc, alpha, // host or device pointer
                                        A, Adesc, B, Bdesc, beta,                            // host or device pointer
                                        C, Cdesc, C, Cdesc, algo, workSpace, workSpaceSize, perfResults[algoCount],
                                        stream, startEvent, stopEvent);
                                    perfResults[algoCount].status = status;
                                    algoCount += (status == CUBLAS_STATUS_SUCCESS);
                                }
                            }
                        }
                        else
                        { // Non-splitK case
                            // if user preference is ok with workspace
                            if (algoCount < kNB_ALGO_COMBINATIONS)
                            {
                                status = customMatmulRun(ltHandle, operationDesc, alpha, // host or device pointer
                                    A, Adesc, B, Bdesc, beta,                            // host or device pointer
                                    C, Cdesc, C, Cdesc, algo, workSpace, workSpaceSize, perfResults[algoCount], stream,
                                    startEvent, stopEvent);
                                perfResults[algoCount].status = status;
                                algoCount += (status == CUBLAS_STATUS_SUCCESS);
                            }
                        }
                    } // end l
                } // end k
            } // end customOption
        } // end tileIdx
    } // end idx

    // Sort the results per run duration
    std::sort(perfResults.begin(), perfResults.end(), timeCompare);

    // Print timing and perf details of the fastest combinations
    for (int32_t i = 0; i < kPRINT_ALGOS && perfResults[i].time != customMatmulPerf_t::kMAX_TIME; i++)
    {
        printPerfStructure(perfResults[i], m, n, k);
    }

    // Descriptors are no longer needed as all GPU work was already enqueued
    PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulPreferenceDestroy(preference));
    PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutDestroy(Cdesc));
    PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutDestroy(Bdesc));
    PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutDestroy(Adesc));
    PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulDescDestroy(operationDesc));
    PLUGIN_CUASSERT(cudaEventDestroy(startEvent));
    PLUGIN_CUASSERT(cudaEventDestroy(stopEvent));
}

FCPluginDynamic::FCPluginDynamic(std::string const name, DataType const type, int32_t const outDim, Weights const& W)
    : mLayerName(name)
    , mType(type)
    , mOutDim(outDim)
    , mNumParams(W.count)
    , mNmax(0)
    , mK(0)
    , mWdev(nullptr)
{
    memset(mAlgo.data, 0, sizeof(mAlgo.data));

    mW.convertAndCopy(W, mType);
    copyToDevice(mW, getWeightsSize(mW, mType), mWdev);
}

FCPluginDynamic::FCPluginDynamic(std::string const name, void const* data, size_t length)
    : mLayerName(name)
    , mWdev(nullptr)
{
    gLogVerbose << "FCPluginDynamic deserialize\n";

    // Deserialize in the same order as serialization
    deserialize_value(&data, &length, &mType);
    deserialize_value(&data, &length, &mOutDim);
    deserialize_value(&data, &length, &mNumParams);
    deserialize_value(&data, &length, &mNmax);
    deserialize_value(&data, &length, &mK);
    deserialize_value(&data, &length, &mAlgo);

    char const* d = static_cast<char const*>(data);

    mW.convertAndCopy(d, mNumParams, mType);
    copyToDevice(mW, getWeightsSize(mW, mType), mWdev);
}

// IPluginV2DynamicExt Methods
IPluginV2DynamicExt* FCPluginDynamic::clone() const noexcept
{
    try
    {
        gLogVerbose << "FCPluginDynamic clone\n";

        auto p = std::make_unique<FCPluginDynamic>(mLayerName, mType, mOutDim, mW);
        p->mAlgo = mAlgo;
        p->setPluginNamespace(mNamespace.c_str());

        return p.release();
    }
    catch (std::exception const& e)
    {
        caughtError(e);
    }
    return nullptr;
}

void FCPluginDynamic::attachToContext(
    cudnnContext* cudnnContext, cublasContext* cublasContext, nvinfer1::IGpuAllocator* gpuAllocator) noexcept
{
    mLtContext.attach();
}

void FCPluginDynamic::detachFromContext() noexcept
{
    mLtContext.detach();
}

DimsExprs FCPluginDynamic::getOutputDimensions(
    int32_t outputIndex, DimsExprs const* inputs, int32_t nbInputs, IExprBuilder& exprBuilder) noexcept
{
    try
    {
        PLUGIN_VALIDATE(nbInputs == 1);
        PLUGIN_VALIDATE(outputIndex == 0);
        PLUGIN_VALIDATE(inputs != nullptr);
        DimsExprs ret;
        ret.nbDims = 5;
        ret.d[0] = inputs[0].d[0];
        ret.d[1] = inputs[0].d[1];
        ret.d[2] = exprBuilder.constant(mOutDim);
        ret.d[3] = exprBuilder.constant(1);
        ret.d[4] = exprBuilder.constant(1);
        return ret;
    }
    catch (std::exception const& e)
    {
        caughtError(e);
    }
    return DimsExprs{};
}

bool FCPluginDynamic::supportsFormatCombination(
    int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
{
    PLUGIN_ASSERT(nbInputs == 1);
    PLUGIN_ASSERT(nbOutputs == 1);
    PLUGIN_ASSERT(inOut != nullptr);

    PluginTensorDesc const& in = inOut[pos];
    if (pos == 0)
    {
        return (in.type == mType) && (in.format == TensorFormat::kLINEAR);
    }
    PluginTensorDesc const& prev = inOut[pos - 1];

    // output
    return in.type == prev.type && in.format == prev.format;
}

void FCPluginDynamic::configurePlugin(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
    DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
{
    try
    {
        // Validate input arguments
        PLUGIN_VALIDATE(nbOutputs == 1);
        PLUGIN_VALIDATE(nbInputs == 1);
        PLUGIN_VALIDATE(inputs != nullptr);
        PLUGIN_VALIDATE(outputs != nullptr);
        PLUGIN_VALIDATE(mType == inputs[0].desc.type);
        auto const& inDims0 = inputs[0].desc.dims;

        PLUGIN_VALIDATE(inDims0.nbDims == 5);
        mK = inDims0.d[HDIM]; // hiddensize
        // PLUGIN_ASSERT(hiddenSize * mOutDim == mNumParams);
        PLUGIN_VALIDATE(inDims0.d[3] == 1);
        PLUGIN_VALIDATE(inDims0.d[4] == 1);

        // m and k are mOutDim
        // n is B*S
        int32_t const S = inputs->max.d[SDIM];
        int32_t const B = inputs->max.d[BDIM];

        mNmax = S * B;

        // Cleanup LtContext descriptors before creating new ones.
        mLtContext.destroy();

        if (mType == DataType::kFLOAT)
        {
            Gemm<float> g(mOutDim, mNmax, mK, false, false);
            mLtContext.create(g, kMAX_WORKSPACE_BYTES);
        }
        else if (mType == DataType::kHALF)
        {
            Gemm<half> g(mOutDim, mNmax, mK, false, false);
            mLtContext.create(g, kMAX_WORKSPACE_BYTES);
        }
        else
        {
            std::string const msg = "Unsupported type error, expected [kHALF,kFLOAT], but received ";
            PLUGIN_VALIDATE(false, (msg + std::to_string(static_cast<int32_t>(mType))).c_str());
        }

        gLogVerbose << "FCPluginDynamic configurePlugin m=" << mOutDim << ", n=" << mNmax << ", k=" << mK << std::endl;

        size_t actualWorkspace = 0;
        if (std::all_of(std::begin(mAlgo.data), std::end(mAlgo.data), [](auto v) { return v == 0; }))
        {
            gLogVerbose << "FCPluginDynamic gemmSearch\n";
            if (mSharedStream == nullptr)
            {
                SharedStream ss{};
                mSharedStream = static_cast<SharedStream*>(
                    getPluginRegistry()->acquirePluginResource(kFCPLUGIN_SHARED_STREAM_KEY, &ss))
                                    ->mStream;
            }
            if (mType == DataType::kFLOAT)
            {
                mAlgo = gemmSearch<float>(mOutDim, mNmax, mK, kMAX_WORKSPACE_BYTES, actualWorkspace, mSharedStream);
            }
            else if (mType == DataType::kHALF)
            {
                mAlgo = gemmSearch<half>(mOutDim, mNmax, mK, kMAX_WORKSPACE_BYTES, actualWorkspace, mSharedStream);
            }
        }

        AlgoProps p;
        p.populate(mAlgo);

        if (mType == DataType::kFLOAT && p.numericImpl == CUBLASLT_NUMERICAL_IMPL_FLAGS_HMMA)
        {
            gLogWarning << "cuBLAS might use mixed precision instead of FP32" << std::endl;
        }

        if (mType == DataType::kHALF && p.numericImpl != CUBLASLT_NUMERICAL_IMPL_FLAGS_HMMA)
        {
            gLogWarning << "TensorCore support was not selected" << std::endl;
        }

        gLogVerbose << "FCPluginDynamic configuration Algo=" << p.algoId << " Tile=" << p.tile << " ("
                    << matmulTileName[p.tile] << ") K=" << p.numSplitsK << " Red.Sch.=" << p.reductionScheme
                    << " Swiz=" << p.swizzle << " Cust=" << p.customOption << " numericImpl=" << p.numericImpl
                    << " ws=" << actualWorkspace << std::endl;
    }
    catch (std::exception const& e)
    {
        caughtError(e);
    }
}

size_t FCPluginDynamic::getWorkspaceSize(
    PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
{
    return kMAX_WORKSPACE_BYTES;
}

int32_t FCPluginDynamic::enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc,
    void const* const* inputs, void* const* outputs, void* workSpace, cudaStream_t stream) noexcept
{
    try
    {
        PLUGIN_VALIDATE(inputDesc != nullptr && outputDesc != nullptr && inputs != nullptr && outputs != nullptr
            && workSpace != nullptr);

        size_t const workspaceSize = getWorkspaceSize(inputDesc, 1, outputDesc, 1);

        int32_t const S = inputDesc->dims.d[SDIM];
        int32_t const B = inputDesc->dims.d[BDIM];
        int32_t const n = S * B;
        PLUGIN_VALIDATE(n >= 0);
        mLtContext.setN(static_cast<uint64_t>(n));

        if (mType == DataType::kFLOAT)
        {
            auto const* const input = static_cast<float const*>(inputs[0]);
            auto* output = static_cast<float*>(outputs[0]);

            Gemm<float> g(mOutDim, n, mK, false, false);
            if (mWdev == nullptr)
            {
                return STATUS_FAILURE;
            }
            g.A = static_cast<float*>(mWdev.get());
            g.B = const_cast<float*>(input);
            g.C = output;

            return cublasLtMatmul(mLtContext, g, mAlgo, workSpace, workspaceSize, stream);
        }
        if (mType == DataType::kHALF)
        {
            auto const* const input = static_cast<half const*>(inputs[0]);
            auto* output = static_cast<half*>(outputs[0]);

            Gemm<half> g(mOutDim, n, mK, false, false);
            if (mWdev == nullptr)
            {
                return STATUS_FAILURE;
            }
            g.A = static_cast<half*>(mWdev.get());
            g.B = const_cast<half*>(input);
            g.C = output;
            return cublasLtMatmul(mLtContext, g, mAlgo, workSpace, workspaceSize, stream);
        }
        else
        {
            gLogError << "Unsupported type error, expected [kHALF,kFLOAT], but received " << static_cast<int32_t>(mType)
                      << std::endl;
            return STATUS_FAILURE;
        }
    }
    catch (std::exception const& e)
    {
        caughtError(e);
    }
    return STATUS_FAILURE;
}

// IPluginV2Ext Methods
DataType FCPluginDynamic::getOutputDataType(int32_t index, DataType const* inputTypes, int32_t nbInputs) const noexcept
{
    PLUGIN_ASSERT(index == 0);
    PLUGIN_ASSERT(nbInputs == 1);
    PLUGIN_ASSERT(inputTypes != nullptr);
    PLUGIN_ASSERT(inputTypes[0] == DataType::kFLOAT || inputTypes[0] == DataType::kHALF);
    return inputTypes[0];
}

// IPluginV2 Methods
char const* FCPluginDynamic::getPluginType() const noexcept
{
    return kFC_NAME;
}

char const* FCPluginDynamic::getPluginVersion() const noexcept
{
    return kFC_VERSION;
}

int32_t FCPluginDynamic::getNbOutputs() const noexcept
{
    return 1;
}

int32_t FCPluginDynamic::initialize() noexcept
{
    gLogVerbose << "FCPluginDynamic initialize\n";
    return 0;
}

void FCPluginDynamic::terminate() noexcept
{
    gLogVerbose << "FCPluginDynamic terminate\n";
    if (mSharedStream)
    {
        TRT_UNUSED(getPluginRegistry()->releasePluginResource(kFCPLUGIN_SHARED_STREAM_KEY));
        mSharedStream = nullptr;
    }
}

size_t FCPluginDynamic::getSerializationSize() const noexcept
{
    size_t wordSize = getElementSize(mType);
    return wordSize * mNumParams + sizeof(mType) + sizeof(mOutDim) + sizeof(mNumParams) + sizeof(mAlgo) + sizeof(mNmax)
        + sizeof(mK);
}

void FCPluginDynamic::serialize(void* buffer) const noexcept
{
    serialize_value(&buffer, mType);
    serialize_value(&buffer, mOutDim);
    serialize_value(&buffer, mNumParams);
    serialize_value(&buffer, mNmax);
    serialize_value(&buffer, mK);
    serialize_value(&buffer, mAlgo);

    size_t wordSize = getElementSize(mType);
    char* d = static_cast<char*>(buffer);
    serFromDev(d, static_cast<char*>(mWdev.get()), mNumParams * wordSize);
}

void FCPluginDynamic::destroy() noexcept
{
    gLogVerbose << "FCPluginDynamic destroy\n";
    // This gets called when the network containing plugin is destroyed
    mLtContext.destroy();
    mWdev.reset(nullptr);
    delete this;
}

void FCPluginDynamic::setPluginNamespace(char const* libNamespace) noexcept
{
    try
    {
        PLUGIN_VALIDATE(libNamespace != nullptr);
        mNamespace = libNamespace;
    }
    catch (std::exception const& e)
    {
        caughtError(e);
    }
}

char const* FCPluginDynamic::getPluginNamespace() const noexcept
{
    return mNamespace.c_str();
}

/////////////////////////////////////////////////////////

FCPluginDynamicCreator::FCPluginDynamicCreator()
{
    mPluginAttributes.clear();
    mPluginAttributes.emplace_back(PluginField("out_dims", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("W", nullptr, PluginFieldType::kFLOAT32, 1));

    mFC.nbFields = mPluginAttributes.size();
    mFC.fields = mPluginAttributes.data();
}

char const* FCPluginDynamicCreator::getPluginName() const noexcept
{
    return kFC_NAME;
}

char const* FCPluginDynamicCreator::getPluginVersion() const noexcept
{
    return kFC_VERSION;
}

PluginFieldCollection const* FCPluginDynamicCreator::getFieldNames() noexcept
{
    return &mFC;
}

IPluginV2* FCPluginDynamicCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept
{
    try
    {
        gLogVerbose << "Creating FCPluginDynamicCreator...\n";
        PLUGIN_VALIDATE(name != nullptr);
        PLUGIN_VALIDATE(fc != nullptr);

        int32_t outDims = 0;
        int32_t typeId = -1;
        Weights W{DataType::kFLOAT, nullptr, 0LL};
        plugin::validateRequiredAttributesExist({"out_dims", "type_id", "W"}, fc);

        for (int32_t i = 0; i < fc->nbFields; i++)
        {
            std::string_view const fieldName = fc->fields[i].name;
            if (fieldName == "out_dims"sv)
            {
                outDims = static_cast<int32_t const*>(fc->fields[i].data)[0];
                gLogVerbose << "Building outDims: " << outDims << std::endl;
            }

            if (fieldName == "type_id"sv)
            {
                typeId = static_cast<int32_t const*>(fc->fields[i].data)[0];
                gLogVerbose << "Building typeId: " << outDims << std::endl;
            }

            if (fieldName == "W"sv)
            {
                gLogVerbose << "Building W...\n";
                W.values = fc->fields[i].data;
                W.count = fc->fields[i].length;
                W.type = fieldTypeToDataType(fc->fields[i].type);
                gLogVerbose << "Is W float32: " << (W.type == DataType::kFLOAT) << std::endl;
            }
        }

        if (outDims <= 0)
        {
            gLogError << "Invalid output dimension" << std::endl;
        }
        if (typeId < 0 || typeId > 1)
        {
            gLogError << "Invalid type id" << typeId << std::endl;
        }
        if (W.count == 0 || W.values == nullptr || W.count < outDims)
        {
            gLogError << "Invalid weights" << std::endl;
        }

        DataType type = static_cast<DataType>(typeId);
        return new FCPluginDynamic(name, type, outDims, W);
    }
    catch (std::exception const& e)
    {
        caughtError(e);
    }
    return nullptr;
}

IPluginV2* FCPluginDynamicCreator::deserializePlugin(
    char const* name, void const* serialData, size_t serialLength) noexcept
{
    // This object will be deleted when the network is destroyed, which will
    // call FCPluginDynamic::destroy()
    try
    {
        return new FCPluginDynamic(name, serialData, serialLength);
    }
    catch (std::exception const& e)
    {
        caughtError(e);
    }
    return nullptr;
}

void FCPluginDynamicCreator::setPluginNamespace(char const* libNamespace) noexcept
{
    try
    {
        PLUGIN_VALIDATE(libNamespace != nullptr);
        mNamespace = libNamespace;
    }
    catch (std::exception const& e)
    {
        caughtError(e);
    }
}

char const* FCPluginDynamicCreator::getPluginNamespace() const noexcept
{
    return mNamespace.c_str();
}

#endif // #if CUDA_VERSION >= 10010