Skip to content

Commit ff4a4e2

Browse files
9prady99prady9
committed
FEATURE: generalized indexing function
Accepts any combination of af_seq and af_array
1 parent 8a608c2 commit ff4a4e2

17 files changed

Lines changed: 877 additions & 14 deletions

File tree

‎include/af/index.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,11 @@ extern "C" {
7474

7575
AFAPI af_err af_flip(af_array *out, const af_array in, const unsigned dim);
7676

77+
// generalized indexing function that accepts either af_array or af_seq
78+
// along a dimension to index the input array and create the corresponding
79+
// output array
80+
AFAPI af_err af_index_gen(af_array *out, const af_array in, const dim_type ndims, const af_index_t* indexers);
81+
7782
#ifdef __cplusplus
7883
}
7984
#endif

‎include/af/seq.h‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,21 @@ af_make_seq(double begin, double end, double step);
2121

2222
static const af_seq af_span = {1, 1, 0};
2323

24+
typedef struct {
25+
// if seq is used for current dimension
26+
// mIsSeq is set to 'true' and mIndexer.seq
27+
// should be used. Otherwise, mIndexer.arr
28+
// should be used.
29+
union {
30+
af_array arr;
31+
af_seq seq;
32+
} mIndexer;
33+
// below variable is used to determine if
34+
// the current dimension is indexed using
35+
// af_array or af_seq
36+
bool mIsSeq;
37+
} af_index_t;
38+
2439
#ifdef __cplusplus
2540
namespace af
2641
{
@@ -68,5 +83,7 @@ class AFAPI seq
6883
extern AFAPI int end;
6984
extern AFAPI seq span;
7085

86+
typedef af_index_t indexType;
87+
7188
}
7289
#endif

‎src/api/c/index.cpp‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,14 @@
1212

1313
#include <af/array.h>
1414
#include <af/index.h>
15+
#include <af/arith.h>
1516
#include <ArrayInfo.hpp>
1617
#include <err_common.hpp>
1718
#include <handle.hpp>
1819
#include <backend.hpp>
1920
#include <Array.hpp>
2021
#include <lookup.hpp>
22+
#include <index.hpp>
2123

2224
using namespace detail;
2325
using std::vector;
@@ -116,3 +118,91 @@ af_make_seq(double begin, double end, double step) {
116118
af_seq seq = {begin, end, step};
117119
return seq;
118120
}
121+
122+
// idxrs parameter to the below static function
123+
// expects 4 values which is handled appropriately
124+
// by the C-API af_index_gen
125+
template<typename T>
126+
static inline
127+
af_array genIndex(const af_array& in, const af_index_t idxrs[])
128+
{
129+
return getHandle<T>(index<T>(getArray<T>(in), idxrs));
130+
}
131+
132+
af_err af_index_gen(af_array *out, const af_array in, const dim_type ndims, const af_index_t* indexers)
133+
{
134+
af_array output = 0;
135+
// spanner is sequence indexer used for indexing along the
136+
// dimensions after ndims
137+
af_index_t spanner;
138+
spanner.mIndexer.seq = af_span;
139+
spanner.mIsSeq = true;
140+
141+
try {
142+
ARG_ASSERT(2, (ndims>0));
143+
ARG_ASSERT(3, (indexers!=NULL));
144+
145+
int track = 0;
146+
af_seq seqs[] = {af_span, af_span, af_span, af_span};
147+
for (dim_type i = 0; i < ndims; i++) {
148+
if (indexers[i].mIsSeq) {
149+
track++;
150+
seqs[i] = indexers[i].mIndexer.seq;
151+
}
152+
}
153+
154+
if (track==ndims) {
155+
// all indexers are sequences, redirecting to af_index
156+
return af_index(out, in, ndims, seqs);
157+
}
158+
159+
af_index_t idxrs[4];
160+
// set all dimensions above ndims to spanner indexer
161+
for (dim_type i=ndims; i<4; ++i) idxrs[i] = spanner;
162+
163+
for (dim_type i=0; i<ndims; ++i) {
164+
if (!indexers[i].mIsSeq) {
165+
// check if all af_arrays have atleast one value
166+
// to enable indexing along that dimension
167+
ArrayInfo idxInfo = getInfo(indexers[i].mIndexer.arr);
168+
af_dtype idxType = idxInfo.getType();
169+
170+
ARG_ASSERT(3, (idxType!=c32));
171+
ARG_ASSERT(3, (idxType!=c64));
172+
ARG_ASSERT(3, (idxType!=b8 ));
173+
174+
idxrs[i].mIndexer.arr = indexers[i].mIndexer.arr;
175+
idxrs[i].mIsSeq = indexers[i].mIsSeq;
176+
} else {
177+
// af_seq is being used for this dimension
178+
// just copy the indexer to local variable
179+
idxrs[i] = indexers[i];
180+
}
181+
}
182+
183+
ArrayInfo iInfo = getInfo(in);
184+
dim4 iDims = iInfo.dims();
185+
186+
ARG_ASSERT(1, (iDims.ndims()>0));
187+
188+
af_dtype inType = getInfo(in).getType();
189+
switch(inType) {
190+
case c64: output = genIndex<cdouble>(in, idxrs); break;
191+
case f64: output = genIndex<double >(in, idxrs); break;
192+
case c32: output = genIndex<cfloat >(in, idxrs); break;
193+
case f32: output = genIndex<float >(in, idxrs); break;
194+
case u64: output = genIndex<uintl >(in, idxrs); break;
195+
case u32: output = genIndex<uint >(in, idxrs); break;
196+
case s64: output = genIndex<intl >(in, idxrs); break;
197+
case s32: output = genIndex<int >(in, idxrs); break;
198+
case u8: output = genIndex<uchar >(in, idxrs); break;
199+
case b8: output = genIndex<char >(in, idxrs); break;
200+
default: TYPE_ERROR(1, inType);
201+
}
202+
}
203+
CATCHALL;
204+
205+
std::swap(*out, output);
206+
207+
return AF_SUCCESS;
208+
}

