| 247 | |
| 248 | template<typename Ti, typename To, af_op_t op> |
| 249 | static void scan_first(Param<To> out, Param<Ti> in, bool inclusive_scan) { |
| 250 | uint threads_x = nextpow2(std::max(32u, (uint)out.info.dims[0])); |
| 251 | threads_x = std::min(threads_x, THREADS_PER_BLOCK); |
| 252 | uint threads_y = THREADS_PER_BLOCK / threads_x; |
| 253 | |
| 254 | uint groups_x = divup(out.info.dims[0], threads_x * REPEAT); |
| 255 | uint groups_y = divup(out.info.dims[1], threads_y); |
| 256 | |
| 257 | if (groups_x == 1) { |
| 258 | scan_first_launcher<Ti, To, op>(out, out, in, groups_x, groups_y, |
| 259 | threads_x, true, inclusive_scan); |
| 260 | } else { |
| 261 | Param<To> tmp = out; |
| 262 | |
| 263 | tmp.info.dims[0] = groups_x; |
| 264 | tmp.info.strides[0] = 1; |
| 265 | for (int k = 1; k < 4; k++) |
| 266 | tmp.info.strides[k] = |
| 267 | tmp.info.strides[k - 1] * tmp.info.dims[k - 1]; |
| 268 | |
| 269 | int tmp_elements = tmp.info.strides[3] * tmp.info.dims[3]; |
| 270 | auto tmp_alloc = memAlloc<To>(tmp_elements); |
| 271 | tmp.data = tmp_alloc.get(); |
| 272 | |
| 273 | scan_first_launcher<Ti, To, op>(out, tmp, in, groups_x, groups_y, |
| 274 | threads_x, false, inclusive_scan); |
| 275 | |
| 276 | // FIXME: Is there an alternative to the if condition ? |
| 277 | if (op == af_notzero_t) { |
| 278 | scan_first_launcher<To, To, af_add_t>(tmp, tmp, tmp, 1, groups_y, |
| 279 | threads_x, true, true); |
| 280 | } else { |
| 281 | scan_first_launcher<To, To, op>(tmp, tmp, tmp, 1, groups_y, |
| 282 | threads_x, true, true); |
| 283 | } |
| 284 | |
| 285 | bcast_first_launcher<To, op>(out, tmp, groups_x, groups_y, threads_x, |
| 286 | inclusive_scan); |
| 287 | } |
| 288 | } |
| 289 | |
| 290 | } // namespace kernel |
| 291 | } // namespace oneapi |