| 261 | |
| 262 | template<typename Ti, typename To, af_op_t op, uint DIMX> |
| 263 | __global__ static void reduce_all_kernel(Param<To> out, |
| 264 | Param<unsigned> retirementCount, |
| 265 | Param<To> tmp, CParam<Ti> in, |
| 266 | uint blocks_x, uint blocks_y, |
| 267 | uint repeat, bool change_nan, |
| 268 | To nanval) { |
| 269 | const uint tidx = threadIdx.x; |
| 270 | const uint tidy = threadIdx.y; |
| 271 | const uint tid = tidy * DIMX + tidx; |
| 272 | |
| 273 | const uint zid = blockIdx.x / blocks_x; |
| 274 | const uint blockIdx_x = blockIdx.x - (blocks_x)*zid; |
| 275 | const uint xid = blockIdx_x * blockDim.x * repeat + tidx; |
| 276 | |
| 277 | const uint wid = (blockIdx.y + blockIdx.z * gridDim.y) / blocks_y; |
| 278 | const uint blockIdx_y = |
| 279 | (blockIdx.y + blockIdx.z * gridDim.y) - (blocks_y)*wid; |
| 280 | const uint yid = blockIdx_y * blockDim.y + tidy; |
| 281 | |
| 282 | common::Binary<compute_t<To>, op> reduce; |
| 283 | common::Transform<Ti, compute_t<To>, op> transform; |
| 284 | |
| 285 | const int nwarps = THREADS_PER_BLOCK / 32; |
| 286 | __shared__ compute_t<To> s_val[nwarps]; |
| 287 | |
| 288 | const data_t<Ti> *const iptr = |
| 289 | in.ptr + |
| 290 | (wid * in.strides[3] + zid * in.strides[2] + yid * in.strides[1]); |
| 291 | |
| 292 | bool cond = yid < in.dims[1] && zid < in.dims[2] && wid < in.dims[3]; |
| 293 | |
| 294 | int lim = min((int)(xid + repeat * DIMX), in.dims[0]); |
| 295 | |
| 296 | compute_t<To> out_val = common::Binary<compute_t<To>, op>::init(); |
| 297 | for (int id = xid; cond && id < lim; id += DIMX) { |
| 298 | compute_t<To> in_val = transform(iptr[id]); |
| 299 | if (change_nan) |
| 300 | in_val = |
| 301 | !IS_NAN(in_val) ? in_val : static_cast<compute_t<To>>(nanval); |
| 302 | out_val = reduce(in_val, out_val); |
| 303 | } |
| 304 | |
| 305 | const int warpid = tid / 32; |
| 306 | const int lid = tid % 32; |
| 307 | |
| 308 | typedef cub::WarpReduce<compute_t<To>> WarpReduce; |
| 309 | __shared__ typename WarpReduce::TempStorage temp_storage[nwarps]; |
| 310 | |
| 311 | out_val = WarpReduce(temp_storage[warpid]).Reduce(out_val, reduce); |
| 312 | |
| 313 | if (cond && lid == 0) { |
| 314 | s_val[warpid] = out_val; |
| 315 | } else if (!cond) { |
| 316 | s_val[warpid] = common::Binary<compute_t<To>, op>::init(); |
| 317 | } |
| 318 | __syncthreads(); |
| 319 | |
| 320 | if (tid < 32) { |