| 189 | |
| 190 | template<typename Ti, typename To, af_op_t op, uint DIMX> |
| 191 | __global__ static void reduce_first_kernel(Param<To> out, CParam<Ti> in, |
| 192 | uint blocks_x, uint blocks_y, |
| 193 | uint repeat, bool change_nan, |
| 194 | To nanval) { |
| 195 | const uint tidx = threadIdx.x; |
| 196 | const uint tidy = threadIdx.y; |
| 197 | const uint tid = tidy * blockDim.x + tidx; |
| 198 | |
| 199 | const uint zid = blockIdx.x / blocks_x; |
| 200 | const uint blockIdx_x = blockIdx.x - (blocks_x)*zid; |
| 201 | const uint xid = blockIdx_x * blockDim.x * repeat + tidx; |
| 202 | |
| 203 | common::Binary<compute_t<To>, op> reduce; |
| 204 | common::Transform<Ti, compute_t<To>, op> transform; |
| 205 | |
| 206 | __shared__ compute_t<To> s_val[THREADS_PER_BLOCK]; |
| 207 | |
| 208 | const uint wid = (blockIdx.y + blockIdx.z * gridDim.y) / blocks_y; |
| 209 | const uint blockIdx_y = |
| 210 | (blockIdx.y + blockIdx.z * gridDim.y) - (blocks_y)*wid; |
| 211 | const uint yid = blockIdx_y * blockDim.y + tidy; |
| 212 | |
| 213 | const data_t<Ti> *const iptr = |
| 214 | in.ptr + |
| 215 | (wid * in.strides[3] + zid * in.strides[2] + yid * in.strides[1]); |
| 216 | |
| 217 | if (yid >= in.dims[1] || zid >= in.dims[2] || wid >= in.dims[3]) return; |
| 218 | |
| 219 | int lim = min((int)(xid + repeat * DIMX), in.dims[0]); |
| 220 | |
| 221 | compute_t<To> out_val = common::Binary<compute_t<To>, op>::init(); |
| 222 | for (int id = xid; id < lim; id += DIMX) { |
| 223 | compute_t<To> in_val = transform(iptr[id]); |
| 224 | if (change_nan) |
| 225 | in_val = |
| 226 | !IS_NAN(in_val) ? in_val : static_cast<compute_t<To>>(nanval); |
| 227 | out_val = reduce(in_val, out_val); |
| 228 | } |
| 229 | |
| 230 | s_val[tid] = out_val; |
| 231 | |
| 232 | __syncthreads(); |
| 233 | compute_t<To> *s_ptr = s_val + tidy * DIMX; |
| 234 | |
| 235 | if (DIMX == 256) { |
| 236 | if (tidx < 128) s_ptr[tidx] = reduce(s_ptr[tidx], s_ptr[tidx + 128]); |
| 237 | __syncthreads(); |
| 238 | } |
| 239 | |
| 240 | if (DIMX >= 128) { |
| 241 | if (tidx < 64) s_ptr[tidx] = reduce(s_ptr[tidx], s_ptr[tidx + 64]); |
| 242 | __syncthreads(); |
| 243 | } |
| 244 | |
| 245 | if (DIMX >= 64) { |
| 246 | if (tidx < 32) s_ptr[tidx] = reduce(s_ptr[tidx], s_ptr[tidx + 32]); |
| 247 | __syncthreads(); |
| 248 | } |