| 223 | } |
| 224 | |
| 225 | af_err af_scan_by_key(af_array* out, const af_array key, const af_array in, |
| 226 | const int dim, af_binary_op op, bool inclusive_scan) { |
| 227 | try { |
| 228 | ARG_ASSERT(2, dim >= 0); |
| 229 | ARG_ASSERT(2, dim < 4); |
| 230 | |
| 231 | const ArrayInfo& in_info = getInfo(in); |
| 232 | const ArrayInfo& key_info = getInfo(key); |
| 233 | |
| 234 | if (dim >= static_cast<int>(in_info.ndims())) { |
| 235 | *out = retain(in); |
| 236 | return AF_SUCCESS; |
| 237 | } |
| 238 | |
| 239 | ARG_ASSERT(2, in_info.dims() == key_info.dims()); |
| 240 | |
| 241 | af_dtype type = in_info.getType(); |
| 242 | af_array res; |
| 243 | |
| 244 | switch (type) { |
| 245 | case f32: |
| 246 | res = scan_op<float, float>(key, in, dim, op, inclusive_scan); |
| 247 | break; |
| 248 | case f64: |
| 249 | res = scan_op<double, double>(key, in, dim, op, inclusive_scan); |
| 250 | break; |
| 251 | case c32: |
| 252 | res = scan_op<cfloat, cfloat>(key, in, dim, op, inclusive_scan); |
| 253 | break; |
| 254 | case c64: |
| 255 | res = |
| 256 | scan_op<cdouble, cdouble>(key, in, dim, op, inclusive_scan); |
| 257 | break; |
| 258 | case s16: |
| 259 | case s32: |
| 260 | case s8: |
| 261 | res = scan_op<int, int>(key, in, dim, op, inclusive_scan); |
| 262 | break; |
| 263 | case u64: |
| 264 | res = scan_op<uintl, uintl>(key, in, dim, op, inclusive_scan); |
| 265 | break; |
| 266 | case s64: |
| 267 | res = scan_op<intl, intl>(key, in, dim, op, inclusive_scan); |
| 268 | break; |
| 269 | case u16: |
| 270 | case u32: |
| 271 | case u8: |
| 272 | case b8: |
| 273 | res = scan_op<uint, uint>(key, in, dim, op, inclusive_scan); |
| 274 | break; |
| 275 | default: TYPE_ERROR(1, type); |
| 276 | } |
| 277 | |
| 278 | std::swap(*out, res); |
| 279 | } |
| 280 | CATCHALL; |
| 281 | |
| 282 | return AF_SUCCESS; |