Unsued variable removed.

Unsued variable removed.

/*******************************************************************************
 *
 * MIT License
 *
 * Copyright (c) 2017 Advanced Micro Devices, Inc.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 *
 *******************************************************************************/

#define PPCAT_NX(A, B) A##B
#define PPCAT(A, B) PPCAT_NX(A, B)
#define TWO 2
#define FOUR 4
#define EIGHT 8

#if MIOPEN_USE_FP16 == 1
#define MIO_BN_NODPP 1
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
#define _FLOAT half
#ifndef HALF_MAX
#define MAX_VAL 65504 /* max value */
#else
#define MAX_VAL HALF_MAX
#endif
#endif
#if MIOPEN_USE_FP32 == 1
#define _FLOAT float
#ifndef FLT_MAX
#define MAX_VAL 3.402823466e+38F /* max value */
#else
#define MAX_VAL FLT_MAX
#endif
#endif

#define _FLOAT2 PPCAT(_FLOAT, TWO)
#define _FLOAT4 PPCAT(_FLOAT, FOUR)
#define _FLOAT8 PPCAT(_FLOAT, EIGHT)
#define _AS_FLOAT PPCAT(as_, _FLOAT)

#ifndef MIO_BN_LDSGCN_SIZE
#define MIO_BN_LDSGCN_SIZE 16
#endif

#ifndef MIO_BN_LDS_SIZE
#define MIO_BN_LDS_SIZE 256
#endif

#ifndef MIO_BN_C
#define MIO_BN_C 1
#endif

#ifndef MIO_BN_N
#define MIO_BN_N 1
#endif

#ifndef MIO_BN_NHW
#define MIO_BN_NHW 1
#endif

#ifndef MIO_BN_INHW
#define MIO_BN_INHW 1
#endif

#ifndef MIO_BN_CHW
#define MIO_BN_CHW 1
#endif

#ifndef MIO_BN_HW
#define MIO_BN_HW 1
#endif

#ifndef MIO_BN_GRP0
#define MIO_BN_GRP0 1
#endif

#ifndef MIO_BN_GRP1
#define MIO_BN_GRP1 1
#endif

#ifndef MIO_BN_GRP2
#define MIO_BN_GRP2 1
#endif

#ifndef MIO_BN_NGRPS
#define MIO_BN_NGRPS 1
#endif

#ifndef MIO_BN_NCHW
#define MIO_BN_NCHW 1
#endif

#ifndef MIO_BN_VARIANT
#define MIO_BN_VARIANT 255
#endif

#define MIO_BN_MAXN 512

#ifndef MIO_BN_NODPP
#define MIO_BN_NODPP 0
#elif(MIO_BN_NODPP == 1)
#undef __AMDGCN__
#endif

#ifdef __AMDGCN__
#undef __AMDGCN__
#endif

#define UNUSED __attribute__((__unused__))

// Disable specific warnings
#ifdef __clang__
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wconditional-uninitialized"
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wsometimes-uninitialized"
#endif

#ifndef __AMDGCN__

static inline void ReduceKernel(__local _FLOAT* lcl_mem,
                                unsigned int sum_stride,
                                unsigned int unit_id,
                                unsigned int unit_len)
{
    _FLOAT sum              = (_FLOAT)0.;
    unsigned int lcl_offset = unit_id * unit_len;

#pragma unroll
    for(unsigned int i = 0; i < unit_len; i += sum_stride)
    {
        sum += lcl_mem[lcl_offset + i];
    }
    lcl_mem[lcl_offset] = sum;
}

static inline void
regLDSreduce(_FLOAT* value, __local _FLOAT* data, unsigned int localID, _FLOAT scale)
{
    data[localID] = *value;
    barrier(CLK_LOCAL_MEM_FENCE);
    if(localID < (MIO_BN_LDS_SIZE >> 2))
        ReduceKernel(data, 1, localID, 4);
    barrier(CLK_LOCAL_MEM_FENCE);
    if(localID < (MIO_BN_LDS_SIZE >> 4))
        ReduceKernel(data, 4, localID, 16);
    barrier(CLK_LOCAL_MEM_FENCE);
    if(localID == 0)
        ReduceKernel(data, 16, localID, MIO_BN_LDS_SIZE);
    barrier(CLK_LOCAL_MEM_FENCE);
    *value = data[0] * scale;
}
#endif

#ifdef __AMDGCN__

static inline void dppSimpleRedNoBcast64(_FLOAT* value)
{
    _FLOAT tmp = (_FLOAT)0.;
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x111, 0xF, 0xF, 0));
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x112, 0xF, 0xF, 0));
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x114, 0xF, 0xF, 0));
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x118, 0xF, 0xF, 0));
    tmp = _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x142, 0xF, 0xF, 0));
    *value += tmp;
    tmp = _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x143, 0xF, 0xF, 0));
    *value += tmp;
}

static inline void dppSimpleRedBcast64(_FLOAT* value)
{
    _FLOAT tmp = (_FLOAT)0.;
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x111, 0xF, 0xF, 0));
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x112, 0xF, 0xF, 0));
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x114, 0xF, 0xF, 0));
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x118, 0xF, 0xF, 0));
    tmp = _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x142, 0xF, 0xF, 0));
    *value += tmp;
    tmp = _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x143, 0xF, 0xF, 0));
    *value += tmp;
    barrier(CLK_LOCAL_MEM_FENCE);
    *value = _AS_FLOAT(__builtin_amdgcn_readlane(as_int(*value), 63));
}

#endif

#if(MIO_BN_VARIANT == 0)

#define MIO_BN_SEGTMP (MIO_BN_HW * (MIO_BN_GRP0 / MIO_BN_HW))
#define MIO_BN_SEGMENT ((MIO_BN_SEGTMP > MIO_BN_NHW) ? (MIO_BN_NHW) : (MIO_BN_SEGTMP))
#define MIO_BN_NLOOP ((MIO_BN_NHW + MIO_BN_SEGMENT - 1) / MIO_BN_SEGMENT)
#define MIO_BN_SEGIHW (MIO_BN_SEGMENT / MIO_BN_HW)
#define MIO_BN_NLOOPM (MIO_BN_NLOOP - 1)
#define MIO_BN_SNHW (MIO_BN_NLOOPM * MIO_BN_SEGIHW)

