/*
 * 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.
 */

#ifndef INSTANCE_NORM_COMMON_H
#define INSTANCE_NORM_COMMON_H

#include "common/plugin.h"
#include <stdint.h>
using namespace nvinfer1::pluginInternal;

#define DEVICE_FUNCTION static inline __device__

template <typename T, int32_t ELEMENTS_PER_LDG>
struct PackedStorage
{
    enum
    {
        PACKED_ELEMENTS_PER_LDG = ELEMENTS_PER_LDG
    };
    typedef T Type;
};

template <int32_t ELEMENTS_PER_LDG>
struct PackedStorage<uint16_t, ELEMENTS_PER_LDG>
{
    enum
    {
        PACKED_ELEMENTS_PER_LDG = ELEMENTS_PER_LDG / 2
    };
    typedef int32_t Type;
};

template <int32_t ELEMENTS_PER_LDG>
struct PackedStorage<int8_t, ELEMENTS_PER_LDG>
{
    enum
    {
        PACKED_ELEMENTS_PER_LDG = ELEMENTS_PER_LDG / 4
    };
    typedef int32_t Type;
};

template <int32_t N>
DEVICE_FUNCTION void fromFloat(int32_t (&dst)[N], float const (&src)[2 * N])
{
#pragma unroll
    for (int32_t i = 0; i < N; ++i)
    {
        uint16_t lo, hi;
        asm volatile("cvt.rn.f16.f32 %0, %1;" : "=h"(lo) : "f"(src[2 * i + 0]));
        asm volatile("cvt.rn.f16.f32 %0, %1;" : "=h"(hi) : "f"(src[2 * i + 1]));
        asm volatile("mov.b32 %0, {%1, %2};" : "=r"(dst[i]) : "h"(lo), "h"(hi));
    }
}

template <int32_t N>
DEVICE_FUNCTION void fromFloat(int32_t (&dst)[N], float const (&src)[4 * N], float scale)
{
    union Pack_t
    {
        int8_t x[4];
        int32_t val;
    };
#pragma unroll
    for (int32_t i = 0; i < N; ++i)
    {
        Pack_t packed;
#pragma unroll
        for (int32_t ii = 0; ii < 4; ii++)
        {
            packed.x[ii] = __float_as_int(min(max(src[4 * i + ii] * scale + 12582912.0F, 12582785.0F), 12583039.0F));
        }
        dst[i] = packed.val;
    }
}

template <int32_t N>
DEVICE_FUNCTION void fromFloat(float (&dst)[N], float const (&src)[N])
{
#pragma unroll
    for (int32_t i = 0; i < N; ++i)
    {
        dst[i] = src[i];
    }
}

template <int32_t N, bool DO_SCALE = false>
DEVICE_FUNCTION void toFloat(float (&dst)[2 * N], int32_t (&src)[N], float scale = 1.f)
{
#pragma unroll
    for (int32_t i = 0; i < N; ++i)
    {
        uint16_t lo, hi;
        asm volatile("mov.b32 {%0, %1}, %2;" : "=h"(lo), "=h"(hi) : "r"(src[i]));
        asm volatile("cvt.f32.f16 %0, %1;" : "=f"(dst[2 * i + 0]) : "h"(lo));
        asm volatile("cvt.f32.f16 %0, %1;" : "=f"(dst[2 * i + 1]) : "h"(hi));
    }
}

template <int32_t N, bool DO_SCALE = false>
DEVICE_FUNCTION void toFloat(float (&dst)[4 * N], int32_t (&src)[N], float scale = 1.f)
{
    union Pack_t
    {
        int8_t x[4];
        int32_t val;
    };
#pragma unroll
    for (int32_t i = 0; i < N; ++i)
    {
        Pack_t packed;
        packed.val = src[i];
#pragma unroll
        for (int32_t ii = 0; ii < 4; ++ii)
        {
            dst[4 * i + ii]
                = (DO_SCALE) ? __int2float_rn((int32_t) packed.x[ii]) * scale : __int2float_rn((int32_t) packed.x[ii]);
        }
    }
}

template <int32_t N, bool DO_SCALE = false>
DEVICE_FUNCTION void toFloat(float (&dst)[N], float (&src)[N], float scale = 1.f)
{
#pragma unroll
    for (int32_t i = 0; i < N; ++i)
    {
        dst[i] = (DO_SCALE) ? src[i] * scale : src[i];
    }
}

