Repository: numba/numba

Issue:
Corruption on ppc64le/arm64 (potential codegen/LLVM opt bug)
## Reporting a bug

- [x] I have tried using the latest released version of Numba (most recent is
 visible in the change log (https://github.com/numba/numba/blob/master/CHANGE_LOG).
- [x] I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'.

The test `numba.tests.test_np_functions.TestNPFunctions.test_take_along_axis` currently fails against e617b39a0b4b23d7b69d16f482fd66b4ac6cc307 (HEAD of current mainline) with `SIGSEGV` on `linux-ppc64le` and `linux-arm64`.

This is a somewhat minimised reproducer:
```python
from numba import njit, types, config
from numba.extending import register_jitable
import numpy as np

if config.DISABLE_JIT:
    def tuple_setitem(tup, idx, val):
        return tup[:idx] + (val,) + tup[idx+1:]
else:
    from numba.cpython.unsafe.tuple import tuple_setitem

@register_jitable
def _tal_impl(arr, indices, axis, indices_broadcast_shape):
    arr_shape = list(arr.shape) # comment just this and a ValueError is raised from the broadcast (ppc64le)
    print(indices_broadcast_shape) # comment just this and it "works" on ppc64le, ValueError in broadcast on arm64
    arr_broadcast_shape = tuple_setitem(indices_broadcast_shape, axis, arr.shape[axis])
    # Call these in either order, effect is the same, corruption
    np.broadcast_to(indices, indices_broadcast_shape)
    np.broadcast_to(arr, arr_broadcast_shape)

a = np.arange(24).reshape(3, 1, 4, 2)

@njit
def impl(a, i):
    return _tal_impl(a, i, -1, (3, 1, 4, 2))

ai = np.argsort(a, axis=-1)
impl(a, ai)
```
Failure table:
| `NUMBA_OPT` level|ARM64       |PPC64LE|
|----------------------------|-------------------|----------|
| 0 | ok | ok |
| 1 | `ValueError` from `np.broadcast_to` impl | ok |
| 2 | `SIGBUS` | `SIGSEGV` |
| 3 | `SIGBUS` | `SIGSEGV` |

Notes:
 * Running the script under `valgrind` on `ppc64le` with `NUMBA_OPT=3` leads to a `SIGBUS`.
 * Running the script under `valgrind` on `x86_64` with `NUMBA_OPT=3` has nothing suspicious.
 * Commenting out various lines in the script (as noted) triggers different behaviour.
 * Running the script with `NUMBA_DISABLE_JIT=1` has no errors.

I think this somewhat hints at corruption from an invalid optimisation.


Discussion:
With thanks to @gmarkall for spotting this, turns out the `-1` in this:
```python
@njit
def impl(a, i):
    return _tal_impl(a, i, -1, (3, 1, 4, 2))
```
ends up as the index in a `getelementptr` on an array type that is 4 elements in size and to make it worse it's specified as `inbounds`, this is all compile time constant and likely triggers UB. When debugging this on power8, there's the tuple `(3, 1, 4, 2)` on the stack and it reads 1 before this and it ends up clobbering a register with value `0x4` which is then read as an address in an `NRT_incref` almost immediately after.


Resolving patch:
diff --git a/numba/np/arrayobj.py b/numba/np/arrayobj.py
--- a/numba/np/arrayobj.py
+++ b/numba/np/arrayobj.py
@@ -5545,6 +5545,11 @@ def _take_along_axis_impl(
     # With addition of pre-broadcasting:
     # https://github.com/numpy/numpy/issues/19704
 
+    # Wrap axis, it's used in tuple_setitem so must be (axis >= 0) to ensure
+    # the GEP is in bounds.
+    if axis < 0:
+        axis = arr.ndim + axis
+
     # Broadcast the two arrays to matching shapes:
     arr_shape = list(arr.shape)
     arr_shape[axis] = 1
@@ -5620,7 +5625,7 @@ def take_along_axis_impl(arr, indices, axis):
         check_is_integer(axis, "axis")
         if not isinstance(axis, types.IntegerLiteral):
             raise ValueError(
-                "axis must be a literal value (i.e. not an argument) for now"
+                "axis must be a literal value (i.e. not an argument)"
             )
         axis = axis.literal_value
         if axis < 0:


Test patch:
diff --git a/numba/tests/test_np_functions.py b/numba/tests/test_np_functions.py
--- a/numba/tests/test_np_functions.py
+++ b/numba/tests/test_np_functions.py
@@ -4461,14 +4461,19 @@ def test_take_along_axis_broadcasting(self):
         arr = np.ones((3, 4, 1))
         ai = np.ones((1, 2, 5), dtype=np.intp)
 
-        @njit
-        def check(a, i):
-            return np.take_along_axis(a, i, axis=1)
+        def gen(axis):
+            @njit
+            def impl(a, i):
+                return np.take_along_axis(a, i, axis)
+            return impl
 
-        expected = check.py_func(arr, ai)
-        actual = check(arr, ai)
-        self.assertPreciseEqual(expected, actual)
-        self.assertEqual(actual.shape, (3, 2, 5))
+        # Check same axis but expressed as positive/negative value
+        for i in (1, -2):
+            check = gen(i)
+            expected = check.py_func(arr, ai)
+            actual = check(arr, ai)
+            self.assertPreciseEqual(expected, actual)
+            self.assertEqual(actual.shape, (3, 2, 5))
 
     def test_take_along_axis_exceptions(self):
         arr2d = np.arange(8).reshape(2, 4)