Repository: numba/numba

Issue:
SSA pass causes invalid compilation (array write elided).
From https://github.com/pymatting/pymatting/issues/15, whilst long, this compiles incorrectly if the SSA passes are enabled. Problem is marked in the source by `PROBLEM IS HERE` comment.
```python
import numpy as np
import llvmlite
import numba
import platform
from numba import njit

print("System:", platform.uname())
print("llvmlite version:", llvmlite.__version__)
print("numba version:", numba.__version__)
print("numpy version:", np.__version__)

# If this should crash, try without this line
@njit("i8(i8[:], i8[:], i8[:], i8[:], i8[:], f4[:, :, :], f4[:], f4[:, :], i8[:], i8)",
      boundscheck=True, debug=True)
def _make_tree(
    i0_inds,
    i1_inds,
    less_inds,
    more_inds,
    split_dims,
    bounds,
    split_values,
    points,
    indices,
    min_leaf_size,
):
    dimension = points.shape[1]
    # Expect log2(len(points) / min_leaf_size) depth, 1000 should be plenty
    stack = np.ones(1000, np.int64) * -99999999999999
    stack_size = 0
    n_nodes = 0
    # min_leaf_size <= leaf_node_size < max_leaf_size
    max_leaf_size = 2 * min_leaf_size

    # Push i0, i1, i_node
    stack[stack_size] = 0
    stack_size += 1
    stack[stack_size] = points.shape[0]
    stack_size += 1
    stack[stack_size] = n_nodes
    n_nodes += 1
    stack_size += 1

    # While there are more tree nodes to process recursively
    while stack_size > 0:
        print("start", stack[:20])
        stack_size -= 1
        i_node = stack[stack_size]
        stack_size -= 1
        i1 = stack[stack_size]
        stack_size -= 1
        i0 = stack[stack_size]

        lo = bounds[i_node, 0]
        hi = bounds[i_node, 1]

        for d in range(dimension):
            lo[d] = points[i0, d]
            hi[d] = points[i0, d]

        # Find lower and upper bounds of points for each dimension
        for i in range(i0 + 1, i1):
            for d in range(dimension):
                lo[d] = min(lo[d], points[i, d])
                hi[d] = max(hi[d], points[i, d])

        # Done if node is small
        if i1 - i0 <= max_leaf_size:
            i0_inds[i_node] = i0
            i1_inds[i_node] = i1
            less_inds[i_node] = -1
            more_inds[i_node] = -1
            split_dims[i_node] = -1
            split_values[i_node] = np.float32(0.0)
        else:
            # Split on largest dimension
            lengths = hi - lo
            split_dim = np.argmax(lengths)
            split_value = lo[split_dim] + np.float32(0.5) * lengths[split_dim]

            # Partition i0:i1 range into points where points[i, split_dim] < split_value
            i = i0
            j = i1 - 1

            while i < j:
                while i < j and points[i, split_dim] < split_value:
                    i += 1
                while i < j and points[j, split_dim] >= split_value:
                    j -= 1

                # Swap points
                if i < j:
                    for d in range(dimension):
                        temp = points[i, d]
                        points[i, d] = points[j, d]
                        points[j, d] = temp

                    temp_i_node = indices[i]
                    indices[i] = indices[j]
                    indices[j] = temp_i_node

            if points[i, split_dim] < split_value:
                i += 1

            i_split = i

            # Now it holds that:
            # for i in range(i0, i_split): assert(points[i, split_dim] < split_value)
            # for i in range(i_split, i1): assert(points[i, split_dim] >= split_value)

            # Ensure that each node has at least min_leaf_size children
            i_split = max(i0 + min_leaf_size, min(i1 - min_leaf_size, i_split))
            print(i_split)

            less = n_nodes
            n_nodes += 1
            more = n_nodes
            n_nodes += 1

            # push i0, i_split, less
            print("i0", i0, "to", stack_size)
            stack[stack_size] = i0
            stack_size += 1
            print("i_split", i_split, "to", stack_size)
            stack[stack_size] = i_split
            stack_size += 1
            print("less", less, "to", stack_size)
            stack[stack_size] = less
            stack_size += 1

            # push i_split, i1, more
            print("i_split", i_split, "to", stack_size)
            stack[stack_size] = i_split # <----- PROBLEM IS HERE, THIS DOESN'T GET WRITTEN
            #stack[stack_size] = np.sqrt(i_split) # DO THIS AND IT WORKS WITH SSA!
            stack_size += 1
            print("i1", i1, "to", stack_size)
            stack[stack_size] = i1
            stack_size += 1
            print("more", more, "to", stack_size)
            stack[stack_size] = more
            stack_size += 1

            i0_inds[i_node] = i0
            i1_inds[i_node] = i1
            less_inds[i_node] = less
            more_inds[i_node] = more
            split_dims[i_node] = split_dim
            split_values[i_node] = split_value
            print("end", stack_size, "stack", stack[:20])
    return n_nodes

def main():
    k = 20
    n_data = 100_000
    n_query = n_data
    dimension = 5
    np.random.seed(0)
    data_points = np.random.rand(n_data, dimension).astype(np.float32)
    min_leaf_size = 8

    n_data, dimension = data_points.shape

    max_nodes = 2 * ((n_data + min_leaf_size - 1) // min_leaf_size)

    i0_inds = np.empty(max_nodes, np.int64)
    i1_inds = np.empty(max_nodes, np.int64)
    less_inds = np.empty(max_nodes, np.int64)
    more_inds = np.empty(max_nodes, np.int64)
    split_dims = np.empty(max_nodes, np.int64)
    bounds = np.empty((max_nodes, 2, dimension), np.float32)
    split_values = np.empty(max_nodes, np.float32)
    shuffled_data_points = data_points.copy()
    shuffled_indices = np.arange(n_data).astype(np.int64)

    _make_tree(
        i0_inds,
        i1_inds,
        less_inds,
        more_inds,
        split_dims,
        bounds,
        split_values,
        shuffled_data_points,
        shuffled_indices,
        min_leaf_size,
    )

main()

Discussion:

Patch:
diff --git a/numba/core/ssa.py b/numba/core/ssa.py
index d38918f8c..a4b30df67 100644
--- a/numba/core/ssa.py
+++ b/numba/core/ssa.py
@@ -430,15 +430,18 @@ class _FixSSAVars(_BaseHandler):
             return self._find_def_from_top(states, label, loc=loc)
 
     def _stmt_index(self, defstmt, block, stop=-1):
-        """Find the postitional index of the statement at ``block``.
+        """Find the positional index of the statement at ``block``.
 
         Assumptions:
         - no two statements can point to the same object.
         """
-        try:
-            return block.body.index(defstmt, 0, stop)
-        except ValueError:
-            return len(block.body)
+        # Compare using id() as IR node equality is for semantic equivalence
+        # opposed to direct equality (the location and scope are not considered
+        # as part of the equality measure, this is important here).
+        for i in range(len(block.body))[:stop]:
+            if block.body[i] is defstmt:
+                return i
+        return len(block.body)
 
 
 def _warn_about_uninitialized_variable(varname, loc):
diff --git a/numba/typed/typedlist.py b/numba/typed/typedlist.py
index 82259b5bc..2d3efd195 100644
--- a/numba/typed/typedlist.py
+++ b/numba/typed/typedlist.py
@@ -355,14 +355,13 @@ class List(MutableSequence):
         return _pop(self, i)
 
     def extend(self, iterable):
-        # Empty iterable, do nothing
-        if len(iterable) == 0:
-            return self
         if not self._typed:
             # Need to get the first element of the iterable to initialise the
             # type of the list. FIXME: this may be a problem if the iterable
             # can not be sliced.
             self._initialise_list(iterable[0])
+            self.append(iterable[0])
+            return _extend(self, iterable[1:])
         return _extend(self, iterable)
 
     def remove(self, item):


Test patch:
diff --git a/numba/tests/test_ssa.py b/numba/tests/test_ssa.py
index 0d3e7bc2d..5df232f34 100644
--- a/numba/tests/test_ssa.py
+++ b/numba/tests/test_ssa.py
@@ -370,3 +370,26 @@ class TestReportedSSAIssues(SSABaseTest):
         res2 = jit(forceobj=True, looplift=False)(foo)(10, 10, data)
         np.testing.assert_array_equal(expect, res1)
         np.testing.assert_array_equal(expect, res2)
+
+    def test_issue5623_equal_statements_in_same_bb(self):
+
+        def foo(pred, stack):
+            i = 0
+            c = 1
+
+            if pred is True:
+                stack[i] = c
+                i += 1
+                stack[i] = c
+                i += 1
+
+        python = np.array([0, 666])
+        foo(True, python)
+
+        nb = np.array([0, 666])
+        njit(foo)(True, nb)
+
+        expect = np.array([1, 1])
+
+        np.testing.assert_array_equal(python, expect)
+        np.testing.assert_array_equal(nb, expect)
diff --git a/numba/tests/test_typedlist.py b/numba/tests/test_typedlist.py
index 3ce886f07..72f17fb5a 100644
--- a/numba/tests/test_typedlist.py
+++ b/numba/tests/test_typedlist.py
@@ -738,33 +738,6 @@ class TestExtend(MemoryLeakMixin, TestCase):
         got = impl()
         self.assertEqual(expected, got)
 
-    def test_extend_single_value_container(self):
-        @njit
-        def impl():
-            l = List()
-            l.extend((100,))
-            return l
-
-        expected = impl.py_func()
-        got = impl()
-        self.assertEqual(expected, got)
-
-    def test_extend_empty_unrefined(self):
-        # Extending an unrefined list with an empty iterable doesn't work in a
-        # jit compiled function as the list remains untyped.
-        l = List()
-        l.extend(tuple())
-        self.assertEqual(len(l), 0)
-        self.assertFalse(l._typed)
-
-    def test_extend_empty_refiend(self):
-        # Extending a refined list with an empty iterable doesn't work in a
-        # jit compiled function as the (empty) argument can't be typed
-        l = List((1,))
-        l.extend(tuple())
-        self.assertEqual(len(l), 1)
-        self.assertTrue(l._typed)
-
 
 @njit
 def cmp(a, b):