template <typename T>
DEVICE_FUNCTION void ldg(int32_t (&dst)[1], T const* gmem)
{
    dst[0] = __ldg((int32_t const*) gmem);
}

template <typename T>
DEVICE_FUNCTION void ldgStream(int32_t (&dst)[1], T const* gmem)
{
    uint32_t tmp;
    asm volatile("ld.global.cs.nc.s32 %0, [%1];" : "=r"(tmp) : "l"((uint32_t const*) gmem));
    dst[0] = tmp;
}

template <typename T>
DEVICE_FUNCTION void ldg(int32_t (&dst)[2], T const* gmem)
{
    int2 tmp = __ldg((int2 const*) gmem);
    dst[0] = tmp.x;
    dst[1] = tmp.y;
}

template <typename T>
DEVICE_FUNCTION void ldgStream(int32_t (&dst)[2], T const* gmem)
{
    int2 tmp;
    asm volatile("ld.global.cs.nc.v2.s32 {%0,%1}, [%2];" : "=r"(tmp.x), "=r"(tmp.y) : "l"((int2 const*) gmem));
    dst[0] = tmp.x;
    dst[1] = tmp.y;
}

DEVICE_FUNCTION void ldg(int32_t (&dst)[2], uint16_t const* gmem)
{
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 320
    int2 tmp = __ldg((int2 const*) gmem);
    dst[0] = tmp.x;
    dst[1] = tmp.y;
#endif
}

DEVICE_FUNCTION void ldgStream(int32_t (&dst)[2], uint16_t const* gmem)
{
    int2 tmp;
    asm volatile("ld.global.cs.nc.v2.s32 {%0,%1}, [%2];" : "=r"(tmp.x), "=r"(tmp.y) : "l"((int2 const*) gmem));
    dst[0] = tmp.x;
    dst[1] = tmp.y;
}

template <int32_t N>
DEVICE_FUNCTION void ldg(float (&dst)[N], uint16_t const* gmem)
{
    int32_t tmp[N / 2];
    ldg(tmp, gmem);
    toFloat(dst, tmp);
}

template <int32_t N>
DEVICE_FUNCTION void ldgStream(float (&dst)[N], uint16_t const* gmem)
{
    int32_t tmp[N / 2];
    ldgStream(tmp, gmem);
    toFloat(dst, tmp);
}

template <int32_t N>
DEVICE_FUNCTION void ldg(float (&dst)[N], int8_t const* gmem)
{
    int32_t tmp[N / 4];
    ldg(tmp, gmem);
    toFloat(dst, tmp);
}

template <int32_t N>
DEVICE_FUNCTION void ldgStream(float (&dst)[N], int8_t const* gmem)
{
    int32_t tmp[N / 4];
    ldgStream(tmp, gmem);
    toFloat(dst, tmp);
}

template <typename T>
DEVICE_FUNCTION void stg(T* gmem, int32_t (&src)[1])
{
    reinterpret_cast<int32_t*>(gmem)[0] = src[0];
}

template <typename T>
DEVICE_FUNCTION void stgStream(T* gmem, int32_t (&src)[1])
{
    uint32_t tmp = src[0];
    asm volatile("st.global.cs.s32 [%0], %1;" ::"l"((uint32_t*) gmem), "r"(tmp));
}

template <typename T>
DEVICE_FUNCTION void stg(T* gmem, int32_t (&src)[2])
{
    reinterpret_cast<int2*>(gmem)[0] = make_int2(src[0], src[1]);
}

template <typename T>
DEVICE_FUNCTION void stgStream(T* gmem, int32_t (&src)[2])
{
    asm volatile("st.global.cs.v2.s32 [%0], {%1,%2};" ::"l"((uint32_t*) gmem), "r"(src[0]), "r"(src[1]));
}

template <int32_t N>
DEVICE_FUNCTION void stg(uint16_t* gmem, float (&src)[N], float scale)
{
    int32_t tmp[N / 2];
    fromFloat(tmp, src);
    stg(gmem, tmp);
}

template <int32_t N>
DEVICE_FUNCTION void stgStream(uint16_t* gmem, float (&src)[N], float scale)
{
    int32_t tmp[N / 2];
    fromFloat(tmp, src);
    stgStream(gmem, tmp);
}

template <int32_t N>
DEVICE_FUNCTION void stg(int8_t* gmem, float (&src)[N], float scale)
{
    int32_t tmp[N / 4];
    fromFloat(tmp, src, scale);
    stg(gmem, tmp);
}

