-
Notifications
You must be signed in to change notification settings - Fork 555
Expand file tree
/
Copy pathindex.cpp
More file actions
94 lines (81 loc) · 2.59 KB
/
Copy pathindex.cpp
File metadata and controls
94 lines (81 loc) · 2.59 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
/*******************************************************
* Copyright (c) 2022, ArrayFire
* All rights reserved.
*
* This file is distributed under 3-clause BSD license.
* The complete license agreement can be obtained at:
* http://arrayfire.com/licenses/BSD-3-Clause
********************************************************/
#include <index.hpp>
#include <Array.hpp>
#include <err_oneapi.hpp>
#include <handle.hpp>
#include <kernel/assign_kernel_param.hpp>
#include <kernel/index.hpp>
#include <memory.hpp>
#include <af/dim4.hpp>
using arrayfire::common::half;
using arrayfire::oneapi::IndexKernelParam;
namespace arrayfire {
namespace oneapi {
template<typename T>
Array<T> index(const Array<T>& in, const af_index_t idxrs[]) {
IndexKernelParam p;
std::vector<af_seq> seqs(4, af_span);
// create seq vector to retrieve output
// dimensions, offsets & offsets
for (dim_t x = 0; x < 4; ++x) {
if (idxrs[x].isSeq) { seqs[x] = idxrs[x].idx.seq; }
}
// retrieve dimensions, strides and offsets
const dim4& iDims = in.dims();
dim4 dDims = in.getDataDims();
dim4 oDims = toDims(seqs, iDims);
dim4 iOffs = toOffset(seqs, dDims);
dim4 iStrds = in.strides();
for (dim_t i = 0; i < 4; ++i) {
p.isSeq[i] = idxrs[i].isSeq;
p.offs[i] = iOffs[i];
p.strds[i] = iStrds[i];
p.steps[i] = 0;
if (idxrs[i].isSeq) {
af_seq seq = idxrs[i].idx.seq;
// The step for af_span used in the kernel must be 1
if (seq.begin == af_span.begin && seq.end == af_span.end &&
seq.step == af_span.step)
p.steps[i] = 1;
else
p.steps[i] = seq.step;
}
}
std::vector<Array<uint>> idxArrs(4, createEmptyArray<uint>(dim4(1)));
// look through indexs to read af_array indexs
for (dim_t x = 0; x < 4; ++x) {
if (!p.isSeq[x]) {
idxArrs[x] = castArray<uint>(idxrs[x].idx.arr);
oDims[x] = idxArrs[x].elements();
}
}
Array<T> out = createEmptyArray<T>(oDims);
if (oDims.elements() == 0) { return out; }
kernel::index<T>(out, in, p, idxArrs);
return out;
}
#define INSTANTIATE(T) \
template Array<T> index<T>(const Array<T>& in, const af_index_t idxrs[]);
INSTANTIATE(cdouble)
INSTANTIATE(double)
INSTANTIATE(cfloat)
INSTANTIATE(float)
INSTANTIATE(int)
INSTANTIATE(uint)
INSTANTIATE(intl)
INSTANTIATE(uintl)
INSTANTIATE(schar)
INSTANTIATE(uchar)
INSTANTIATE(char)
INSTANTIATE(short)
INSTANTIATE(ushort)
INSTANTIATE(half)
} // namespace oneapi
} // namespace arrayfire