| 18 | |
| 19 | template<typename T> |
| 20 | void index(Param<T> out, CParam<T> in, const af::dim4 dDims, |
| 21 | std::vector<bool> const isSeq, std::vector<af_seq> const seqs, |
| 22 | std::vector<CParam<uint>> idxArrs) { |
| 23 | const af::dim4 iDims = in.dims(); |
| 24 | const af::dim4 iOffs = toOffset(seqs, dDims); |
| 25 | const af::dim4 iStrds = in.strides(); |
| 26 | const af::dim4 oDims = out.dims(); |
| 27 | const af::dim4 oStrides = out.strides(); |
| 28 | const T* src = in.get(); |
| 29 | T* dst = out.get(); |
| 30 | const uint* ptr0 = idxArrs[0].get(); |
| 31 | const uint* ptr1 = idxArrs[1].get(); |
| 32 | const uint* ptr2 = idxArrs[2].get(); |
| 33 | const uint* ptr3 = idxArrs[3].get(); |
| 34 | |
| 35 | for (dim_t l = 0; l < oDims[3]; ++l) { |
| 36 | dim_t lOff = l * oStrides[3]; |
| 37 | dim_t inIdx3 = trimIndex( |
| 38 | isSeq[3] ? l * seqs[3].step + iOffs[3] : ptr3[l], iDims[3]); |
| 39 | dim_t inOff3 = inIdx3 * iStrds[3]; |
| 40 | |
| 41 | for (dim_t k = 0; k < oDims[2]; ++k) { |
| 42 | dim_t kOff = k * oStrides[2]; |
| 43 | dim_t inIdx2 = trimIndex( |
| 44 | isSeq[2] ? k * seqs[2].step + iOffs[2] : ptr2[k], iDims[2]); |
| 45 | dim_t inOff2 = inIdx2 * iStrds[2]; |
| 46 | |
| 47 | for (dim_t j = 0; j < oDims[1]; ++j) { |
| 48 | dim_t jOff = j * oStrides[1]; |
| 49 | dim_t inIdx1 = trimIndex( |
| 50 | isSeq[1] ? j * seqs[1].step + iOffs[1] : ptr1[j], iDims[1]); |
| 51 | dim_t inOff1 = inIdx1 * iStrds[1]; |
| 52 | |
| 53 | for (dim_t i = 0; i < oDims[0]; ++i) { |
| 54 | dim_t iOff = i * oStrides[0]; |
| 55 | dim_t inIdx0 = trimIndex( |
| 56 | isSeq[0] ? i * seqs[0].step + iOffs[0] : ptr0[i], |
| 57 | iDims[0]); |
| 58 | dim_t inOff0 = inIdx0 * iStrds[0]; |
| 59 | |
| 60 | dst[lOff + kOff + jOff + iOff] = |
| 61 | src[inOff3 + inOff2 + inOff1 + inOff0]; |
| 62 | } |
| 63 | } |
| 64 | } |
| 65 | } |
| 66 | } |
| 67 | |
| 68 | } // namespace kernel |
| 69 | } // namespace cpu |