Function parallelSums_8x4

plugin/instanceNormalizationPlugin/instanceNormCommon.h:550–611

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)
        {