Function parallelSums

plugin/instanceNormalizationPlugin/instanceNormCommon.h:614–732

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);
        }
    }