| 195 | |
| 196 | template<af_op_t op, typename Ti, typename Tk, typename To> |
| 197 | void reduce_by_key_first(Array<Tk> &keys_out, Array<To> &vals_out, |
| 198 | const Array<Tk> &keys, const Array<Ti> &vals, |
| 199 | bool change_nan, double nanval) { |
| 200 | dim4 kdims = keys.dims(); |
| 201 | dim4 odims = vals.dims(); |
| 202 | |
| 203 | // allocate space for output and temporary working arrays |
| 204 | Array<Tk> reduced_keys = createEmptyArray<Tk>(kdims); |
| 205 | Array<To> reduced_vals = createEmptyArray<To>(odims); |
| 206 | Array<Tk> t_reduced_keys = createEmptyArray<Tk>(kdims); |
| 207 | Array<To> t_reduced_vals = createEmptyArray<To>(odims); |
| 208 | |
| 209 | // flags determining more reduction is necessary |
| 210 | auto needs_another_reduction = memAlloc<int>(1); |
| 211 | auto needs_block_boundary_reduction = memAlloc<int>(1); |
| 212 | |
| 213 | // reset flags |
| 214 | CUDA_CHECK(cudaMemsetAsync(needs_another_reduction.get(), 0, sizeof(int), |
| 215 | getActiveStream())); |
| 216 | CUDA_CHECK(cudaMemsetAsync(needs_block_boundary_reduction.get(), 0, |
| 217 | sizeof(int), getActiveStream())); |
| 218 | |
| 219 | int nelems = kdims[0]; |
| 220 | |
| 221 | const unsigned int numThreads = 128; |
| 222 | int numBlocksD0 = divup(nelems, numThreads); |
| 223 | |
| 224 | auto reduced_block_sizes = memAlloc<int>(numBlocksD0); |
| 225 | |
| 226 | size_t temp_storage_bytes = 0; |
| 227 | cub::DeviceScan::InclusiveSum( |
| 228 | NULL, temp_storage_bytes, reduced_block_sizes.get(), |
| 229 | reduced_block_sizes.get(), numBlocksD0, getActiveStream()); |
| 230 | auto d_temp_storage = memAlloc<char>(temp_storage_bytes); |
| 231 | |
| 232 | int n_reduced_host = nelems; |
| 233 | int needs_another_reduction_host; |
| 234 | int needs_block_boundary_reduction_host; |
| 235 | |
| 236 | bool first_pass = true; |
| 237 | do { |
| 238 | numBlocksD0 = divup(n_reduced_host, numThreads); |
| 239 | dim3 blocks(numBlocksD0, odims[1], odims[2] * odims[3]); |
| 240 | |
| 241 | if (first_pass) { |
| 242 | CUDA_LAUNCH( |
| 243 | (kernel::reduce_blocks_by_key<Ti, Tk, To, op, numThreads>), |
| 244 | blocks, numThreads, reduced_block_sizes.get(), reduced_keys, |
| 245 | reduced_vals, keys, vals, nelems, change_nan, |
| 246 | scalar<To>(nanval), odims[2]); |
| 247 | POST_LAUNCH_CHECK(); |
| 248 | first_pass = false; |
| 249 | } else { |
| 250 | constexpr af_op_t op2 = op == af_notzero_t ? af_add_t : op; |
| 251 | CUDA_LAUNCH( |
| 252 | (kernel::reduce_blocks_by_key<To, Tk, To, op2, numThreads>), |
| 253 | blocks, numThreads, reduced_block_sizes.get(), reduced_keys, |
| 254 | reduced_vals, t_reduced_keys, t_reduced_vals, n_reduced_host, |