wehub-resource-sync
541 lines
20 KiB
C++
541 lines
20 KiB
C++
/*
|
|
* 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.
|
|
*/
|
|
|
|
#include <cuda.h>
|
|
#if CUDA_VERSION >= 10010
|
|
|
|
#ifndef BERT_COMMON_H
|
|
#define BERT_COMMON_H
|
|
|
|
#include "NvInfer.h"
|
|
#include "NvInferRuntimeCommon.h"
|
|
#include "common/checkMacrosPlugin.h"
|
|
#include "common/cublasWrapper.h"
|
|
#include "common/plugin.h"
|
|
#include <cuda_fp16.h>
|
|
|
|
#include <algorithm>
|
|
#include <cassert>
|
|
#include <cuda_runtime_api.h>
|
|
#include <memory>
|
|
#include <numeric>
|
|
#include <stdexcept>
|
|
#include <vector>
|
|
|
|
#define TRT_UNUSED (void)
|
|
|
|
#define BERT_PRINT_DEBUG_MSG 0
|
|
|
|
#if BERT_PRINT_DEBUG_MSG
|
|
#define BERT_DEBUG_MSG(msg) (gLogVerbose << (msg) << std::endl)
|
|
#define BERT_DEBUG_VALUE(key, value) (gLogVerbose << key << value << std::endl)
|
|
#else
|
|
#define BERT_DEBUG_MSG(msg) TRT_UNUSED(msg)
|
|
#define BERT_DEBUG_VALUE(key, value) \
|
|
TRT_UNUSED(key); \
|
|
TRT_UNUSED(value)
|
|
#endif
|
|
|
|
using half = __half;
|
|
|
|
constexpr uint32_t BDIM = 1; // batch dimension
|
|
constexpr uint32_t SDIM = 0; // seq len dimension
|
|
constexpr uint32_t HDIM = 2; // hidden dimension
|
|
|
|
constexpr int32_t kSM_75 = 75;
|
|
constexpr int32_t kSM_80 = 80;
|
|
constexpr int32_t kSM_86 = 86;
|
|
constexpr int32_t kSM_87 = 87;
|
|
constexpr int32_t kSM_89 = 89;
|
|
constexpr int32_t kSM_90 = 90;
|
|
constexpr int32_t kSM_100 = 100;
|
|
constexpr int32_t kSM_120 = 120;
|
|
|
|
// For full mask mode, we must produce the compressed mask format expected by the fused attention path. Currently, only
|
|
// two sequence lengths are supported. We hard code the sizes here.
|
|
// The number of threads per CTA: warps_m * warps_n * warps_k * 32;
|
|
constexpr size_t threadsPerCta128 = 2 * 2 * 32;
|
|
constexpr size_t threadsPerCta384 = 1 * 8 * 32;
|
|
|
|
// The number of xmmas in the M dimension. We use one uint32_t per XMMA in the M dimension: (s + 16*warps_m - 1)
|
|
// / (16*warps_m);
|
|
constexpr size_t xmmasM128 = 4;
|
|
constexpr size_t xmmasM384 = 24;
|
|
|
|
// Packed mask size per batch. Layout is XMMAS_M * THREADS_PER_CTA.
|
|
constexpr size_t unfusedMaskSize = 1;
|
|
constexpr size_t packedMaskSize64 = xmmasM128 * threadsPerCta128;
|
|
constexpr size_t packedMaskSize96 = xmmasM128 * threadsPerCta128;
|
|
constexpr size_t packedMaskSize128 = xmmasM128 * threadsPerCta128;
|
|
constexpr size_t packedMaskSize384 = xmmasM384 * threadsPerCta384;
|
|
|
|
namespace nvinfer1
|
|
{
|
|
namespace pluginInternal
|
|
{
|
|
template <typename T>
|
|
struct CudaDeleter
|
|
{
|
|
void operator()(T* buf)
|
|
{
|
|
PLUGIN_CUASSERT(cudaFree(buf));
|
|
}
|
|
};
|
|
|
|
} // namespace pluginInternal
|
|
namespace plugin
|
|
{
|
|
|
|
namespace bert
|
|
{
|
|
|
|
//! \brief Checks if the first argument matches any of the list items.
|
|
//! \return True if v is a member of list.
|
|
template <typename TElem, typename Container = std::initializer_list<TElem>>
|
|
bool elem(TElem const& v, Container const& list)
|
|
{
|
|
return std::any_of(std::begin(list), std::end(list), [&v](TElem const& t) { return t == v; });
|
|
}
|
|
|
|
inline int32_t getMHAMaskPackedSize(int32_t smVersion, nvinfer1::DataType dataType, int32_t sequenceLength)
|
|
{
|
|
// this code must match EmbLayerNormPluginDynamic::getOutputDimensions in embLayerNormPlugin.cpp
|
|
int32_t packedSize = unfusedMaskSize;
|
|
bool const isSmOK = elem(smVersion, {kSM_75, kSM_80, kSM_86, kSM_87, kSM_89, kSM_90, kSM_100, kSM_120});
|
|
bool isPrecisionOK = (dataType == nvinfer1::DataType::kINT8 || dataType == nvinfer1::DataType::kHALF);
|
|
if (isSmOK && isPrecisionOK)
|
|
{
|
|
if (sequenceLength == 64)
|
|
{
|
|
packedSize = packedMaskSize64;
|
|
}
|
|
else if (sequenceLength == 96)
|
|
{
|
|
packedSize = packedMaskSize96;
|
|
}
|
|
else if (sequenceLength == 128)
|
|
{
|
|
packedSize = packedMaskSize128;
|
|
}
|
|
else if (sequenceLength == 384)
|
|
{
|
|
packedSize = packedMaskSize384;
|
|
}
|
|
}
|
|
return packedSize;
|
|
}
|
|
|
|
inline uint32_t getElementSize(nvinfer1::DataType t)
|
|
{
|
|
switch (t)
|
|
{
|
|
case nvinfer1::DataType::kINT64: return 8;
|
|
case nvinfer1::DataType::kINT32:
|
|
case nvinfer1::DataType::kFLOAT: return 4;
|
|
case nvinfer1::DataType::kBF16:
|
|
case nvinfer1::DataType::kHALF: return 2;
|
|
case nvinfer1::DataType::kBOOL:
|
|
case nvinfer1::DataType::kUINT8:
|
|
case nvinfer1::DataType::kINT8:
|
|
case nvinfer1::DataType::kFP8: return 1;
|
|
case nvinfer1::DataType::kINT4:
|
|
case nvinfer1::DataType::kFP4:
|
|
case nvinfer1::DataType::kE8M0: PLUGIN_FAIL("Element size is not implemented for sub-byte data-types");
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
inline int64_t getWeightsSize(nvinfer1::Weights const& w, nvinfer1::DataType type)
|
|
{
|
|
return w.count * getElementSize(type);
|
|
}
|
|
|
|
inline int64_t volume(nvinfer1::Dims const& d)
|
|
{
|
|
return std::accumulate(d.d, d.d + d.nbDims, int64_t{1}, std::multiplies<int64_t>{});
|
|
}
|
|
|
|
//! Check if the hardware supports BERT Multi-Head Attention plugins
|
|
//! The plugin calls precompiled cubins (compiled from fmha_v2/xmma kernels)
|
|
//! that are SM-specific.
|
|
inline bool doesHwSupportBertMHAPlugin() noexcept
|
|
{
|
|
int32_t device{-1};
|
|
cudaGetDevice(&device);
|
|
int32_t smMajor{0};
|
|
int32_t smMinor{0};
|
|
cudaDeviceGetAttribute(&smMajor, cudaDevAttrComputeCapabilityMajor, device);
|
|
cudaDeviceGetAttribute(&smMinor, cudaDevAttrComputeCapabilityMinor, device);
|
|
int32_t smVersion = (smMajor << 4) | (smMinor);
|
|
// Turing and above
|
|
static constexpr int32_t kSM_TURING_HEX{0x75};
|
|
static constexpr int32_t kSM_BLACKWELL_100_HEX{0xA0};
|
|
static constexpr int32_t kSM_BLACKWELL_120_HEX{0xC0};
|
|
static constexpr int32_t kSM_ORIN_HEX{0x87};
|
|
bool isAuto = smVersion == kSM_ORIN_HEX;
|
|
bool isSm100OrLower = smVersion >= kSM_TURING_HEX && smVersion <= kSM_BLACKWELL_100_HEX;
|
|
bool isHardwareSupported = (isSm100OrLower || smVersion == kSM_BLACKWELL_120_HEX) && !isAuto;
|
|
|
|
return isHardwareSupported;
|
|
}
|
|
|
|
template <typename IntType>
|
|
constexpr IntType ceildiv(IntType a, IntType b)
|
|
{
|
|
return (a + b - 1) / b;
|
|
}
|
|
template <typename IntType>
|
|
constexpr IntType alignTo(IntType a, IntType b)
|
|
{
|
|
return ceildiv(a, b) * b;
|
|
}
|
|
|
|
template <typename T>
|
|
inline T* deserToDev(char const*& buffer, size_t nbElem)
|
|
{
|
|
void* dev{nullptr};
|
|
const size_t len = sizeof(T) * nbElem;
|
|
PLUGIN_CUASSERT(cudaMalloc(&dev, len));
|
|
PLUGIN_CUASSERT(cudaMemcpy(dev, buffer, len, cudaMemcpyHostToDevice));
|
|
|
|
buffer += len;
|
|
return static_cast<T*>(dev);
|
|
}
|
|
|
|
template <typename T>
|
|
inline void serFromDev(char*& buffer, T const* data, size_t nbElem)
|
|
{
|
|
const size_t len = sizeof(T) * nbElem;
|
|
PLUGIN_CUASSERT(cudaMemcpy(buffer, static_cast<void const*>(data), len, cudaMemcpyDeviceToHost));
|
|
buffer += len;
|
|
}
|
|
|
|
template <typename T>
|
|
inline T* devToDev(T const* data, size_t nbElem)
|
|
{
|
|
void* dev{nullptr};
|
|
const size_t len = sizeof(T) * nbElem;
|
|
PLUGIN_CUASSERT(cudaMalloc(&dev, len));
|
|
PLUGIN_CUASSERT(cudaMemcpy(dev, static_cast<void const*>(data), len, cudaMemcpyDeviceToDevice));
|
|
return static_cast<T*>(dev);
|
|
}
|
|
|
|
template <typename T>
|
|
nvinfer1::pluginInternal::cublasStatus_t inline cublasGemm(nvinfer1::pluginInternal::cublasHandle_t handle,
|
|
nvinfer1::pluginInternal::cublasOperation_t transa, nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m,
|
|
int32_t n, int32_t k, const T alpha, T const* A, int32_t lda, T const* B, int32_t ldb, const T beta, T* C,
|
|
int32_t ldc);
|
|
|
|
template <>
|
|
nvinfer1::pluginInternal::cublasStatus_t inline cublasGemm(nvinfer1::pluginInternal::cublasHandle_t handle,
|
|
nvinfer1::pluginInternal::cublasOperation_t transa, nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m,
|
|
int32_t n, int32_t k, float const alpha, float const* A, int32_t lda, float const* B, int32_t ldb, float const beta,
|
|
float* C, int32_t ldc)
|
|
{
|
|
nvinfer1::pluginInternal::CublasWrapper& wrapper = nvinfer1::pluginInternal::getCublasWrapper();
|
|
return wrapper.cublasSgemm(handle, transa, transb, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
|
|
}
|
|
|
|
template <>
|
|
nvinfer1::pluginInternal::cublasStatus_t inline cublasGemm(nvinfer1::pluginInternal::cublasHandle_t handle,
|
|
nvinfer1::pluginInternal::cublasOperation_t transa, nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m,
|
|
int32_t n, int32_t k, const half alpha, half const* A, int32_t lda, half const* B, int32_t ldb, const half beta,
|
|
half* C, int32_t ldc)
|
|
{
|
|
nvinfer1::pluginInternal::CublasWrapper& wrapper = nvinfer1::pluginInternal::getCublasWrapper();
|
|
return wrapper.cublasHgemm(handle, transa, transb, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
|
|
}
|
|
|
|
template <typename T>
|
|
nvinfer1::pluginInternal::cublasStatus_t inline cublasGemmStridedBatchedEx(
|
|
nvinfer1::pluginInternal::cublasHandle_t handle, nvinfer1::pluginInternal::cublasOperation_t transa,
|
|
nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m, int32_t n, int32_t k, const T alpha, T const* A,
|
|
int32_t lda, int64_t strideA, T const* B, int32_t ldb, int64_t strideB, const T beta, T* C, int32_t ldc,
|
|
int64_t strideC, int32_t batchCount, nvinfer1::pluginInternal::cublasGemmAlgo_t algo);
|
|
|
|
template <>
|
|
nvinfer1::pluginInternal::cublasStatus_t inline cublasGemmStridedBatchedEx(
|
|
nvinfer1::pluginInternal::cublasHandle_t handle, nvinfer1::pluginInternal::cublasOperation_t transa,
|
|
nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m, int32_t n, int32_t k, float const alpha,
|
|
float const* A, int32_t lda, int64_t strideA, float const* B, int32_t ldb, int64_t strideB, float const beta,
|
|
float* C, int32_t ldc, int64_t strideC, int32_t batchCount, nvinfer1::pluginInternal::cublasGemmAlgo_t algo)
|
|
{
|
|
nvinfer1::pluginInternal::CublasWrapper& wrapper = nvinfer1::pluginInternal::getCublasWrapper();
|
|
return wrapper.cublasGemmStridedBatchedEx(handle, transa, transb, m, n, k, &alpha, A, CUDA_R_32F, lda, strideA, B,
|
|
CUDA_R_32F, ldb, strideB, &beta, C, CUDA_R_32F, ldc, strideC, batchCount, CUDA_R_32F, algo);
|
|
}
|
|
|
|
template <>
|
|
nvinfer1::pluginInternal::cublasStatus_t inline cublasGemmStridedBatchedEx(
|
|
nvinfer1::pluginInternal::cublasHandle_t handle, nvinfer1::pluginInternal::cublasOperation_t transa,
|
|
nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m, int32_t n, int32_t k, const half alpha,
|
|
half const* A, int32_t lda, int64_t strideA, half const* B, int32_t ldb, int64_t strideB, const half beta, half* C,
|
|
int32_t ldc, int64_t strideC, int32_t batchCount, nvinfer1::pluginInternal::cublasGemmAlgo_t algo)
|
|
{
|
|
nvinfer1::pluginInternal::CublasWrapper& wrapper = nvinfer1::pluginInternal::getCublasWrapper();
|
|
return wrapper.cublasGemmStridedBatchedEx(handle, transa, transb, m, n, k, &alpha, A, CUDA_R_16F, lda, strideA, B,
|
|
CUDA_R_16F, ldb, strideB, &beta, C, CUDA_R_16F, ldc, strideC, batchCount, CUDA_R_16F, algo);
|
|
}
|
|
|
|
template <typename T>
|
|
nvinfer1::pluginInternal::cublasStatus_t inline cublasGemmStridedBatched(
|
|
nvinfer1::pluginInternal::cublasHandle_t handle, nvinfer1::pluginInternal::cublasOperation_t transa,
|
|
nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m, int32_t n, int32_t k, const T alpha, T const* A,
|
|
int32_t lda, int64_t strideA, T const* B, int32_t ldb, int64_t strideB, const T beta, T* C, int32_t ldc,
|
|
int64_t strideC, int32_t batchCount);
|
|
|
|
template <>
|
|
nvinfer1::pluginInternal::cublasStatus_t inline cublasGemmStridedBatched(
|
|
nvinfer1::pluginInternal::cublasHandle_t handle, nvinfer1::pluginInternal::cublasOperation_t transa,
|
|
nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m, int32_t n, int32_t k, float const alpha,
|
|
float const* A, int32_t lda, int64_t strideA, float const* B, int32_t ldb, int64_t strideB, float const beta,
|
|
float* C, int32_t ldc, int64_t strideC, int32_t batchCount)
|
|
{
|
|
nvinfer1::pluginInternal::CublasWrapper& wrapper = nvinfer1::pluginInternal::getCublasWrapper();
|
|
return wrapper.cublasSgemmStridedBatched(
|
|
handle, transa, transb, m, n, k, &alpha, A, lda, strideA, B, ldb, strideB, &beta, C, ldc, strideC, batchCount);
|
|
}
|
|
|
|
template <>
|
|
nvinfer1::pluginInternal::cublasStatus_t inline cublasGemmStridedBatched(
|
|
nvinfer1::pluginInternal::cublasHandle_t handle, nvinfer1::pluginInternal::cublasOperation_t transa,
|
|
nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m, int32_t n, int32_t k, const half alpha,
|
|
half const* A, int32_t lda, int64_t strideA, half const* B, int32_t ldb, int64_t strideB, const half beta, half* C,
|
|
int32_t ldc, int64_t strideC, int32_t batchCount)
|
|
{
|
|
nvinfer1::pluginInternal::CublasWrapper& wrapper = nvinfer1::pluginInternal::getCublasWrapper();
|
|
return wrapper.cublasHgemmStridedBatched(
|
|
handle, transa, transb, m, n, k, &alpha, A, lda, strideA, B, ldb, strideB, &beta, C, ldc, strideC, batchCount);
|
|
}
|
|
|
|
struct CublasConfigHelper
|
|
{
|
|
nvinfer1::pluginInternal::cublasPointerMode_t pm;
|
|
nvinfer1::pluginInternal::cublasMath_t mm;
|
|
nvinfer1::pluginInternal::cublasHandle_t cublas;
|
|
nvinfer1::pluginInternal::CublasWrapper& wrapper = nvinfer1::pluginInternal::getCublasWrapper();
|
|
CublasConfigHelper(nvinfer1::pluginInternal::cublasHandle_t cublas_)
|
|
: cublas(cublas_)
|
|
{
|
|
PLUGIN_CUBLASASSERT(wrapper.cublasGetPointerMode(cublas, &pm));
|
|
PLUGIN_CUBLASASSERT(wrapper.cublasGetMathMode(cublas, &mm));
|
|
PLUGIN_CUBLASASSERT(wrapper.cublasSetPointerMode(cublas, nvinfer1::pluginInternal::CUBLAS_POINTER_MODE_HOST));
|
|
PLUGIN_CUBLASASSERT(wrapper.cublasSetMathMode(cublas, nvinfer1::pluginInternal::CUBLAS_TENSOR_OP_MATH));
|
|
}
|
|
~CublasConfigHelper()
|
|
{
|
|
wrapper.cublasSetMathMode(cublas, mm);
|
|
wrapper.cublasSetPointerMode(cublas, pm);
|
|
}
|
|
};
|
|
|
|
template <typename T>
|
|
using cuda_unique_ptr = std::unique_ptr<T, pluginInternal::CudaDeleter<T>>;
|
|
|
|
template <typename T>
|
|
using cuda_shared_ptr = std::shared_ptr<T>;
|
|
|
|
template <typename T>
|
|
void make_cuda_shared(cuda_shared_ptr<T>& ptr, void* cudaMem)
|
|
{
|
|
ptr.reset(static_cast<T*>(cudaMem), pluginInternal::CudaDeleter<T>());
|
|
}
|
|
|
|
struct WeightsWithOwnership : public nvinfer1::Weights
|
|
{
|
|
WeightsWithOwnership()
|
|
{
|
|
values = nullptr;
|
|
count = 0;
|
|
}
|
|
~WeightsWithOwnership()
|
|
{
|
|
operator delete[](const_cast<void*>(values));
|
|
}
|
|
|
|
WeightsWithOwnership(WeightsWithOwnership const&) = delete;
|
|
WeightsWithOwnership operator=(WeightsWithOwnership const&) = delete;
|
|
WeightsWithOwnership(WeightsWithOwnership const&&) = delete;
|
|
WeightsWithOwnership operator=(WeightsWithOwnership const&&) = delete;
|
|
|
|
void convertAndCopy(nvinfer1::Weights const& src, nvinfer1::DataType type)
|
|
{
|
|
this->type = type;
|
|
this->count = src.count;
|
|
|
|
if (type == nvinfer1::DataType::kFLOAT)
|
|
{
|
|
auto destBuf = new float[src.count];
|
|
this->values = destBuf;
|
|
|
|
if (src.type == nvinfer1::DataType::kFLOAT)
|
|
{
|
|
BERT_DEBUG_MSG("Float Weights(Host) => Float Array(Host)");
|
|
std::copy_n(static_cast<float const*>(src.values), src.count, destBuf);
|
|
}
|
|
else
|
|
{
|
|
PLUGIN_ASSERT(src.type == nvinfer1::DataType::kHALF);
|
|
|
|
BERT_DEBUG_MSG("Half Weights(Host) => Float Array(Host)");
|
|
auto const s = static_cast<half const*>(src.values);
|
|
auto d = static_cast<float*>(const_cast<void*>(this->values));
|
|
|
|
for (auto it = 0; it < src.count; it++)
|
|
{
|
|
d[it] = __half2float(s[it]);
|
|
}
|
|
}
|
|
}
|
|
else if (type == nvinfer1::DataType::kHALF)
|
|
{
|
|
auto destBuf = new half[src.count];
|
|
this->values = destBuf;
|
|
|
|
if (src.type == nvinfer1::DataType::kHALF)
|
|
{
|
|
BERT_DEBUG_MSG("Half Weights(Host) => Half Array(Host)");
|
|
std::copy_n(static_cast<half const*>(src.values), src.count, destBuf);
|
|
}
|
|
else
|
|
{
|
|
PLUGIN_ASSERT(src.type == nvinfer1::DataType::kFLOAT);
|
|
|
|
BERT_DEBUG_MSG("Float Weights(Host) => Half Array(Host)");
|
|
auto const s = static_cast<float const*>(src.values);
|
|
auto d = static_cast<half*>(const_cast<void*>(this->values));
|
|
|
|
for (auto it = 0; it < src.count; it++)
|
|
{
|
|
d[it] = __float2half(s[it]);
|
|
}
|
|
}
|
|
}
|
|
else
|
|
{
|
|
throw std::runtime_error("Unsupported DataType specified for plugin.");
|
|
}
|
|
}
|
|
|
|
void convertAndCopy(char const*& srcBuf, size_t count, nvinfer1::DataType type) noexcept
|
|
{
|
|
this->type = type;
|
|
this->count = count;
|
|
auto const nbBytes = getWeightsSize(*this, type);
|
|
auto destBuf = new char[nbBytes];
|
|
this->values = destBuf;
|
|
|
|
std::copy_n(srcBuf, nbBytes, destBuf);
|
|
srcBuf += nbBytes;
|
|
}
|
|
};
|
|
|
|
template <typename T>
|
|
inline void copyToDevice(WeightsWithOwnership& hostWeights, size_t nbBytes, cuda_unique_ptr<T>& cudaWeights)
|
|
{
|
|
if (hostWeights.values)
|
|
{
|
|
void* cudaMem{nullptr};
|
|
PLUGIN_CUASSERT(cudaMalloc(&cudaMem, nbBytes));
|
|
PLUGIN_CUASSERT(cudaMemcpy(cudaMem, hostWeights.values, nbBytes, cudaMemcpyHostToDevice));
|
|
cudaWeights.reset(static_cast<T*>(cudaMem));
|
|
}
|
|
}
|
|
|
|
inline void convertAndCopyToDevice(nvinfer1::Weights const& src, float* destDev)
|
|
{
|
|
|
|
size_t wordSize = sizeof(float);
|
|
size_t nbBytes = src.count * wordSize;
|
|
if (src.type == nvinfer1::DataType::kFLOAT)
|
|
{
|
|
BERT_DEBUG_MSG("Float Weights(Host) => Float Array(Device)");
|
|
PLUGIN_CUASSERT(cudaMemcpy(destDev, src.values, nbBytes, cudaMemcpyHostToDevice));
|
|
}
|
|
else
|
|
{
|
|
BERT_DEBUG_MSG("Half Weights(Host) => Float Array(Device)");
|
|
std::vector<float> tmp(src.count);
|
|
half const* values = reinterpret_cast<half const*>(src.values);
|
|
|
|
for (size_t it = 0; it < tmp.size(); it++)
|
|
{
|
|
tmp[it] = __half2float(values[it]);
|
|
}
|
|
|
|
PLUGIN_CUASSERT(cudaMemcpy(destDev, tmp.data(), nbBytes, cudaMemcpyHostToDevice));
|
|
}
|
|
}
|
|
|
|
inline void convertAndCopyToDevice(nvinfer1::Weights const& src, half* destDev)
|
|
{
|
|
size_t wordSize = sizeof(half);
|
|
size_t nbBytes = src.count * wordSize;
|
|
if (src.type == nvinfer1::DataType::kHALF)
|
|
{
|
|
BERT_DEBUG_MSG("Half Weights(Host) => Half Array(Device)");
|
|
PLUGIN_CUASSERT(cudaMemcpy(destDev, src.values, nbBytes, cudaMemcpyHostToDevice));
|
|
}
|
|
else
|
|
{
|
|
BERT_DEBUG_MSG("Float Weights(Host) => Half Array(Device)");
|
|
std::vector<half> tmp(src.count);
|
|
float const* values = reinterpret_cast<float const*>(src.values);
|
|
|
|
for (size_t it = 0; it < tmp.size(); it++)
|
|
{
|
|
tmp[it] = __float2half(values[it]);
|
|
}
|
|
PLUGIN_CUASSERT(cudaMemcpy(destDev, tmp.data(), nbBytes, cudaMemcpyHostToDevice));
|
|
}
|
|
}
|
|
|
|
inline nvinfer1::DataType fieldTypeToDataType(const nvinfer1::PluginFieldType ftype)
|
|
{
|
|
switch (ftype)
|
|
{
|
|
case nvinfer1::PluginFieldType::kFLOAT32:
|
|
{
|
|
BERT_DEBUG_MSG("PluginFieldType is Float32");
|
|
return nvinfer1::DataType::kFLOAT;
|
|
}
|
|
case nvinfer1::PluginFieldType::kFLOAT16:
|
|
{
|
|
BERT_DEBUG_MSG("PluginFieldType is Float16");
|
|
return nvinfer1::DataType::kHALF;
|
|
}
|
|
case nvinfer1::PluginFieldType::kINT32:
|
|
{
|
|
BERT_DEBUG_MSG("PluginFieldType is Int32");
|
|
return nvinfer1::DataType::kINT32;
|
|
}
|
|
case nvinfer1::PluginFieldType::kINT8:
|
|
{
|
|
BERT_DEBUG_MSG("PluginFieldType is Int8");
|
|
return nvinfer1::DataType::kINT8;
|
|
}
|
|
default: throw std::invalid_argument("No corresponding datatype for plugin field type");
|
|
}
|
|
}
|
|
|
|
} // namespace bert
|
|
} // namespace plugin
|
|
} // namespace nvinfer1
|
|
#endif // BERT_COMMON_H
|
|
|
|
#endif // CUDA_VERSION >= 10010
|