| 21 | namespace cpu { |
| 22 | template<af_op_t op, typename Ti, typename Tk, typename To> |
| 23 | Array<To> scan(const Array<Tk>& key, const Array<Ti>& in, const int dim, |
| 24 | bool inclusive_scan) { |
| 25 | const dim4& dims = in.dims(); |
| 26 | Array<To> out = createEmptyArray<To>(dims); |
| 27 | kernel::scan_dim_by_key<op, Ti, Tk, To, 1> func1(inclusive_scan); |
| 28 | kernel::scan_dim_by_key<op, Ti, Tk, To, 2> func2(inclusive_scan); |
| 29 | kernel::scan_dim_by_key<op, Ti, Tk, To, 3> func3(inclusive_scan); |
| 30 | kernel::scan_dim_by_key<op, Ti, Tk, To, 4> func4(inclusive_scan); |
| 31 | |
| 32 | switch (in.ndims()) { |
| 33 | case 1: getQueue().enqueue(func1, out, 0, key, 0, in, 0, dim); break; |
| 34 | case 2: getQueue().enqueue(func2, out, 0, key, 0, in, 0, dim); break; |
| 35 | case 3: getQueue().enqueue(func3, out, 0, key, 0, in, 0, dim); break; |
| 36 | case 4: getQueue().enqueue(func4, out, 0, key, 0, in, 0, dim); break; |
| 37 | } |
| 38 | |
| 39 | return out; |
| 40 | } |
| 41 | |
| 42 | #define INSTANTIATE_SCAN_BY_KEY(ROp, Ti, Tk, To) \ |
| 43 | template Array<To> scan<ROp, Ti, Tk, To>( \ |