| 289 | |
| 290 | template<typename T, af_op_t op, int dim> |
| 291 | void ireduce_dim(Param<T> out, Param<uint> oloc, Param<T> in, |
| 292 | Param<uint> rlen) { |
| 293 | uint threads_y = std::min(creduce::THREADS_Y, nextpow2(in.info.dims[dim])); |
| 294 | uint threads_x = creduce::THREADS_X; |
| 295 | |
| 296 | dim_t blocks_dim[] = {divup(in.info.dims[0], threads_x), in.info.dims[1], |
| 297 | in.info.dims[2], in.info.dims[3]}; |
| 298 | |
| 299 | blocks_dim[dim] = divup(in.info.dims[dim], threads_y * creduce::REPEAT); |
| 300 | |
| 301 | Param<T> tmp = out; |
| 302 | Param<uint> tlptr = oloc; |
| 303 | bufptr<T> tmp_alloc; |
| 304 | bufptr<uint> tlptr_alloc; |
| 305 | |
| 306 | if (blocks_dim[dim] > 1) { |
| 307 | int tmp_elements = 1; |
| 308 | tmp.info.dims[dim] = blocks_dim[dim]; |
| 309 | |
| 310 | for (int k = 0; k < 4; k++) tmp_elements *= tmp.info.dims[k]; |
| 311 | tmp_alloc = memAlloc<T>(tmp_elements); |
| 312 | tlptr_alloc = memAlloc<uint>(tmp_elements); |
| 313 | tmp.data = tmp_alloc.get(); |
| 314 | tlptr.data = tlptr_alloc.get(); |
| 315 | |
| 316 | for (int k = dim + 1; k < 4; k++) |
| 317 | tmp.info.strides[k] *= blocks_dim[dim]; |
| 318 | } |
| 319 | |
| 320 | Param<uint> nullparam; |
| 321 | ireduce_dim_launcher<T, op, dim, true>(tmp, tlptr, in, nullparam, threads_y, |
| 322 | blocks_dim, rlen); |
| 323 | |
| 324 | if (blocks_dim[dim] > 1) { |
| 325 | blocks_dim[dim] = 1; |
| 326 | |
| 327 | ireduce_dim_launcher<T, op, dim, false>(out, oloc, tmp, tlptr, |
| 328 | threads_y, blocks_dim, rlen); |
| 329 | } |
| 330 | } |
| 331 | |
| 332 | template<typename T, af_op_t op, bool is_first, uint DIMX> |
| 333 | class ireduceFirstKernelSMEM { |