/*******************************************************************************
* 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/include/post_ops.h"
#include "gpu/intel/include/types.h"
#define offset6D(d0, d1, d2, d3, d4, d5, s0, s1, s2, s3, s4, s5) \
((d0) * (s0) + (d1) * (s1) + (d2) * (s2) + (d3) * (s3) + (d4) * (s4) \
+ (d5) * (s5))
__kernel void ref_matmul(__global SRC_DATA_T *A, __global WEI_DATA_T *B,
__global DST_DATA_T *C, __global BIA_DATA_T *bia,
#if WITH_HOST_SRC_ZP
SRC_ZP_DATA_T a0_value,
#else
__global SRC_ZP_DATA_T *a0,
#endif
long src_zp_stride_k, long src_zp_stride_m, long src_zp_group_k,
#if WITH_HOST_WEI_ZP
WEI_ZP_DATA_T b0_value,
#else
__global WEI_ZP_DATA_T *b0,
#endif
long wei_zp_stride_n, long wei_zp_stride_k, long wei_zp_stride_d0,
long wei_zp_stride_d1, long wei_zp_group_n, long wei_zp_group_k,
#if WITH_HOST_DST_ZP
int c0_value,
#else
__global int *c0,
#endif
#if WITH_HOST_SRC_SCALE
SRC_SCALES_DATA_T src_scale_value,
#else
__global SRC_SCALES_DATA_T *src_scales,
#endif
long src_scale_stride_k, long src_scale_stride_m,
long src_scale_stride_d0, long src_scale_stride_d1,
long src_scale_group_m, long src_scale_group_k,
#if WITH_HOST_WEI_SCALE
WEI_SCALES_DATA_T wei_scale_value,
#else
__global WEI_SCALES_DATA_T *wei_scales,
#endif
long wei_scale_stride_n, long wei_scale_stride_k,
long wei_scale_stride_d0, long wei_scale_stride_d1,
long wei_scale_group_n, long wei_scale_group_k,
#if WITH_HOST_DST_SCALE
DST_SCALES_DATA_T dst_scale_value,
#else
__global DST_SCALES_DATA_T *dst_scales,
#endif
__global SRC_GS_DATA_T *ag, long src_gs_stride_k, long src_gs_stride_m,
long src_gs_stride_d0, long src_gs_stride_d1, long src_gs_group_k,
long group_K, long K, long N, long M, long D0, long D1, long D2,
long bia_stride_d3, long bia_stride_d2, long bia_stride_d1,
long bia_stride_d0, long bia_stride_m, long bia_stride_n,
long a_stride_d3, long a_stride_d2, long a_stride_d1, long a_stride_d0,
long a_stride_m, long a_stride_k, long b_stride_d3, long b_stride_d2,
long b_stride_d1, long b_stride_d0, long b_stride_k, long b_stride_n,
long c_stride_d3, long c_stride_d2, long c_stride_d1, long c_stride_d0,
long c_stride_m, long c_stride_n
#if WITH_DROPOUT
,
__global uchar *dropout_mask_buf,
#if USE_HOST_SCALARS
long dropout_seed, long dropout_offset, float dropout_p
#else
__global SEED_DATA_T *dropout_seed_buf,
__global long *dropout_offset_buf, __global float *dropout_p_buf
#endif
#endif
#if WITH_SROUND
,
__global uint *sround_seed_buf
#endif
POST_OP_ARGS) {
#if WITH_HOST_SRC_ZP
SRC_ZP_DATA_T *a0 = &a0_value#endif
#if WITH_HOST_WEI_ZP
WEI_ZP_DATA_T *b0 = &b0_value#endif
#if WITH_HOST_DST_ZP
int *c0 = &c0_value#endif
#if WITH_HOST_SRC_SCALE
SRC_SCALES_DATA_T *src_scales = &src_scale_value#endif
#if WITH_HOST_WEI_SCALE
WEI_SCALES_DATA_T *wei_scales = &wei_scale_value#endif
#if WITH_HOST_DST_SCALE
DST_SCALES_DATA_T *dst_scales = &dst_scale_value#endif
#if WITH_DROPOUT
#if !USE_HOST_SCALARS
SEED_DATA_T 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#endif
#if WITH_SROUND
uint sround_seed = sround_seed_buf[0]#endif
long m = get_global_id(0) long n = get_global_id(1) int mb = get_global_id(2)
#if WITH_DST_ZPOINTS
int dst_zp = c0[0]#else
int dst_zp = 0#endif
// NOTE: non-standard notation
// In matmul all dimensions above 2 are considered as batch
// In below code that deals with broadcast of batch dimensions,
// d0 means lowest batch dimension and d3 the highest
// decompose mb into batch dimensions (d0..d3)
long d3 = mb / D0 / D1 / D2 long d2 = (mb / D0 / D1) % D2 long d1 = (mb / D0) % D1 long d0 = mb % D0
// With groups, compute `k` over each group, and iterate over k_groups.
// Inside each group, compute acc as `ACC_DATA_T` but once reduction
// happens, convert to float and apply scales.
long n_groups_k = K / group_K
FLT_ACC_DATA_T acc = 0.f for (long g = 0 ACC_DATA_T acc_g = 0 for (long k_g = 0 auto k = k_g + g * group_K#if RUNTIME_DIMS
long src_off
= offset6D(m, k, d0, d1, d2, d3, a_stride_m, a_stride_k,
a_stride_d0, a_stride_d1, a_stride_d2, a_stride_d3) long wei_off
= offset6D(k, n, d0, d1, d2, d3, b_stride_k, b_stride_n,
b_stride_d0, b_stride_d1, b_stride_d2, b_stride_d3)#else
#if NDIMS == 5
long src_off = SRC_OFF(d2 % SRC_D0, d1 % SRC_D1, d0 % SRC_D2, m, k) long wei_off
= WEI_OFF(0, d2 % WEI_D0, d1 % WEI_D1, d0 % WEI_D2, k, n)#elif NDIMS == 4
long src_off = SRC_OFF(d1 % SRC_D0, d0 % SRC_D1, 0, m, k) long wei_off = WEI_OFF(0, d1 % WEI_D0, d0 % WEI_D1, 0, k, n)#elif NDIMS == 3
long src_off = SRC_OFF(d0 % SRC_D0, m, 0, 0, k) long wei_off = WEI_OFF(0, d0 % WEI_D0, k, 0, 0, n)#else
long src_off = SRC_OFF(m, k, 0, 0, 0) long wei_off = WEI_OFF(0, k, n, 0, 0, 0)#endif
#endif
int wei_zp = 0#if WITH_WEI_ZPOINTS && !WITH_SRC_GROUP_SUMS
long wei_zp_off = wei_zp_stride_n * (n / wei_zp_group_n)
+ wei_zp_stride_k * (k / wei_zp_group_k)
+ wei_zp_stride_d0 * d0 + wei_zp_stride_d1 * d1 wei_zp = WEI_ZP_TO_REF(b0, wei_zp_off)#endif
int src_zp = 0#if WITH_SRC_ZPOINTS && !WITH_SRC_GROUP_SUMS
long src_zp_off = src_zp_stride_k * (k / src_zp_group_k)
+ src_zp_stride_m * m src_zp = SRC_ZP_TO_REF(a0, src_zp_off)#endif
#if SRC_DT_F4_E2M1 || SRC_DT_F4_E3M0
ACC_DATA_T s
= TO_ACC(SRC_TO_REF(GET_HALF_BYTE(A, src_off)) - src_zp)#else
ACC_DATA_T s = TO_ACC(SRC_TO_REF(A[src_off]) - src_zp)#endif
#if WEI_DT_S4 || WEI_DT_U4 || WEI_DT_F4_E2M1 || WEI_DT_F4_E3M0
ACC_DATA_T w_raw = WEI_TO_REF(GET_HALF_BYTE(B, wei_off))#else
ACC_DATA_T w_raw = WEI_TO_REF(B[wei_off])#endif
ACC_DATA_T w = TO_ACC(w_raw - wei_zp) acc_g += s * w }
FLT_ACC_DATA_T src_scale = 1.f FLT_ACC_DATA_T wei_scale = 1.f#if WITH_SRC_SCALES
long src_scale_off = src_scale_stride_m * (m / src_scale_group_m)
+ src_scale_stride_k * (g * group_K / src_scale_group_k)
+ src_scale_stride_d0 * d0 + src_scale_stride_d1 * d1 src_scale = SRC_SCALES_TO_REF(src_scales[src_scale_off])#endif
#if WITH_WEI_SCALES
long wei_scale_off = wei_scale_stride_n * (n / wei_scale_group_n)
+ wei_scale_stride_k * (g * group_K / wei_scale_group_k)
+ wei_scale_stride_d0 * d0 + wei_scale_stride_d1 * d1 wei_scale = WEI_SCALES_TO_REF(wei_scales[wei_scale_off])#endif
FLT_ACC_DATA_T acc_g_to_f = ACC_TO_REF(acc_g) * src_scale * wei_scale acc += acc_g_to_f }
#if WITH_SRC_GROUP_SUMS
for (int g = 0, gend = K / src_gs_group_k FLT_ACC_DATA_T src_scale = 1.f FLT_ACC_DATA_T wei_scale = 1.f#if WITH_SRC_SCALES
long src_scale_g = g * src_gs_group_k / src_scale_group_k long src_scale_off = src_scale_stride_m * (m / wei_scale_group_n)
+ src_scale_stride_k * src_scale_g + src_scale_stride_d0 * d0
+ src_scale_stride_d1 * d1 src_scale = SRC_SCALES_TO_REF(src_scales[src_scale_off])#endif
#if WITH_WEI_SCALES
long wei_scale_g = g * src_gs_group_k / wei_scale_group_k long wei_scale_off = wei_scale_stride_n * (n / wei_scale_group_n)
+ wei_scale_stride_k * wei_scale_g + wei_scale_stride_d0 * d0
+ wei_scale_stride_d1 * d1 wei_scale = WEI_SCALES_TO_REF(wei_scales[wei_scale_off])#endif
long src_gs_off = src_gs_stride_m * m + src_gs_stride_k * g
+ src_gs_stride_d0 * d0 + src_gs_stride_d1 * d1 int src_gs = SRC_ZP_TO_REF(ag, src_gs_off) long wei_zp_off = wei_zp_stride_n * (n / wei_zp_group_n)
+ wei_zp_stride_k * (g * src_gs_group_k / wei_zp_group_k)
+ wei_zp_stride_d0 * d0 + wei_zp_stride_d1 * d1 int wei_zp = WEI_ZP_TO_REF(b0, wei_zp_off) acc -= src_scale * wei_scale * TO_ACC(src_gs) * TO_ACC(wei_zp) }
#endif
#if RUNTIME_DIMS
long dst_off = offset6D(m, n, d0, d1, d2, d3, c_stride_m, c_stride_n,
c_stride_d0, c_stride_d1, c_stride_d2, c_stride_d3)#else
#if NDIMS == 5
long dst_off = DST_OFF(d2 % DST_D0, d1 % DST_D1, d0 % DST_D2, m, n)#elif NDIMS == 4
long dst_off = DST_OFF(d1 % DST_D0, d0 % DST_D1, 0, m, n)#elif NDIMS == 3
long dst_off = DST_OFF(d0 % DST_D0, m, 0, 0, n)#else
long dst_off = DST_OFF(m, n, 0, 0, 0)#endif
#endif
#if WITH_BIAS || WITH_DROPOUT || NON_DEFAULT_ATTRS
POST_OP_DATA_T temp = (POST_OP_DATA_T)acc#if WITH_BIAS
long bia_off = offset6D(m, n, d0, d1, d2, d3, bia_stride_m, bia_stride_n,
bia_stride_d0, bia_stride_d1, bia_stride_d2, bia_stride_d3) temp += BIA_TO_REF(bia[bia_off])#endif // WITH_BIAS
float dst_data#if WITH_SUM
dst_data = convert_float(DATA_TO_REF(C[dst_off]))#endif // WITH_SUM
float po_acc = convert_float(temp)
#if WITH_DROPOUT
#if WITH_SEED_S64 && USE_OFFSET
uint res = philox_4x32_s64(
dst_off, (ulong)dropout_seed, (ulong)dropout_offset)#else
uint res = philox_4x32((uint)dst_off, (uint)dropout_seed)#endif
uchar dropout = res > dropout_threshold po_acc = (dropout) ? po_acc * dropout_inv_q : 0#if HAS_OUTPUT_MASK
dropout_mask_buf[dst_off] = dropout#endif
#endif
#if WITH_SROUND
po_acc = stochastic_round_fwd(po_acc, dst_off, sround_seed)#endif
if (DST_NDIMS == 2)
APPLY_POST_OPS_SERIAL(po_acc, dst_data, m, n, 0, 0, 0, 0) if (DST_NDIMS == 3)
APPLY_POST_OPS_SERIAL(po_acc, dst_data, d0, m, n, 0, 0, 0) if (DST_NDIMS == 4)
APPLY_POST_OPS_SERIAL(po_acc, dst_data, d1, d0, m, n, 0, 0) if (DST_NDIMS == 5)
APPLY_POST_OPS_SERIAL(po_acc, dst_data, d2, d1, d0, m, n, 0) if (DST_NDIMS == 6)
APPLY_POST_OPS_SERIAL(po_acc, dst_data, d3, d2, d1, d0, m, n)
#if WITH_DST_SCALES
#if DST_SCALES_MASK == 0
po_acc /= DST_SCALES_TO_REF(dst_scales[0])#elif WITH_DYN_DST_SCALE == 0
po_acc /= DST_SCALES_TO_REF(dst_scales[n])#endif
#endif
po_acc += dst_zp
#if WITH_DYN_DST_SCALE
((__global ACC_DATA_T *)C)[dst_off] = po_acc#else
C[dst_off] = TO_DST(po_acc)#endif
#else // WITH_BIAS || NON_DEFAULT_ATTRS
#if WITH_DYN_DST_SCALE
((__global ACC_DATA_T *)C)[dst_off] = acc#else
C[dst_off] = TO_DST(acc)#endif
#endif // WITH_BIAS || NON_DEFAULT_ATTRS
}