Added support for multidimensional array indexing

[kc611]

This PR builds on top of #8238

As titled, this PR adds support for multidimensional indices while indexing NumPy arrays in Numba JIT functions:

import numba
import numpy as np

@numba.njit
def func(x, idx):
    return x[:, 2, idx, :, :]

a = np.random.randint(0, 100, (10, 11, 12, 13, 14))
b = np.random.randint(0, 10, (4, 5)) # Previously, these were limited to 1-D.

print(func(a, b).shape) 
# (10, 4, 5, 13, 14)
print(np.allclose(func(a, b), func.py_func(a, b)))
# True

[kc611]

The memory leak I was facing in this issue has been isolated:

import numpy as np
import unittest
from numba import njit
from numba.tests.support import MemoryLeakMixin, TestCase

class TestFancyIndexingMultiDim(MemoryLeakMixin, TestCase):
    shape = (4, 5, 6)
    
    def check_setitem_indices(self, arr_shape, index):
        @njit     
        def set_item(array, idx, item):
            array[idx] = item

        arr = np.random.randint(0, 11, size=arr_shape)
        src = arr[index]
        expected = np.zeros_like(arr)
        got = np.zeros_like(arr)

        set_item.py_func(expected, index, src)
        set_item(got, index, src)

        np.testing.assert_equal(got, expected)

    def test_setitem_with_tuple(self):
        # When index is an array within a tuple
        idx = (np.array([1,2,3]),)
        # No memory leak
        self.check_setitem_indices(self.shape, idx)

    def test_setitem_without_tuple(self):
        # When index itself is an array, it is
        # supposed to be implicitly cast to same form as
        # the example above i.e. both should and do 
        # give same answer, except for the leak
        idx = np.array([1,2,3])
        # Memory leaks
        self.check_setitem_indices(self.shape, idx)
    
if __name__ == '__main__':
    unittest.main()

This I presume is a consequence of a fix for memory leak that was happening earlier and was 'fixed' by having the following decrefs in place

https://github.com/numba/numba/pull/8491/files#diff-d0aefd08783016c4e2a19e29d4725c7870806ea2aa1ab4caa6e8f0adc08c6e64R1717-R1721

I heavily suspect that for some reason these decrefs aren't being triggered in the particular case given above.

[stuartarchibald]

@kc611 please could you resolve conflicts against main? As #8238 is merged the diff will hopefully shrink. Many thanks.

[stuartarchibald]

Thanks for the PR @kc611, great to see this implemented. I've given this an initial review and have also done some initial manual testing. Once the comments are address I'll look at the implementation details more closely, but from a cursory inspection the approach seems like it should work. Thanks again!

[stuartarchibald]

I don't think that this path in its present form can be supported on CUDA as the flat_imp* function compiled below relies on:

• NumPy functions
• CPU only implementations for reshape
• returns an array
• needs the NRT

I'm not sure of the best way to "detect" this, immediate thoughts are to either require the NRT and raise if the context doesn't have it, or alternatively declare these functions as overloads targetting CPU only which will then fail for CUDA users. CC @gmarkall do you have an opinion on this?

[stuartarchibald]

As per a later comment, this requires the presence of the NRT which isn't guaranteed and needs guarding against.

[stuartarchibald]

Resolving this can be postponed as the 3 points in #8491 (review) will cover it.

[stuartarchibald]

Given this check is removed, does the proposed implementation correctly handle "overlap"?

[kc611]

Define overlap in context of arrays ?

[stuartarchibald]

I think in this case it would be if the source and destination share any of the same memory location.

[kc611]

With the current implementation, I don't think that's possible since fancy indexing always creates a copy of the data. i.e. In fancy indexing the source and destination array never share memory. This is stated in the NumPy documentation as well.

I'm in-fact curious as of why it was the case over here that they did.

[stuartarchibald]

