chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,222 @@
|
||||
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
||||
// SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
|
||||
// AFFILIATES. 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cublas_v2.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime_api.h>
|
||||
#include <algorithm>
|
||||
#include <cassert>
|
||||
#include <memory>
|
||||
#include <numeric>
|
||||
#include <stdexcept>
|
||||
#include <vector>
|
||||
#include "NvInfer.h"
|
||||
#include "NvInferRuntimeCommon.h"
|
||||
#include "paddle/fluid/inference/tensorrt/plugin/common/plugin.h"
|
||||
|
||||
#define TRT_UNUSED (void)
|
||||
#define BERT_PRINT_DEBUG_MSG 0
|
||||
#if BERT_PRINT_DEBUG_MSG
|
||||
#define TRANSFORMER_DEBUG_MSG(msg) (gLogVerbose << (msg) << std::endl)
|
||||
#define BERT_DEBUG_VALUE(key, value) (gLogVerbose << key << value << std::endl)
|
||||
#else
|
||||
#define TRANSFORMER_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_53 = 53;
|
||||
constexpr int32_t kSM_70 = 70;
|
||||
constexpr int32_t kSM_72 = 72;
|
||||
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 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 paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
namespace plugin {
|
||||
|
||||
inline uint32_t getElementSize(nvinfer1::DataType t) noexcept {
|
||||
switch (t) {
|
||||
case nvinfer1::DataType::kINT32:
|
||||
return 4;
|
||||
case nvinfer1::DataType::kFLOAT:
|
||||
return 4;
|
||||
case nvinfer1::DataType::kHALF:
|
||||
return 2;
|
||||
case nvinfer1::DataType::kBOOL:
|
||||
case nvinfer1::DataType::kINT8:
|
||||
return 1;
|
||||
default:
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
|
||||
inline int64_t getWeightsSize(const nvinfer1::Weights& w,
|
||||
nvinfer1::DataType type) {
|
||||
return w.count * getElementSize(type);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline void serFromDev(char** buffer, const T* data, size_t nbElem) {
|
||||
const size_t len = sizeof(T) * nbElem;
|
||||
cudaMemcpy(
|
||||
*buffer, static_cast<const void*>(data), len, cudaMemcpyDeviceToHost);
|
||||
*buffer += len;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct CudaDeleter {
|
||||
void operator()(T* buf) { cudaFree(buf); }
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
using cuda_unique_ptr = std::unique_ptr<T, 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), CudaDeleter<T>());
|
||||
}
|
||||
|
||||
struct WeightsWithOwnership : public nvinfer1::Weights {
|
||||
WeightsWithOwnership() {
|
||||
values = nullptr;
|
||||
count = 0;
|
||||
}
|
||||
~WeightsWithOwnership() { operator delete[](const_cast<void*>(values)); }
|
||||
|
||||
WeightsWithOwnership(const WeightsWithOwnership&) = delete;
|
||||
WeightsWithOwnership operator=(const WeightsWithOwnership&) = delete;
|
||||
WeightsWithOwnership(const WeightsWithOwnership&&) = delete;
|
||||
WeightsWithOwnership operator=(const WeightsWithOwnership&&) = delete;
|
||||
|
||||
void convertAndCopy(const nvinfer1::Weights& 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) {
|
||||
TRANSFORMER_DEBUG_MSG("Float Weights(Host) => Float Array(Host)");
|
||||
std::copy_n(static_cast<const float*>(src.values), src.count, destBuf);
|
||||
} else {
|
||||
assert(src.type == nvinfer1::DataType::kHALF);
|
||||
TRANSFORMER_DEBUG_MSG("Half Weights(Host) => Float Array(Host)");
|
||||
const auto s = static_cast<const half*>(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) {
|
||||
TRANSFORMER_DEBUG_MSG("Half Weights(Host) => Half Array(Host)");
|
||||
std::copy_n(static_cast<const half*>(src.values), src.count, destBuf);
|
||||
} else {
|
||||
assert(src.type == nvinfer1::DataType::kFLOAT);
|
||||
TRANSFORMER_DEBUG_MSG("Float Weights(Host) => Half Array(Host)");
|
||||
const auto s = static_cast<const float*>(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(const char** srcBuf,
|
||||
size_t count,
|
||||
nvinfer1::DataType type) noexcept {
|
||||
this->type = type;
|
||||
this->count = count;
|
||||
const auto 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};
|
||||
cudaMalloc(&cudaMem, nbBytes);
|
||||
cudaMemcpy(cudaMem, hostWeights->values, nbBytes, cudaMemcpyHostToDevice);
|
||||
cudaWeights->reset(static_cast<T*>(cudaMem));
|
||||
}
|
||||
}
|
||||
|
||||
inline nvinfer1::DataType fieldTypeToDataType(
|
||||
const nvinfer1::PluginFieldType ftype) {
|
||||
switch (ftype) {
|
||||
case nvinfer1::PluginFieldType::kFLOAT32: {
|
||||
TRANSFORMER_DEBUG_MSG("PluginFieldType is Float32");
|
||||
return nvinfer1::DataType::kFLOAT;
|
||||
}
|
||||
case nvinfer1::PluginFieldType::kFLOAT16: {
|
||||
TRANSFORMER_DEBUG_MSG("PluginFieldType is Float16");
|
||||
return nvinfer1::DataType::kHALF;
|
||||
}
|
||||
case nvinfer1::PluginFieldType::kINT32: {
|
||||
TRANSFORMER_DEBUG_MSG("PluginFieldType is Int32");
|
||||
return nvinfer1::DataType::kINT32;
|
||||
}
|
||||
case nvinfer1::PluginFieldType::kINT8: {
|
||||
TRANSFORMER_DEBUG_MSG("PluginFieldType is Int8");
|
||||
return nvinfer1::DataType::kINT8;
|
||||
}
|
||||
default:
|
||||
TRANSFORMER_DEBUG_MSG("PluginFieldType is Float32");
|
||||
return nvinfer1::DataType::kFLOAT;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace plugin
|
||||
} // namespace tensorrt
|
||||
} // namespace inference
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,301 @@
|
||||
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
||||
// SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
|
||||
// AFFILIATES. 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cublas_v2.h>
|
||||
#include <cub/cub.cuh>
|
||||
#include "paddle/phi/core/platform/device_context.h"
|
||||
|
||||
using kv_float = cub::KeyValuePair<float, float>;
|
||||
using kv_half = cub::KeyValuePair<half, half>;
|
||||
using kv_half2 = cub::KeyValuePair<half2, half2>;
|
||||
|
||||
template <typename T>
|
||||
__device__ inline T rsqrt(const T& x);
|
||||
|
||||
template <>
|
||||
__device__ inline float rsqrt(const float& x) {
|
||||
return rsqrtf(x);
|
||||
}
|
||||
|
||||
namespace cub {
|
||||
__host__ __device__ inline kv_float operator+(const kv_float& a,
|
||||
const kv_float& b) {
|
||||
return kv_float(a.key + b.key, a.value + b.value);
|
||||
}
|
||||
} // namespace cub
|
||||
|
||||
// Half Operations
|
||||
__device__ inline half2 __hadd2_with_fallback(const half2 a, const half2 b) {
|
||||
#if __CUDA_ARCH__ >= 530
|
||||
return __hadd2(a, b);
|
||||
#else
|
||||
float2 out{};
|
||||
out.x = __half2float(a.x) + __half2float(b.x);
|
||||
out.y = __half2float(a.y) + __half2float(b.y);
|
||||
return __float22half2_rn(out);
|
||||
#endif
|
||||
}
|
||||
#if __CUDA_ARCH__ < 530
|
||||
template <typename T>
|
||||
__device__ inline T operator+(const T& a, const T& b);
|
||||
template <typename T>
|
||||
__device__ inline T operator*(const T& a, const T& b);
|
||||
template <>
|
||||
__device__ inline half2 operator+(const half2& a, const half2& b) {
|
||||
return __hadd2_with_fallback(a, b);
|
||||
}
|
||||
template <>
|
||||
__device__ inline half2 operator*(const half2& a, const half2& b) {
|
||||
float2 out{};
|
||||
out.x = __half2float(a.x) * __half2float(b.x);
|
||||
out.y = __half2float(a.y) * __half2float(b.y);
|
||||
return __float22half2_rn(out);
|
||||
}
|
||||
template <typename T>
|
||||
__device__ inline T operator+(const T& a, const T& b);
|
||||
template <typename T>
|
||||
__device__ inline T operator/(const T& a, const T& b);
|
||||
template <typename T>
|
||||
__device__ inline T& operator+=(T& a, const T& b); // NOLINT
|
||||
template <typename T>
|
||||
__device__ inline T operator-(const T& a, const T& b);
|
||||
template <typename T>
|
||||
__device__ inline T operator*(const T& a, const T& b);
|
||||
template <>
|
||||
__device__ inline half operator+(const half& a, const half& b) {
|
||||
return __float2half(__half2float(a) + __half2float(b));
|
||||
}
|
||||
template <>
|
||||
__device__ inline half& operator+=(half& a, const half& b) { // NOLINT
|
||||
a = __float2half(__half2float(a) + __half2float(b));
|
||||
return a;
|
||||
}
|
||||
template <>
|
||||
__device__ inline half operator-(const half& a, const half& b) {
|
||||
return __float2half(__half2float(a) - __half2float(b));
|
||||
}
|
||||
template <>
|
||||
__device__ inline half operator*(const half& a, const half& b) {
|
||||
return __float2half(__half2float(a) * __half2float(b));
|
||||
}
|
||||
template <>
|
||||
__device__ inline half operator/(const half& a, const half& b) {
|
||||
return __float2half(__half2float(a) / __half2float(b));
|
||||
}
|
||||
#endif
|
||||
|
||||
template <>
|
||||
__device__ inline half rsqrt(const half& x) {
|
||||
#if __CUDA_ARCH__ >= 530
|
||||
return hrsqrt(x);
|
||||
#else
|
||||
return __float2half(rsqrt(__half2float(x)));
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ inline kv_half operator+(const kv_half& a, const kv_half& b) {
|
||||
const half2 a2 = __halves2half2(a.key, a.value);
|
||||
const half2 b2 = __halves2half2(b.key, b.value);
|
||||
const half2 res = __hadd2_with_fallback(a2, b2);
|
||||
return kv_half(res.x, res.y);
|
||||
}
|
||||
|
||||
__device__ inline kv_half2 operator+(const kv_half2& a, const kv_half2& b) {
|
||||
return kv_half2(__hadd2_with_fallback(a.key, b.key),
|
||||
__hadd2_with_fallback(a.value, b.value));
|
||||
}
|
||||
// Helper Functions
|
||||
template <typename T>
|
||||
using kvp = cub::KeyValuePair<T, T>;
|
||||
template <typename T, typename R, typename P, int TPB>
|
||||
__device__ inline void layerNorm(const kvp<R>& threadData,
|
||||
const int ld,
|
||||
const int offset,
|
||||
const P* beta,
|
||||
const P* gamma,
|
||||
T* output) {
|
||||
// Assuming threadData is already divided by ld
|
||||
using BlockReduce = cub::BlockReduce<kvp<R>, TPB>;
|
||||
__shared__ typename BlockReduce::TempStorage temp_storage;
|
||||
__shared__ R mu; // mean
|
||||
__shared__ R rsigma; // 1 / std.dev.
|
||||
const auto sumKV = BlockReduce(temp_storage).Reduce(threadData, cub::Sum());
|
||||
if (threadIdx.x == 0) {
|
||||
mu = sumKV.key;
|
||||
rsigma = rsqrt(sumKV.value - mu * mu);
|
||||
}
|
||||
__syncthreads();
|
||||
for (int i = threadIdx.x; i < ld; i += TPB) {
|
||||
const int idx = offset + i;
|
||||
const R val = output[idx];
|
||||
const R g(gamma[i]);
|
||||
const R b(beta[i]);
|
||||
output[idx] = g * (val - mu) * rsigma + b;
|
||||
}
|
||||
}
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
namespace plugin {
|
||||
// Helper Functions for multihead related plugins
|
||||
template <typename T>
|
||||
__global__ void transpose(T* src,
|
||||
T* dst,
|
||||
const int batch_size,
|
||||
const int seq_len,
|
||||
const int head_num,
|
||||
const int size_per_head) {
|
||||
int batch_id = blockIdx.x / (head_num * seq_len);
|
||||
int seq_id = blockIdx.x % seq_len;
|
||||
int head_id = (blockIdx.x % (head_num * seq_len)) / seq_len;
|
||||
dst[batch_id * (head_num * seq_len * size_per_head) +
|
||||
seq_id * head_num * size_per_head + head_id * size_per_head +
|
||||
threadIdx.x] = src[blockIdx.x * size_per_head + threadIdx.x];
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void TransposeQkvKernel(const int H, const T* input, T* output) {
|
||||
// Input: BxSx3xNxH
|
||||
// Bias: 3xSxB
|
||||
// Output: 3xBxNxSxH
|
||||
int n = threadIdx.y;
|
||||
int s = blockIdx.x;
|
||||
int b = blockIdx.y;
|
||||
int m = blockIdx.z;
|
||||
|
||||
const int N = blockDim.y;
|
||||
const int S = gridDim.x;
|
||||
const int B = gridDim.y;
|
||||
|
||||
const int NH = N * H;
|
||||
const int NHS = NH * S;
|
||||
const int in_offset = n * H + m * NH + s * 3 * NH + b * NHS * 3;
|
||||
const int out_offset = s * H + n * S * H + b * NHS + m * NHS * B;
|
||||
|
||||
const int i = threadIdx.x;
|
||||
output[out_offset + i] = input[in_offset + i];
|
||||
}
|
||||
|
||||
inline void TransposeQKV(const int batch,
|
||||
const int seq_len,
|
||||
const int head_size,
|
||||
const int head_num,
|
||||
const float* input,
|
||||
float* output,
|
||||
cudaStream_t stream) {
|
||||
int scratch_size = batch * head_num * seq_len * seq_len;
|
||||
const dim3 grid(seq_len, batch, 3);
|
||||
if (head_size % 4 == 0 && scratch_size % 4 == 0) {
|
||||
const int h = head_size / 4;
|
||||
const float4* input4 = reinterpret_cast<const float4*>(input);
|
||||
float4* output4 = reinterpret_cast<float4*>(output);
|
||||
const dim3 block(h, head_num, 1);
|
||||
// limit h * head_num to max block size(1024).
|
||||
PADDLE_ENFORCE_LE(h * head_num,
|
||||
1024,
|
||||
common::errors::InvalidArgument(
|
||||
"head_num (%d) * head_size (%d) should <= %d",
|
||||
head_num,
|
||||
head_size,
|
||||
1024 * 4));
|
||||
TransposeQkvKernel<float4><<<grid, block, 0, stream>>>(h, input4, output4);
|
||||
} else if (head_size % 2 == 0 && scratch_size % 2 == 0) {
|
||||
const int h = head_size / 2;
|
||||
const float2* input2 = reinterpret_cast<const float2*>(input);
|
||||
float2* output2 = reinterpret_cast<float2*>(output);
|
||||
const dim3 block(h, head_num, 1);
|
||||
// limit h * head_num to max block size(1024).
|
||||
PADDLE_ENFORCE_LE(h * head_num,
|
||||
1024,
|
||||
common::errors::InvalidArgument(
|
||||
"head_num (%d) * head_size (%d) should <= %d",
|
||||
head_num,
|
||||
head_size,
|
||||
1024 * 2));
|
||||
TransposeQkvKernel<float2><<<grid, block, 0, stream>>>(h, input2, output2);
|
||||
} else {
|
||||
const dim3 block(head_size, head_num, 1);
|
||||
// limit head_size * head_num to max block size(1024).
|
||||
PADDLE_ENFORCE_LE(head_size * head_num,
|
||||
1024,
|
||||
common::errors::InvalidArgument(
|
||||
"head_num (%d) * head_size (%d) should <= %d",
|
||||
head_num,
|
||||
head_size,
|
||||
1024));
|
||||
TransposeQkvKernel<float>
|
||||
<<<grid, block, 0, stream>>>(head_size, input, output);
|
||||
}
|
||||
}
|
||||
|
||||
inline void TransposeQKV(const int batch,
|
||||
const int seq_len,
|
||||
const int head_size,
|
||||
const int head_num,
|
||||
const half* input,
|
||||
half* output,
|
||||
cudaStream_t stream) {
|
||||
int scratch_size = batch * head_num * seq_len * seq_len;
|
||||
const dim3 grid(seq_len, batch, 3);
|
||||
if (head_size % 8 == 0 && scratch_size % 8 == 0) {
|
||||
int h = head_size / 8;
|
||||
const int4* input4 = reinterpret_cast<const int4*>(input);
|
||||
int4* output4 = reinterpret_cast<int4*>(output);
|
||||
dim3 block(h, head_num, 1);
|
||||
// limit h * head_num to max block size(1024).
|
||||
PADDLE_ENFORCE_LE(h * head_num,
|
||||
1024,
|
||||
common::errors::InvalidArgument(
|
||||
"head_num (%d) * head_size (%d) should <= %d",
|
||||
head_num,
|
||||
head_size,
|
||||
1024 * 8));
|
||||
TransposeQkvKernel<int4><<<grid, block, 0, stream>>>(h, input4, output4);
|
||||
} else if (head_size % 2 == 0 && scratch_size % 2 == 0) {
|
||||
const int h = head_size / 2;
|
||||
const half2* input2 = reinterpret_cast<const half2*>(input);
|
||||
half2* output2 = reinterpret_cast<half2*>(output);
|
||||
const dim3 block(h, head_num, 1);
|
||||
// limit h * head_num to max block size(1024).
|
||||
PADDLE_ENFORCE_LE(h * head_num,
|
||||
1024,
|
||||
common::errors::InvalidArgument(
|
||||
"head_num (%d) * head_size (%d) should <= %d",
|
||||
head_num,
|
||||
head_size,
|
||||
1024 * 2));
|
||||
TransposeQkvKernel<half2><<<grid, block, 0, stream>>>(h, input2, output2);
|
||||
} else {
|
||||
const dim3 block(head_size, head_num, 1);
|
||||
// limit head_size * head_num to max block size(1024).
|
||||
PADDLE_ENFORCE_LE(head_size * head_num,
|
||||
1024,
|
||||
common::errors::InvalidArgument(
|
||||
"head_num (%d) * head_size (%d) should <= %d",
|
||||
head_num,
|
||||
head_size,
|
||||
1024));
|
||||
TransposeQkvKernel<half>
|
||||
<<<grid, block, 0, stream>>>(head_size, input, output);
|
||||
}
|
||||
}
|
||||
} // namespace plugin
|
||||
} // namespace tensorrt
|
||||
} // namespace inference
|
||||
} // namespace paddle
|
||||
@@ -0,0 +1,62 @@
|
||||
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
||||
// SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
|
||||
// AFFILIATES. 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cstring>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include "NvInfer.h"
|
||||
#include "NvInferPlugin.h"
|
||||
|
||||
typedef enum {
|
||||
STATUS_SUCCESS = 0,
|
||||
STATUS_FAILURE = 1,
|
||||
STATUS_BAD_PARAM = 2,
|
||||
STATUS_NOT_SUPPORTED = 3,
|
||||
STATUS_NOT_INITIALIZED = 4
|
||||
} pluginStatus_t;
|
||||
|
||||
namespace nvinfer1 {
|
||||
|
||||
class BasePlugin : public IPluginV2 {
|
||||
protected:
|
||||
void setPluginNamespace(const char* libNamespace) noexcept override {
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
const char* getPluginNamespace() const noexcept override {
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
std::string mNamespace;
|
||||
};
|
||||
|
||||
class BaseCreator : public IPluginCreator {
|
||||
public:
|
||||
void setPluginNamespace(const char* libNamespace) noexcept override {
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
const char* getPluginNamespace() const noexcept override {
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
|
||||
protected:
|
||||
std::string mNamespace;
|
||||
};
|
||||
|
||||
} // namespace nvinfer1
|
||||
@@ -0,0 +1,119 @@
|
||||
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
||||
// SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
|
||||
// AFFILIATES. 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cassert>
|
||||
#include <cstring>
|
||||
#include <type_traits>
|
||||
#include <vector>
|
||||
|
||||
#include <iostream>
|
||||
using std::cerr;
|
||||
using std::cout;
|
||||
using std::endl;
|
||||
|
||||
template <typename T>
|
||||
inline void serialize_value(void** buffer, T const& value);
|
||||
|
||||
template <typename T>
|
||||
inline void deserialize_value(void const** buffer,
|
||||
size_t* buffer_size,
|
||||
T* value);
|
||||
|
||||
template <typename T, class Enable = void>
|
||||
struct Serializer {};
|
||||
|
||||
template <typename T>
|
||||
struct Serializer<T,
|
||||
typename std::enable_if<std::is_arithmetic<T>::value ||
|
||||
std::is_enum<T>::value ||
|
||||
std::is_pod<T>::value>::type> {
|
||||
static size_t serialized_size(T const&) { return sizeof(T); }
|
||||
static void serialize(void** buffer, T const& value) {
|
||||
::memcpy(*buffer, &value, sizeof(T));
|
||||
reinterpret_cast<char*&>(*buffer) += sizeof(T);
|
||||
}
|
||||
static void deserialize(void const** buffer, size_t* buffer_size, T* value) {
|
||||
assert(*buffer_size >= sizeof(T));
|
||||
::memcpy(value, *buffer, sizeof(T));
|
||||
reinterpret_cast<char const*&>(*buffer) += sizeof(T);
|
||||
*buffer_size -= sizeof(T);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Serializer<const char*> {
|
||||
static size_t serialized_size(const char* value) { return strlen(value) + 1; }
|
||||
static void serialize(void** buffer, const char* value) {
|
||||
::strcpy(static_cast<char*>(*buffer), value); // NOLINT
|
||||
reinterpret_cast<char*&>(*buffer) += strlen(value) + 1;
|
||||
}
|
||||
static void deserialize(void const** buffer,
|
||||
size_t* buffer_size,
|
||||
const char** value) {
|
||||
*value = static_cast<char const*>(*buffer);
|
||||
size_t data_size = strnlen(*value, *buffer_size) + 1;
|
||||
assert(*buffer_size >= data_size);
|
||||
reinterpret_cast<char const*&>(*buffer) += data_size;
|
||||
*buffer_size -= data_size;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct Serializer<std::vector<T>,
|
||||
typename std::enable_if<std::is_arithmetic<T>::value ||
|
||||
std::is_enum<T>::value ||
|
||||
std::is_pod<T>::value>::type> {
|
||||
static size_t serialized_size(std::vector<T> const& value) {
|
||||
return sizeof(value.size()) + value.size() * sizeof(T);
|
||||
}
|
||||
static void serialize(void** buffer, std::vector<T> const& value) {
|
||||
serialize_value(buffer, value.size());
|
||||
size_t nbyte = value.size() * sizeof(T);
|
||||
::memcpy(*buffer, value.data(), nbyte);
|
||||
reinterpret_cast<char*&>(*buffer) += nbyte;
|
||||
}
|
||||
static void deserialize(void const** buffer,
|
||||
size_t* buffer_size,
|
||||
std::vector<T>* value) {
|
||||
size_t size;
|
||||
deserialize_value(buffer, buffer_size, &size);
|
||||
value->resize(size);
|
||||
size_t nbyte = value->size() * sizeof(T);
|
||||
assert(*buffer_size >= nbyte);
|
||||
::memcpy(value->data(), *buffer, nbyte);
|
||||
reinterpret_cast<char const*&>(*buffer) += nbyte;
|
||||
*buffer_size -= nbyte;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
inline size_t serialized_size(T const& value) {
|
||||
return Serializer<T>::serialized_size(value);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline void serialize_value(void** buffer, T const& value) {
|
||||
return Serializer<T>::serialize(buffer, value);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline void deserialize_value(void const** buffer,
|
||||
size_t* buffer_size,
|
||||
T* value) {
|
||||
return Serializer<T>::deserialize(buffer, buffer_size, value);
|
||||
}
|
||||
Reference in New Issue
Block a user