| 41 | |
| 42 | template<af_op_t op, typename Ti, typename Tk, typename To> |
| 43 | void reduce_by_key_dim(Array<Tk> &keys_out, Array<To> &vals_out, |
| 44 | const Array<Tk> &keys, const Array<Ti> &vals, |
| 45 | bool change_nan, double nanval, const int dim) { |
| 46 | std::vector<int> dim_ordering = {dim}; |
| 47 | for (int i = 0; i < 4; ++i) { |
| 48 | if (i != dim) { dim_ordering.push_back(i); } |
| 49 | } |
| 50 | |
| 51 | dim4 kdims = keys.dims(); |
| 52 | dim4 odims = vals.dims(); |
| 53 | |
| 54 | // allocate space for output and temporary working arrays |
| 55 | Array<Tk> reduced_keys = createEmptyArray<Tk>(kdims); |
| 56 | Array<To> reduced_vals = createEmptyArray<To>(odims); |
| 57 | Array<Tk> t_reduced_keys = createEmptyArray<Tk>(kdims); |
| 58 | Array<To> t_reduced_vals = createEmptyArray<To>(odims); |
| 59 | |
| 60 | // flags determining more reduction is necessary |
| 61 | auto needs_another_reduction = memAlloc<int>(1); |
| 62 | auto needs_block_boundary_reduction = memAlloc<int>(1); |
| 63 | |
| 64 | // reset flags |
| 65 | CUDA_CHECK(cudaMemsetAsync(needs_another_reduction.get(), 0, sizeof(int), |
| 66 | getActiveStream())); |
| 67 | CUDA_CHECK(cudaMemsetAsync(needs_block_boundary_reduction.get(), 0, |
| 68 | sizeof(int), getActiveStream())); |
| 69 | |
| 70 | int nelems = kdims[0]; |
| 71 | |
| 72 | const unsigned int numThreads = 128; |
| 73 | int numBlocksD0 = divup(nelems, numThreads); |
| 74 | |
| 75 | auto reduced_block_sizes = memAlloc<int>(numBlocksD0); |
| 76 | |
| 77 | size_t temp_storage_bytes = 0; |
| 78 | cub::DeviceScan::InclusiveSum( |
| 79 | NULL, temp_storage_bytes, reduced_block_sizes.get(), |
| 80 | reduced_block_sizes.get(), numBlocksD0, getActiveStream()); |
| 81 | auto d_temp_storage = memAlloc<char>(temp_storage_bytes); |
| 82 | |
| 83 | int n_reduced_host = nelems; |
| 84 | int needs_another_reduction_host; |
| 85 | int needs_block_boundary_reduction_host; |
| 86 | |
| 87 | bool first_pass = true; |
| 88 | do { |
| 89 | numBlocksD0 = divup(n_reduced_host, numThreads); |
| 90 | dim3 blocks(numBlocksD0, odims[dim_ordering[1]], |
| 91 | odims[dim_ordering[2]] * odims[dim_ordering[3]]); |
| 92 | |
| 93 | int folded_dim_sz = odims[dim_ordering[2]]; |
| 94 | if (first_pass) { |
| 95 | CUDA_LAUNCH( |
| 96 | (kernel::reduce_blocks_dim_by_key<Ti, Tk, To, op, numThreads>), |
| 97 | blocks, numThreads, reduced_block_sizes.get(), reduced_keys, |
| 98 | reduced_vals, keys, vals, nelems, change_nan, |
| 99 | scalar<To>(nanval), dim, folded_dim_sz); |
| 100 | POST_LAUNCH_CHECK(); |