__attribute__((reqd_work_group_size(MIO_BN_GRP0, MIO_BN_GRP1, MIO_BN_GRP2))) __kernel void
BatchNormFwdTrainSpatial(const __global _FLOAT* __restrict in,
                         __global _FLOAT* __restrict out,
                         __constant _FLOAT* __restrict scale,
                         __constant _FLOAT* __restrict bias,
                         _FLOAT INHW,
#if(MIO_RUNNING_RESULT == 1)
                         double expAvgFactor,
                         __global _FLOAT* __restrict resultRunningMean,
                         __global _FLOAT* __restrict resultRunningVariance,
#endif
                         double epsilon
#if(MIO_SAVE_MEAN_VARIANCE == 1)
                         ,
                         __global _FLOAT* __restrict resultSaveMean,
                         __global _FLOAT* __restrict resultSaveInvVariance
#endif
                         )
{

    // SPATIAL
    _FLOAT mean        = (_FLOAT)0.;
    _FLOAT variance    = (_FLOAT)0.;
    _FLOAT invVariance = (_FLOAT)0.;
    _FLOAT pvscale     = (_FLOAT)0.;
    _FLOAT pvbias      = (_FLOAT)0.;
    _FLOAT batchvalues[MIO_BN_NLOOP];

    __local _FLOAT lcl_bias;
    __local _FLOAT lcl_scale;

    unsigned int index  = 0;
    unsigned int lid    = get_local_id(0);
    unsigned int grpid  = get_group_id(0);
    unsigned int chwid  = grpid * MIO_BN_HW + (lid % MIO_BN_HW);
    unsigned int lidihw = lid / MIO_BN_HW;
    unsigned int nid    = 0;

    if(lid == 0)
    {
        lcl_scale = *(scale + grpid);
        lcl_bias  = *(bias + grpid);
    }
    barrier(CLK_LOCAL_MEM_FENCE);

    if(lid < MIO_BN_SEGMENT)
    {
#pragma unroll
        for(unsigned int n = 0; n < MIO_BN_NLOOPM; ++n)
        {
            nid            = n * MIO_BN_SEGIHW + lidihw;
            index          = nid * MIO_BN_CHW + chwid;
            batchvalues[n] = *(in + index);
            mean += batchvalues[n];
            variance = mad(batchvalues[n], batchvalues[n], variance);
        }
        nid                        = MIO_BN_SNHW + lidihw;
        index                      = nid * MIO_BN_CHW + chwid;
        batchvalues[MIO_BN_NLOOPM] = (index < MIO_BN_NCHW) ? *(in + index) : (_FLOAT)0.;
        mean += batchvalues[MIO_BN_NLOOPM];
        variance = mad(batchvalues[MIO_BN_NLOOPM], batchvalues[MIO_BN_NLOOPM], variance);
    }

#ifndef __AMDGCN__
    __local _FLOAT lcl_data[MIO_BN_LDS_SIZE];

    // Reduce mean
    lcl_data[lid] = mean;
    barrier(CLK_LOCAL_MEM_FENCE);
    for(unsigned int red = (MIO_BN_GRP0 >> 1); red > 256; red >>= 1)
    {
        if(lid < red)
            lcl_data[lid] += lcl_data[lid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&mean, lcl_data, lid, (_FLOAT)INHW);
    barrier(CLK_LOCAL_MEM_FENCE);

    // Reduce variance
    lcl_data[lid] = variance;
    barrier(CLK_LOCAL_MEM_FENCE);
    for(unsigned int red = (MIO_BN_GRP0 >> 1); red > 256; red >>= 1)
    {
        if(lid < red)
            lcl_data[lid] += lcl_data[lid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&variance, lcl_data, lid, (_FLOAT)INHW);
    barrier(CLK_LOCAL_MEM_FENCE);
#else

    unsigned int ldsidx = lid >> 6;
    __local _FLOAT lcl_mean[MIO_BN_LDSGCN_SIZE];
    __local _FLOAT lcl_variance[MIO_BN_LDSGCN_SIZE];

    dppSimpleRedNoBcast64(&mean);
    dppSimpleRedNoBcast64(&variance);

    if((lid % 64) == 63)
    {
        lcl_mean[ldsidx]     = mean;
        lcl_variance[ldsidx] = variance;
    }
    barrier(CLK_LOCAL_MEM_FENCE);
    mean = variance = 0.;
#pragma unroll
    for(uint i = 0; i < MIO_BN_LDSGCN_SIZE; i++)
    {
        mean += lcl_mean[i];
        variance += lcl_variance[i];
    }
    mean *= (_FLOAT)INHW;
    variance *= (_FLOAT)INHW;

#endif

    variance    = mad(-mean, mean, variance);
    invVariance = rsqrt(variance + epsilon);
    pvscale     = lcl_scale;
    pvbias      = lcl_bias;

    if(lid < MIO_BN_SEGMENT)
    {
        //==== CALC NORM =======================
        _FLOAT inhat = 0.;

#pragma unroll
        for(unsigned int n = 0; n < MIO_BN_NLOOPM; n++)
        { // apply normalization
            inhat      = (batchvalues[n] - mean) * invVariance;
            nid        = n * MIO_BN_SEGIHW + lidihw;
            index      = nid * MIO_BN_CHW + chwid;
            out[index] = mad(pvscale, inhat, pvbias);
        } // end for

        // Tail of loop
        inhat = (batchvalues[MIO_BN_NLOOPM] - mean) * invVariance;
        nid   = MIO_BN_SNHW + lidihw;
        index = nid * MIO_BN_CHW + chwid;
        if(index < MIO_BN_NCHW)
            out[index] = mad(pvscale, inhat, pvbias);
    }

#if(MIO_SAVE_MEAN_VARIANCE == 1 || MIO_RUNNING_RESULT == 1)
    if(lid == 0)
    {

// Save mean and calculate and save running mean
#if(MIO_SAVE_MEAN_VARIANCE == 1)
        *(resultSaveMean + grpid)        = mean;
        *(resultSaveInvVariance + grpid) = invVariance;
#endif

#if(MIO_RUNNING_RESULT == 1)
        _FLOAT pvt_runMean = *(resultRunningMean + grpid);
        _FLOAT pvt_newRunMean =
            mad((_FLOAT)-expAvgFactor, pvt_runMean, pvt_runMean); // tmp = oldRunMean*(1-factor)
        resultRunningMean[grpid] =
            mad(mean, (_FLOAT)expAvgFactor, pvt_newRunMean); // newMean*factor + tmp
        const _FLOAT adjust = (MIO_BN_NHW == 1)
                                  ? variance
                                  : variance * ((_FLOAT)MIO_BN_NHW / (_FLOAT)(MIO_BN_NHW - 1.0));
        resultRunningVariance[grpid] =
            (1 - (_FLOAT)expAvgFactor) * *(resultRunningVariance + grpid) +
            (_FLOAT)expAvgFactor * adjust;
#endif
    }
#endif
} // end spatial norm

#elif(MIO_BN_VARIANT == 1)

#define MIO_BN_REM (MIO_BN_NHW - ((MIO_BN_NHW / MIO_BN_GRP1) * MIO_BN_GRP1))
#define MIO_BN_LESS (MIO_BN_NHW - MIO_BN_REM)

__attribute__((reqd_work_group_size(MIO_BN_GRP0, MIO_BN_GRP1, MIO_BN_GRP2))) __kernel void
BatchNormFwdTrainSpatial(const __global _FLOAT* __restrict in,
                         __global _FLOAT* __restrict out,
                         __constant _FLOAT* __restrict scale,
                         __constant _FLOAT* __restrict bias,
                         _FLOAT INHW,
#if(MIO_RUNNING_RESULT == 1)
                         double expAvgFactor,
                         __global _FLOAT* __restrict resultRunningMean,
                         __global _FLOAT* __restrict resultRunningVariance,
#endif
                         double epsilon
#if(MIO_SAVE_MEAN_VARIANCE == 1)
                         ,
                         __global _FLOAT* __restrict resultSaveMean,
                         __global _FLOAT* __restrict resultSaveInvVariance
#endif
                         )
{

    // SPATIAL
    _FLOAT mean        = (_FLOAT)0.;
    _FLOAT variance    = (_FLOAT)0.;
    _FLOAT invVariance = (_FLOAT)0.;
    _FLOAT pvscale, pvbias;

    __local _FLOAT lcl_bias;
    __local _FLOAT lcl_scale;

    int index = 0;
    int lid   = get_local_id(1);
    int xgid  = get_global_id(0);
    int grpid = get_group_id(0);
    int chwid = grpid * MIO_BN_HW;
    int nidx  = 0;
    int hwidx = 0;

#if(MIO_BN_NHW < MIO_BN_MAXN)
    _FLOAT input[MIO_BN_NHW];
#endif

    if(lid == 0)
    {
        lcl_scale = *(scale + xgid);
        lcl_bias  = *(bias + xgid);
    }
    barrier(CLK_LOCAL_MEM_FENCE);

#pragma unroll
    for(unsigned int k = lid, lesskey = 0; k < MIO_BN_LESS; k += MIO_BN_GRP1, ++lesskey)
    {
        nidx           = k / MIO_BN_HW;
        hwidx          = k - (nidx * MIO_BN_HW);
        index          = nidx * MIO_BN_CHW + chwid + hwidx;
#if(MIO_BN_NHW < MIO_BN_MAXN)
        input[lesskey] = *(in + index);
        mean += input[lesskey];
        variance = mad(input[lesskey], input[lesskey], variance);
#else
        _FLOAT xin = *(in + index);
        mean += xin;
        variance        = mad(xin, xin, variance);
#endif
    }
#if(MIO_BN_REM)
    unsigned int remkey = lid + MIO_BN_LESS;
    nidx                = remkey / MIO_BN_HW;
    hwidx               = remkey - (nidx * MIO_BN_HW);
    index               = nidx * MIO_BN_CHW + chwid + hwidx;
#if(MIO_BN_NHW < MIO_BN_MAXN)
    input[remkey]       = (index < MIO_BN_NCHW) ? *(in + index) : 0.;
    mean += input[remkey];
    variance = mad(input[remkey], input[remkey], variance);
#else
    _FLOAT xin = (index < MIO_BN_NCHW) ? *(in + index) : 0.;
    mean += xin;
    variance           = mad(xin, xin, variance);
#endif
#endif

// REDUCE MEAN AND VARIANCE -----------------------
#ifndef __AMDGCN__
    local _FLOAT lcl_data[MIO_BN_LDS_SIZE];
    lcl_data[lid] = mean;
    barrier(CLK_LOCAL_MEM_FENCE);

    for(unsigned int red = (MIO_BN_GRP1 >> 1); red > 256; red >>= 1)
    {
        if(lid < red)
            lcl_data[lid] += lcl_data[lid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&mean, lcl_data, lid, (_FLOAT)INHW);

    barrier(CLK_LOCAL_MEM_FENCE);
    lcl_data[lid] = variance;
    barrier(CLK_LOCAL_MEM_FENCE);

    for(unsigned int red = (MIO_BN_GRP1 >> 1); red > 256; red >>= 1)
    {
        if(lid < red)
            lcl_data[lid] += lcl_data[lid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&variance, lcl_data, lid, (_FLOAT)INHW);
    barrier(CLK_LOCAL_MEM_FENCE);

#else
    unsigned int ldsidx = lid >> 6;
    __local _FLOAT lcl_mean[MIO_BN_LDSGCN_SIZE];
    __local _FLOAT lcl_variance[MIO_BN_LDSGCN_SIZE];

    dppSimpleRedNoBcast64(&mean);
    dppSimpleRedNoBcast64(&variance);
    if((lid % 64) == 63)
    {
        lcl_mean[ldsidx]     = mean;
        lcl_variance[ldsidx] = variance;
    }
    barrier(CLK_LOCAL_MEM_FENCE);
    mean = variance = 0.;
#pragma unroll
    for(uint i = 0; i < MIO_BN_LDSGCN_SIZE; i++)
    {
        mean += lcl_mean[i];
        variance += lcl_variance[i];
    }
    mean *= (_FLOAT)INHW;
    variance *= (_FLOAT)INHW;
#endif
    // REDUCTION COMPLETE ---------------------------

    variance    = mad(-mean, mean, variance);
    invVariance = rsqrt(variance + epsilon);

    pvscale = lcl_scale;
    pvbias  = lcl_bias;

#pragma unroll
    for(unsigned int k = lid, lesskey = 0; k < MIO_BN_LESS; k += MIO_BN_GRP1, ++lesskey)
    {
        nidx       = k / MIO_BN_HW;
        hwidx      = k - (nidx * MIO_BN_HW);
        index      = nidx * MIO_BN_CHW + chwid + hwidx;

#if(MIO_BN_NHW < MIO_BN_MAXN)
        out[index] = mad(pvscale, (input[lesskey] - mean) * invVariance, pvbias);
#else
        out[index] = mad(pvscale, (*(in + index) - mean) * invVariance, pvbias);
#endif
    } // end for
#if(MIO_BN_REM)
    nidx  = remkey / MIO_BN_HW;
    hwidx = remkey - (nidx * MIO_BN_HW);
    index = nidx * MIO_BN_CHW + chwid + hwidx;
    if(index < MIO_BN_NCHW)
    {
#if(MIO_BN_NHW < MIO_BN_MAXN)
        *(out + index) = mad(pvscale, (input[remkey] - mean) * invVariance, pvbias);
#else
        *(out + index) = mad(pvscale, (*(in + index) - mean) * invVariance, pvbias);
#endif
    }
#endif

#if(MIO_SAVE_MEAN_VARIANCE == 1 || MIO_RUNNING_RESULT == 1)
    if(get_global_id(1) == 0)
    {

// Save mean and calculate and save running mean
#if(MIO_SAVE_MEAN_VARIANCE == 1)
        *(resultSaveMean + xgid)        = mean;
        *(resultSaveInvVariance + xgid) = invVariance;
#endif

#if(MIO_RUNNING_RESULT == 1)
        _FLOAT pvt_runMean              = *(resultRunningMean + xgid);
        _FLOAT pvt_newRunMean =
            mad((_FLOAT)-expAvgFactor, pvt_runMean, pvt_runMean); // tmp = oldRunMean*(1-factor)
        resultRunningMean[xgid] =
            mad(mean, (_FLOAT)expAvgFactor, pvt_newRunMean); // newMean*factor + tmp
        const _FLOAT adjust = (MIO_BN_NHW == 1)
                                  ? variance
                                  : variance * ((_FLOAT)MIO_BN_NHW / (_FLOAT)(MIO_BN_NHW - 1.0));
        resultRunningVariance[xgid] = (1 - (_FLOAT)expAvgFactor) * *(resultRunningVariance + xgid) +
                                      (_FLOAT)expAvgFactor * adjust;
#endif
    }
#endif
} // end spatial norm

#elif(MIO_BN_VARIANT == 2) // MULTI-KERNEL reduction for > 33M elements

__attribute__((reqd_work_group_size(MIO_BN_GRP0, MIO_BN_GRP1, MIO_BN_GRP2))) __kernel void
BatchNormFwdTrainSpatialNorm(const __global _FLOAT* __restrict in,
                             __global _FLOAT* __restrict out,
                             const __global _FLOAT* __restrict scale,
                             const __global _FLOAT* __restrict bias)
{

    // SPATIAL
    _FLOAT mean        = (_FLOAT)0.;
    _FLOAT invVariance = (_FLOAT)0.;
    _FLOAT inhat       = (_FLOAT)0.;
    _FLOAT pvt_scale   = (_FLOAT)0.;
    _FLOAT pvt_bias    = (_FLOAT)0.;
    __local _FLOAT lcl_mean, lcl_ivar, lcl_scale, lcl_bias;

    unsigned int ygrp_id = get_group_id(1);
    unsigned int xgid    = get_global_id(0);
    unsigned int ygid    = get_global_id(1);
    unsigned int ygrp_sz = get_local_size(1);
    unsigned int index;
    unsigned int cidx           = xgid * MIO_BN_HW;
    unsigned int meanstashindex = cidx + ygrp_sz * ygrp_id + 1;
    unsigned int varstashindex  = cidx + ygrp_sz * ygrp_id + 3;

    // #4 apply the normalization :: x_hat = (x_i - mean) / sqrt(variance_accum + epsilon)
    if(get_local_id(1) == 0)
    {
        lcl_scale = *(scale + xgid);
        lcl_bias  = *(bias + xgid);
        lcl_mean  = *(out + meanstashindex); // load stashed mean
        lcl_ivar  = *(out + varstashindex);
    }
    barrier(CLK_LOCAL_MEM_FENCE);

    if(ygid < MIO_BN_HW)
    {
        mean        = lcl_mean;
        invVariance = lcl_ivar;
        pvt_scale   = lcl_scale;
        pvt_bias    = lcl_bias;
#pragma unroll
        for(unsigned int n = 0; n < MIO_BN_N; n++)
        { // apply normalization
            index        = n * MIO_BN_CHW + cidx + ygid;
            _FLOAT inhat = (*(in + index) - mean) * invVariance;
            // #5 Gamma and Beta adjust :: y_i = gamma*x_hat + beta
            out[index] = mad(pvt_scale, inhat, pvt_bias);
        } // end for(n)
    }     // end if(inImgIndex)
} // end spatial norm

__attribute__((reqd_work_group_size(MIO_BN_GRP0, MIO_BN_GRP1, MIO_BN_GRP2))) __kernel void
BatchNormFwdTrainSpatialFinalMeanVariance(
    __global _FLOAT* __restrict meanvarbuff,
    _FLOAT INHW
#if(MIO_RUNNING_RESULT == 1)
    ,
    double expAvgFactor /* input momentum */
    ,
    __global _FLOAT* __restrict resultRunningMean, /*input and output*/
    __global _FLOAT* __restrict resultRunningVariance
#endif
    ,
    double epsilon
#if(MIO_SAVE_MEAN_VARIANCE == 1)
    ,
    __global _FLOAT* __restrict resultSaveMean /*output only*/
    ,
    __global _FLOAT* __restrict resultSaveInvVariance
#endif
    )
{
    _FLOAT variance             = (_FLOAT)0.;
    _FLOAT invVariance          = (_FLOAT)0.;
    _FLOAT mean                 = (_FLOAT)0.;
    unsigned int lid            = get_local_id(1);
    unsigned int ygrp_id        = get_group_id(1);
    unsigned int xgid           = get_global_id(0);
    unsigned int ygrp_sz        = get_local_size(1);
    unsigned int yngrps         = get_num_groups(1);
    unsigned int cidx           = xgid * MIO_BN_HW;
    unsigned int meanstashindex = cidx + ygrp_sz * ygrp_id + 1;
    unsigned int varstashindex  = cidx + ygrp_sz * ygrp_id + 3;
    unsigned int commitID       = 0;

    for(int gn = 0; gn < yngrps; gn++)
    {
        unsigned int offset    = gn * ygrp_sz + lid;
        unsigned int meanindex = cidx + ygrp_sz * offset;
        unsigned int varindex  = cidx + ygrp_sz * offset + 2;
        if(offset < yngrps)
        { // modify to span larger number of groups
            mean += *(meanvarbuff + meanindex);
            variance += *(meanvarbuff + varindex); // load per group variance
        }
    }
#ifndef __AMDGCN__
    __local _FLOAT lcl_data[MIO_BN_LDS_SIZE];
    lcl_data[lid] = mean;
    barrier(CLK_LOCAL_MEM_FENCE);

#if(MIO_BN_NGRPS > 64)
    for(unsigned int red = (MIO_BN_GRP1 >> 1); red > 256; red >>= 1)
    {
        if(lid < red)
            lcl_data[lid] += lcl_data[lid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&mean, lcl_data, lid, (_FLOAT)INHW);
#elif(MIO_BN_NGRPS <= 64)
    regLDSreduce(&mean, lcl_data, lid, (_FLOAT)INHW);
    commitID = 0;
#else
    mean = (_FLOAT)0.;
#pragma unroll
    for(int i = 0; i < MIO_BN_NGRPS; i++)
    {
        mean += lcl_data[i];
    }

#endif

    lcl_data[lid] = variance;
    barrier(CLK_LOCAL_MEM_FENCE);

#if(MIO_BN_NGRPS > 256)
    for(unsigned int red = (MIO_BN_GRP1 >> 1); red > 256; red >>= 1)
    {
        if(lid < red)
            lcl_data[lid] += lcl_data[lid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&variance, lcl_data, lid, (_FLOAT)INHW);
#elif(MIO_BN_NGRPS > 64)
    regLDSreduce(&variance, lcl_data, lid, (_FLOAT)INHW);
#elif(MIO_BN_NGRPS > 16)
    regLDSreduce(&variance, lcl_data, lid, (_FLOAT)INHW);
#else //(MIO_BN_NGRPS <= 16)
    variance = (_FLOAT)0.;
#pragma unroll
    for(int i = 0; i < MIO_BN_NGRPS; i++)
    {
        variance += lcl_data[i];
    }
#endif

#else
    commitID            = 64;
    unsigned int ldsidx = lid >> 6;
    __local _FLOAT lcl_mean[MIO_BN_LDSGCN_SIZE];
    __local _FLOAT lcl_variance[MIO_BN_LDSGCN_SIZE];

    dppSimpleRedNoBcast64(&mean);
    if((lid % 64) == 63)
    {
        lcl_mean[ldsidx] = mean;
    }
    barrier(CLK_LOCAL_MEM_FENCE | CLK_GLOBAL_MEM_FENCE);
    mean = 0.;

#pragma unroll
    for(uint i = 0; i < MIO_BN_LDSGCN_SIZE; i++)
    {
        mean += lcl_mean[i];
    }
    mean *= (_FLOAT)INHW;

    dppSimpleRedNoBcast64(&variance);
    if((lid % 64) == 63)
    {
        lcl_variance[ldsidx] = variance;
    }
    barrier(CLK_LOCAL_MEM_FENCE | CLK_GLOBAL_MEM_FENCE);
    variance = 0.;
#pragma unroll
    for(uint i = 0; i < MIO_BN_LDSGCN_SIZE; i++)
    {
        variance += lcl_variance[i];
    }
    variance *= (_FLOAT)INHW;

#endif
    invVariance = rsqrt(variance + epsilon);
    if(lid == commitID)
    {
        meanvarbuff[meanstashindex] = mean;        // stash mean
        meanvarbuff[varstashindex]  = invVariance; // stash mean
    }

// Save mean and calculate and save running mean
#if(MIO_SAVE_MEAN_VARIANCE == 1 || MIO_RUNNING_RESULT == 1)
    unsigned int ygid = get_global_id(1);
    if(ygid == commitID)
    {
#if(MIO_SAVE_MEAN_VARIANCE == 1)
        resultSaveMean[xgid]        = mean;
        resultSaveInvVariance[xgid] = invVariance;
#endif
#if(MIO_RUNNING_RESULT == 1)
        _FLOAT pvt_runMean          = resultRunningMean[xgid];
        _FLOAT pvt_newRunMean =
            mad((_FLOAT)-expAvgFactor, pvt_runMean, pvt_runMean); // tmp = oldRunMean*(1-factor)
        resultRunningMean[xgid] =
            mad(mean, (_FLOAT)expAvgFactor, pvt_newRunMean); // newMean*factor + tmp

        _FLOAT NHW                  = (_FLOAT)MIO_BN_NHW;
        const _FLOAT adjust         = (MIO_BN_NHW == 1) ? variance : variance * (NHW / (NHW - 1));
        resultRunningVariance[xgid] = (1 - (_FLOAT)expAvgFactor) * *(resultRunningVariance + xgid) +
                                      (_FLOAT)expAvgFactor * adjust;
#endif
    }
#endif
}

__attribute__((reqd_work_group_size(MIO_BN_GRP0, MIO_BN_GRP1, MIO_BN_GRP2))) __kernel void
BatchNormFwdTrainSpatialMeanVariance(const __global _FLOAT* __restrict in,
                                     __global _FLOAT* __restrict mvbuff)
{

    __local _FLOAT lcl_data[MIO_BN_LDS_SIZE];

    unsigned int ylid    = get_local_id(1);
    unsigned int ygrp_id = get_group_id(1);
    unsigned int xgid    = get_global_id(0);
    unsigned int ygid    = get_global_id(1);
    unsigned int ygrp_sz = get_local_size(1);
    unsigned int index;
    unsigned int cidx      = xgid * MIO_BN_HW;
    unsigned int meanindex = cidx + ygrp_sz * ygrp_id;
    unsigned int varindex  = meanindex + 2;
    _FLOAT mean            = (_FLOAT)0.;
    _FLOAT variance        = (_FLOAT)0.;
    _FLOAT value           = (_FLOAT)0.;

    if(ygid < MIO_BN_HW)
    {
#pragma unroll
        for(unsigned int n = 0; n < MIO_BN_N; n++)
        {
            index = n * MIO_BN_CHW + cidx + ygid;
            value = *(in + index);
            mean += value;
            variance = mad(value, value, variance);
        }
    }

#ifdef __AMDGCN__
    unsigned int ldsidx = ylid >> 6;
    __local _FLOAT lcl_mean[MIO_BN_LDSGCN_SIZE];
    __local _FLOAT lcl_variance[MIO_BN_LDSGCN_SIZE];

    dppSimpleRedNoBcast64(&mean);
    dppSimpleRedNoBcast64(&variance);

    if((ylid % 64) == 63)
    {
        lcl_mean[ldsidx]     = mean;
        lcl_variance[ldsidx] = variance;
    }
    barrier(CLK_LOCAL_MEM_FENCE | CLK_GLOBAL_MEM_FENCE);
    mean = variance = 0.;

#pragma unroll
    for(uint i = 0; i < MIO_BN_LDSGCN_SIZE; i++)
    {
        mean += lcl_mean[i];
        variance += lcl_variance[i];
    }

#else
    lcl_data[ylid] = mean;
    barrier(CLK_LOCAL_MEM_FENCE);

    for(unsigned int red = (MIO_BN_GRP1 >> 1); red > 256; red >>= 1)
    {
        if(ylid < red)
            lcl_data[ylid] += lcl_data[ylid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&mean, lcl_data, ylid, 1);
    barrier(CLK_LOCAL_MEM_FENCE);

    lcl_data[ylid] = variance;
    barrier(CLK_LOCAL_MEM_FENCE);

    for(unsigned int red = (MIO_BN_GRP1 >> 1); red > 256; red >>= 1)
    {
        if(ylid < red)
            lcl_data[ylid] += lcl_data[ylid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&variance, lcl_data, ylid, 1);
    barrier(CLK_LOCAL_MEM_FENCE);
#endif

    if(ylid == 0)
    {
        mvbuff[meanindex] = mean;
        mvbuff[varindex]  = variance;
    }
} // end spatial mean kernel

#endif

// Restore warnings
#ifdef __clang__
#pragma clang diagnostic pop
#pragma clang diagnostic pop
#endif

/*******************************************************************************
 *
 * MIT License
 *
 * Copyright (c) 2017 Advanced Micro Devices, Inc.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 *
 *******************************************************************************/

#define PPCAT_NX(A, B) A##B
#define PPCAT(A, B) PPCAT_NX(A, B)
#define TWO 2
#define FOUR 4
#define EIGHT 8

#if MIOPEN_USE_FP16 == 1
#define MIO_BN_NODPP 1
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
#define _FLOAT half
#ifndef HALF_MAX
#define MAX_VAL 65504 /* max value */
#else
#define MAX_VAL HALF_MAX
#endif
#endif
#if MIOPEN_USE_FP32 == 1
#define _FLOAT float
#ifndef FLT_MAX
#define MAX_VAL 3.402823466e+38F /* max value */
#else
#define MAX_VAL FLT_MAX
#endif
#endif

#define _FLOAT2 PPCAT(_FLOAT, TWO)
#define _FLOAT4 PPCAT(_FLOAT, FOUR)
#define _FLOAT8 PPCAT(_FLOAT, EIGHT)
#define _AS_FLOAT PPCAT(as_, _FLOAT)

#ifndef MIO_BN_LDSGCN_SIZE
#define MIO_BN_LDSGCN_SIZE 16
#endif

#ifndef MIO_BN_LDS_SIZE
#define MIO_BN_LDS_SIZE 256
#endif

#ifndef MIO_BN_C
#define MIO_BN_C 1
#endif

#ifndef MIO_BN_N
#define MIO_BN_N 1
#endif

#ifndef MIO_BN_NHW
#define MIO_BN_NHW 1
#endif

#ifndef MIO_BN_INHW
#define MIO_BN_INHW 1
#endif

#ifndef MIO_BN_CHW
#define MIO_BN_CHW 1
#endif

#ifndef MIO_BN_HW
#define MIO_BN_HW 1
#endif

#ifndef MIO_BN_GRP0
#define MIO_BN_GRP0 1
#endif

#ifndef MIO_BN_GRP1
#define MIO_BN_GRP1 1
#endif

#ifndef MIO_BN_GRP2
#define MIO_BN_GRP2 1
#endif

#ifndef MIO_BN_NGRPS
#define MIO_BN_NGRPS 1
#endif

#ifndef MIO_BN_NCHW
#define MIO_BN_NCHW 1
#endif

#ifndef MIO_BN_VARIANT
#define MIO_BN_VARIANT 255
#endif

#define MIO_BN_MAXN 512

#ifndef MIO_BN_NODPP
#define MIO_BN_NODPP 0
#elif(MIO_BN_NODPP == 1)
#undef __AMDGCN__
#endif

#ifdef __AMDGCN__
#undef __AMDGCN__
#endif

#define UNUSED __attribute__((__unused__))

// Disable specific warnings
#ifdef __clang__
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wconditional-uninitialized"
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wsometimes-uninitialized"
#endif

#ifndef __AMDGCN__

static inline void ReduceKernel(__local _FLOAT* lcl_mem,
                                unsigned int sum_stride,
                                unsigned int unit_id,
                                unsigned int unit_len)
{
    _FLOAT sum              = (_FLOAT)0.;
    unsigned int lcl_offset = unit_id * unit_len;

#pragma unroll
    for(unsigned int i = 0; i < unit_len; i += sum_stride)
    {
        sum += lcl_mem[lcl_offset + i];
    }
    lcl_mem[lcl_offset] = sum;
}

static inline void
regLDSreduce(_FLOAT* value, __local _FLOAT* data, unsigned int localID, _FLOAT scale)
{
    data[localID] = *value;
    barrier(CLK_LOCAL_MEM_FENCE);
    if(localID < (MIO_BN_LDS_SIZE >> 2))
        ReduceKernel(data, 1, localID, 4);
    barrier(CLK_LOCAL_MEM_FENCE);
    if(localID < (MIO_BN_LDS_SIZE >> 4))
        ReduceKernel(data, 4, localID, 16);
    barrier(CLK_LOCAL_MEM_FENCE);
    if(localID == 0)
        ReduceKernel(data, 16, localID, MIO_BN_LDS_SIZE);
    barrier(CLK_LOCAL_MEM_FENCE);
    *value = data[0] * scale;
}
#endif

#ifdef __AMDGCN__

static inline void dppSimpleRedNoBcast64(_FLOAT* value)
{
    _FLOAT tmp = (_FLOAT)0.;
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x111, 0xF, 0xF, 0));
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x112, 0xF, 0xF, 0));
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x114, 0xF, 0xF, 0));
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x118, 0xF, 0xF, 0));
    tmp = _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x142, 0xF, 0xF, 0));
    *value += tmp;
    tmp = _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x143, 0xF, 0xF, 0));
    *value += tmp;
}

