| 86 | |
| 87 | template<typename Ti, typename Tk, typename To, af_op_t op> |
| 88 | void scan_first_by_key(Param<To> out, CParam<Ti> in, CParam<Tk> key, |
| 89 | bool inclusive_scan) { |
| 90 | uint threads_x = nextpow2(std::max(32u, (uint)out.dims[0])); |
| 91 | threads_x = std::min(threads_x, THREADS_PER_BLOCK); |
| 92 | uint threads_y = THREADS_PER_BLOCK / threads_x; |
| 93 | |
| 94 | uint blocks_x = static_cast<uint>(divup(out.dims[0], threads_x * REPEAT)); |
| 95 | uint blocks_y = static_cast<uint>(divup(out.dims[1], threads_y)); |
| 96 | |
| 97 | if (blocks_x == 1) { |
| 98 | scan_final_launcher<Ti, Tk, To, op>(out, in, key, blocks_x, blocks_y, |
| 99 | threads_x, true, inclusive_scan); |
| 100 | } else { |
| 101 | Param<char> tmpflg; |
| 102 | Param<int> tmpid; |
| 103 | Param<To> tmp = out; |
| 104 | tmp.dims[0] = blocks_x; |
| 105 | tmp.strides[0] = 1; |
| 106 | for (int k = 1; k < AF_MAX_DIMS; k++) { |
| 107 | tmp.strides[k] = tmp.strides[k - 1] * tmp.dims[k - 1]; |
| 108 | } |
| 109 | for (int k = 0; k < AF_MAX_DIMS; k++) { |
| 110 | tmpflg.dims[k] = tmp.dims[k]; |
| 111 | tmpflg.strides[k] = tmp.strides[k]; |
| 112 | tmpid.dims[k] = tmp.dims[k]; |
| 113 | tmpid.strides[k] = tmp.strides[k]; |
| 114 | } |
| 115 | |
| 116 | int tmp_elements = tmp.strides[3] * tmp.dims[3]; |
| 117 | auto tmp_alloc = memAlloc<To>(tmp_elements); |
| 118 | auto tmpflg_alloc = memAlloc<char>(tmp_elements); |
| 119 | auto tmpid_alloc = memAlloc<int>(tmp_elements); |
| 120 | tmp.ptr = tmp_alloc.get(); |
| 121 | tmpflg.ptr = tmpflg_alloc.get(); |
| 122 | tmpid.ptr = tmpid_alloc.get(); |
| 123 | |
| 124 | scan_nonfinal_launcher<Ti, Tk, To, op>(out, tmp, tmpflg, tmpid, in, key, |
| 125 | blocks_x, blocks_y, threads_x, |
| 126 | inclusive_scan); |
| 127 | |
| 128 | scan_final_launcher<To, char, To, op>(tmp, tmp, tmpflg, 1, blocks_y, |
| 129 | threads_x, false, true); |
| 130 | |
| 131 | bcast_first_launcher<To, op>(out, tmp, tmpid, blocks_x, blocks_y, |
| 132 | threads_x); |
| 133 | } |
| 134 | } |
| 135 | } // namespace kernel |
| 136 | |
| 137 | #define INSTANTIATE_SCAN_FIRST_BY_KEY(ROp, Ti, Tk, To) \ |