| 28 | |
| 29 | template<typename T> |
| 30 | Array<T> shift(const Array<T> &in, const int sdims[4]) { |
| 31 | // Shift should only be the first node in the JIT tree. |
| 32 | // Force input to be evaluated so that in is always a buffer. |
| 33 | in.eval(); |
| 34 | |
| 35 | string name_str("Sh"); |
| 36 | name_str += shortname<T>(true); |
| 37 | const dim4 &iDims = in.dims(); |
| 38 | dim4 oDims = iDims; |
| 39 | |
| 40 | array<int, 4> shifts{}; |
| 41 | for (int i = 0; i < 4; i++) { |
| 42 | // sdims_[i] will always be positive and always [0, oDims[i]]. |
| 43 | // Negative shifts are converted to position by going the other way |
| 44 | // round |
| 45 | shifts[i] = -(sdims[i] % static_cast<int>(oDims[i])) + |
| 46 | oDims[i] * (sdims[i] > 0); |
| 47 | assert(shifts[i] >= 0 && shifts[i] <= oDims[i]); |
| 48 | } |
| 49 | |
| 50 | auto node = make_shared<ShiftNode<T>>( |
| 51 | static_cast<af::dtype>(dtype_traits<T>::af_type), |
| 52 | static_pointer_cast<jit::BufferNode<T>>(in.getNode()), shifts); |
| 53 | return createNodeArray<T>(oDims, common::Node_ptr(node)); |
| 54 | } |
| 55 | |
| 56 | #define INSTANTIATE(T) \ |
| 57 | template Array<T> shift<T>(const Array<T> &in, const int sdims[4]); |