I suppose the question is... in Numba, what if in practice they do share memory? Assessing and resolving this can be deferred to a subsequent PR. @DrTodd13 was asking about functions for assessing whether arrays overlap at the Numba public meeting last week. The associated functions and algorithms in NumPy look like they are something Numba could support, but are out of scope for implementing in this PR.

[stuartarchibald]

It would be good for the testing to now also include indexing into F-order and A-order arrays?

[kc611]

Alright will add that.

[sklam]

Bad syntax

[stuartarchibald]

Thanks for the updates @kc611. I've given them a review and I think there's a few things to be resolved but otherwise looks good. The unit tests that are failing CI are failing due to the use of compile_isolated and Flags(). To fix, this:

numba/numba/tests/test_indexing.py

Lines 15 to 16 in 74cb715

	enable_pyobj_flags = Flags()
	enable_pyobj_flags.enable_pyobject = True

could have:

enable_pyobj_flags.nrt = True

adding. Or alternatively, as object mode fallback is deprecated, forced object mode testing could be achieved by using alternative flags such as those available through importing these from testing support:

numba/numba/tests/support.py

Lines 70 to 71 in 74cb715

	force_pyobj_flags = Flags()
	force_pyobj_flags.force_pyobject = True

Hope this helps?

[stuartarchibald]

I think the second part of this still needs doing, i.e. move the "flat" implementations to module level so as to make use of the compilation cache.

[stuartarchibald]

To be consistent, I think retty should be derived from the callsig, specifically, it should be callsig.return_type?

[stuartarchibald]

Please could this be done throughout? Numba uses spaces after commas in this context.

[stuartarchibald]

Thanks for the updates @kc611, few minor things to resolve now else looks good. OOB I spoke with @gmarkall about the impact of the use of the NRT on the CUDA target. The conclusion is as follows:

1. It's probably best to concentrate on getting this and the next related PR ([WIP] Advanced Indexing #3: Added support for multiple multidimensional Indices #8912) merged as the implementation needs to be figured out to start with and it's almost certainly easier to do this on the CPU with the NRT present.
2. Go back to e.g. the release0.57 branch and see what sort of "fancy indexing" actually worked on CUDA, write tests for this.
3. Get the tests in 2. working on main once these PRs are merged. I suspect it'll only be the case where the index is a non-contiguous array that's a problem as the copy is currently needed as part of "flattening" it. Once we have the tests I expect the options and limitations will become more clear.

[gmarkall]

gpuci run tests

[stuartarchibald]

Many thanks for all your efforts on this @kc611. It's great to see this feature implemented. There's a few outstanding issues that I've left as a "review", but these should be captured in a new issue for completion in subsequent work. Most of the outstanding items relate to use of the Numba runtime which is going to take careful assessment. As alluded to previously, this will manifest in having to write a number of tests for the CUDA target to ensure that no regressions are introduced. Thanks again for work on this challenging implementation and feature, patch is approved!

[stuartarchibald]

Ideally this would be tested, but I'm inclined to leave this to a subsequent PR as I am relatively sure it's going to need to be moved.

[stuartarchibald]

As above, ideally this would be tested, but I'm inclined to leave this to a subsequent PR as I am relatively sure it's going to need to be moved.

[stuartarchibald]

Resolving this can be postponed as the 3 points in #8491 (review) will cover it.

[stuartarchibald]

I suppose the question is... in Numba, what if in practice they do share memory? Assessing and resolving this can be deferred to a subsequent PR. @DrTodd13 was asking about functions for assessing whether arrays overlap at the Numba public meeting last week. The associated functions and algorithms in NumPy look like they are something Numba could support, but are out of scope for implementing in this PR.

[stuartarchibald]

This needs moving to a module global in a subsequent PR.

[stuartarchibald]

Use of NRT needs assessing, defer to later PR.

[stuartarchibald]

Setting order='A' in the NumPy constructor won't have the effect of it being A ordered within Numba's type system (assuming that was the intent?) Suggest deferring fixing this to a later PR.