| 291 | |
| 292 | template<typename Ti, typename To, af_op_t op, int dim> |
| 293 | static void scan_dim(Param<To> out, Param<Ti> in, bool inclusive_scan) { |
| 294 | uint threads_y = std::min(THREADS_Y, nextpow2(out.info.dims[dim])); |
| 295 | uint threads_x = THREADS_X; |
| 296 | |
| 297 | dim_t blocks_all[] = {divup(out.info.dims[0], threads_x), out.info.dims[1], |
| 298 | out.info.dims[2], out.info.dims[3]}; |
| 299 | |
| 300 | blocks_all[dim] = divup(out.info.dims[dim], threads_y * REPEAT); |
| 301 | |
| 302 | if (blocks_all[dim] == 1) { |
| 303 | scan_dim_launcher<Ti, To, op, dim>(out, out, in, threads_y, blocks_all, |
| 304 | true, inclusive_scan); |
| 305 | } else { |
| 306 | Param<To> tmp = out; |
| 307 | |
| 308 | tmp.info.dims[dim] = blocks_all[dim]; |
| 309 | tmp.info.strides[0] = 1; |
| 310 | for (int k = 1; k < 4; k++) |
| 311 | tmp.info.strides[k] = |
| 312 | tmp.info.strides[k - 1] * tmp.info.dims[k - 1]; |
| 313 | |
| 314 | int tmp_elements = tmp.info.strides[3] * tmp.info.dims[3]; |
| 315 | auto tmp_alloc = memAlloc<To>(tmp_elements); |
| 316 | tmp.data = tmp_alloc.get(); |
| 317 | |
| 318 | scan_dim_launcher<Ti, To, op, dim>(out, tmp, in, threads_y, blocks_all, |
| 319 | false, inclusive_scan); |
| 320 | |
| 321 | int bdim = blocks_all[dim]; |
| 322 | blocks_all[dim] = 1; |
| 323 | |
| 324 | // FIXME: Is there an alternative to the if condition ? |
| 325 | if (op == af_notzero_t) { |
| 326 | scan_dim_launcher<To, To, af_add_t, dim>(tmp, tmp, tmp, threads_y, |
| 327 | blocks_all, true, true); |
| 328 | } else { |
| 329 | scan_dim_launcher<To, To, op, dim>(tmp, tmp, tmp, threads_y, |
| 330 | blocks_all, true, true); |
| 331 | } |
| 332 | |
| 333 | blocks_all[dim] = bdim; |
| 334 | bcast_dim_launcher<To, op, dim>(out, tmp, threads_y, blocks_all, |
| 335 | inclusive_scan); |
| 336 | } |
| 337 | } |
| 338 | |
| 339 | } // namespace kernel |
| 340 | } // namespace oneapi |