wehub-resource-sync
308 lines
11 KiB
Plaintext
308 lines
11 KiB
Plaintext
/*
|
|
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 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
|
|
|
|
#include "NvInfer.h"
|
|
#include "common/bertCommon.h"
|
|
#include "common/common.cuh"
|
|
#include "common/serialize.hpp"
|
|
#include "skipLayerNormPlugin.h"
|
|
#include "skipLayerNormPluginLegacy.h"
|
|
|
|
#include <cassert>
|
|
#include <cstring>
|
|
#include <limits>
|
|
#include <vector>
|
|
|
|
using namespace nvinfer1;
|
|
using namespace nvinfer1::plugin;
|
|
|
|
namespace nvinfer1
|
|
{
|
|
namespace plugin
|
|
{
|
|
namespace bert
|
|
{
|
|
|
|
template <int32_t TPB, int32_t VPT, bool hasBias>
|
|
__global__ void skiplnDQQ(int32_t const ld, int8_t const* input, int8_t const* skip, int8_t* output, __half const* beta,
|
|
__half const* gamma, __half const* bias, float const dqScaleIn, float const dqScaleSkip, float const qScale)
|
|
{
|
|
int32_t const idx = ld * blockIdx.x + threadIdx.x * VPT;
|
|
// 4 * 1024 * 4 * 2 Bytes = 16KB per block
|
|
int8_t inLocal[VPT];
|
|
int8_t skipLocal[VPT];
|
|
|
|
__half inLocalDQ[VPT]; // dequantized input + skip + bias
|
|
__half biasLocal[VPT]; // bias and beta
|
|
__half gammaLocal[VPT];
|
|
copy<sizeof(int8_t) * VPT>(&input[idx], inLocal);
|
|
copy<sizeof(int8_t) * VPT>(&skip[idx], skipLocal);
|
|
copy<sizeof(__half) * VPT>(&bias[threadIdx.x * VPT], biasLocal);
|
|
__half2 loc = __floats2half2_rn(0.f, 0.f); // accumulator
|
|
|
|
const __half rld = __half(1) / __half(ld);
|
|
#pragma unroll
|
|
for (int32_t it = 0; it < VPT; it++)
|
|
{
|
|
// DQ input and skip
|
|
float const tmpIn = inLocal[it];
|
|
float const tmpSkip = skipLocal[it];
|
|
inLocalDQ[it] = dqScaleIn * tmpIn + dqScaleSkip * tmpSkip;
|
|
|
|
if (hasBias)
|
|
inLocalDQ[it] += biasLocal[it];
|
|
const __half tmp = rld * inLocalDQ[it];
|
|
const __half2 tmp2 = __halves2half2(tmp, tmp * inLocalDQ[it]);
|
|
loc = loc + tmp2;
|
|
}
|
|
// load parameters
|
|
copy<sizeof(__half) * VPT>(&beta[threadIdx.x * VPT], biasLocal);
|
|
copy<sizeof(__half) * VPT>(&gamma[threadIdx.x * VPT], gammaLocal);
|
|
|
|
using BlockReduce = cub::BlockReduce<__half2, TPB>;
|
|
__shared__ typename BlockReduce::TempStorage tempStorage;
|
|
__shared__ __half mu; // mean
|
|
__shared__ __half rsigma; // 1 / std.dev.
|
|
|
|
const __half2 sum2 = BlockReduce(tempStorage).Reduce(loc, [](auto const& lhs, auto const& rhs){return lhs + rhs;});
|
|
|
|
if (threadIdx.x == 0)
|
|
{
|
|
mu = __low2half(sum2);
|
|
rsigma = rsqrt(__high2half(sum2) - mu * mu + std::numeric_limits<half>::epsilon());
|
|
}
|
|
__syncthreads();
|
|
|
|
static_assert(VPT % 4 == 0, "");
|
|
uint32_t outLocal[VPT / 4U];
|
|
#pragma unroll
|
|
for (int32_t it = 0; it < VPT / 4U; it++)
|
|
{
|
|
float const tmp0 = gammaLocal[it * 4 + 0] * (inLocalDQ[it * 4 + 0] - mu) * rsigma + biasLocal[it * 4 + 0];
|
|
float const tmp1 = gammaLocal[it * 4 + 1] * (inLocalDQ[it * 4 + 1] - mu) * rsigma + biasLocal[it * 4 + 1];
|
|
float const tmp2 = gammaLocal[it * 4 + 2] * (inLocalDQ[it * 4 + 2] - mu) * rsigma + biasLocal[it * 4 + 2];
|
|
float const tmp3 = gammaLocal[it * 4 + 3] * (inLocalDQ[it * 4 + 3] - mu) * rsigma + biasLocal[it * 4 + 3];
|
|
outLocal[it] = float4_to_char4(tmp0 * qScale, tmp1 * qScale, tmp2 * qScale, tmp3 * qScale);
|
|
}
|
|
|
|
copy<sizeof(int8_t) * VPT>(outLocal, &output[idx]);
|
|
}
|
|
|
|
template <typename T, int32_t TPB, int32_t VPT, bool hasBias>
|
|
__global__ void skipln_vec(
|
|
int32_t const ld, const T* input, const T* skip, T* output, const T* beta, const T* gamma, const T* bias)
|
|
{
|
|
int32_t const idx = ld * blockIdx.x + threadIdx.x * VPT;
|
|
// 4 * 1024 * 4 * 2 Bytes = 16KB per block
|
|
T inLocal[VPT];
|
|
T skipLocal[VPT];
|
|
T biasLocal[VPT];
|
|
// T gammaLocal[VPT];
|
|
copy<sizeof(T) * VPT>(&input[idx], inLocal);
|
|
copy<sizeof(T) * VPT>(&skip[idx], skipLocal);
|
|
copy<sizeof(T) * VPT>(&bias[threadIdx.x * VPT], biasLocal);
|
|
T local = 0.f;
|
|
T local2 = 0.f;
|
|
|
|
const T rld = T(1) / T(ld);
|
|
#pragma unroll
|
|
for (int32_t it = 0; it < VPT; it++)
|
|
{
|
|
inLocal[it] += skipLocal[it];
|
|
if (hasBias)
|
|
inLocal[it] += biasLocal[it];
|
|
const T tmp = rld * inLocal[it];
|
|
local += tmp;
|
|
local2 += tmp * inLocal[it];
|
|
}
|
|
|
|
copy<sizeof(T) * VPT>(&beta[threadIdx.x * VPT], biasLocal);
|
|
copy<sizeof(T) * VPT>(&gamma[threadIdx.x * VPT], skipLocal);
|
|
|
|
using BlockReduce = cub::BlockReduce<kvp<T>, TPB>;
|
|
__shared__ typename BlockReduce::TempStorage tempStorage;
|
|
__shared__ T mu; // mean
|
|
__shared__ T rsigma; // 1 / std.dev.
|
|
|
|
auto const sumKV = BlockReduce(tempStorage).Reduce(kvp<T>(local, local2), [](auto const& lhs, auto const& rhs){return lhs + rhs;});
|
|
|
|
if (threadIdx.x == 0)
|
|
{
|
|
mu = sumKV.key;
|
|
rsigma = rsqrt(sumKV.value - mu * mu + std::numeric_limits<T>::epsilon());
|
|
}
|
|
__syncthreads();
|
|
///*
|
|
#pragma unroll
|
|
for (int32_t it = 0; it < VPT; it++)
|
|
{
|
|
inLocal[it] = skipLocal[it] * (inLocal[it] - mu) * rsigma + biasLocal[it];
|
|
}
|
|
/* */
|
|
|
|
copy<sizeof(T) * VPT>(inLocal, &output[idx]);
|
|
}
|
|
|
|
template <typename T, unsigned TPB, bool hasBias>
|
|
__global__ void skipLayerNormKernelSmall(
|
|
int32_t const ld, const T* input, const T* skip, const T* beta, const T* gamma, T* output, const T* bias)
|
|
{
|
|
|
|
const T rld = T(1) / T(ld);
|
|
int32_t const offset = blockIdx.x * ld;
|
|
|
|
// reduce x and x^2
|
|
kvp<T> threadData(0, 0);
|
|
int32_t const idx = offset + threadIdx.x;
|
|
T val = 0;
|
|
|
|
if (threadIdx.x < ld)
|
|
{
|
|
|
|
val = input[idx] + skip[idx];
|
|
if (hasBias)
|
|
{
|
|
val += bias[threadIdx.x];
|
|
}
|
|
|
|
const T rldval = rld * val;
|
|
threadData = threadData + kvp<T>(rldval, rldval * val);
|
|
}
|
|
|
|
layerNormSmall<T, T, TPB>(val, threadData, ld, idx, beta, gamma, output);
|
|
}
|
|
|
|
template <typename T, unsigned TPB, bool hasBias>
|
|
__global__ void skipLayerNormKernel(
|
|
int32_t const ld, const T* input, const T* skip, const T* beta, const T* gamma, T* output, const T* bias)
|
|
{
|
|
const T rld = T(1) / T(ld);
|
|
int32_t const offset = blockIdx.x * ld;
|
|
|
|
// reduce x and x^2
|
|
kvp<T> threadData(0, 0);
|
|
|
|
for (int32_t i = threadIdx.x; i < ld; i += TPB)
|
|
{
|
|
int32_t const idx = offset + i;
|
|
T val = T(input[idx]) + T(skip[idx]);
|
|
|
|
if (hasBias)
|
|
{
|
|
val += T(bias[i]);
|
|
}
|
|
const T rldval = rld * val;
|
|
threadData = threadData + kvp<T>(rldval, rldval * val);
|
|
output[idx] = val;
|
|
}
|
|
|
|
layerNorm<T, T, T, TPB>(threadData, ld, offset, beta, gamma, output);
|
|
}
|
|
|
|
template <bool hasBias>
|
|
int32_t computeSkipLayerNormDQQ(cudaStream_t stream, int32_t const ld, int32_t const n, int8_t const* input,
|
|
int8_t const* skip, __half const* beta, __half const* gamma, int8_t* output, __half const* bias,
|
|
float const dqScaleIn, float const dqScaleSkip, float const qScale)
|
|
{
|
|
// this must be true because n is the total size of the tensor
|
|
PLUGIN_VALIDATE(n % ld == 0);
|
|
|
|
int32_t const gridSize = n / ld;
|
|
// we're limited by the size of the parameters, i.e. 8-wide instead of 16
|
|
constexpr int32_t VPT = 16 / sizeof(__half);
|
|
if (ld == 768)
|
|
{
|
|
constexpr int32_t TPB = 768 / VPT;
|
|
skiplnDQQ<TPB, VPT, hasBias>
|
|
<<<gridSize, TPB, 0, stream>>>(ld, input, skip, output, beta, gamma, bias, dqScaleIn, dqScaleSkip, qScale);
|
|
}
|
|
else if (ld == 1024)
|
|
{
|
|
constexpr int32_t TPB = 1024 / VPT;
|
|
skiplnDQQ<TPB, VPT, hasBias>
|
|
<<<gridSize, TPB, 0, stream>>>(ld, input, skip, output, beta, gamma, bias, dqScaleIn, dqScaleSkip, qScale);
|
|
}
|
|
else
|
|
{
|
|
// TODO need to implement this
|
|
PLUGIN_ERROR(("SkipLayerNormDQQ - FATAL: unsupported hidden layer size: " + std::to_string(ld)).c_str());
|
|
}
|
|
PLUGIN_CHECK(cudaPeekAtLastError());
|
|
|
|
return 0;
|
|
}
|
|
|
|
template <typename T, bool hasBias>
|
|
int32_t computeSkipLayerNorm(cudaStream_t stream, int32_t const ld, int32_t const n, const T* input, const T* skip,
|
|
const T* beta, const T* gamma, T* output, const T* bias)
|
|
{
|
|
|
|
// this must be true because n is the total size of the tensor
|
|
PLUGIN_VALIDATE(n % ld == 0);
|
|
int32_t const gridSize = n / ld;
|
|
constexpr int32_t VPT = 16 / sizeof(T);
|
|
if (ld <= 32)
|
|
{
|
|
constexpr int32_t blockSize = 32;
|
|
skipLayerNormKernelSmall<T, blockSize, hasBias>
|
|
<<<gridSize, blockSize, 0, stream>>>(ld, input, skip, beta, gamma, output, bias);
|
|
}
|
|
else if (ld == 768)
|
|
{
|
|
constexpr int32_t TPB = 768 / VPT;
|
|
skipln_vec<T, TPB, VPT, hasBias><<<gridSize, TPB, 0, stream>>>(ld, input, skip, output, beta, gamma, bias);
|
|
}
|
|
else if (ld == 1024)
|
|
{
|
|
constexpr int32_t TPB = 1024 / VPT;
|
|
skipln_vec<T, TPB, VPT, hasBias><<<gridSize, TPB, 0, stream>>>(ld, input, skip, output, beta, gamma, bias);
|
|
}
|
|
else
|
|
{
|
|
constexpr int32_t blockSize = 256;
|
|
skipLayerNormKernel<T, blockSize, hasBias>
|
|
<<<gridSize, blockSize, 0, stream>>>(ld, input, skip, beta, gamma, output, bias);
|
|
}
|
|
PLUGIN_CHECK(cudaPeekAtLastError());
|
|
|
|
return 0;
|
|
}
|
|
|
|
template int32_t computeSkipLayerNormDQQ<true>(cudaStream_t stream, int32_t const ld, int32_t const n,
|
|
int8_t const* input, int8_t const* skip, __half const* beta, __half const* gamma, int8_t* output,
|
|
__half const* bias, float const dqScaleIn, float const dqScaleSkip, float const qScale);
|
|
template int32_t computeSkipLayerNormDQQ<false>(cudaStream_t stream, int32_t const ld, int32_t const n,
|
|
int8_t const* input, int8_t const* skip, __half const* beta, __half const* gamma, int8_t* output,
|
|
__half const* bias, float const dqScaleIn, float const dqScaleSkip, float const qScale);
|
|
|
|
template int32_t computeSkipLayerNorm<float, true>(cudaStream_t, int32_t const, int32_t const, float const*,
|
|
float const*, float const*, float const*, float*, float const*);
|
|
template int32_t computeSkipLayerNorm<float, false>(cudaStream_t, int32_t const, int32_t const, float const*,
|
|
float const*, float const*, float const*, float*, float const*);
|
|
template int32_t computeSkipLayerNorm<half, true>(
|
|
cudaStream_t, int32_t const, int32_t const, half const*, half const*, half const*, half const*, half*, half const*);
|
|
template int32_t computeSkipLayerNorm<half, false>(
|
|
cudaStream_t, int32_t const, int32_t const, half const*, half const*, half const*, half const*, half*, half const*);
|
|
|
|
} // namespace bert
|
|
} // namespace plugin
|
|
} // namespace nvinfer1
|
|
#endif // CUDA_VERSION >= 10010
|