‎src/backend/cpu/index.cpp‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
/*******************************************************
2+
* Copyright (c) 2014, ArrayFire
3+
* All rights reserved.
4+
*
5+
* This file is distributed under 3-clause BSD license.
6+
* The complete license agreement can be obtained at:
7+
* http://arrayfire.com/licenses/BSD-3-Clause
8+
********************************************************/
9+
10+
#include <af/dim4.hpp>
11+
#include <af/defines.h>
12+
#include <ArrayInfo.hpp>
13+
#include <Array.hpp>
14+
#include <index.hpp>
15+
#include <handle.hpp>
16+
#include <err_cpu.hpp>
17+
#include <vector>
18+
19+
using af::dim4;
20+
21+
namespace cpu
22+
{
23+
24+
static inline
25+
dim_type trimIndex(dim_type idx, const dim_type &len)
26+
{
27+
dim_type ret_val = idx;
28+
dim_type offset = abs(ret_val)%len;
29+
if (ret_val<0) {
30+
ret_val = offset-1;
31+
} else if (ret_val>=len) {
32+
ret_val = len-offset-1;
33+
}
34+
return ret_val;
35+
}
36+
37+
template<typename T>
38+
Array<T> index(const Array<T>& in, const af_index_t idxrs[])
39+
{
40+
bool isSeq[4];
41+
std::vector<af_seq> seqs(4, af_span);
42+
// create seq vector to retrieve output
43+
// dimensions, offsets & offsets
44+
for (dim_type x=0; x<4; ++x) {
45+
if (idxrs[x].mIsSeq) {
46+
seqs[x] = idxrs[x].mIndexer.seq;
47+
}
48+
isSeq[x] = idxrs[x].mIsSeq;
49+
}
50+
51+
// rettrieve
52+
dim4 iDims = in.dims();
53+
dim4 dDims = in.getDataDims();
54+
dim4 oDims = af::toDims (seqs, iDims);
55+
dim4 iOffs = af::toOffset(seqs, dDims);
56+
dim4 iStrds= af::toStride(seqs, dDims);
57+
58+
std::vector< Array<uint> > idxArrs(4, createEmptyArray<uint>(dim4()));
59+
// look through indexers to read af_array indexers
60+
for (dim_type x=0; x<4; ++x) {
61+
if (!isSeq[x]) {
62+
idxArrs[x] = castArray<uint>(idxrs[x].mIndexer.arr);
63+
// set output array ith dimension value
64+
oDims[x] = idxArrs[x].elements();
65+
}
66+
}
67+
68+
Array<T> out = createEmptyArray<T>(oDims);
69+
dim4 oStrides= out.strides();
70+
71+
const T *src = in.get();
72+
T *dst = out.get();
73+
74+
const uint* ptr0 = idxArrs[0].get();
75+
const uint* ptr1 = idxArrs[1].get();
76+
const uint* ptr2 = idxArrs[2].get();
77+
const uint* ptr3 = idxArrs[3].get();
78+
79+
for (dim_type l=0; l<oDims[3]; ++l) {
80+
81+
dim_type lOff = l*oStrides[3];
82+
dim_type inIdx3 = trimIndex(isSeq[3] ? l+iOffs[3] : ptr3[l], iDims[3]);
83+
dim_type inOff3 = inIdx3*iStrds[3];
84+
85+
for (dim_type k=0; k<oDims[2]; ++k) {
86+
87+
dim_type kOff = k*oStrides[2];
88+
dim_type inIdx2 = trimIndex(isSeq[2] ? k+iOffs[2] : ptr2[k], iDims[2]);
89+
dim_type inOff2 = inIdx2*iStrds[2];
90+
91+
for (dim_type j=0; j<oDims[1]; ++j) {
92+
93+
dim_type jOff = j*oStrides[1];
94+
dim_type inIdx1 = trimIndex(isSeq[1] ? j+iOffs[1] : ptr1[j], iDims[1]);
95+
dim_type inOff1 = inIdx1*iStrds[1];
96+
97+
for (dim_type i=0; i<oDims[0]; ++i) {
98+
99+
dim_type iOff = i*oStrides[0];
100+
dim_type inIdx0 = trimIndex(isSeq[0] ? i+iOffs[0] : ptr0[i], iDims[0]);
101+
dim_type inOff0 = inIdx0*iStrds[0];
102+
103+
dst[lOff+kOff+jOff+iOff] = src[inOff3+inOff2+inOff1+inOff0];
104+
}
105+
}
106+
}
107+
}
108+
109+
return out;
110+
}
111+
112+
#define INSTANTIATE(T) \
113+
template Array<T> index<T>(const Array<T>& in, const af_index_t idxrs[]);
114+
115+
INSTANTIATE(cdouble)
116+
INSTANTIATE(double )
117+
INSTANTIATE(cfloat )
118+
INSTANTIATE(float )
119+
INSTANTIATE(uintl )
120+
INSTANTIATE(uint )
121+
INSTANTIATE(intl )
122+
INSTANTIATE(int )
123+
INSTANTIATE(uchar )
124+
INSTANTIATE(char )
125+
126+
}