static inline void dppSimpleRedBcast64(_FLOAT* value)
{
    _FLOAT tmp = (_FLOAT)0.;
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x111, 0xF, 0xF, 0));
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x112, 0xF, 0xF, 0));
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x114, 0xF, 0xF, 0));
    *value += _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x118, 0xF, 0xF, 0));
    tmp = _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x142, 0xF, 0xF, 0));
    *value += tmp;
    tmp = _AS_FLOAT(__builtin_amdgcn_mov_dpp(as_int(*value), 0x143, 0xF, 0xF, 0));
    *value += tmp;
    barrier(CLK_LOCAL_MEM_FENCE);
    *value = _AS_FLOAT(__builtin_amdgcn_readlane(as_int(*value), 63));
}

#endif

#if(MIO_BN_VARIANT == 0)

#define MIO_BN_SEGTMP (MIO_BN_HW * (MIO_BN_GRP0 / MIO_BN_HW))
#define MIO_BN_SEGMENT ((MIO_BN_SEGTMP > MIO_BN_NHW) ? (MIO_BN_NHW) : (MIO_BN_SEGTMP))
#define MIO_BN_NLOOP ((MIO_BN_NHW + MIO_BN_SEGMENT - 1) / MIO_BN_SEGMENT)
#define MIO_BN_SEGIHW (MIO_BN_SEGMENT / MIO_BN_HW)
#define MIO_BN_NLOOPM (MIO_BN_NLOOP - 1)
#define MIO_BN_SNHW (MIO_BN_NLOOPM * MIO_BN_SEGIHW)

