wehub-resource-sync
103 lines
2.7 KiB
Metal
103 lines
2.7 KiB
Metal
// DS4 Metal row-sum kernel.
|
|
|
|
#define FC_SUM_ROWS 1400
|
|
|
|
#define OP_SUM_ROWS_NUM_SUM_ROWS 10
|
|
#define OP_SUM_ROWS_NUM_MEAN 11
|
|
|
|
struct ds4_metal_args_sum_rows {
|
|
int64_t ne00;
|
|
int64_t ne01;
|
|
int64_t ne02;
|
|
int64_t ne03;
|
|
uint64_t nb00;
|
|
uint64_t nb01;
|
|
uint64_t nb02;
|
|
uint64_t nb03;
|
|
int64_t ne0;
|
|
int64_t ne1;
|
|
int64_t ne2;
|
|
int64_t ne3;
|
|
uint64_t nb0;
|
|
uint64_t nb1;
|
|
uint64_t nb2;
|
|
uint64_t nb3;
|
|
};
|
|
|
|
static inline float sum(float x) {
|
|
return x;
|
|
}
|
|
|
|
static inline float sum(float4 x) {
|
|
return x[0] + x[1] + x[2] + x[3];
|
|
}
|
|
|
|
constant short FC_sum_rows_op [[function_constant(FC_SUM_ROWS + 0)]];
|
|
|
|
// Reduces each row to a sum or mean. DS4 mainly uses the sum form to preserve
|
|
// the compressor-pooling graph boundary in the single-compressor-row case.
|
|
template <typename T0, typename T>
|
|
kernel void kernel_sum_rows_impl(
|
|
constant ds4_metal_args_sum_rows & args,
|
|
device const char * src0,
|
|
device char * dst,
|
|
threadgroup char * shmem [[threadgroup(0)]],
|
|
uint3 tgpig[[threadgroup_position_in_grid]],
|
|
ushort3 tpitg[[thread_position_in_threadgroup]],
|
|
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
|
ushort tiisg[[thread_index_in_simdgroup]],
|
|
ushort3 ntg[[threads_per_threadgroup]]) {
|
|
#define FC_OP FC_sum_rows_op
|
|
|
|
const int i3 = tgpig.z;
|
|
const int i2 = tgpig.y;
|
|
const int i1 = tgpig.x;
|
|
|
|
threadgroup T0 * shmem_t = (threadgroup T0 *) shmem;
|
|
|
|
if (sgitg == 0) {
|
|
shmem_t[tiisg] = 0.0f;
|
|
}
|
|
|
|
device const T0 * src_row = (device const T0 *) (src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03);
|
|
device T * dst_row = (device T *) (dst + i1*args.nb1 + i2*args.nb2 + i3*args.nb3);
|
|
|
|
T0 sumf = T0(0.0f);
|
|
|
|
for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) {
|
|
sumf += src_row[i0];
|
|
}
|
|
|
|
sumf = simd_sum(sumf);
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
if (tiisg == 0) {
|
|
shmem_t[sgitg] = sumf;
|
|
}
|
|
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
|
|
sumf = shmem_t[tiisg];
|
|
sumf = simd_sum(sumf);
|
|
|
|
if (tpitg.x == 0) {
|
|
if (FC_OP == OP_SUM_ROWS_NUM_MEAN) {
|
|
if (is_same<float4, T0>::value) {
|
|
dst_row[0] = sum(sumf) / (4*args.ne00);
|
|
} else {
|
|
dst_row[0] = sum(sumf) / args.ne00;
|
|
}
|
|
} else {
|
|
dst_row[0] = sum(sumf);
|
|
}
|
|
}
|
|
|
|
#undef FC_OP
|
|
}
|
|
|
|
typedef decltype(kernel_sum_rows_impl<float, float>) kernel_sum_rows_t;
|
|
|
|
// Host-visible F32 row reduction used by compressor pooling.
|
|
template [[host_name("kernel_sum_rows_f32_f32")]] kernel kernel_sum_rows_t kernel_sum_rows_impl<float, float>;
|