‎src/backend/cpu/index.hpp‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
/*******************************************************
2+
* Copyright (c) 2014, ArrayFire
3+
* All rights reserved.
4+
*
5+
* This file is distributed under 3-clause BSD license.
6+
* The complete license agreement can be obtained at:
7+
* http://arrayfire.com/licenses/BSD-3-Clause
8+
********************************************************/
9+
10+
#include <Array.hpp>
11+
12+
namespace cpu
13+
{
14+
15+
template<typename T>
16+
Array<T> index(const Array<T>& in, const af_index_t idxrs[]);
17+
18+
}

‎src/backend/cpu/lookup.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
namespace cpu
1515
{
1616

17+
static inline
1718
dim_type trimIndex(dim_type idx, const dim_type &len)
1819
{
1920
dim_type ret_val = idx;

‎src/backend/cuda/index.cu‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
/*******************************************************
2+
* Copyright (c) 2014, ArrayFire
3+
* All rights reserved.
4+
*
5+
* This file is distributed under 3-clause BSD license.
6+
* The complete license agreement can be obtained at:
7+
* http://arrayfire.com/licenses/BSD-3-Clause
8+
********************************************************/
9+
10+
#include <af/dim4.hpp>
11+
#include <af/defines.h>
12+
#include <ArrayInfo.hpp>
13+
#include <Array.hpp>
14+
#include <handle.hpp>
15+
#include <index.hpp>
16+
#include <kernel/index.hpp>
17+
#include <err_cuda.hpp>
18+
19+
using af::dim4;
20+
21+
namespace cuda
22+
{
23+
24+
template<typename T>
25+
Array<T> index(const Array<T>& in, const af_index_t idxrs[])
26+
{
27+
kernel::IndexKernelParam_t p;
28+
std::vector<af_seq> seqs(4, af_span);
29+
// create seq vector to retrieve output
30+
// dimensions, offsets & offsets
31+
for (dim_type x=0; x<4; ++x) {
32+
if (idxrs[x].mIsSeq) {
33+
seqs[x] = idxrs[x].mIndexer.seq;
34+
}
35+
}
36+
37+
// retrieve dimensions, strides and offsets
38+
dim4 iDims = in.dims();
39+
dim4 dDims = in.getDataDims();
40+
dim4 oDims = af::toDims (seqs, iDims);
41+
dim4 iOffs = af::toOffset(seqs, dDims);
42+
dim4 iStrds= af::toStride(seqs, dDims);
43+
44+
for (dim_type i=0; i<4; ++i) {
45+
p.isSeq[i] = idxrs[i].mIsSeq;
46+
p.offs[i] = iOffs[i];
47+
p.strds[i] = iStrds[i];
48+
}
49+
50+
std::vector< Array<uint> > idxArrs(4, createEmptyArray<uint>(dim4()));
51+
// look through indexers to read af_array indexers
52+
for (dim_type x=0; x<4; ++x) {
53+
// set idxPtrs to null
54+
p.ptr[x] = 0;
55+
// set index pointers were applicable
56+
if (!p.isSeq[x]) {
57+
idxArrs[x] = castArray<uint>(idxrs[x].mIndexer.arr);
58+
p.ptr[x] = idxArrs[x].get();
59+
// set output array ith dimension value
60+
oDims[x] = idxArrs[x].elements();
61+
}
62+
}
63+
64+
Array<T> out = createEmptyArray<T>(oDims);
65+
66+
kernel::index<T>(out, in, p);
67+
68+
return out;
69+
}
70+
71+
#define INSTANTIATE(T) \
72+
template Array<T> index<T>(const Array<T>& in, const af_index_t idxrs[]);
73+
74+
INSTANTIATE(cdouble)
75+
INSTANTIATE(double )
76+
INSTANTIATE(cfloat )
77+
INSTANTIATE(float )
78+
INSTANTIATE(uintl )
79+
INSTANTIATE(uint )
80+
INSTANTIATE(intl )
81+
INSTANTIATE(int )
82+
INSTANTIATE(uchar )
83+
INSTANTIATE(char )
84+
85+
}

0 commit comments

Comments
 (0)