__attribute__((reqd_work_group_size(MIO_BN_GRP0, MIO_BN_GRP1, MIO_BN_GRP2))) __kernel void
BatchNormFwdTrainSpatial(const __global _FLOAT* __restrict in,
                         __global _FLOAT* __restrict out,
                         __constant _FLOAT* __restrict scale,
                         __constant _FLOAT* __restrict bias,
                         _FLOAT INHW,
#if(MIO_RUNNING_RESULT == 1)
                         double expAvgFactor,
                         __global _FLOAT* __restrict resultRunningMean,
                         __global _FLOAT* __restrict resultRunningVariance,
#endif
                         double epsilon
#if(MIO_SAVE_MEAN_VARIANCE == 1)
                         ,
                         __global _FLOAT* __restrict resultSaveMean,
                         __global _FLOAT* __restrict resultSaveInvVariance
#endif
                         )
{

    // SPATIAL
    _FLOAT mean        = (_FLOAT)0.;
    _FLOAT variance    = (_FLOAT)0.;
    _FLOAT invVariance = (_FLOAT)0.;
    _FLOAT pvscale     = (_FLOAT)0.;
    _FLOAT pvbias      = (_FLOAT)0.;
    _FLOAT batchvalues[MIO_BN_NLOOP];

    __local _FLOAT lcl_bias;
    __local _FLOAT lcl_scale;

    unsigned int index  = 0;
    unsigned int lid    = get_local_id(0);
    unsigned int grpid  = get_group_id(0);
    unsigned int chwid  = grpid * MIO_BN_HW + (lid % MIO_BN_HW);
    unsigned int lidihw = lid / MIO_BN_HW;
    unsigned int nid    = 0;

    if(lid == 0)
    {
        lcl_scale = *(scale + grpid);
        lcl_bias  = *(bias + grpid);
    }
    barrier(CLK_LOCAL_MEM_FENCE);

    if(lid < MIO_BN_SEGMENT)
    {
#pragma unroll
        for(unsigned int n = 0; n < MIO_BN_NLOOPM; ++n)
        {
            nid            = n * MIO_BN_SEGIHW + lidihw;
            index          = nid * MIO_BN_CHW + chwid;
            batchvalues[n] = *(in + index);
            mean += batchvalues[n];
            variance = mad(batchvalues[n], batchvalues[n], variance);
        }
        nid                        = MIO_BN_SNHW + lidihw;
        index                      = nid * MIO_BN_CHW + chwid;
        batchvalues[MIO_BN_NLOOPM] = (index < MIO_BN_NCHW) ? *(in + index) : (_FLOAT)0.;
        mean += batchvalues[MIO_BN_NLOOPM];
        variance = mad(batchvalues[MIO_BN_NLOOPM], batchvalues[MIO_BN_NLOOPM], variance);
    }

#ifndef __AMDGCN__
    __local _FLOAT lcl_data[MIO_BN_LDS_SIZE];

    // Reduce mean
    lcl_data[lid] = mean;
    barrier(CLK_LOCAL_MEM_FENCE);
    for(unsigned int red = (MIO_BN_GRP0 >> 1); red > 256; red >>= 1)
    {
        if(lid < red)
            lcl_data[lid] += lcl_data[lid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&mean, lcl_data, lid, (_FLOAT)INHW);
    barrier(CLK_LOCAL_MEM_FENCE);

    // Reduce variance
    lcl_data[lid] = variance;
    barrier(CLK_LOCAL_MEM_FENCE);
    for(unsigned int red = (MIO_BN_GRP0 >> 1); red > 256; red >>= 1)
    {
        if(lid < red)
            lcl_data[lid] += lcl_data[lid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&variance, lcl_data, lid, (_FLOAT)INHW);
    barrier(CLK_LOCAL_MEM_FENCE);
#else

    unsigned int ldsidx = lid >> 6;
    __local _FLOAT lcl_mean[MIO_BN_LDSGCN_SIZE];
    __local _FLOAT lcl_variance[MIO_BN_LDSGCN_SIZE];

    dppSimpleRedNoBcast64(&mean);
    dppSimpleRedNoBcast64(&variance);

    if((lid % 64) == 63)
    {
        lcl_mean[ldsidx]     = mean;
        lcl_variance[ldsidx] = variance;
    }
    barrier(CLK_LOCAL_MEM_FENCE);
    mean = variance = 0.;
#pragma unroll
    for(uint i = 0; i < MIO_BN_LDSGCN_SIZE; i++)
    {
        mean += lcl_mean[i];
        variance += lcl_variance[i];
    }
    mean *= (_FLOAT)INHW;
    variance *= (_FLOAT)INHW;

#endif

    variance    = mad(-mean, mean, variance);
    invVariance = rsqrt(variance + epsilon);
    pvscale     = lcl_scale;
    pvbias      = lcl_bias;

    if(lid < MIO_BN_SEGMENT)
    {
        //==== CALC NORM =======================
        _FLOAT inhat = 0.;

#pragma unroll
        for(unsigned int n = 0; n < MIO_BN_NLOOPM; n++)
        { // apply normalization
            inhat      = (batchvalues[n] - mean) * invVariance;
            nid        = n * MIO_BN_SEGIHW + lidihw;
            index      = nid * MIO_BN_CHW + chwid;
            out[index] = mad(pvscale, inhat, pvbias);
        } // end for

        // Tail of loop
        inhat = (batchvalues[MIO_BN_NLOOPM] - mean) * invVariance;
        nid   = MIO_BN_SNHW + lidihw;
        index = nid * MIO_BN_CHW + chwid;
        if(index < MIO_BN_NCHW)
            out[index] = mad(pvscale, inhat, pvbias);
    }

#if(MIO_SAVE_MEAN_VARIANCE == 1 || MIO_RUNNING_RESULT == 1)
    if(lid == 0)
    {

// Save mean and calculate and save running mean
#if(MIO_SAVE_MEAN_VARIANCE == 1)
        *(resultSaveMean + grpid)        = mean;
        *(resultSaveInvVariance + grpid) = invVariance;
#endif

#if(MIO_RUNNING_RESULT == 1)
        _FLOAT pvt_runMean = *(resultRunningMean + grpid);
        _FLOAT pvt_newRunMean =
            mad((_FLOAT)-expAvgFactor, pvt_runMean, pvt_runMean); // tmp = oldRunMean*(1-factor)
        resultRunningMean[grpid] =
            mad(mean, (_FLOAT)expAvgFactor, pvt_newRunMean); // newMean*factor + tmp
        const _FLOAT adjust = (MIO_BN_NHW == 1)
                                  ? variance
                                  : variance * ((_FLOAT)MIO_BN_NHW / (_FLOAT)(MIO_BN_NHW - 1.0));
        resultRunningVariance[grpid] =
            (1 - (_FLOAT)expAvgFactor) * *(resultRunningVariance + grpid) +
            (_FLOAT)expAvgFactor * adjust;
#endif
    }
#endif
} // end spatial norm

#elif(MIO_BN_VARIANT == 1)

#define MIO_BN_REM (MIO_BN_NHW - ((MIO_BN_NHW / MIO_BN_GRP1) * MIO_BN_GRP1))
#define MIO_BN_LESS (MIO_BN_NHW - MIO_BN_REM)

__attribute__((reqd_work_group_size(MIO_BN_GRP0, MIO_BN_GRP1, MIO_BN_GRP2))) __kernel void
BatchNormFwdTrainSpatial(const __global _FLOAT* __restrict in,
                         __global _FLOAT* __restrict out,
                         __constant _FLOAT* __restrict scale,
                         __constant _FLOAT* __restrict bias,
                         _FLOAT INHW,
#if(MIO_RUNNING_RESULT == 1)
                         double expAvgFactor,
                         __global _FLOAT* __restrict resultRunningMean,
                         __global _FLOAT* __restrict resultRunningVariance,
#endif
                         double epsilon
#if(MIO_SAVE_MEAN_VARIANCE == 1)
                         ,
                         __global _FLOAT* __restrict resultSaveMean,
                         __global _FLOAT* __restrict resultSaveInvVariance
#endif
                         )
{

    // SPATIAL
    _FLOAT mean     = (_FLOAT)0.;
    _FLOAT variance = (_FLOAT)0.;
    _FLOAT invVariance, inhat;
    _FLOAT pvscale, pvbias;

    __local _FLOAT lcl_bias;
    __local _FLOAT lcl_scale;

    int index = 0;
    int lid   = get_local_id(1);
    int xgid  = get_global_id(0);
    int grpid = get_group_id(0);
    int chwid = grpid * MIO_BN_HW;
    int nidx  = 0;
    int hwidx = 0;

#if(MIO_BN_NHW < MIO_BN_MAXN)
    _FLOAT input[MIO_BN_NHW];
#endif

    if(lid == 0)
    {
        lcl_scale = *(scale + xgid);
        lcl_bias  = *(bias + xgid);
    }
    barrier(CLK_LOCAL_MEM_FENCE);

#pragma unroll
    for(unsigned int k = lid, lesskey = 0; k < MIO_BN_LESS; k += MIO_BN_GRP1, ++lesskey)
    {
        nidx           = k / MIO_BN_HW;
        hwidx          = k - (nidx * MIO_BN_HW);
        index          = nidx * MIO_BN_CHW + chwid + hwidx;
#if(MIO_BN_NHW < MIO_BN_MAXN)
        input[lesskey] = *(in + index);
        mean += input[lesskey];
        variance = mad(input[lesskey], input[lesskey], variance);
#else
        _FLOAT xin = *(in + index);
        mean += xin;
        variance        = mad(xin, xin, variance);
#endif
    }
#if(MIO_BN_REM)
    unsigned int remkey = lid + MIO_BN_LESS;
    nidx                = remkey / MIO_BN_HW;
    hwidx               = remkey - (nidx * MIO_BN_HW);
    index               = nidx * MIO_BN_CHW + chwid + hwidx;
#if(MIO_BN_NHW < MIO_BN_MAXN)
    input[remkey]       = (index < MIO_BN_NCHW) ? *(in + index) : 0.;
    mean += input[remkey];
    variance = mad(input[remkey], input[remkey], variance);
#else
    _FLOAT xin = (index < MIO_BN_NCHW) ? *(in + index) : 0.;
    mean += xin;
    variance           = mad(xin, xin, variance);
#endif
#endif

// REDUCE MEAN AND VARIANCE -----------------------
#ifndef __AMDGCN__
    local _FLOAT lcl_data[MIO_BN_LDS_SIZE];
    lcl_data[lid] = mean;
    barrier(CLK_LOCAL_MEM_FENCE);

    for(unsigned int red = (MIO_BN_GRP1 >> 1); red > 256; red >>= 1)
    {
        if(lid < red)
            lcl_data[lid] += lcl_data[lid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&mean, lcl_data, lid, (_FLOAT)INHW);

    barrier(CLK_LOCAL_MEM_FENCE);
    lcl_data[lid] = variance;
    barrier(CLK_LOCAL_MEM_FENCE);

    for(unsigned int red = (MIO_BN_GRP1 >> 1); red > 256; red >>= 1)
    {
        if(lid < red)
            lcl_data[lid] += lcl_data[lid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&variance, lcl_data, lid, (_FLOAT)INHW);
    barrier(CLK_LOCAL_MEM_FENCE);

#else
    unsigned int ldsidx = lid >> 6;
    __local _FLOAT lcl_mean[MIO_BN_LDSGCN_SIZE];
    __local _FLOAT lcl_variance[MIO_BN_LDSGCN_SIZE];

    dppSimpleRedNoBcast64(&mean);
    dppSimpleRedNoBcast64(&variance);
    if((lid % 64) == 63)
    {
        lcl_mean[ldsidx]     = mean;
        lcl_variance[ldsidx] = variance;
    }
    barrier(CLK_LOCAL_MEM_FENCE);
    mean = variance = 0.;
#pragma unroll
    for(uint i = 0; i < MIO_BN_LDSGCN_SIZE; i++)
    {
        mean += lcl_mean[i];
        variance += lcl_variance[i];
    }
    mean *= (_FLOAT)INHW;
    variance *= (_FLOAT)INHW;
#endif
    // REDUCTION COMPLETE ---------------------------

    variance    = mad(-mean, mean, variance);
    invVariance = rsqrt(variance + epsilon);

    pvscale = lcl_scale;
    pvbias  = lcl_bias;

#pragma unroll
    for(unsigned int k = lid, lesskey = 0; k < MIO_BN_LESS; k += MIO_BN_GRP1, ++lesskey)
    {
        nidx       = k / MIO_BN_HW;
        hwidx      = k - (nidx * MIO_BN_HW);
        index      = nidx * MIO_BN_CHW + chwid + hwidx;

#if(MIO_BN_NHW < MIO_BN_MAXN)
        out[index] = mad(pvscale, (input[lesskey] - mean) * invVariance, pvbias);
#else
        out[index] = mad(pvscale, (*(in + index) - mean) * invVariance, pvbias);
#endif
    } // end for
#if(MIO_BN_REM)
    nidx  = remkey / MIO_BN_HW;
    hwidx = remkey - (nidx * MIO_BN_HW);
    index = nidx * MIO_BN_CHW + chwid + hwidx;
    if(index < MIO_BN_NCHW)
    {
#if(MIO_BN_NHW < MIO_BN_MAXN)
        *(out + index) = mad(pvscale, (input[remkey] - mean) * invVariance, pvbias);
#else
        *(out + index) = mad(pvscale, (*(in + index) - mean) * invVariance, pvbias);
#endif
    }
#endif

#if(MIO_SAVE_MEAN_VARIANCE == 1 || MIO_RUNNING_RESULT == 1)
    if(get_global_id(1) == 0)
    {

// Save mean and calculate and save running mean
#if(MIO_SAVE_MEAN_VARIANCE == 1)
        *(resultSaveMean + xgid)        = mean;
        *(resultSaveInvVariance + xgid) = invVariance;
#endif

#if(MIO_RUNNING_RESULT == 1)
        _FLOAT pvt_runMean              = *(resultRunningMean + xgid);
        _FLOAT pvt_newRunMean =
            mad((_FLOAT)-expAvgFactor, pvt_runMean, pvt_runMean); // tmp = oldRunMean*(1-factor)
        resultRunningMean[xgid] =
            mad(mean, (_FLOAT)expAvgFactor, pvt_newRunMean); // newMean*factor + tmp
        const _FLOAT adjust = (MIO_BN_NHW == 1)
                                  ? variance
                                  : variance * ((_FLOAT)MIO_BN_NHW / (_FLOAT)(MIO_BN_NHW - 1.0));
        resultRunningVariance[xgid] = (1 - (_FLOAT)expAvgFactor) * *(resultRunningVariance + xgid) +
                                      (_FLOAT)expAvgFactor * adjust;
#endif
    }
#endif
} // end spatial norm

#elif(MIO_BN_VARIANT == 2) // MULTI-KERNEL reduction for > 33M elements

__attribute__((reqd_work_group_size(MIO_BN_GRP0, MIO_BN_GRP1, MIO_BN_GRP2))) __kernel void
BatchNormFwdTrainSpatialNorm(const __global _FLOAT* __restrict in,
                             __global _FLOAT* __restrict out,
                             const __global _FLOAT* __restrict scale,
                             const __global _FLOAT* __restrict bias)
{

    // SPATIAL
    _FLOAT mean        = (_FLOAT)0.;
    _FLOAT invVariance = (_FLOAT)0.;
    _FLOAT inhat       = (_FLOAT)0.;
    _FLOAT pvt_scale   = (_FLOAT)0.;
    _FLOAT pvt_bias    = (_FLOAT)0.;
    __local _FLOAT lcl_mean, lcl_ivar, lcl_scale, lcl_bias;

    unsigned int ygrp_id = get_group_id(1);
    unsigned int xgid    = get_global_id(0);
    unsigned int ygid    = get_global_id(1);
    unsigned int ygrp_sz = get_local_size(1);
    unsigned int index;
    unsigned int cidx           = xgid * MIO_BN_HW;
    unsigned int meanstashindex = cidx + ygrp_sz * ygrp_id + 1;
    unsigned int varstashindex  = cidx + ygrp_sz * ygrp_id + 3;

    // #4 apply the normalization :: x_hat = (x_i - mean) / sqrt(variance_accum + epsilon)
    if(get_local_id(1) == 0)
    {
        lcl_scale = *(scale + xgid);
        lcl_bias  = *(bias + xgid);
        lcl_mean  = *(out + meanstashindex); // load stashed mean
        lcl_ivar  = *(out + varstashindex);
    }
    barrier(CLK_LOCAL_MEM_FENCE);

    if(ygid < MIO_BN_HW)
    {
        mean        = lcl_mean;
        invVariance = lcl_ivar;
        pvt_scale   = lcl_scale;
        pvt_bias    = lcl_bias;
#pragma unroll
        for(unsigned int n = 0; n < MIO_BN_N; n++)
        { // apply normalization
            index        = n * MIO_BN_CHW + cidx + ygid;
            _FLOAT inhat = (*(in + index) - mean) * invVariance;
            // #5 Gamma and Beta adjust :: y_i = gamma*x_hat + beta
            out[index] = mad(pvt_scale, inhat, pvt_bias);
        } // end for(n)
    }     // end if(inImgIndex)
} // end spatial norm

__attribute__((reqd_work_group_size(MIO_BN_GRP0, MIO_BN_GRP1, MIO_BN_GRP2))) __kernel void
BatchNormFwdTrainSpatialFinalMeanVariance(
    __global _FLOAT* __restrict meanvarbuff,
    _FLOAT INHW
#if(MIO_RUNNING_RESULT == 1)
    ,
    double expAvgFactor /* input momentum */
    ,
    __global _FLOAT* __restrict resultRunningMean, /*input and output*/
    __global _FLOAT* __restrict resultRunningVariance
#endif
    ,
    double epsilon
#if(MIO_SAVE_MEAN_VARIANCE == 1)
    ,
    __global _FLOAT* __restrict resultSaveMean /*output only*/
    ,
    __global _FLOAT* __restrict resultSaveInvVariance
#endif
    )
{
    _FLOAT variance             = (_FLOAT)0.;
    _FLOAT invVariance          = (_FLOAT)0.;
    _FLOAT mean                 = (_FLOAT)0.;
    unsigned int lid            = get_local_id(1);
    unsigned int ygrp_id        = get_group_id(1);
    unsigned int xgid           = get_global_id(0);
    unsigned int ygrp_sz        = get_local_size(1);
    unsigned int yngrps         = get_num_groups(1);
    unsigned int cidx           = xgid * MIO_BN_HW;
    unsigned int meanstashindex = cidx + ygrp_sz * ygrp_id + 1;
    unsigned int varstashindex  = cidx + ygrp_sz * ygrp_id + 3;
    unsigned int commitID       = 0;

    for(int gn = 0; gn < yngrps; gn++)
    {
        unsigned int offset    = gn * ygrp_sz + lid;
        unsigned int meanindex = cidx + ygrp_sz * offset;
        unsigned int varindex  = cidx + ygrp_sz * offset + 2;
        if(offset < yngrps)
        { // modify to span larger number of groups
            mean += *(meanvarbuff + meanindex);
            variance += *(meanvarbuff + varindex); // load per group variance
        }
    }
#ifndef __AMDGCN__
    __local _FLOAT lcl_data[MIO_BN_LDS_SIZE];
    lcl_data[lid] = mean;
    barrier(CLK_LOCAL_MEM_FENCE);

#if(MIO_BN_NGRPS > 64)
    for(unsigned int red = (MIO_BN_GRP1 >> 1); red > 256; red >>= 1)
    {
        if(lid < red)
            lcl_data[lid] += lcl_data[lid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&mean, lcl_data, lid, (_FLOAT)INHW);
#elif(MIO_BN_NGRPS <= 64)
    regLDSreduce(&mean, lcl_data, lid, (_FLOAT)INHW);
    commitID = 0;
#else
    mean = (_FLOAT)0.;
#pragma unroll
    for(int i = 0; i < MIO_BN_NGRPS; i++)
    {
        mean += lcl_data[i];
    }

#endif

    lcl_data[lid] = variance;
    barrier(CLK_LOCAL_MEM_FENCE);

#if(MIO_BN_NGRPS > 256)
    for(unsigned int red = (MIO_BN_GRP1 >> 1); red > 256; red >>= 1)
    {
        if(lid < red)
            lcl_data[lid] += lcl_data[lid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&variance, lcl_data, lid, (_FLOAT)INHW);
#elif(MIO_BN_NGRPS > 64)
    regLDSreduce(&variance, lcl_data, lid, (_FLOAT)INHW);
#elif(MIO_BN_NGRPS > 16)
    regLDSreduce(&variance, lcl_data, lid, (_FLOAT)INHW);
#else //(MIO_BN_NGRPS <= 16)
    variance = (_FLOAT)0.;
#pragma unroll
    for(int i = 0; i < MIO_BN_NGRPS; i++)
    {
        variance += lcl_data[i];
    }
#endif

#else
    commitID            = 64;
    unsigned int ldsidx = lid >> 6;
    __local _FLOAT lcl_mean[MIO_BN_LDSGCN_SIZE];
    __local _FLOAT lcl_variance[MIO_BN_LDSGCN_SIZE];

    dppSimpleRedNoBcast64(&mean);
    if((lid % 64) == 63)
    {
        lcl_mean[ldsidx] = mean;
    }
    barrier(CLK_LOCAL_MEM_FENCE | CLK_GLOBAL_MEM_FENCE);
    mean = 0.;

#pragma unroll
    for(uint i = 0; i < MIO_BN_LDSGCN_SIZE; i++)
    {
        mean += lcl_mean[i];
    }
    mean *= (_FLOAT)INHW;

    dppSimpleRedNoBcast64(&variance);
    if((lid % 64) == 63)
    {
        lcl_variance[ldsidx] = variance;
    }
    barrier(CLK_LOCAL_MEM_FENCE | CLK_GLOBAL_MEM_FENCE);
    variance = 0.;
#pragma unroll
    for(uint i = 0; i < MIO_BN_LDSGCN_SIZE; i++)
    {
        variance += lcl_variance[i];
    }
    variance *= (_FLOAT)INHW;

#endif
    invVariance = rsqrt(variance + epsilon);
    if(lid == commitID)
    {
        meanvarbuff[meanstashindex] = mean;        // stash mean
        meanvarbuff[varstashindex]  = invVariance; // stash mean
    }

// Save mean and calculate and save running mean
#if(MIO_SAVE_MEAN_VARIANCE == 1 || MIO_RUNNING_RESULT == 1)
    unsigned int ygid = get_global_id(1);
    if(ygid == commitID)
    {
#if(MIO_SAVE_MEAN_VARIANCE == 1)
        resultSaveMean[xgid]        = mean;
        resultSaveInvVariance[xgid] = invVariance;
#endif
#if(MIO_RUNNING_RESULT == 1)
        _FLOAT pvt_runMean          = resultRunningMean[xgid];
        _FLOAT pvt_newRunMean =
            mad((_FLOAT)-expAvgFactor, pvt_runMean, pvt_runMean); // tmp = oldRunMean*(1-factor)
        resultRunningMean[xgid] =
            mad(mean, (_FLOAT)expAvgFactor, pvt_newRunMean); // newMean*factor + tmp

        _FLOAT NHW                  = (_FLOAT)MIO_BN_NHW;
        const _FLOAT adjust         = (MIO_BN_NHW == 1) ? variance : variance * (NHW / (NHW - 1));
        resultRunningVariance[xgid] = (1 - (_FLOAT)expAvgFactor) * *(resultRunningVariance + xgid) +
                                      (_FLOAT)expAvgFactor * adjust;
#endif
    }
#endif
}

__attribute__((reqd_work_group_size(MIO_BN_GRP0, MIO_BN_GRP1, MIO_BN_GRP2))) __kernel void
BatchNormFwdTrainSpatialMeanVariance(const __global _FLOAT* __restrict in,
                                     __global _FLOAT* __restrict mvbuff)
{

    __local _FLOAT lcl_data[MIO_BN_LDS_SIZE];

    unsigned int ylid    = get_local_id(1);
    unsigned int ygrp_id = get_group_id(1);
    unsigned int xgid    = get_global_id(0);
    unsigned int ygid    = get_global_id(1);
    unsigned int ygrp_sz = get_local_size(1);
    unsigned int index;
    unsigned int cidx      = xgid * MIO_BN_HW;
    unsigned int meanindex = cidx + ygrp_sz * ygrp_id;
    unsigned int varindex  = meanindex + 2;
    _FLOAT mean            = (_FLOAT)0.;
    _FLOAT variance        = (_FLOAT)0.;
    _FLOAT value           = (_FLOAT)0.;

    if(ygid < MIO_BN_HW)
    {
#pragma unroll
        for(unsigned int n = 0; n < MIO_BN_N; n++)
        {
            index = n * MIO_BN_CHW + cidx + ygid;
            value = *(in + index);
            mean += value;
            variance = mad(value, value, variance);
        }
    }

#ifdef __AMDGCN__
    unsigned int ldsidx = ylid >> 6;
    __local _FLOAT lcl_mean[MIO_BN_LDSGCN_SIZE];
    __local _FLOAT lcl_variance[MIO_BN_LDSGCN_SIZE];

    dppSimpleRedNoBcast64(&mean);
    dppSimpleRedNoBcast64(&variance);

    if((ylid % 64) == 63)
    {
        lcl_mean[ldsidx]     = mean;
        lcl_variance[ldsidx] = variance;
    }
    barrier(CLK_LOCAL_MEM_FENCE | CLK_GLOBAL_MEM_FENCE);
    mean = variance = 0.;

#pragma unroll
    for(uint i = 0; i < MIO_BN_LDSGCN_SIZE; i++)
    {
        mean += lcl_mean[i];
        variance += lcl_variance[i];
    }

#else
    lcl_data[ylid] = mean;
    barrier(CLK_LOCAL_MEM_FENCE);

    for(unsigned int red = (MIO_BN_GRP1 >> 1); red > 256; red >>= 1)
    {
        if(ylid < red)
            lcl_data[ylid] += lcl_data[ylid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&mean, lcl_data, ylid, 1);
    barrier(CLK_LOCAL_MEM_FENCE);

    lcl_data[ylid] = variance;
    barrier(CLK_LOCAL_MEM_FENCE);

    for(unsigned int red = (MIO_BN_GRP1 >> 1); red > 256; red >>= 1)
    {
        if(ylid < red)
            lcl_data[ylid] += lcl_data[ylid + red];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    regLDSreduce(&variance, lcl_data, ylid, 1);
    barrier(CLK_LOCAL_MEM_FENCE);
#endif

    if(ylid == 0)
    {
        mvbuff[meanindex] = mean;
        mvbuff[varindex]  = variance;
    }
} // end spatial mean kernel

#endif

// Restore warnings
#ifdef __clang__
#pragma clang diagnostic pop
#pragma clang diagnostic pop
#endif


updated pooling

updated pooling

/*******************************************************************************
 *
 * MIT License
 *
 * Copyright (c) 2017 Advanced Micro Devices, Inc.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 *
 *******************************************************************************/

#define _FLOAT float
#define _FLOAT2 float2
#define _FLOAT4 float4
#define _FLOAT8 float8
#define _INT_MASK_GLOBAL uchar
#define _INT_MASK_LOCAL uchar

#ifndef FLT_MAX
#define FLT_MAX 3.402823466e+38F /* max value */
#endif

#define UNUSED __attribute__((__unused__))

#define MLO_POOLING_OP_AVE 0
#define MLO_POOLING_OP_MAX 1
#define MLO_POOLING_OP_STC 2

#define MLO_POOLING_GROUP_SZ2 1

#ifndef MLO_POOLING_OP_ID
#define MLO_POOLING_OP_ID 0
#endif
// max
#if MLO_POOLING_OP_ID == MLO_POOLING_OP_MAX
#define MLO_POOLING_OP(A, B) fmax(A, B);
#elif MLO_POOLING_OP_ID == MLO_POOLING_OP_AVE
#define MLO_POOLING_OP(A, B) (A + B);
#endif

/*********************************************************************************

**********************************************************************************/

#define MLO_BOT_DATA_SZ0 \
    (MLO_POOLING_N_HORIZ_OUT_PIX * MLO_POOLING_STRIDE0 + MLO_POOLING_PAD0 + MLO_POOLING_PAD0)
#define MLO_BOT_DATA_SZ1 \
    (MLO_POOLING_N_VERT_OUT_PIX * MLO_POOLING_STRIDE1 + MLO_POOLING_PAD1 + MLO_POOLING_PAD1)

__attribute__((reqd_work_group_size(MLO_POOLING_GROUP_SZ0,
                                    MLO_POOLING_GROUP_SZ1,
                                    MLO_POOLING_GROUP_SZ2))) __kernel void
mloPoolingG(const __global _FLOAT* bot,
            __global _FLOAT* top,
#if !defined(MLO_POOLING_DO_BACKWARD) || MLO_POOLING_OP_ID != MLO_POOLING_OP_MAX
            UNUSED
#endif
                __global _INT_MASK_GLOBAL* mask)
{

    uint x       = get_group_id(0) * MLO_POOLING_GROUP_SZ0 * MLO_POOLING_N_HORIZ_OUT_PIX;
    uint y       = get_group_id(1) * MLO_POOLING_GROUP_SZ1 * MLO_POOLING_N_VERT_OUT_PIX;
    uint lcl_id0 = get_local_id(0);
    uint lcl_id1 = get_local_id(1);
    //		int lcl_id = (lcl_id1 << MLO_POOLING_GROUP_LG2SZ0) + lcl_id0;
    uint ob      = get_global_id(2); // output * batch_sz
    uint b       = ob / MLO_POOLING_N_OUTPUTS;
    uint o       = ob - b * MLO_POOLING_N_OUTPUTS;
    uint bot_x   = (x + lcl_id0 * MLO_POOLING_N_HORIZ_OUT_PIX) * MLO_POOLING_STRIDE0;
    uint bot_y   = (y + lcl_id1 * MLO_POOLING_N_VERT_OUT_PIX) * MLO_POOLING_STRIDE1;
    uint bot_off = b * MLO_POOLING_BOT_BATCH_STRIDE + o * MLO_POOLING_BOT_CHANNEL_STRIDE;

    _FLOAT bot_data[MLO_BOT_DATA_SZ1][MLO_BOT_DATA_SZ0];
    _FLOAT res[MLO_POOLING_N_VERT_OUT_PIX][MLO_POOLING_N_HORIZ_OUT_PIX];
#if defined(MLO_POOLING_DO_BACKWARD) && MLO_POOLING_OP_ID == MLO_POOLING_OP_MAX
    _INT_MASK_LOCAL mask_private[MLO_POOLING_N_VERT_OUT_PIX][MLO_POOLING_N_HORIZ_OUT_PIX];
#endif
    for(int k = 0; k < MLO_POOLING_N_VERT_OUT_PIX; k++)
    {
        for(int l = 0; l < MLO_POOLING_N_HORIZ_OUT_PIX; l++)
        {
#if MLO_POOLING_OP_ID == MLO_POOLING_OP_MAX
            res[k][l] = -FLT_MAX;
#elif MLO_POOLING_OP_ID == MLO_POOLING_OP_AVE
            res[k][l] = 0;
#endif
        }
    }

    for(uint j = 0; j < MLO_BOT_DATA_SZ1; ++j)
    {
        int run_y = (int)bot_y + j - MLO_POOLING_PAD1;

        for(uint i = 0; i < MLO_BOT_DATA_SZ0; ++i)
        {
            int run_x        = (int)bot_x + i - MLO_POOLING_PAD0;
            uint bot_gbl_off = bot_off + (uint)run_y * MLO_POOLING_BOT_STRIDE + (uint)run_x;
            bool vis         = ((run_y >= 0 && run_y < MLO_POOLING_BOT_HEIGHT) &&
                        (run_x >= 0 && run_x < MLO_POOLING_BOT_WIDTH))
                           ? true
                           : false;
            bot_data[j][i] = (vis) ? bot[bot_gbl_off] :
#if MLO_POOLING_OP_ID == MLO_POOLING_OP_MAX
                                   -FLT_MAX
#elif MLO_POOLING_OP_ID == MLO_POOLING_OP_AVE
                                   0
#endif
                ;
        }
    }

    for(uint k = 0; k < MLO_POOLING_N_VERT_OUT_PIX; k++)
    {
#if MLO_POOLING_OP_ID == MLO_POOLING_OP_AVE
        uint y_dst = y + lcl_id1 * MLO_POOLING_N_VERT_OUT_PIX + k;
        int hstart = (int)y_dst * MLO_POOLING_STRIDE1 - MLO_POOLING_PAD1;
        int hend   = min((hstart + MLO_POOLING_KERNEL_SZ1),
                       (int)(MLO_POOLING_BOT_HEIGHT + MLO_POOLING_PAD1));
#endif
        for(uint l = 0; l < MLO_POOLING_N_HORIZ_OUT_PIX; l++)
        {

#if MLO_POOLING_OP_ID == MLO_POOLING_OP_AVE
            uint x_dst = x + lcl_id0 * MLO_POOLING_N_HORIZ_OUT_PIX + l;
            int wstart = (int)x_dst * MLO_POOLING_STRIDE0 - MLO_POOLING_PAD0;
            int wend   = min((wstart + MLO_POOLING_KERNEL_SZ0),
                           (int)(MLO_POOLING_BOT_WIDTH + MLO_POOLING_PAD0));
            uint pool_size = (hend - hstart) * (wend - wstart);
#endif
#if defined(MLO_POOLING_DO_BACKWARD) && MLO_POOLING_OP_ID == MLO_POOLING_OP_MAX
            mask_private[k][l] = 0xFF;
#endif

            for(uint j = 0; j < MLO_POOLING_KERNEL_SZ1; j++)
            {
                for(uint i = 0; i < MLO_POOLING_KERNEL_SZ0; i++)
                {

                    _FLOAT bot_val =
                        bot_data[j + k * MLO_POOLING_STRIDE1][i + l * MLO_POOLING_STRIDE0];
#if 0
						if (y_dst == 0 && x_dst == 6)
						{

							printf("k: %d %f %f\n",
								lcl_off + (k * MLO_POOLING_STRIDE1+j)*MLO_POOLING_LCL_DATA_WIDTH + (l * MLO_POOLING_STRIDE0+i),
								res[k][l],
								bot_val
								);
						}
#endif
#if defined(MLO_POOLING_DO_BACKWARD) && MLO_POOLING_OP_ID == MLO_POOLING_OP_MAX
                    if(bot_val > res[k][l])
                    {
                        res[k][l]          = bot_val;
                        mask_private[k][l] = i + MLO_POOLING_KERNEL_SZ0 * j;
                    }
#else
                    res[k][l] = MLO_POOLING_OP(res[k][l], bot_val);
#endif
                }
            }

#if MLO_POOLING_OP_ID == MLO_POOLING_OP_AVE
            res[k][l] *= (_FLOAT)1.f / (_FLOAT)pool_size;
#endif
        }
    }

    uint top_y   = (y + lcl_id1 * MLO_POOLING_N_VERT_OUT_PIX);
    uint top_x   = (x + lcl_id0 * MLO_POOLING_N_HORIZ_OUT_PIX);
    uint top_off = b * MLO_POOLING_TOP_BATCH_STRIDE + o * MLO_POOLING_TOP_CHANNEL_STRIDE +
                   top_y * MLO_POOLING_TOP_STRIDE + top_x;
    for(uint k = 0; k < MLO_POOLING_N_VERT_OUT_PIX; k++)
    {
        for(uint l = 0; l < MLO_POOLING_N_HORIZ_OUT_PIX; l++)
        {
            if(top_y + k < MLO_POOLING_TOP_HEIGHT && top_x + l < MLO_POOLING_TOP_WIDTH)
            {
                top[top_off + k * MLO_POOLING_TOP_STRIDE + l] = res[k][l];
#if defined(MLO_POOLING_DO_BACKWARD) && MLO_POOLING_OP_ID == MLO_POOLING_OP_MAX
                mask[top_off + k * MLO_POOLING_TOP_STRIDE + l] = mask_private[k][l];
#endif
            }
        }
    }
}

/*******************************************************************************
 *
 * MIT License
 *
 * Copyright (c) 2017 Advanced Micro Devices, Inc.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 *
 *******************************************************************************/

#define _FLOAT float
#define _FLOAT2 float2
#define _FLOAT4 float4
#define _FLOAT8 float8
#define _INT_MASK_GLOBAL uchar
#define _INT_MASK_LOCAL uchar

#ifndef FLT_MAX
#define FLT_MAX 3.402823466e+38F /* max value */
#endif

#define UNUSED __attribute__((__unused__))

#define MLO_POOLING_OP_AVE 0
#define MLO_POOLING_OP_MAX 1
#define MLO_POOLING_OP_STC 2

#define MLO_POOLING_GROUP_SZ2 1

#ifndef MLO_POOLING_OP_ID
#define MLO_POOLING_OP_ID 0
#endif
// max
#if MLO_POOLING_OP_ID == MLO_POOLING_OP_MAX
#define MLO_POOLING_OP(A, B) fmax(A, B);
#elif MLO_POOLING_OP_ID == MLO_POOLING_OP_AVE
#define MLO_POOLING_OP(A, B) (A + B);
#endif

/*********************************************************************************

**********************************************************************************/

#define MLO_BOT_DATA_SZ0 \
    (MLO_POOLING_N_HORIZ_OUT_PIX * MLO_POOLING_STRIDE0 + MLO_POOLING_KERNEL_SZ0 - 1)
#define MLO_BOT_DATA_SZ1 \
    (MLO_POOLING_N_VERT_OUT_PIX * MLO_POOLING_STRIDE1 + MLO_POOLING_KERNEL_SZ1 - 1)

__attribute__((reqd_work_group_size(MLO_POOLING_GROUP_SZ0,
                                    MLO_POOLING_GROUP_SZ1,
                                    MLO_POOLING_GROUP_SZ2))) __kernel void
mloPoolingG(const __global _FLOAT* bot,
            __global _FLOAT* top,
#if !defined(MLO_POOLING_DO_BACKWARD) || MLO_POOLING_OP_ID != MLO_POOLING_OP_MAX
            UNUSED
#endif
                __global _INT_MASK_GLOBAL* mask)
{

    uint x       = get_group_id(0) * MLO_POOLING_GROUP_SZ0 * MLO_POOLING_N_HORIZ_OUT_PIX;
    uint y       = get_group_id(1) * MLO_POOLING_GROUP_SZ1 * MLO_POOLING_N_VERT_OUT_PIX;
    uint lcl_id0 = get_local_id(0);
    uint lcl_id1 = get_local_id(1);
    //		int lcl_id = (lcl_id1 << MLO_POOLING_GROUP_LG2SZ0) + lcl_id0;
    uint ob      = get_global_id(2); // output * batch_sz
    uint b       = ob / MLO_POOLING_N_OUTPUTS;
    uint o       = ob - b * MLO_POOLING_N_OUTPUTS;
    uint bot_x   = (x + lcl_id0 * MLO_POOLING_N_HORIZ_OUT_PIX) * MLO_POOLING_STRIDE0;
    uint bot_y   = (y + lcl_id1 * MLO_POOLING_N_VERT_OUT_PIX) * MLO_POOLING_STRIDE1;
    uint bot_off = b * MLO_POOLING_BOT_BATCH_STRIDE + o * MLO_POOLING_BOT_CHANNEL_STRIDE;

    _FLOAT bot_data[MLO_BOT_DATA_SZ1][MLO_BOT_DATA_SZ0];
    _FLOAT res[MLO_POOLING_N_VERT_OUT_PIX][MLO_POOLING_N_HORIZ_OUT_PIX];
#if defined(MLO_POOLING_DO_BACKWARD) && MLO_POOLING_OP_ID == MLO_POOLING_OP_MAX
    _INT_MASK_LOCAL mask_private[MLO_POOLING_N_VERT_OUT_PIX][MLO_POOLING_N_HORIZ_OUT_PIX];
#endif
    for(int k = 0; k < MLO_POOLING_N_VERT_OUT_PIX; k++)
    {
        for(int l = 0; l < MLO_POOLING_N_HORIZ_OUT_PIX; l++)
        {
#if MLO_POOLING_OP_ID == MLO_POOLING_OP_MAX
            res[k][l] = -FLT_MAX;
#elif MLO_POOLING_OP_ID == MLO_POOLING_OP_AVE
            res[k][l]         = 0;
#endif
        }
    }

    for(uint j = 0; j < MLO_BOT_DATA_SZ1; ++j)
    {
        int run_y = (int)bot_y + j - MLO_POOLING_PAD1;

        for(uint i = 0; i < MLO_BOT_DATA_SZ0; ++i)
        {
            int run_x        = (int)bot_x + i - MLO_POOLING_PAD0;
            uint bot_gbl_off = bot_off + (uint)run_y * MLO_POOLING_BOT_STRIDE + (uint)run_x;
            bool vis         = ((run_y >= 0 && run_y < MLO_POOLING_BOT_HEIGHT) &&
                        (run_x >= 0 && run_x < MLO_POOLING_BOT_WIDTH))
                           ? true
                           : false;
            bot_data[j][i] = (vis) ? bot[bot_gbl_off] : 0;
        }
    }

    for(uint k = 0; k < MLO_POOLING_N_VERT_OUT_PIX; k++)
    {
#if MLO_POOLING_OP_ID == MLO_POOLING_OP_AVE
        uint y_dst  = y + lcl_id1 * MLO_POOLING_N_VERT_OUT_PIX + k;
        uint hstart = y_dst * MLO_POOLING_STRIDE1 - MLO_POOLING_PAD1;
        uint hend   = min((uint)(hstart + MLO_POOLING_KERNEL_SZ1),
                        (uint)(MLO_POOLING_BOT_HEIGHT + MLO_POOLING_PAD1));
#endif
        for(uint l = 0; l < MLO_POOLING_N_HORIZ_OUT_PIX; l++)
        {

#if MLO_POOLING_OP_ID == MLO_POOLING_OP_AVE
            uint x_dst  = x + lcl_id0 * MLO_POOLING_N_HORIZ_OUT_PIX + l;
            uint wstart = x_dst * MLO_POOLING_STRIDE0 - MLO_POOLING_PAD0;
            uint wend   = min((uint)(wstart + MLO_POOLING_KERNEL_SZ0),
                            (uint)(MLO_POOLING_BOT_WIDTH + MLO_POOLING_PAD0));
            uint pool_size = (hend - hstart) * (wend - wstart);
#endif
#if defined(MLO_POOLING_DO_BACKWARD) && MLO_POOLING_OP_ID == MLO_POOLING_OP_MAX
            mask_private[k][l] = 0xFF;
#endif

            for(uint j = 0; j < MLO_POOLING_KERNEL_SZ1; j++)
            {
                for(uint i = 0; i < MLO_POOLING_KERNEL_SZ0; i++)
                {

                    _FLOAT bot_val =
                        bot_data[j + k * MLO_POOLING_STRIDE1][i + l * MLO_POOLING_STRIDE0];
#if 0
						if (y_dst == 0 && x_dst == 6)
						{

							printf("k: %d %f %f\n",
								lcl_off + (k * MLO_POOLING_STRIDE1+j)*MLO_POOLING_LCL_DATA_WIDTH + (l * MLO_POOLING_STRIDE0+i),
								res[k][l],
								bot_val
								);
						}
#endif
#if defined(MLO_POOLING_DO_BACKWARD) && MLO_POOLING_OP_ID == MLO_POOLING_OP_MAX
                    if(bot_val > res[k][l])
                    {
                        res[k][l]          = bot_val;
                        mask_private[k][l] = i + MLO_POOLING_KERNEL_SZ0 * j;
                    }
#else
                    res[k][l] = MLO_POOLING_OP(res[k][l], bot_val);
#endif
                }
            }

#if MLO_POOLING_OP_ID == MLO_POOLING_OP_AVE
            res[k][l] *= (_FLOAT)1.f / (_FLOAT)pool_size;
#endif
        }
    }

    uint top_y   = (y + lcl_id1 * MLO_POOLING_N_VERT_OUT_PIX);
    uint top_x   = (x + lcl_id0 * MLO_POOLING_N_HORIZ_OUT_PIX);
    uint top_off = b * MLO_POOLING_TOP_BATCH_STRIDE + o * MLO_POOLING_TOP_CHANNEL_STRIDE +
                   top_y * MLO_POOLING_TOP_STRIDE + top_x;
    for(uint k = 0; k < MLO_POOLING_N_VERT_OUT_PIX; k++)
    {
        for(uint l = 0; l < MLO_POOLING_N_HORIZ_OUT_PIX; l++)
        {
            if(top_y + k < MLO_POOLING_TOP_HEIGHT && top_x + l < MLO_POOLING_TOP_WIDTH)
            {
                top[top_off + k * MLO_POOLING_TOP_STRIDE + l] = res[k][l];
#if defined(MLO_POOLING_DO_BACKWARD) && MLO_POOLING_OP_ID == MLO_POOLING_OP_MAX
                mask[top_off + k * MLO_POOLING_TOP_STRIDE + l] = mask_private[k][l];
#endif
            }
        }
    }
}


Minor changes

Minor changes

/******************************************************************************
 * function: merge_lit_by_val_chids                                          *
 * input:                                                                     *
 * output:                                                                    *
 ******************************************************************************/

/**
 * High level function:
 * size_t merge_lit_by_val_chids(global uint* input, global uint* chids,
 *                                global uint* lits, size_t li, size_t work_size);
 */

// prototypes
uint sumReduce128(local uint* arr);
uint compactSIMDPrefixSum(local const uint* dsData, local const uint* dsValid,
                         local uint* dsCompact,    local uint* dsLocalIndex);
uint exclusivePrescan128(local const uint* in, local uint* outAndTemp);

/**
 * Steps
 * - create heads
 * - lazy_segmented_scan
 * - stream compaction to reduce all that have a 0 in heads?
 * (There is probably a more efficient solution)
 */

kernel void create_heads(global uint* keys_high, global uint* keys_low,
                         global uint* heads) {
  uint idx = get_global_id(0);
  heads[idx] = 0;
  if (idx > 0) {
    if ((keys_high[idx] != keys_high[idx - 1]) ||
        (keys_low [idx] != keys_low [idx - 1]))
    heads[idx] = 1;
  } else {
    heads[idx] = 1;
  }
}

// somewhat inefficient segmented scan (not really a scan as it stores the
// accumulated value at the beginning of the segment)
kernel void lazy_segmented_scan(global uint* heads, global uint* data) {
  uint idx = get_global_id(0);
  uint maximum = get_global_size(0);
  if (heads[idx] != 0) {
    uint val = data[idx];
    uint curr = idx + 1;
    while (heads[curr] == 0 && curr < maximum) {
      val |= data[curr]; // OR operation
      curr += 1;
    }
    data[idx] = val;
  }
}

/*
// exclusive scan on ...
kernel void exclusive_scan() {

}

// compact to keep only the
kernel void stream_compaction(global uint* config, global uint* heads,
                              global uint* keys_low, global uint* heads,
                              global uint* values) {

}
*/

/*
 *  Code for CUDA stream compaction. Roughly based on:
 *  Billeter M, Olsson O, Assarsson U.
 *  Efficient Stream Compaction on Wide SIMD Many-Core Architectures.
 *      High Performance Graphics 2009.
 *
 *  Project: https://github.com/pandegroup/openmm
 *  File:    platforms/opencl/src/kernels/compact.cl
 *  Notes:
 *    - paper recommends 128 threads/block, so this is hard coded.
 *    - I only implement the prefix-sum based compact primitive,
 *      and not the POPC one, as that is more complicated and performs poorly
 *      on current hardware
 *    - I only implement the scattered- and staged-write variant of phase III
 *      as it they have reasonable performance across most of the tested
 *      workloads in the paper. The selective variant is not implemented.
 *    - The prefix sum of per-block element counts (phase II) is not done in
 *      a particularly efficient manner. It is, however, done in a very easy
 *      to program manner, and integrated into the top of phase III, reducing
 *      the number of kernel invocations required. If one wanted to use
 *      existing code, it'd be easy to take the CUDA SDK scanLargeArray
 *      sample, and do a prefix sum over dgBlockCounts in a phase II kernel.
 *      You could also adapt the existing prescan128 to take an initial value,
 *      and scan dgBlockCounts in stages.
 *
 * Date:         23 Aug 2009
 * Author:       CUDA version by Imran Haque (ihaque@cs.stanford.edu),
 *               converted to OpenCL by Peter Eastman
 * Affiliation:  Stanford University
 * License:      Public Domain
 */

#define N 128

// Phase 1: Count valid elements per thread block
// Hard-code 128 thd/blk
uint sumReduce128(local uint* arr) {
  // Parallel reduce element counts, assumes 128 thd/block
  const uint idx = get_local_id(0);
  if (idx < 64) arr[idx] += arr[idx + 64];
  barrier(CLK_LOCAL_MEM_FENCE);
  if (idx < 32) arr[idx] += arr[idx + 32];
  barrier(CLK_LOCAL_MEM_FENCE);
  if (idx < 16) arr[idx] += arr[idx + 16];
  barrier(CLK_LOCAL_MEM_FENCE);
  if (idx <  8) arr[idx] += arr[idx +  8];
  barrier(CLK_LOCAL_MEM_FENCE);
  if (idx <  4) arr[idx] += arr[idx +  4];
  barrier(CLK_LOCAL_MEM_FENCE);
  if (idx <  2) arr[idx] += arr[idx +  2];
  barrier(CLK_LOCAL_MEM_FENCE);
  if (idx <  1) arr[idx] += arr[idx +  1];
  barrier(CLK_LOCAL_MEM_FENCE);
  return arr[0];
}

// len in config
kernel void countElts(global uint* restrict config,
                      global uint* restrict dgBlockCounts,
                      global uint* restrict dgValid,
                      local  uint* restrict dsCount) {
  const uint len = config[0];
  const uint idx  = get_local_id(0);
  const uint gidx = get_group_id(0);
  const uint ngrps = get_num_groups(0);
  const uint lsize = get_local_size(0);
  const uint epb = len / ngrps + ((len % ngrps) ? 1 : 0); // epb: eltsPerBlock
  const uint ub  = (len < (gidx + 1) * epb) ? len : ((gidx + 1) * epb);
  dsCount[idx] = 0;
  for (uint base = gidx * epb; base < (gidx + 1) * epb; base += lsize) {
    if ((base + idx) < ub && dgValid[base + idx])
      dsCount[idx]++;
  }
  barrier(CLK_LOCAL_MEM_FENCE);
  uint blockCount = sumReduce128(dsCount);
  if (idx == 0)
      dgBlockCounts[gidx] = blockCount;
  return;
}

// Phase 2/3: Move valid elements using SIMD compaction. Phase 2 is done
// implicitly at top of global__ method. Exclusive prefix scan over 128
// elements, assumes 128 threads. Taken from cuda SDK "scan" sample for
// naive scan, with small modifications.
uint exclusivePrescan128(local const uint* in, local uint* outAndTemp) {
  const uint idx = get_local_id(0);
  int pout = 1;
  int pin  = 0;

  // load input into outAndTemp
  // This is exclusive scan, so shift right by one and set first elt to 0
  outAndTemp[pout * N + idx] = (idx > 0) ? in[idx - 1] : 0;
  barrier(CLK_LOCAL_MEM_FENCE);
  for (uint offset = 1; offset < N; offset *= 2) {
    pout = 1 - pout; // swap double buffer indices
    pin  = 1 - pout;
    barrier(CLK_LOCAL_MEM_FENCE);
    outAndTemp[pout * N + idx] = outAndTemp[pin * N + idx];
    if (idx >= offset)
      outAndTemp[pout * N + idx] += outAndTemp[pin * N + idx - offset];
  }
  barrier(CLK_LOCAL_MEM_FENCE);
  return outAndTemp[N - 1] + in[N - 1]; // Return sum of all elements
}

uint compactSIMDPrefixSum(local const uint* dsData, local const uint* dsValid,
                         local uint* dsCompact,    local uint* dsLocalIndex) {
  uint idx = get_local_id(0);
  uint numValid = exclusivePrescan128(dsValid,dsLocalIndex);
  if (dsValid[idx])
    dsCompact[dsLocalIndex[idx]] = dsData[idx];
  return numValid;
}

kernel void moveValidElementsStaged(global       uint* restrict config,
                                    global const uint* restrict dgData,
                                    global       uint* restrict dgCompact,
                                    global const uint* restrict dgValid,
                                    global const uint* restrict dgBlockCounts,
                                    local        uint* restrict inBlock,
                                    local        uint* restrict validBlock,
                                    local        uint* restrict compactBlock) {
  uint len = config[0];
  // dNumValidElements = config[1]
  uint idx = get_local_id(0);
  uint gidx = get_group_id(0);
  uint ngrps = get_num_groups(0);
  uint lsize = get_local_size(0);
  local uint dsLocalIndex[256];
  int blockOutOffset=0;
  // Sum up the blockCounts before us to find our offset
  // This is totally inefficient - lots of repeated work b/w blocks,
  // and uneven balancing. Paper implements this as a prefix sum kernel
  // in phase II. May still be faster than an extra kernel invocation?
  for (uint base = 0; base < gidx; base += lsize) {
    // Load up the count of valid elements for each block before us in batches of 128
    validBlock[idx] = 0;
    if ((base + idx) < gidx)
        validBlock[idx] = dgBlockCounts[base + idx];
    barrier(CLK_LOCAL_MEM_FENCE);
    // Parallel reduce these counts
    // Accumulate in the final offset variable
    blockOutOffset += sumReduce128(validBlock);
    barrier(CLK_LOCAL_MEM_FENCE);
  }
  // epb = eltsPerBlock
  const uint epb = len / ngrps + ((len % ngrps) ? 1 : 0);
  uint ub = (len < (gidx + 1) * epb) ? len : ((gidx + 1) * epb);
  for (uint base = gidx * epb; base < (gidx + 1) * epb; base += lsize) {
    validBlock[idx] = 0;
    if ((base + idx) < ub) {
      validBlock[idx] = dgValid[base + idx];
      inBlock[idx] = dgData[base + idx];
    }
    barrier(CLK_LOCAL_MEM_FENCE);
    uint numValidBlock
      = compactSIMDPrefixSum(inBlock, validBlock, compactBlock, dsLocalIndex);
    barrier(CLK_LOCAL_MEM_FENCE);
    if (idx < numValidBlock) {
      dgCompact[blockOutOffset + idx] = compactBlock[idx];
    }
    blockOutOffset += numValidBlock;
  }
  if (gidx == (ngrps - 1) && idx == 0)
    config[1] = blockOutOffset;
}

/******************************************************************************
 * function: merge_lit_by_val_chids                                          *
 * input:                                                                     *
 * output:                                                                    *
 ******************************************************************************/

/**
 * High level function:
 * size_t merge_lit_by_val_chids(global uint* input, global uint* chids,
 *                                global uint* lits, size_t li, size_t work_size);
 */

// prototypes
uint sumReduce128(local uint* arr);
int compactSIMDPrefixSum(local const uint* dsData, local const uint* dsValid,
                         local uint* dsCompact,    local uint* dsLocalIndex);
int exclusivePrescan128(local const uint* in, local uint* outAndTemp);

/**
 * Steps
 * - create heads
 * - lazy_segmented_scan
 * - stream compaction to reduce all that have a 0 in heads?
 * (There is probably a more efficient solution)
 */

kernel void create_heads(global uint* keys_high, global uint* keys_low,
                         global uint* heads) {
  uint idx = get_global_id(0);
  heads[idx] = 0;
  if (idx > 0) {
    if ((keys_high[idx] != keys_high[idx - 1]) ||
        (keys_low [idx] != keys_low [idx - 1]))
    heads[idx] = 1;
  } else {
    heads[idx] = 1;
  }
}

// somewhat inefficient segmented scan (not really a scan as it stores the
// accumulated value at the beginning of the segment)
kernel void lazy_segmented_scan(global uint* heads, global uint* data) {
  uint idx = get_global_id(0);
  uint maximum = get_global_size(0);
  if (heads[idx] != 0) {
    uint val = data[idx];
    uint curr = idx + 1;
    while (heads[curr] == 0 && curr < maximum) {
      val |= data[curr]; // OR operation
      curr += 1;
    }
    data[idx] = val;
  }
}

/*
// exclusive scan on ...
kernel void exclusive_scan() {

}

// compact to keep only the
kernel void stream_compaction(global uint* config, global uint* heads,
                              global uint* keys_low, global uint* heads,
                              global uint* values) {

}
*/

/*
 *  Code for CUDA stream compaction. Roughly based on:
 *  Billeter M, Olsson O, Assarsson U.
 *  Efficient Stream Compaction on Wide SIMD Many-Core Architectures.
 *      High Performance Graphics 2009.
 *
 *  Project: https://github.com/pandegroup/openmm
 *  File:    platforms/opencl/src/kernels/compact.cl
 *  Notes:
 *    - paper recommends 128 threads/block, so this is hard coded.
 *    - I only implement the prefix-sum based compact primitive,
 *      and not the POPC one, as that is more complicated and performs poorly
 *      on current hardware
 *    - I only implement the scattered- and staged-write variant of phase III
 *      as it they have reasonable performance across most of the tested
 *      workloads in the paper. The selective variant is not implemented.
 *    - The prefix sum of per-block element counts (phase II) is not done in
 *      a particularly efficient manner. It is, however, done in a very easy
 *      to program manner, and integrated into the top of phase III, reducing
 *      the number of kernel invocations required. If one wanted to use
 *      existing code, it'd be easy to take the CUDA SDK scanLargeArray
 *      sample, and do a prefix sum over dgBlockCounts in a phase II kernel.
 *      You could also adapt the existing prescan128 to take an initial value,
 *      and scan dgBlockCounts in stages.
 *
 * Date:         23 Aug 2009
 * Author:       CUDA version by Imran Haque (ihaque@cs.stanford.edu),
 *               converted to OpenCL by Peter Eastman
 * Affiliation:  Stanford University
 * License:      Public Domain
 */

// Phase 1: Count valid elements per thread block
// Hard-code 128 thd/blk
uint sumReduce128(local uint* arr) {
  // Parallel reduce element counts, assumes 128 thd/block
  const uint idx = get_local_id(0);
  if (idx < 64) arr[idx] += arr[idx + 64];
  barrier(CLK_LOCAL_MEM_FENCE);
  if (idx < 32) arr[idx] += arr[idx + 32];
  barrier(CLK_LOCAL_MEM_FENCE);
  if (idx < 16) arr[idx] += arr[idx + 16];
  barrier(CLK_LOCAL_MEM_FENCE);
  if (idx <  8) arr[idx] += arr[idx +  8];
  barrier(CLK_LOCAL_MEM_FENCE);
  if (idx <  4) arr[idx] += arr[idx +  4];
  barrier(CLK_LOCAL_MEM_FENCE);
  if (idx <  2) arr[idx] += arr[idx +  2];
  barrier(CLK_LOCAL_MEM_FENCE);
  if (idx <  1) arr[idx] += arr[idx +  1];
  barrier(CLK_LOCAL_MEM_FENCE);
  return arr[0];
}

// len in config
kernel void countElts(global uint* restrict config,
                      global uint* restrict dgBlockCounts,
                      global uint* restrict dgValid,
                      local  uint* restrict dsCount) {
  const uint len = config[0];
  const uint idx  = get_local_id(0);
  const uint gidx = get_group_id(0);
  const uint ngrps = get_num_groups(0);
  const uint lsize = get_local_size(0);
  const uint epb = len / ngrps + ((len % ngrps) ? 1 : 0); // epb: eltsPerBlock
  const uint ub  = (len < (gidx + 1) * epb) ? len : ((gidx + 1) * epb);
  dsCount[idx] = 0;
  for (uint base = gidx * epb; base < (gidx + 1) * epb; base += lsize) {
    if ((base + idx) < ub && dgValid[base + idx])
      dsCount[idx]++;
  }
  barrier(CLK_LOCAL_MEM_FENCE);
  uint blockCount = sumReduce128(dsCount);
  if (idx == 0)
      dgBlockCounts[gidx] = blockCount;
  return;
}

// Phase 2/3: Move valid elements using SIMD compaction. Phase 2 is done
// implicitly at top of global__ method. Exclusive prefix scan over 128
// elements, assumes 128 threads. Taken from cuda SDK "scan" sample for
// naive scan, with small modifications.
int exclusivePrescan128(local const uint* in, local uint* outAndTemp) {
  const uint idx = get_local_id(0);
  const uint n = 128;
  //TODO: this temp storage could be reduced since we write to shared
  //      memory in out anyway, and n is hardcoded
  //__shared__ int temp[2*n];
  local uint* temp = outAndTemp;
  int pout = 1;
  int pin  = 0;

  // load input into temp
  // This is exclusive scan, so shift right by one and set first elt to 0
  temp[pout * n + idx] = (idx > 0) ? in[idx - 1] : 0;
  barrier(CLK_LOCAL_MEM_FENCE);
  for (uint offset = 1; offset < n; offset *= 2) {
    pout = 1 - pout; // swap double buffer indices
    pin  = 1 - pout;
    barrier(CLK_LOCAL_MEM_FENCE);
    temp[pout * n + idx] = temp[pin * n + idx];
    if (idx >= offset)
      temp[pout * n + idx] += temp[pin * n + idx - offset];
  }
  //out[idx] = temp[pout*n+idx]; // write output
  barrier(CLK_LOCAL_MEM_FENCE);
  return outAndTemp[127] + in[127]; // Return sum of all elements
}

int compactSIMDPrefixSum(local const uint* dsData, local const uint* dsValid,
                         local uint* dsCompact,    local uint* dsLocalIndex) {
  uint idx = get_local_id(0);
  int numValid = exclusivePrescan128(dsValid,dsLocalIndex);
  if (dsValid[idx])
    dsCompact[dsLocalIndex[idx]] = dsData[idx];
  return numValid;
}

kernel void moveValidElementsStaged(global       uint* restrict config,
                                    global const uint* restrict dgData,
                                    global       uint* restrict dgCompact,
                                    global const uint* restrict dgValid,
                                    global const uint* restrict dgBlockCounts,
                                    local        uint* restrict inBlock,
                                    local        uint* restrict validBlock,
                                    local        uint* restrict compactBlock) {
  uint len = config[0];
  // dNumValidElements = config[1]
  uint idx = get_local_id(0);
  uint gidx = get_group_id(0);
  uint ngrps = get_num_groups(0);
  uint lsize = get_local_size(0);
  local uint dsLocalIndex[256];
  int blockOutOffset=0;
  // Sum up the blockCounts before us to find our offset
  // This is totally inefficient - lots of repeated work b/w blocks,
  // and uneven balancing. Paper implements this as a prefix sum kernel
  // in phase II. May still be faster than an extra kernel invocation?
  for (uint base = 0; base < gidx; base += lsize) {
    // Load up the count of valid elements for each block before us in batches of 128
    validBlock[idx] = 0;
    if ((base + idx) < gidx)
        validBlock[idx] = dgBlockCounts[base + idx];
    barrier(CLK_LOCAL_MEM_FENCE);
    // Parallel reduce these counts
    // Accumulate in the final offset variable
    blockOutOffset += sumReduce128(validBlock);
    barrier(CLK_LOCAL_MEM_FENCE);
  }
  // epb = eltsPerBlock
  const uint epb = len / ngrps + ((len % ngrps) ? 1 : 0);
  uint ub = (len < (gidx + 1) * epb) ? len : ((gidx + 1) * epb);
  for (uint base = gidx * epb; base < (gidx + 1) * epb; base += lsize) {
    validBlock[idx] = 0;
    if ((base + idx) < ub) {
      validBlock[idx] = dgValid[base + idx];
      inBlock[idx] = dgData[base + idx];
    }
    barrier(CLK_LOCAL_MEM_FENCE);
    uint numValidBlock
      = compactSIMDPrefixSum(inBlock, validBlock, compactBlock, dsLocalIndex);
    barrier(CLK_LOCAL_MEM_FENCE);
    if (idx < numValidBlock) {
      dgCompact[blockOutOffset + idx] = compactBlock[idx];
    }
    blockOutOffset += numValidBlock;
  }
  if (gidx == (ngrps - 1) && idx == 0)
    config[1] = blockOutOffset;
}