| 25 | |
| 26 | template<typename T> |
| 27 | void assign(Param<T> out, af::dim4 dDims, CParam<T> rhs, |
| 28 | std::vector<bool> const isSeq, std::vector<af_seq> const seqs, |
| 29 | std::vector<CParam<uint>> idxArrs) { |
| 30 | af::dim4 pDims = out.dims(); |
| 31 | // retrieve dimensions & strides for array to which rhs is being copied to |
| 32 | af::dim4 dst_offsets = toOffset(seqs, dDims); |
| 33 | af::dim4 dst_strides = toStride(seqs, dDims); |
| 34 | // retrieve rhs array dimenesions & strides |
| 35 | af::dim4 src_dims = rhs.dims(); |
| 36 | af::dim4 src_strides = rhs.strides(); |
| 37 | // declare pointers to af_array index data |
| 38 | uint const* const ptr0 = idxArrs[0].get(); |
| 39 | uint const* const ptr1 = idxArrs[1].get(); |
| 40 | uint const* const ptr2 = idxArrs[2].get(); |
| 41 | uint const* const ptr3 = idxArrs[3].get(); |
| 42 | |
| 43 | const T* src = rhs.get(); |
| 44 | T* dst = out.get(); |
| 45 | |
| 46 | for (dim_t l = 0; l < src_dims[3]; ++l) { |
| 47 | dim_t src_loff = l * src_strides[3]; |
| 48 | |
| 49 | dim_t dst_lIdx = |
| 50 | trimIndex(isSeq[3] ? l + dst_offsets[3] : ptr3[l], pDims[3]); |
| 51 | dim_t dst_loff = dst_lIdx * dst_strides[3]; |
| 52 | |
| 53 | for (dim_t k = 0; k < src_dims[2]; ++k) { |
| 54 | dim_t src_koff = k * src_strides[2]; |
| 55 | |
| 56 | dim_t dst_kIdx = |
| 57 | trimIndex(isSeq[2] ? k + dst_offsets[2] : ptr2[k], pDims[2]); |
| 58 | dim_t dst_koff = dst_kIdx * dst_strides[2]; |
| 59 | |
| 60 | for (dim_t j = 0; j < src_dims[1]; ++j) { |
| 61 | dim_t src_joff = j * src_strides[1]; |
| 62 | |
| 63 | dim_t dst_jIdx = trimIndex( |
| 64 | isSeq[1] ? j + dst_offsets[1] : ptr1[j], pDims[1]); |
| 65 | dim_t dst_joff = dst_jIdx * dst_strides[1]; |
| 66 | |
| 67 | for (dim_t i = 0; i < src_dims[0]; ++i) { |
| 68 | dim_t src_ioff = i * src_strides[0]; |
| 69 | dim_t src_idx = src_ioff + src_joff + src_koff + src_loff; |
| 70 | |
| 71 | dim_t dst_iIdx = trimIndex( |
| 72 | isSeq[0] ? i + dst_offsets[0] : ptr0[i], pDims[0]); |
| 73 | dim_t dst_ioff = dst_iIdx * dst_strides[0]; |
| 74 | dim_t dst_idx = dst_ioff + dst_joff + dst_koff + dst_loff; |
| 75 | |
| 76 | dst[dst_idx] = src[src_idx]; |
| 77 | } |
| 78 | } |
| 79 | } |
| 80 | } |
| 81 | } |
| 82 | |
| 83 | } // namespace kernel |
| 84 | } // namespace cpu |