/*******************************************************************************
* Copyright 2019 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/include/philox.h"
#include "gpu/intel/softmax/simple.h"
#if IS_FWD
#if SUB_GROUP_SIZE == 16
__attribute__((reqd_work_group_size(GROUP_SIZE, 1, 1)))
__attribute__((intel_reqd_sub_group_size(SUB_GROUP_SIZE)))
#endif
__kernel void
simple_softmax_fwd_generic(__global SRC_DATA_T *src, __global DATA_T *dst,
__global float *src_scale, __global float *dst_scale
#if WITH_DROPOUT
,
__global uchar *dropout_mask_buf,
#if USE_HOST_SCALARS
long dropout_seed, long dropout_offset,
float dropout_p
#else
__global long *dropout_seed_buf, __global long *dropout_offset_buf,
__global float *dropout_p_buf
#endif
#endif
POST_OP_ARGS) {
const off_t dim[] = {
(get_global_id(0) / GROUP_SIZE) % BLOCK_0,
get_global_id(1) % BLOCK_1,
get_global_id(2) % BLOCK_2,
(get_global_id(0) / GROUP_SIZE) / BLOCK_0,
get_global_id(1) / BLOCK_1,
get_global_id(2) / BLOCK_2,
}
float scale = 1.0f#if WITH_SRC_SCALES
scale *= src_scale[0]#endif
#if SUB_GROUP_SIZE == 16
const int local_id = get_local_id(0)
// SOFTMAX_AXIS is the size of axis around which softmax operation is
// performed
// begin and end indices calculated according to thread's id
const int begin = local_id * (SOFTMAX_AXIS / GROUP_SIZE) const int end = (local_id == GROUP_SIZE - 1)
? SOFTMAX_AXIS
: (local_id + 1) * (SOFTMAX_AXIS / GROUP_SIZE)#if SOFTMAX_AXIS - (GROUP_SIZE - 1) * (SOFTMAX_AXIS / GROUP_SIZE) \
> SOFTMAX_AXIS / GROUP_SIZE
const int buf_size
= SOFTMAX_AXIS - (GROUP_SIZE - 1) * (SOFTMAX_AXIS / GROUP_SIZE)#else
const int buf_size = SOFTMAX_AXIS / GROUP_SIZE#endif
#else
const int begin = 0 const int end = SOFTMAX_AXIS const int buf_size = SOFTMAX_AXIS#endif
#if IS_HALF(SRC_DATA_T) == 1 && IS_HALF(DST_DATA_TYPE) == 1
typedef half acc_t const acc_t acc_max = HALF_MAX const acc_t acc_zero = 0.h#elif DT_F64 || DST_DT_F64
typedef double acc_t const acc_t acc_max = DBL_MAX const acc_t acc_zero = 0.0#else
typedef float acc_t const acc_t acc_max = FLT_MAX const acc_t acc_zero = 0.f#endif
acc_t d[buf_size] acc_t max_ = -acc_max acc_t denom_ = acc_zero
// finding max value for each sub_group
if (!(NEEDS_PADDING(dim[0], dim[1], dim[2], dim[3], dim[4], begin))) {
for (int i = begin off_t data_off
= DATA_OFF(dim[0], dim[1], dim[2], dim[3], dim[4], i) d[i - begin] = SRC_TO_REF(src[data_off]) max_ = max(max_, d[i - begin]) }
}
#if SUB_GROUP_SIZE == 16
// reduce using work_group_reduce if no. of subgroups > 1, for e.g.
// if group_size is 32, there will be 2 sub-groups (size of each sub-group
// is 16 which is an optimal value)
#if GROUP_SIZE == SUB_GROUP_SIZE
max_ = sub_group_reduce_max(max_)#else
max_ = work_group_reduce_max(max_)#endif
#endif
// updating dst tensor and accumulating denom for last step
if (!(NEEDS_PADDING(dim[0], dim[1], dim[2], dim[3], dim[4], begin))) {
for (int i = begin#if LOGSOFTMAX
denom_ += exp(d[i - begin] - max_)#else
d[i - begin] = exp(d[i - begin] - max_) denom_ += d[i - begin]#endif
}
}
#if SUB_GROUP_SIZE == 16
#if GROUP_SIZE == SUB_GROUP_SIZE
denom_ = sub_group_reduce_add(denom_)#else
denom_ = work_group_reduce_add(denom_)#endif
#endif
#if LOGSOFTMAX
denom_ = log(denom_)#else
denom_ = (SOFTMAX_INF_AS_ZERO && denom_ == 0.f) ? 1.0f : 1.0f / denom_#endif
for (int i = begin off_t data_off = DATA_OFF(dim[0], dim[1], dim[2], dim[3], dim[4], i)
POST_OP_DATA_T tmp if (NEEDS_PADDING(dim[0], dim[1], dim[2], dim[3], dim[4], i)) {
tmp = (POST_OP_DATA_T)(acc_zero) } else {
float unscaled#if LOGSOFTMAX
unscaled = d[i - begin] - max_ - denom_#else
unscaled = d[i - begin] * denom_#endif
tmp = (POST_OP_DATA_T)(scale * unscaled) }
// post op service
POST_OP_DATA_T sum_src#if WITH_SUM
sum_src = (POST_OP_DATA_T)DATA_TO_REF(dst[data_off])#endif
#define GET_DATA_IDX(dim_idx) \
(SOFTMAX_AXIS_IDX > (dim_idx)) ? dim[(dim_idx)] \
: (SOFTMAX_AXIS_IDX < (dim_idx)) ? dim[(dim_idx) - 1] \
: i
const unsigned po_d0 = SOFTMAX_AXIS_IDX > 0 ? dim[0] : i const unsigned po_d1 = GET_DATA_IDX(1) const unsigned po_d2 = GET_DATA_IDX(2) const unsigned po_d3 = GET_DATA_IDX(3) const unsigned po_d4 = GET_DATA_IDX(4)
APPLY_POST_OPS_SERIAL(
tmp, sum_src, po_d0, po_d1, po_d2, po_d3, po_d4, 0)#undef GET_DATA_IDX
#if WITH_DST_SCALES
tmp /= dst_scale[0]#endif
#if WITH_DROPOUT
#if !USE_HOST_SCALARS
long dropout_seed = dropout_seed_buf[0] long dropout_offset = USE_OFFSET ? dropout_offset_buf[0] : 0 float dropout_p = dropout_p_buf[0]#endif
uint dropout_threshold = get_dropout_threshold(dropout_p) float dropout_inv_q
= (dropout_p != 1.f) ? 1.f / (1.f - dropout_p) : 0.f#if USE_OFFSET
uint res = philox_4x32_s64(
(ulong)data_off, (ulong)dropout_seed, (ulong)dropout_offset)#else
uint res = philox_4x32(data_off, (uint)dropout_seed)#endif
uchar dropout = res > dropout_threshold tmp = (dropout) ? tmp * dropout_inv_q : 0#if HAS_OUTPUT_MASK
dropout_mask_buf[data_off] = dropout#endif
#endif
dst[data_off] = TO_DST(tmp) }
}
#endif
#if IS_BWD
__attribute__((reqd_work_group_size(GROUP_SIZE, 1, 1)))
__attribute__((intel_reqd_sub_group_size(SUB_GROUP_SIZE)))
__kernel void
simple_softmax_bwd_generic(__global DST_DATA_T *dst,
__global SRC_DATA_T *diff_src, __global DST_DATA_T *diff_dst) {
const off_t dim[] = {
(get_global_id(0) / GROUP_SIZE) % BLOCK_0,
get_global_id(1) % BLOCK_1,
get_global_id(2) % BLOCK_2,
(get_global_id(0) / GROUP_SIZE) / BLOCK_0,
get_global_id(1) / BLOCK_1,
get_global_id(2) / BLOCK_2,
}
const int local_id = get_local_id(0)
// begin and end indices calculated according to thread's id
const int begin = local_id * (SOFTMAX_AXIS / GROUP_SIZE) const int end = (local_id == GROUP_SIZE - 1)
? SOFTMAX_AXIS
: (local_id + 1) * (SOFTMAX_AXIS / GROUP_SIZE)#if SOFTMAX_AXIS - (GROUP_SIZE - 1) * (SOFTMAX_AXIS / GROUP_SIZE) \
> SOFTMAX_AXIS / GROUP_SIZE
const int buf_size
= SOFTMAX_AXIS - (GROUP_SIZE - 1) * (SOFTMAX_AXIS / GROUP_SIZE)#else
const int buf_size = SOFTMAX_AXIS / GROUP_SIZE#endif
#if IS_HALF(SRC_DATA_T) == 1 && IS_HALF(DST_DATA_TYPE) == 1
typedef half acc_t const acc_t acc_zero = 0.h#elif DT_F64 || DST_DT_F64
typedef double acc_t const acc_t acc_zero = 0.0#else
typedef float acc_t const acc_t acc_zero = 0.f#endif
acc_t diff_d[buf_size] acc_t d[buf_size] acc_t sbr = acc_zero
if (!(NEEDS_PADDING(dim[0], dim[1], dim[2], dim[3], dim[4], begin))) {
for (int i = begin off_t idx = DATA_OFF(dim[0], dim[1], dim[2], dim[3], dim[4], i) diff_d[i - begin] = DST_TO_REF(diff_dst[idx]) d[i - begin] = DST_TO_REF(dst[idx])#if LOGSOFTMAX
sbr += diff_d[i - begin]#else
sbr += diff_d[i - begin] * d[i - begin]#endif
}
}
#if GROUP_SIZE == SUB_GROUP_SIZE
sbr = sub_group_reduce_add(sbr)#else
sbr = work_group_reduce_add(sbr)#endif
for (int i = begin off_t idx = DATA_OFF(dim[0], dim[1], dim[2], dim[3], dim[4], i)
if (NEEDS_PADDING(dim[0], dim[1], dim[2], dim[3], dim[4], i)) {
diff_src[idx] = REF_TO_SRC(acc_zero) } else {
#if LOGSOFTMAX
diff_src[idx]
= REF_TO_SRC(diff_d[i - begin] - exp(d[i - begin]) * sbr)#else
acc_t inner_data = diff_d[i - begin] - sbr diff_src[idx] = REF_TO_SRC(d[i - begin] * inner_data)#endif
}
}
}
#endif