template <int32_t N>
DEVICE_FUNCTION void stgStream(int8_t* gmem, float (&src)[N], float scale)
{
    int32_t tmp[N / 4];
    fromFloat(tmp, src, scale);
    stg(gmem, tmp);
}

DEVICE_FUNCTION void readFromGmem(float (&dst)[2], float const* gmem, int32_t idx)
{
    float2 tmp = __ldg((float2*) &gmem[2 * idx]);
    dst[0] = tmp.x;
    dst[1] = tmp.y;
}

DEVICE_FUNCTION void readFromGmem(float (&dst)[4], float const* gmem, int32_t idx)
{
    float4 tmp = __ldg((float4*) &gmem[4 * idx]);
    dst[0] = tmp.x;
    dst[1] = tmp.y;
    dst[2] = tmp.z;
    dst[3] = tmp.w;
}

template <int32_t N>
DEVICE_FUNCTION void readFromGmem(float (&dst)[N], __half const* gmem, int32_t idx)
{
    int32_t ival[N / 2];
    if (N == 4)
        reinterpret_cast<int2*>(ival)[0] = __ldg((int2*) &gmem[4 * idx]);
    else
        reinterpret_cast<int32_t*>(ival)[0] = __ldg((int32_t*) &gmem[2 * idx]);
#pragma unroll
    for (int32_t i = 0; i < N / 2; ++i)
    {
        uint16_t lo, hi;
        asm volatile("mov.b32 {%0, %1}, %2;" : "=h"(lo), "=h"(hi) : "r"(ival[i]));
        asm volatile("cvt.f32.f16 %0, %1;" : "=f"(dst[2 * i + 0]) : "h"(lo));
        asm volatile("cvt.f32.f16 %0, %1;" : "=f"(dst[2 * i + 1]) : "h"(hi));
    }
}

DEVICE_FUNCTION void readFromSmem(float (&x)[2], float const* smem, int32_t idx)
{
    float2 tmp = *(float2 const*) &smem[2 * idx];
    x[0] = tmp.x;
    x[1] = tmp.y;
}

DEVICE_FUNCTION void readFromSmem(float (&x)[4], float const* smem, int32_t idx)
{
    float4 tmp = *(float4 const*) &smem[4 * idx];
    x[0] = tmp.x;
    x[1] = tmp.y;
    x[2] = tmp.z;
    x[3] = tmp.w;
}

DEVICE_FUNCTION void readFromSmem(int32_t (&x)[1], int32_t const* smem, int32_t idx)
{
    x[0] = smem[idx];
}

DEVICE_FUNCTION void readFromSmem(int32_t (&x)[2], int32_t const* smem, int32_t idx)
{
    int2 tmp = *(int2 const*) &smem[2 * idx];
    x[0] = tmp.x;
    x[1] = tmp.y;
}

DEVICE_FUNCTION void writeToGmem(float* gmem, int32_t idx, float const (&src)[2])
{
    reinterpret_cast<float2*>(&gmem[2 * idx])[0] = make_float2(src[0], src[1]);
}

DEVICE_FUNCTION void writeToGmem(float* gmem, int32_t idx, float const (&src)[4])
{
    reinterpret_cast<float4*>(&gmem[4 * idx])[0] = make_float4(src[0], src[1], src[2], src[3]);
}

template <int32_t N>
DEVICE_FUNCTION void writeToGmem(__half* gmem, int32_t idx, float const (&src)[N])
{
    int32_t ival[N / 2];
#pragma unroll
    for (int32_t i = 0; i < N / 2; ++i)
    {
        uint16_t lo;
        uint16_t hi;
        asm volatile("cvt.rn.f16.f32 %0, %1;" : "=h"(lo) : "f"(src[2 * i + 0]));
        asm volatile("cvt.rn.f16.f32 %0, %1;" : "=h"(hi) : "f"(src[2 * i + 1]));
        asm volatile("mov.b32 %0, {%1, %2};" : "=r"(ival[i]) : "h"(lo), "h"(hi));
    }
    if (N == 4)
    {
        reinterpret_cast<int2*>(&gmem[4 * idx])[0] = make_int2(ival[0], ival[1]);
    }
    else
    {
        reinterpret_cast<int32_t*>(&gmem[2 * idx])[0] = ival[0];
    }
}

DEVICE_FUNCTION void writeToSmem(float* smem, int32_t idx, float const (&x)[2])
{
    reinterpret_cast<float2*>(&smem[2 * idx])[0] = make_float2(x[0], x[1]);
}

