chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,540 @@
|
||||
/*
|
||||
* 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
|
||||
Reference in New Issue
Block a user