Skip to content

Fix for lookup offset in opencl and non-column indices #3650

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
edwinsolisf merged 2 commits into from
Apr 4, 2025
Merged

Fix for lookup offset in opencl and non-column indices #3650

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e14e2f2
Fixed missing offset handling in lookup
willyborn Oct 1, 2024
7daaf3f
Added tests for lookup with indices with offsets and lying in differe…
edwinsolisf Mar 31, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 22 additions & 10 deletions src/api/c/index.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,22 +178,34 @@ af_err af_lookup(af_array* out, const af_array in, const af_array indices,
ARG_ASSERT(2, (idxType != b8));

af_array output = 0;
af_array idx = 0;

if (!idxInfo.isColumn()) {
// Force a deep copy to flatten the array and handle subarrays of not column vector arrays correctly
AF_CHECK(af_copy_array(&idx, indices));
} else {
idx = indices;
}

switch (idxType) {
case f32: output = lookup<float>(in, indices, dim); break;
case f64: output = lookup<double>(in, indices, dim); break;
case f32: output = lookup<float>(in, idx, dim); break;
case f64: output = lookup<double>(in, idx, dim); break;
case s32: output = lookup<int>(in, indices, dim); break;
case u32: output = lookup<unsigned>(in, indices, dim); break;
case s16: output = lookup<short>(in, indices, dim); break;
case u16: output = lookup<ushort>(in, indices, dim); break;
case s64: output = lookup<intl>(in, indices, dim); break;
case u64: output = lookup<uintl>(in, indices, dim); break;
case s8: output = lookup<schar>(in, indices, dim); break;
case u8: output = lookup<uchar>(in, indices, dim); break;
case f16: output = lookup<half>(in, indices, dim); break;
case u32: output = lookup<unsigned>(in, idx, dim); break;
case s16: output = lookup<short>(in, idx, dim); break;
case u16: output = lookup<ushort>(in, idx, dim); break;
case s64: output = lookup<intl>(in, idx, dim); break;
case u64: output = lookup<uintl>(in, idx, dim); break;
case s8: output = lookup<schar>(in, idx, dim); break;
case u8: output = lookup<uchar>(in, idx, dim); break;
case f16: output = lookup<half>(in, idx, dim); break;
default: TYPE_ERROR(1, idxType);
}
std::swap(*out, output);

if (idx != indices) {
AF_CHECK(af_release_array(idx)); // Release indices array if a copy has been made
}
}
CATCHALL;
return AF_SUCCESS;
Expand Down
2 changes: 1 addition & 1 deletion src/backend/opencl/kernel/lookup.cl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ kernel void lookupND(global in_t *out, KParam oInfo, global const in_t *in,
int gx = get_local_size(0) * (get_group_id(0) - gz * nBBS0) + lx;
int gy = get_local_size(1) * (get_group_id(1) - gw * nBBS1) + ly;

global const idx_t *idxPtr = indices;
global const idx_t *idxPtr = indices + idxInfo.offset;

int i = iInfo.strides[0] *
(DIM == 0 ? trimIndex((int)idxPtr[gx], iInfo.dims[0]) : gx);
Expand Down
4 changes: 2 additions & 2 deletions src/backend/opencl/lookup.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,8 @@ Array<in_t> lookup(const Array<in_t> &input, const Array<idx_t> &indices,
const dim4 &iDims = input.dims();

dim4 oDims(1);
for (int d = 0; d < 4; ++d) {
oDims[d] = (d == int(dim) ? indices.elements() : iDims[d]);
for (dim_t d = 0; d < 4; ++d) {
oDims[d] = (d == dim ? indices.elements() : iDims[d]);
}

Array<in_t> out = createEmptyArray<in_t>(oDims);
Expand Down
121 changes: 121 additions & 0 deletions test/index.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -809,6 +809,127 @@ TEST(lookup, Issue2009) {
ASSERT_ARRAYS_EQ(a, b);
}

TEST(lookup, Issue3613_FirstDimLookupWithOffset) {
dim4 dims(1);
const int selected_dim = 0; // selected span dimension
dims[selected_dim] = 125; // input size

array a = iota(dims);
array idxs = iota(dim4(5, 4, 3, 2));
array selected_idx = idxs(af::span, 3, 2, 1); // Offsets in second, third, & fourth dimension

array expected_selected_idx = range(dim4(5)) * 1 + 3 * 5 + 2 * (5 * 4) + 1 * (5 * 4 * 3);
ASSERT_ARRAYS_EQ(expected_selected_idx, selected_idx);

array b = af::lookup(a, selected_idx, selected_dim);
dim4 output_dims(1);
output_dims[selected_dim] = 5; // output size
ASSERT_ARRAYS_EQ(af::moddims(expected_selected_idx, output_dims), b); // lookup output should be the same as looked up indices
}

TEST(lookup, Issue3613_SecondDimLookupWithOffset) {
dim4 dims(1);
const int selected_dim = 1; // selected span dimension
dims[selected_dim] = 125; // input size

array a = iota(dims);
array idxs = iota(dim4(5, 4, 3, 2));
array selected_idx = idxs(af::span, 3, 2, 1); // Offsets in second, third, & fourth dimension

array expected_selected_idx = range(dim4(5)) * 1 + 3 * 5 + 2 * (5 * 4) + 1 * (5 * 4 * 3);
ASSERT_ARRAYS_EQ(expected_selected_idx, selected_idx);

array b = af::lookup(a, selected_idx, selected_dim);
dim4 output_dims(1);
output_dims[selected_dim] = 5; // output size
ASSERT_ARRAYS_EQ(af::moddims(expected_selected_idx, output_dims), b); // lookup output should be the same as looked up indices
}


TEST(lookup, Issue3613_ThirdDimLookupWithOffset) {
dim4 dims(1);
const int selected_dim = 2; // selected span dimension
dims[selected_dim] = 125; // input size

array a = iota(dims);
array idxs = iota(dim4(5, 4, 3, 2));
array selected_idx = idxs(af::span, 3, 2, 1); // Offsets in second, third, & fourth dimension

array expected_selected_idx = range(dim4(5)) * 1 + 3 * 5 + 2 * (5 * 4) + 1 * (5 * 4 * 3);
ASSERT_ARRAYS_EQ(expected_selected_idx, selected_idx);

array b = af::lookup(a, selected_idx, selected_dim);
dim4 output_dims(1);
output_dims[selected_dim] = 5; // output size
ASSERT_ARRAYS_EQ(af::moddims(expected_selected_idx, output_dims), b); // lookup output should be the same as looked up indices
}

TEST(lookup, Issue3613_FourthDimLookupWithOffset) {
dim4 dims(1);
const int selected_dim = 3; // selected span dimension
dims[selected_dim] = 125; // input size

array a = iota(dims);
array idxs = iota(dim4(5, 4, 3, 2));
array selected_idx = idxs(af::span, 3, 2, 1); // Offsets in second, third, & fourth dimension

array expected_selected_idx = range(dim4(5)) * 1 + 3 * 5 + 2 * (5 * 4) + 1 * (5 * 4 * 3);
ASSERT_ARRAYS_EQ(expected_selected_idx, selected_idx);

array b = af::lookup(a, selected_idx, selected_dim);
dim4 output_dims(1);
output_dims[selected_dim] = 5; // output size
ASSERT_ARRAYS_EQ(af::moddims(expected_selected_idx, output_dims), b); // lookup output should be the same as looked up indices
}

TEST(lookup, IndicesInSecondDimension) {
const int selected_dim = 1; // selected span dimension
dim4 dims(1);
dims[selected_dim] = 3;

array a = iota(dim4(100));
array idxs = iota(dim4(3, 3, 3, 3));
array selected_idx = idxs(0, af::span, 0, 0); // Indices along the second dimension

array expected_selected_idx = iota(dims) * pow(3, selected_dim);
ASSERT_ARRAYS_EQ(expected_selected_idx, selected_idx);

array b = af::lookup(a, selected_idx);
ASSERT_ARRAYS_EQ(af::moddims(expected_selected_idx, dim4(3)), b);
}

TEST(lookup, IndicesInThirdDimension) {
const int selected_dim = 2; // selected span dimension
dim4 dims(1);
dims[selected_dim] = 3;

array a = iota(dim4(100));
array idxs = iota(dim4(3, 3, 3, 3));
array selected_idx = idxs(0, 0, af::span, 0); // Indices along the third dimension

array expected_selected_idx = iota(dims) * pow(3, selected_dim);
ASSERT_ARRAYS_EQ(expected_selected_idx, selected_idx);

array b = af::lookup(a, selected_idx);
ASSERT_ARRAYS_EQ(af::moddims(expected_selected_idx, dim4(3)), b);
}

TEST(lookup, IndicesInFourthDimension) {
const int selected_dim = 3; // selected span dimension
dim4 dims(1);
dims[selected_dim] = 3;

array a = iota(dim4(100));
array idxs = iota(dim4(3, 3, 3, 3));
array selected_idx = idxs(0, 0, 0, af::span); // Indices along the fourth dimension

array expected_selected_idx = iota(dims) * pow(3, selected_dim);
ASSERT_ARRAYS_EQ(expected_selected_idx, selected_idx);

array b = af::lookup(a, selected_idx);
ASSERT_ARRAYS_EQ(af::moddims(expected_selected_idx, dim4(3)), b);
}

TEST(lookup, SNIPPET_lookup1d) {
//! [ex_index_lookup1d]

Expand Down