DEVICE_FUNCTION void writeToSmem(float* smem, int32_t idx, float const (&x)[4])
{
    reinterpret_cast<float4*>(&smem[4 * idx])[0] = make_float4(x[0], x[1], x[2], x[3]);
}

DEVICE_FUNCTION void writeToSmem(int32_t* smem, int32_t idx, int32_t const (&x)[1])
{
    smem[idx] = x[0];
}

static inline __device__ void writeToSmem(int32_t* smem, int32_t idx, int32_t const (&x)[2])
{
    reinterpret_cast<int2*>(&smem[2 * idx])[0] = make_int2(x[0], x[1]);
}

template <int32_t N>
DEVICE_FUNCTION void zero(int32_t (&dst)[N])
{
#pragma unroll
    for (int32_t i = 0; i < N; ++i)
    {
        dst[i] = 0;
    }
}

template <int32_t N>
DEVICE_FUNCTION void zero(float (&dst)[N])
{
#pragma unroll
    for (int32_t i = 0; i < N; ++i)
    {
        dst[i] = 0.f;
    }
}

template <int32_t N>
DEVICE_FUNCTION void add(float (&x)[N], float const (&y)[N])
{
#pragma unroll
    for (int32_t i = 0; i < N; ++i)
    {
        x[i] += y[i];
    }
}

template <int32_t N>
DEVICE_FUNCTION void normalize(float (&x)[N], float const (&bias)[N], float const (&scale)[N], float const (&m1)[N])
{
#pragma unroll
    for (int32_t i = 0; i < N; ++i)
    {
        x[i] = bias[i] + scale[i] * (x[i] - m1[i]);
    }
}

template <typename Storage>
DEVICE_FUNCTION Storage relu(Storage in, Storage alpha)
{
    Storage zero = (Storage) 0.f;
    return (in < zero) ? in * alpha : in;
}

template <int32_t N>
DEVICE_FUNCTION void reluActivation(float (&x)[N], float alpha)
{
#pragma unroll
    for (int32_t i = 0; i < N; ++i)
    {
        x[i] = relu(x[i], alpha);
    }
}

template <int32_t THREADS_PER_CTA>
DEVICE_FUNCTION void parallelSums_16x2(float* smem, float (&x)[4], int32_t nhw)
{

    // The size of a warp.
    int32_t const THREADS_PER_WARP = 32;
    // The number of warps in a CTA.
    int32_t const WARPS_PER_CTA = THREADS_PER_CTA / THREADS_PER_WARP;
    // The number of threads per pixel.
    int32_t const THREADS_PER_PIXEL = 16;
    // The number of elements per ldg.
    int32_t const ELEMENTS_PER_LDG = 4;
    // The warp decomposition.
    int32_t const warp_id = threadIdx.x / THREADS_PER_WARP;
    int32_t const lane_id = threadIdx.x % THREADS_PER_WARP;

    // Store the values to shared memory.
    writeToSmem(smem, threadIdx.x, x);

    // Compute the parallel sum inside the warp. Use SHFL and reduce the amount of SMEM by 2x?
    __syncwarp();

    // Read the running sum from the other thread in the warp.
    float y[ELEMENTS_PER_LDG];
    if (lane_id < THREADS_PER_PIXEL)
    {
        readFromSmem(y, smem, threadIdx.x + THREADS_PER_PIXEL);
    }

    // Compute the updated sum.
    add(x, y);

    // The data is in SMEM. Do the final reduction.
    __syncthreads();

    // The warp leaders, write to SMEM.
    if (lane_id < THREADS_PER_PIXEL)
    {
        writeToSmem(smem, warp_id * THREADS_PER_PIXEL + lane_id, x);
    }

    // The data is in SMEM. Do the final reduction.
    __syncthreads();

    // The 1st warp does all the work.
    if (warp_id == 0)
    {
        readFromSmem(x, smem, threadIdx.x);
    }

// We do the final reduction each half-warp sequentially reduces the final values.
#pragma unroll
    for (int32_t offset = 1; offset < WARPS_PER_CTA / 2; ++offset)
    {

        // Read the mean and variance from the other pixel.
        if (warp_id == 0)
        {
            readFromSmem(y, smem, threadIdx.x + offset * THREADS_PER_WARP);
        }

        // Compute the updated sum.
        add(x, y);
    }

    // Make sure the data is in SMEM.
    __syncwarp();

    // Store the mean/var for the different pixels. TODO: Use SHFL?
    if (warp_id == 0)
    {
        writeToSmem(smem, threadIdx.x, x);
    }

    // Make sure the data is in SMEM.
    __syncwarp();

    // The first half warp finishes the work.
    if (threadIdx.x < THREADS_PER_PIXEL)
    {
        readFromSmem(y, smem, threadIdx.x + THREADS_PER_PIXEL);
    }

    // Compute the updated sum.
    add(x, y);

    // Make sure the data was read from SMEM.
    __syncwarp();

    // Store the final values.
    if (threadIdx.x < THREADS_PER_PIXEL)
    {
        writeToSmem(smem, threadIdx.x, x);
    }
}

