| 18 | |
| 19 | template<typename InT, typename IndexT> |
| 20 | void lookup(Param<InT> out, CParam<InT> input, CParam<IndexT> indices, |
| 21 | unsigned const dim) { |
| 22 | const af::dim4 iDims = input.dims(); |
| 23 | const af::dim4 oDims = out.dims(); |
| 24 | const af::dim4 iStrides = input.strides(); |
| 25 | const af::dim4 oStrides = out.strides(); |
| 26 | const InT *inPtr = input.get(); |
| 27 | const IndexT *idxPtr = indices.get(); |
| 28 | |
| 29 | InT *outPtr = out.get(); |
| 30 | |
| 31 | for (dim_t l = 0; l < oDims[3]; ++l) { |
| 32 | dim_t iLOff = iStrides[3] * |
| 33 | (dim == 3 ? trimIndex((dim_t)idxPtr[l], iDims[3]) : l); |
| 34 | dim_t oLOff = l * oStrides[3]; |
| 35 | |
| 36 | for (dim_t k = 0; k < oDims[2]; ++k) { |
| 37 | dim_t iKOff = |
| 38 | iStrides[2] * |
| 39 | (dim == 2 ? trimIndex((dim_t)idxPtr[k], iDims[2]) : k); |
| 40 | dim_t oKOff = k * oStrides[2]; |
| 41 | |
| 42 | for (dim_t j = 0; j < oDims[1]; ++j) { |
| 43 | dim_t iJOff = |
| 44 | iStrides[1] * |
| 45 | (dim == 1 ? trimIndex((dim_t)idxPtr[j], iDims[1]) : j); |
| 46 | dim_t oJOff = j * oStrides[1]; |
| 47 | |
| 48 | for (dim_t i = 0; i < oDims[0]; ++i) { |
| 49 | dim_t iIOff = |
| 50 | iStrides[0] * |
| 51 | (dim == 0 ? trimIndex((dim_t)idxPtr[i], iDims[0]) : i); |
| 52 | dim_t oIOff = i * oStrides[0]; |
| 53 | |
| 54 | outPtr[oLOff + oKOff + oJOff + oIOff] = |
| 55 | inPtr[iLOff + iKOff + iJOff + iIOff]; |
| 56 | } |
| 57 | } |
| 58 | } |
| 59 | } |
| 60 | } |
| 61 | |
| 62 | } // namespace kernel |
| 63 | } // namespace cpu |