| 23 | |
| 24 | template<typename T> |
| 25 | Array<T> index(const Array<T>& in, const af_index_t idxrs[]) { |
| 26 | kernel::IndexKernelParam_t p; |
| 27 | std::vector<af_seq> seqs(4, af_span); |
| 28 | // create seq vector to retrieve output |
| 29 | // dimensions, offsets & offsets |
| 30 | for (dim_t x = 0; x < 4; ++x) { |
| 31 | if (idxrs[x].isSeq) { seqs[x] = idxrs[x].idx.seq; } |
| 32 | } |
| 33 | |
| 34 | // retrieve dimensions, strides and offsets |
| 35 | const dim4& iDims = in.dims(); |
| 36 | dim4 dDims = in.getDataDims(); |
| 37 | dim4 oDims = toDims(seqs, iDims); |
| 38 | dim4 iOffs = toOffset(seqs, dDims); |
| 39 | dim4 iStrds = in.strides(); |
| 40 | |
| 41 | for (dim_t i = 0; i < 4; ++i) { |
| 42 | p.isSeq[i] = idxrs[i].isSeq ? 1 : 0; |
| 43 | p.offs[i] = iOffs[i]; |
| 44 | p.strds[i] = iStrds[i]; |
| 45 | p.steps[i] = 0; |
| 46 | if (idxrs[i].isSeq) { |
| 47 | af_seq seq = idxrs[i].idx.seq; |
| 48 | // The step for af_span used in the kernel must be 1 |
| 49 | if (seq.begin == af_span.begin && seq.end == af_span.end && |
| 50 | seq.step == af_span.step) |
| 51 | p.steps[i] = 1; |
| 52 | else |
| 53 | p.steps[i] = seq.step; |
| 54 | } |
| 55 | } |
| 56 | |
| 57 | cl::Buffer* bPtrs[4]; |
| 58 | |
| 59 | auto buf = cl::Buffer(); |
| 60 | std::vector<Array<uint>> idxArrs(4, createEmptyArray<uint>(dim4())); |
| 61 | // look through indexs to read af_array indexs |
| 62 | for (dim_t x = 0; x < 4; ++x) { |
| 63 | // set index pointers were applicable |
| 64 | if (!p.isSeq[x]) { |
| 65 | idxArrs[x] = castArray<uint>(idxrs[x].idx.arr); |
| 66 | bPtrs[x] = idxArrs[x].get(); |
| 67 | // set output array ith dimension value |
| 68 | oDims[x] = idxArrs[x].elements(); |
| 69 | } else { |
| 70 | // alloc an 1-element buffer to avoid OpenCL from failing |
| 71 | bPtrs[x] = &buf; |
| 72 | } |
| 73 | } |
| 74 | |
| 75 | Array<T> out = createEmptyArray<T>(oDims); |
| 76 | if (oDims.elements() == 0) { return out; } |
| 77 | |
| 78 | kernel::index<T>(out, in, p, bPtrs); |
| 79 | |
| 80 | return out; |
| 81 | } |
| 82 | |