template <int32_t THREADS_PER_CTA>
static inline __device__ void parallelSums_8x4(float* smem, float (&x)[4], int32_t nhw)
{
    // The size of a warp.
    int32_t const THREADS_PER_WARP = 32;
    // The number of warps in a CTA.
    int32_t const WARPS_PER_CTA = THREADS_PER_CTA / THREADS_PER_WARP;
    // The number of threads per pixel.
    int32_t const THREADS_PER_PIXEL = 8;
    // The number of elements per ldg.
    int32_t const ELEMENTS_PER_LDG = 4;
    // The warp decomposition.
    int32_t const warp_id = threadIdx.x / THREADS_PER_WARP;
    int32_t const lane_id = threadIdx.x % THREADS_PER_WARP;

#pragma unroll
    for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
    {
        x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL + lane_id);
        x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL * 2 + lane_id);
    }

    // The warp leaders, write to SMEM.
    if (lane_id < THREADS_PER_PIXEL)
    {
        writeToSmem(smem, warp_id * THREADS_PER_PIXEL + lane_id, x);
    }

    // The data is in SMEM. Do the final reduction.
    __syncthreads();

    // The 1st warp does all the work.
    // We do the final reduction each half-warp sequentially reduces the final values.
    if (warp_id == 0)
    {
        readFromSmem(x, smem, threadIdx.x);

#pragma unroll
        for (int32_t offset = 1; offset < WARPS_PER_CTA / (THREADS_PER_WARP / THREADS_PER_PIXEL); ++offset)
        {
            float y[ELEMENTS_PER_LDG];
            // Read the mean and variance from the other pixel.
            readFromSmem(y, smem, threadIdx.x + offset * THREADS_PER_WARP);
            // Compute the updated sum.
            add(x, y);
        }

        for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
        {
            x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL + lane_id);
            x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL * 2 + lane_id);
        }

        // Make sure the data was read from SMEM.
        __syncwarp();

        // Store the final values.
        if (threadIdx.x < THREADS_PER_PIXEL)
        {
            writeToSmem(smem, threadIdx.x, x);
        }
    }
}

