/*******************************************************************************
* Copyright 2020 Intel Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License")* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*******************************************************************************/
#include "gpu/intel/bnorm/xe.h"
#include "gpu/intel/bnorm/xe_reduce.h"
// BWD kernels that are that are specially optimized for NHWC layout
// (NHWC_OPTIMIZED definition).
// These kernels not supported IC tail processing.
// Two types of reduction are implemented:
// 1) Reduction over scratchpad (reduce_temp) with SLM use, implemented by
// xe_reduce_stats_nhwc kernel.
// 2) Atomics-based reduction with SLM use (FUSED_ATOMICS_REDUCTION definition),
// implemented as part of calc kernels, see xe_*_fused_reduction()
// functions in xe.h. This reduction implementation requires zeroing and
// finalization steps, see xe_fused_reduce_* kernels in xe_reduce.cl
NAMED_KERNEL_ATTR(CALC)
__kernel void xe_calculate_stats_nhwc(__global DATA_T *src,
__global float *mean, __global DATA_T *diff_dst, __global char *ws,
__global float *temp_reduce, volatile __global atomic_float *diff_scale,
volatile __global atomic_float *diff_shift) {
const int mb = GWS_GET_STAT_MB() const int c = GWS_GET_STAT_IC() const int sp_block_idx = GWS_GET_STAT_SP() const int ic_block_offset = (c / SG_SIZE) * IC_BLOCK const int offset = ic_block_offset + sp_block_idx * STAT_SP_BLOCK * IC
mean += ic_block_offset src += offset diff_dst += offset ws += offset
float v_mean[IC_BLOCK_SGROUPS] for (int sg = 0 v_mean[sg] = as_float(intel_sub_group_block_read(
(const __global uint *)(&mean[(sg * SG_SIZE)]))) }
VECT_FLOAT_T diff_gamma[IC_BLOCK_SGROUPS / VECT_SIZE] = {0.0f} VECT_FLOAT_T diff_beta[IC_BLOCK_SGROUPS / VECT_SIZE] = {0.0f}#if HAS_IC_VECT_TAIL
float diff_gamma_tail[IC_TAIL_SGROUPS] = {0.0f} float diff_beta_tail[IC_TAIL_SGROUPS] = {0.0f}#else
float *diff_gamma_tail = NULL float *diff_beta_tail = NULL#endif
#if USE_WORKAROUND
for (int sp = 0 if (sp_block_idx * STAT_SP_BLOCK + sp >= SP) break#else // issue
#if HAS_STAT_SP_BLOCK_TAIL
for (int sp = 0 ++sp) {
#else
for (int sp = 0#endif
#endif // USE_WORKAROUND
for (int sg = 0 const int sg_idx = sg * SG_SIZE * VECT_SIZE#if FUSE_BN_RELU
VECT_CHAR_T ws_vect = LOAD_VECT_CHAR(&ws[sg_idx])#endif
VECT_FLOAT_T src_vect = LOAD_VECT_DATA(&src[sg_idx]) VECT_FLOAT_T dd_vect = LOAD_VECT_DATA(&diff_dst[sg_idx]) VECT_FLOAT_T v0 for (int vect = 0 int sg_idx = sg * VECT_SIZE + vect#if VECT_SIZE > 1
v0[vect] = src_vect[vect] - v_mean[sg_idx]#else
v0 = src_vect - v_mean[sg_idx]#endif
}
#if FUSE_BN_RELU
dd_vect = select(
(VECT_FLOAT_T)0.0f, dd_vect, CONVERT_VECT_INT_T(ws_vect))#endif
diff_gamma[sg] = fma(v0, dd_vect, diff_gamma[sg]) diff_beta[sg] += dd_vect }
#if HAS_IC_VECT_TAIL
for (int sg = 0 const int sg_idx = IC_VECT_SGROUPS + sg#if FUSE_BN_RELU
char ws_tail = LOAD_CHAR_1x16(&ws[sg_idx * SG_SIZE])#endif
float src_tail = LOAD_DATA_1x16(&src[sg_idx * SG_SIZE]) float dd_tail = LOAD_DATA_1x16(&diff_dst[sg_idx * SG_SIZE]) float v0 = src_tail - v_mean[sg_idx]#if FUSE_BN_RELU
dd_tail = select(0.0f, dd_tail, convert_int(ws_tail))#endif
diff_gamma_tail[sg] = fma(v0, dd_tail, diff_gamma_tail[sg]) diff_beta_tail[sg] += dd_tail }
#endif
src += IC diff_dst += IC#if FUSE_BN_RELU
ws += IC#endif
}
#if FUSED_ATOMICS_REDUCTION
__local float local_gamma[2 * CALC_SLM_SIZE] __local float *local_beta = local_gamma + CALC_SLM_SIZE xe_calc_fused_reduction(diff_scale, diff_shift, ic_block_offset, diff_gamma,
diff_beta, diff_gamma_tail, diff_beta_tail, local_gamma,
local_beta)#else
// scratchpad layout:
// PADDED_IC - diff_gamma reduction, wrote by xe_reduce_stats kernel
// REDUCE_STAT_NBLOCKS * PADDED_IC - diff_gamma stats calculated by this kernel
// PADDED_IC - diff_beta reduction, wrote by xe_reduce_stats kernel
// REDUCE_STAT_NBLOCKS * PADDED_IC - diff_beta stats calculated by this kernel
for (int sg = 0 const int reduce_off = sp_block_idx * SG_SIZE
+ REDUCE_STAT_NBLOCKS * SG_SIZE
* (sg + (int)(c / SG_SIZE) * (IC_BLOCK / SG_SIZE))
const int diff_gamma_offset = PADDED_IC + reduce_off const int diff_beta_offset
= 2 * PADDED_IC + REDUCE_STAT_NBLOCKS * PADDED_IC + reduce_off
#if HAS_IC_VECT_TAIL
if (sg >= IC_VECT_SGROUPS) {
STORE_FLOAT_1x16(&temp_reduce[diff_gamma_offset],
diff_gamma_tail[sg - IC_VECT_SGROUPS]) STORE_FLOAT_1x16(&temp_reduce[diff_beta_offset],
diff_beta_tail[sg - IC_VECT_SGROUPS]) } else
#endif
{
#if VECT_SIZE > 1
STORE_FLOAT_1x16(&temp_reduce[diff_gamma_offset],
diff_gamma[sg / VECT_SIZE][sg % VECT_SIZE]) STORE_FLOAT_1x16(&temp_reduce[diff_beta_offset],
diff_beta[sg / VECT_SIZE][sg % VECT_SIZE])#else
STORE_FLOAT_1x16(&temp_reduce[diff_gamma_offset], diff_gamma[sg]) STORE_FLOAT_1x16(&temp_reduce[diff_beta_offset], diff_beta[sg])#endif
}
}
#endif // FUSED_ATOMICS_REDUCTION
}
KERNEL_ATTR
__kernel void xe_bnorm_bwd_nhwc(__global DATA_T *src, __global float *mean,
__global float *variance, __global DATA_T *diff_dst,
__global float *scaleshift, __global char *ws,
__global DATA_T *diff_src, __global float *diff_scale,
__global float *diff_shift, float eps, __global DATA_T *diff_src_add) {
const int c = GWS_GET_IC() const int ic_block_offset = (c / SG_SIZE) * IC_BLOCK
variance += ic_block_offset mean += ic_block_offset diff_scale += ic_block_offset diff_shift += ic_block_offset scaleshift += ic_block_offset
VECT_FLOAT_T v_variance[IC_BLOCK_SGROUPS / VECT_SIZE],
v_mean[IC_BLOCK_SGROUPS / VECT_SIZE],
diff_gamma[IC_BLOCK_SGROUPS / VECT_SIZE],
diff_beta[IC_BLOCK_SGROUPS / VECT_SIZE],
sqrt_variance[IC_BLOCK_SGROUPS / VECT_SIZE],
gamma[IC_BLOCK_SGROUPS / VECT_SIZE] for (int sg = 0 const int sg_idx = sg * SG_SIZE * VECT_SIZE v_variance[sg] = LOAD_VECT_FLOAT(&variance[sg_idx])#if CALCULATE_DIFF_STATS == 1
v_mean[sg] = LOAD_VECT_FLOAT(&mean[sg_idx]) diff_gamma[sg] = LOAD_VECT_FLOAT(&diff_scale[sg_idx])#if DIFF_SHIFT == 1
diff_beta[sg] = LOAD_VECT_FLOAT(&diff_shift[sg_idx])#else
diff_beta[sg] = LOAD_VECT_FLOAT(
&diff_shift[IC + REDUCE_STAT_NBLOCKS * IC + sg_idx])#endif // #if DIFF_SHIFT == 1
#endif // #if CALCULATE_DIFF_STATS == 1
#if USE_SCALE == 1
gamma[sg] = LOAD_VECT_FLOAT(&scaleshift[sg_idx])#else
gamma[sg] = (VECT_FLOAT_T)1.0f#endif
sqrt_variance[sg]
= (VECT_FLOAT_T)1.0f / sqrt(v_variance[sg] + (VECT_FLOAT_T)eps) }
#if HAS_IC_VECT_TAIL
float v_variance_tail[IC_TAIL_SGROUPS], v_mean_tail[IC_TAIL_SGROUPS],
diff_gamma_tail[IC_TAIL_SGROUPS], diff_beta_tail[IC_TAIL_SGROUPS],
sqrt_variance_tail[IC_TAIL_SGROUPS], gamma_tail[IC_TAIL_SGROUPS] for (int sg = 0 const int sg_idx = (IC_VECT_SGROUPS + sg) * SG_SIZE v_variance_tail[sg] = LOAD_FLOAT_1x16(&variance[sg_idx])
#if CALCULATE_DIFF_STATS == 1
v_mean_tail[sg] = LOAD_FLOAT_1x16(&mean[sg_idx]) diff_gamma_tail[sg] = LOAD_FLOAT_1x16(&diff_scale[sg_idx])#if DIFF_SHIFT == 1
diff_beta_tail[sg] = LOAD_FLOAT_1x16(&diff_shift[sg_idx])#else
diff_beta_tail[sg] = LOAD_FLOAT_1x16(
&diff_shift[IC + REDUCE_STAT_NBLOCKS * IC + sg_idx])#endif // #if DIFF_SHIFT == 1
#endif // #if CALCULATE_DIFF_STATS == 1
#if USE_SCALE == 1
gamma_tail[sg] = LOAD_FLOAT_1x16(&scaleshift[sg_idx])#else
gamma_tail[sg] = 1.0f#endif
sqrt_variance_tail[sg] = 1.0f / sqrt(v_variance_tail[sg] + eps) }
#endif
const int sp_block_idx = GWS_GET_SP() const int offset = ic_block_offset + sp_block_idx * UPDATE_SP_BLOCK * IC
src += offset diff_dst += offset ws += offset diff_src += offset#if FUSE_BN_ADD_RELU
diff_src_add += offset#endif
#if HAS_UPDATE_SP_BLOCK_TAIL
for (int sp = 0 sp < min(UPDATE_SP_BLOCK, SP - sp_block_idx * UPDATE_SP_BLOCK) sp += UPDATE_SP_UNROLL) {
#else
for (int sp = 0#endif
for (int sg = 0 const int sg_idx = sg * SG_SIZE * VECT_SIZE
VECT_FLOAT_T src_vect[UPDATE_SP_UNROLL] unroll_for(int i = 0 src_vect[i] = LOAD_VECT_DATA(&src[sg_idx + IC * i]) }
VECT_FLOAT_T dd_vect[UPDATE_SP_UNROLL] unroll_for(int i = 0 dd_vect[i] = LOAD_VECT_DATA(&diff_dst[sg_idx + IC * i]) }
#if FUSE_BN_RELU
VECT_CHAR_T ws_vect[UPDATE_SP_UNROLL] unroll_for(int i = 0 ws_vect[i] = LOAD_VECT_CHAR(&ws[sg_idx + IC * i]) dd_vect[i] = select((VECT_FLOAT_T)0.0f, dd_vect[i],
CONVERT_VECT_INT_T(ws_vect[i])) }
#if FUSE_BN_ADD_RELU
unroll_for(int i = 0 STORE_VECT_DATA(&diff_src_add[sg_idx + IC * i], dd_vect[i]) }
#endif
#endif
#if CALCULATE_DIFF_STATS == 1
unroll_for(int i = 0 dd_vect[i]
-= (diff_beta[sg]
+ (src_vect[i] - v_mean[sg]) * diff_gamma[sg]
* sqrt_variance[sg])
/ (MB * ID * IH * IW) }
#endif // #if CALCULATE_DIFF_STATS == 1
unroll_for(int i = 0 dd_vect[i] *= gamma[sg] * sqrt_variance[sg] }
unroll_for(int i = 0 STORE_VECT_DATA(&diff_src[sg_idx + IC * i], dd_vect[i]) }
}
#if HAS_IC_VECT_TAIL
for (int sg = 0 const int sg_idx = (IC_VECT_SGROUPS + sg) * SG_SIZE float src_tail[UPDATE_SP_UNROLL] unroll_for(int i = 0 src_tail[i] = LOAD_DATA_1x16(&src[sg_idx + IC * i]) }
float dd_tail[UPDATE_SP_UNROLL] unroll_for(int i = 0 dd_tail[i] = LOAD_DATA_1x16(&diff_dst[sg_idx + IC * i]) }
#if FUSE_BN_RELU
char ws_tail[UPDATE_SP_UNROLL] unroll_for(int i = 0 ws_tail[i] = LOAD_CHAR_1x16(&ws[sg_idx + IC * i]) dd_tail[i] = select(0.0f, dd_tail[i], convert_int(ws_tail[i])) }
#if FUSE_BN_ADD_RELU
unroll_for(int i = 0 STORE_DATA_1x16(&diff_src_add[sg_idx + IC * i], dd_tail[i]) }
#endif
#endif
#if CALCULATE_DIFF_STATS == 1
unroll_for(int i = 0 dd_tail[i] -= (diff_beta_tail[sg]
+ (src_tail[i] - v_mean_tail[sg])
* diff_gamma_tail[sg]
* sqrt_variance_tail[sg])
/ (MB * ID * IH * IW) }
#endif // #if CALCULATE_DIFF_STATS == 1
unroll_for(int i = 0 dd_tail[i] *= gamma_tail[sg] * sqrt_variance_tail[sg] STORE_DATA_1x16(&diff_src[sg_idx + IC * i], dd_tail[i]) }
}
#endif
src += IC * UPDATE_SP_UNROLL diff_dst += IC * UPDATE_SP_UNROLL diff_src += IC * UPDATE_SP_UNROLL#if FUSE_BN_RELU
#if FUSE_BN_ADD_RELU
diff_src_add += IC * UPDATE_SP_UNROLL#endif
ws += IC * UPDATE_SP_UNROLL#endif
}
}