| 135 | |
| 136 | template<typename Ti, typename Tk, typename To, af_op_t op> |
| 137 | void scanFirstByKey(Param &out, const Param &in, const Param &key, |
| 138 | const bool inclusiveScan) { |
| 139 | uint threads_x = nextpow2(std::max(32u, (uint)out.info.dims[0])); |
| 140 | threads_x = std::min(threads_x, THREADS_PER_GROUP); |
| 141 | uint threads_y = THREADS_PER_GROUP / threads_x; |
| 142 | |
| 143 | uint groups_x = divup(out.info.dims[0], threads_x * REPEAT); |
| 144 | uint groups_y = divup(out.info.dims[1], threads_y); |
| 145 | |
| 146 | if (groups_x == 1) { |
| 147 | scanFirstByKeyFinalLauncher<Ti, Tk, To, op>( |
| 148 | out, in, key, true, groups_x, groups_y, threads_x, inclusiveScan); |
| 149 | |
| 150 | } else { |
| 151 | Param tmp = out; |
| 152 | tmp.info.dims[0] = groups_x; |
| 153 | tmp.info.strides[0] = 1; |
| 154 | for (int k = 1; k < 4; k++) { |
| 155 | tmp.info.strides[k] = |
| 156 | tmp.info.strides[k - 1] * tmp.info.dims[k - 1]; |
| 157 | } |
| 158 | Param tmpflg = tmp; |
| 159 | Param tmpid = tmp; |
| 160 | |
| 161 | int tmp_elements = tmp.info.strides[3] * tmp.info.dims[3]; |
| 162 | |
| 163 | tmp.data = bufferAlloc(tmp_elements * sizeof(To)); |
| 164 | tmpflg.data = bufferAlloc(tmp_elements * sizeof(char)); |
| 165 | tmpid.data = bufferAlloc(tmp_elements * sizeof(int)); |
| 166 | |
| 167 | scanFirstByKeyNonfinalLauncher<Ti, Tk, To, op>( |
| 168 | out, tmp, tmpflg, tmpid, in, key, groups_x, groups_y, threads_x, |
| 169 | inclusiveScan); |
| 170 | |
| 171 | if (op == af_notzero_t) { |
| 172 | scanFirstByKeyFinalLauncher<To, char, To, af_add_t>( |
| 173 | tmp, tmp, tmpflg, false, 1, groups_y, threads_x, true); |
| 174 | } else { |
| 175 | scanFirstByKeyFinalLauncher<To, char, To, op>( |
| 176 | tmp, tmp, tmpflg, false, 1, groups_y, threads_x, true); |
| 177 | } |
| 178 | |
| 179 | bcastFirstByKeyLauncher<To, Tk, To, op>( |
| 180 | out, tmp, tmpid, groups_x, groups_y, threads_x, inclusiveScan); |
| 181 | |
| 182 | bufferFree(tmp.data); |
| 183 | bufferFree(tmpflg.data); |
| 184 | bufferFree(tmpid.data); |
| 185 | } |
| 186 | } |
| 187 | } // namespace kernel |
| 188 | |
| 189 | #define INSTANTIATE_SCAN_FIRST_BY_KEY(ROp, Ti, Tk, To) \ |
nothing calls this directly
no test coverage detected