chore: import upstream snapshot with attribution

plugin/skipLayerNormPlugin/skipLayerNormKernel.cu

@@ -0,0 +1,307 @@
+/*
+ * 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