| 32 | |
| 33 | template<typename Ti, typename To, af_op_t op, uint dim, uint DIMY> |
| 34 | __global__ static void reduce_dim_kernel(Param<To> out, CParam<Ti> in, |
| 35 | uint blocks_x, uint blocks_y, |
| 36 | uint offset_dim, bool change_nan, |
| 37 | To nanval) { |
| 38 | const uint tidx = threadIdx.x; |
| 39 | const uint tidy = threadIdx.y; |
| 40 | const uint tid = tidy * THREADS_X + tidx; |
| 41 | |
| 42 | const uint zid = blockIdx.x / blocks_x; |
| 43 | const uint blockIdx_x = blockIdx.x - (blocks_x)*zid; |
| 44 | const uint xid = blockIdx_x * blockDim.x + tidx; |
| 45 | |
| 46 | __shared__ compute_t<To> s_val[THREADS_X * DIMY]; |
| 47 | |
| 48 | const uint wid = (blockIdx.y + blockIdx.z * gridDim.y) / blocks_y; |
| 49 | const uint blockIdx_y = |
| 50 | (blockIdx.y + blockIdx.z * gridDim.y) - (blocks_y)*wid; |
| 51 | const uint yid = blockIdx_y; // yid of output. updated for input later. |
| 52 | |
| 53 | uint ids[4] = {xid, yid, zid, wid}; |
| 54 | |
| 55 | // There is only one element per block for out |
| 56 | // There are blockDim.y elements per block for in |
| 57 | // Hence increment ids[dim] just after offseting out and before offsetting |
| 58 | // in |
| 59 | data_t<To> *const optr = out.ptr + ids[3] * out.strides[3] + |
| 60 | ids[2] * out.strides[2] + ids[1] * out.strides[1] + |
| 61 | ids[0]; |
| 62 | |
| 63 | const uint blockIdx_dim = ids[dim]; |
| 64 | ids[dim] = ids[dim] * blockDim.y + tidy; |
| 65 | |
| 66 | const data_t<Ti> *iptr = in.ptr + ids[3] * in.strides[3] + |
| 67 | ids[2] * in.strides[2] + ids[1] * in.strides[1] + |
| 68 | ids[0]; |
| 69 | |
| 70 | const uint id_dim_in = ids[dim]; |
| 71 | const uint istride_dim = in.strides[dim]; |
| 72 | |
| 73 | bool is_valid = (ids[0] < in.dims[0]) && (ids[1] < in.dims[1]) && |
| 74 | (ids[2] < in.dims[2]) && (ids[3] < in.dims[3]); |
| 75 | |
| 76 | common::Transform<Ti, compute_t<To>, op> transform; |
| 77 | common::Binary<compute_t<To>, op> reduce; |
| 78 | compute_t<To> out_val = common::Binary<compute_t<To>, op>::init(); |
| 79 | for (int id = id_dim_in; is_valid && (id < in.dims[dim]); |
| 80 | id += offset_dim * blockDim.y) { |
| 81 | compute_t<To> in_val = transform(*iptr); |
| 82 | if (change_nan) |
| 83 | in_val = !IS_NAN(in_val) ? in_val : compute_t<To>(nanval); |
| 84 | out_val = reduce(in_val, out_val); |
| 85 | iptr = iptr + offset_dim * blockDim.y * istride_dim; |
| 86 | } |
| 87 | |
| 88 | s_val[tid] = out_val; |
| 89 | |
| 90 | compute_t<To> *s_ptr = s_val + tid; |
| 91 | __syncthreads(); |