template <int32_t THREADS_PER_CTA, int32_t THREADS_PER_PIXEL, int32_t ELEMENTS_PER_LDG>
DEVICE_FUNCTION void parallelSums(float* smem, float (&x)[ELEMENTS_PER_LDG], int32_t nhw)
{

    // The size of a warp.
    int32_t const THREADS_PER_WARP = 32;
    // The number of warps in a CTA.
    int32_t const WARPS_PER_CTA = THREADS_PER_CTA / THREADS_PER_WARP;
    // The number of pixels computed by a single warp.
    int32_t const PIXELS_PER_WARP = THREADS_PER_WARP / THREADS_PER_PIXEL;

    // The position in the warp.
    int32_t const nhw_in_warp = nhw % PIXELS_PER_WARP;
    // The C in the warp.
    int32_t const c_in_warp = threadIdx.x % THREADS_PER_PIXEL;

    // Store the values to shared memory.
    writeToSmem(smem, threadIdx.x, x);

    // Compute the parallel sums.
    for (int32_t offset = PIXELS_PER_WARP / 2; offset > 0; offset /= 2)
    {

        if ((WARPS_PER_CTA * THREADS_PER_WARP) / THREADS_PER_PIXEL > THREADS_PER_WARP)
        {
            __syncthreads();
        }
        else
        {
            // NOP.
            __syncwarp();
        }

        // Read the running sum from the other thread.
        float y[ELEMENTS_PER_LDG];
        if (nhw_in_warp < offset)
        {
            readFromSmem(y, smem, threadIdx.x + offset * THREADS_PER_PIXEL);
        }

        // Compute the updated sum.
        add(x, y);

        if ((WARPS_PER_CTA * THREADS_PER_WARP) / THREADS_PER_PIXEL > THREADS_PER_WARP)
        {
            __syncthreads();
        }
        else
        {
            // NOP.
            __syncwarp();
        }

        // Update the sum in SMEM.
        if (offset > 1 && nhw_in_warp < offset)
        {
            writeToSmem(smem, threadIdx.x, x);
        }
    }

    // The warps are done. Do the final reduction at the CTA level.
    __syncthreads();

    // The warp leaders, write to SMEM.
    int32_t const idx = (threadIdx.x / THREADS_PER_WARP) * THREADS_PER_PIXEL + c_in_warp;
    if (nhw_in_warp == 0)
    {
        writeToSmem(smem, idx, x);
    }

    // The data is in SMEM. Do the final reduction.
    __syncthreads();

    // Read the 1st element to prepare the work.
    if (nhw < WARPS_PER_CTA / 2)
    {
        readFromSmem(x, smem, threadIdx.x);
    }

    // We have the running mean and running m2. Let's build the mean/var of the CTA.
    for (int32_t offset = WARPS_PER_CTA / 2; offset > 0; offset /= 2)
    {

        if ((WARPS_PER_CTA * THREADS_PER_WARP) / THREADS_PER_PIXEL > THREADS_PER_WARP)
        {
            __syncthreads();
        }
        else
        {
            // NOP.
            __syncwarp();
        }

        // Read the mean and variance from the other pixel.
        float y[ELEMENTS_PER_LDG];
        if (nhw < offset)
        {
            readFromSmem(y, smem, threadIdx.x + offset * THREADS_PER_PIXEL);
        }

        // Compute the updated sum.
        add(x, y);

        if ((WARPS_PER_CTA * THREADS_PER_WARP) / THREADS_PER_PIXEL > THREADS_PER_WARP)
        {
            __syncthreads();
        }
        else
        {
            // NOP.
            __syncwarp();
        }

        // Store the mean/var for the different pixels.
        if (nhw < offset)
        {
            writeToSmem(smem, threadIdx.x, x);
        }
    }
}

template <int32_t THREADS_PER_PIXEL, int32_t ELEMENTS_PER_LDG>
struct ParallelSums
{
    template <int32_t THREADS_PER_CTA>
    DEVICE_FUNCTION void dispatch(float* smem, float (&x)[ELEMENTS_PER_LDG], int32_t nhw)
    {
        parallelSums<THREADS_PER_CTA, THREADS_PER_PIXEL, ELEMENTS_PER_LDG>(smem, x, nhw);
    }
};

template <>
struct ParallelSums<16, 4>
{
    template <int32_t THREADS_PER_CTA>
    DEVICE_FUNCTION void dispatch(float* smem, float (&x)[4], int32_t nhw)
    {
        parallelSums_16x2<THREADS_PER_CTA>(smem, x, nhw);
    }
};

template <>
struct ParallelSums<8, 4>
{
    template <int32_t THREADS_PER_CTA>
    static inline __device__ void dispatch(float* smem, float (&x)[4], int32_t nhw)
    {
        parallelSums_8x4<THREADS_PER_CTA>(smem, x, nhw);
    }
};

namespace
{
int32_t divUp(int32_t m, int32_t n)
{
    PLUGIN_ASSERT(m >= 0);
    PLUGIN_ASSERT(n > 0);
    // Use unsigned arithmetic to preclude overflow.
    auto const mu = static_cast<uint32_t>(m);
    auto const nu = static_cast<uint32_t>(n);
    return (mu + nu - 1U) / nu;
}

cudnnStatus_t convertTrt2cudnnDtype(nvinfer1::DataType trt_dtype, cudnnDataType_t* cudnn_dtype)
{
    switch (trt_dtype)
    {
    case nvinfer1::DataType::kFLOAT: *cudnn_dtype = CUDNN_DATA_FLOAT; break;
    case nvinfer1::DataType::kHALF: *cudnn_dtype = CUDNN_DATA_HALF; break;
    default: return CUDNN_STATUS_BAD_PARAM;
    }
    return CUDNN_STATUS_SUCCESS;
}

} // namespace
template <typename T, int32_t THREADS_PER_CTA>
__global__ __launch_bounds__(THREADS_PER_CTA) void in3dReluActivation(T* dst, T const* src, T alpha, int32_t count)
{
    int32_t idx = blockIdx.x * THREADS_PER_CTA + threadIdx.x;
    if (idx >= count)
    {
        return;
    }

    T val = src[idx];
    dst[idx] = (val < static_cast<T>(0.F)) ? val * alpha : val;
}

#endif // INSTANCE_NORM_COMMON_H