| 48 | |
| 49 | template<typename Ti, typename Tk, typename To, af_op_t op> |
| 50 | void reduceBlocksByKey(sycl::buffer<int> &reduced_block_sizes, |
| 51 | Array<Tk> keys_out, Array<To> vals_out, |
| 52 | const Array<Tk> keys, const Array<Ti> vals, |
| 53 | int change_nan, double nanval, const int n, |
| 54 | const int threads_x) { |
| 55 | int numBlocks = divup(n, threads_x); |
| 56 | |
| 57 | sycl::range<3> local(threads_x, 1, 1); |
| 58 | sycl::range<3> global(local[0] * numBlocks, vals_out.dims()[1], |
| 59 | vals_out.dims()[2] * vals_out.dims()[3]); |
| 60 | |
| 61 | auto keys_out_get = keys_out.get(); |
| 62 | auto vals_out_get = vals_out.get(); |
| 63 | auto keys_get = keys.get(); |
| 64 | auto vals_get = vals.get(); |
| 65 | getQueue().submit([&](sycl::handler &h) { |
| 66 | sycl::accessor<int> reduced_block_sizes_acc{reduced_block_sizes, h}; |
| 67 | write_accessor<Tk> keys_out_acc{*keys_out_get, h}; |
| 68 | write_accessor<To> vals_out_acc{*vals_out_get, h}; |
| 69 | read_accessor<Tk> keys_acc{*keys_get, h}; |
| 70 | read_accessor<Ti> vals_acc{*vals_get, h}; |
| 71 | |
| 72 | auto l_keys = sycl::local_accessor<Tk>(threads_x, h); |
| 73 | auto l_vals = sycl::local_accessor<compute_t<To>>(threads_x, h); |
| 74 | auto l_reduced_keys = sycl::local_accessor<Tk>(threads_x, h); |
| 75 | auto l_reduced_vals = sycl::local_accessor<compute_t<To>>(threads_x, h); |
| 76 | auto l_unique_ids = sycl::local_accessor<int>(threads_x, h); |
| 77 | auto l_wq_temp = sycl::local_accessor<int>(threads_x, h); |
| 78 | auto l_unique_flags = sycl::local_accessor<int>(threads_x, h); |
| 79 | auto l_reduced_block_size = sycl::local_accessor<int>(1, h); |
| 80 | |
| 81 | h.parallel_for( |
| 82 | sycl::nd_range<3>(global, local), |
| 83 | kernel::reduceBlocksByKeyKernel<Ti, Tk, To, op>( |
| 84 | reduced_block_sizes_acc, keys_out_acc, keys_out, vals_out_acc, |
| 85 | vals_out, keys_acc, keys, vals_acc, vals, change_nan, |
| 86 | scalar<To>(nanval), n, static_cast<int>(vals_out.dims()[2]), |
| 87 | threads_x, l_keys, l_vals, l_reduced_keys, l_reduced_vals, |
| 88 | l_unique_ids, l_wq_temp, l_unique_flags, l_reduced_block_size)); |
| 89 | }); |
| 90 | ONEAPI_DEBUG_FINISH(getQueue()); |
| 91 | } |
| 92 | |
| 93 | template<typename Ti, typename Tk, typename To, af_op_t op> |
| 94 | void reduceBlocksByKeyDim(sycl::buffer<int> &reduced_block_sizes, |