| 92 | |
| 93 | template<typename Ti, typename Tk, typename To, af_op_t op> |
| 94 | void reduceBlocksByKeyDim(sycl::buffer<int> &reduced_block_sizes, |
| 95 | Array<Tk> keys_out, Array<To> vals_out, |
| 96 | const Array<Tk> keys, const Array<Ti> vals, |
| 97 | int change_nan, double nanval, const int n, |
| 98 | const int threads_x, const int dim, |
| 99 | std::vector<int> dim_ordering) { |
| 100 | int numBlocks = divup(n, threads_x); |
| 101 | |
| 102 | sycl::range<3> local(threads_x, 1, 1); |
| 103 | sycl::range<3> global( |
| 104 | local[0] * numBlocks, vals_out.dims()[dim_ordering[1]], |
| 105 | vals_out.dims()[dim_ordering[2]] * vals_out.dims()[dim_ordering[3]]); |
| 106 | |
| 107 | auto keys_out_get = keys_out.get(); |
| 108 | auto vals_out_get = vals_out.get(); |
| 109 | auto keys_get = keys.get(); |
| 110 | auto vals_get = vals.get(); |
| 111 | getQueue().submit([&](sycl::handler &h) { |
| 112 | sycl::accessor<int> reduced_block_sizes_acc{reduced_block_sizes, h}; |
| 113 | write_accessor<Tk> keys_out_acc{*keys_out_get, h}; |
| 114 | write_accessor<To> vals_out_acc{*vals_out_get, h}; |
| 115 | read_accessor<Tk> keys_acc{*keys_get, h}; |
| 116 | read_accessor<Ti> vals_acc{*vals_get, h}; |
| 117 | |
| 118 | auto l_keys = sycl::local_accessor<Tk>(threads_x, h); |
| 119 | auto l_vals = sycl::local_accessor<compute_t<To>>(threads_x, h); |
| 120 | auto l_reduced_keys = sycl::local_accessor<Tk>(threads_x, h); |
| 121 | auto l_reduced_vals = sycl::local_accessor<compute_t<To>>(threads_x, h); |
| 122 | auto l_unique_ids = sycl::local_accessor<int>(threads_x, h); |
| 123 | auto l_wq_temp = sycl::local_accessor<int>(threads_x, h); |
| 124 | auto l_unique_flags = sycl::local_accessor<int>(threads_x, h); |
| 125 | auto l_reduced_block_size = sycl::local_accessor<int>(1, h); |
| 126 | |
| 127 | h.parallel_for( |
| 128 | sycl::nd_range<3>(global, local), |
| 129 | kernel::reduceBlocksByKeyDimKernel<Ti, Tk, To, op>( |
| 130 | reduced_block_sizes_acc, keys_out_acc, keys_out, vals_out_acc, |
| 131 | vals_out, keys_acc, keys, vals_acc, vals, change_nan, |
| 132 | scalar<To>(nanval), n, static_cast<int>(vals_out.dims()[2]), |
| 133 | threads_x, dim, l_keys, l_vals, l_reduced_keys, l_reduced_vals, |
| 134 | l_unique_ids, l_wq_temp, l_unique_flags, l_reduced_block_size)); |
| 135 | }); |
| 136 | ONEAPI_DEBUG_FINISH(getQueue()); |
| 137 | } |
| 138 | |
| 139 | template<typename Tk, typename To, af_op_t op> |
| 140 | void finalBoundaryReduce(sycl::buffer<int> &reduced_block_sizes, Array<Tk> keys, |