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