Repository: numba/numba

Issue:
Invalid copy of Array analysis object for IR extension handlers (regression in 0.49rc1)
<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

For more general "how do I do X?" type questions, please speak to us in real
time on https://gitter.im/numba/numba or post to the Numba mailing list
https://groups.google.com/a/continuum.io/forum/#!forum/numba-users.

-->

## Reporting a bug

<!--

Before submitting a bug report please ensure that you can check off these boxes:

-->

- [ ] I am using the latest released version of Numba (most recent is visible in
 the change log (https://github.com/numba/numba/blob/master/CHANGE_LOG).
- [ ] I have included below a minimal working reproducer (if you are unsure how
 to write one see http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports).

<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->
The commit eae2d2bc341df8032879d1fd4c4ec7a2d9dd401d in array analysis refactoring PR (#5320)  avoids mutation of array analysis objects when passed to IR extension handlers. However, this defeats the whole purpose of these handlers, which are supposed to add more information. See handler of parfors [here](https://github.com/numba/numba/blob/1ea770564cb3c0c6cb9d8ab92e7faf23cd4c4c19/numba/parfors/parfor.py#L646).

@sklam was there a problem with these handlers that copying solves?
Invalid copy of Array analysis object for IR extension handlers (regression in 0.49rc1)
<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

For more general "how do I do X?" type questions, please speak to us in real
time on https://gitter.im/numba/numba or post to the Numba mailing list
https://groups.google.com/a/continuum.io/forum/#!forum/numba-users.

-->

## Reporting a bug

<!--

Before submitting a bug report please ensure that you can check off these boxes:

-->

- [ ] I am using the latest released version of Numba (most recent is visible in
 the change log (https://github.com/numba/numba/blob/master/CHANGE_LOG).
- [ ] I have included below a minimal working reproducer (if you are unsure how
 to write one see http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports).

<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->
The commit eae2d2bc341df8032879d1fd4c4ec7a2d9dd401d in array analysis refactoring PR (#5320)  avoids mutation of array analysis objects when passed to IR extension handlers. However, this defeats the whole purpose of these handlers, which are supposed to add more information. See handler of parfors [here](https://github.com/numba/numba/blob/1ea770564cb3c0c6cb9d8ab92e7faf23cd4c4c19/numba/parfors/parfor.py#L646).

@sklam was there a problem with these handlers that copying solves?


Discussion:
Thanks for the issue report - marked to triage by @sklam.
Thanks for the issue report - marked to triage by @sklam.

Resolving patch:
diff --git a/numba/parfors/array_analysis.py b/numba/parfors/array_analysis.py
--- a/numba/parfors/array_analysis.py
+++ b/numba/parfors/array_analysis.py
@@ -1058,13 +1058,6 @@ def __init__(self, context, func_ir, typemap, calltypes):
         # keep track of pruned precessors when branch degenerates to jump
         self.pruned_predecessors = {}
 
-    def copy(self):
-        """Returns a fresh copy with shared references to ``context``,
-        ``func_ir``, ``typemap``, ``calltypes``.
-        """
-        return ArrayAnalysis(
-            self.context, self.func_ir, self.typemap, self.calltypes)
-
     def get_equiv_set(self, block_label):
         """Return the equiv_set object of an block given its label.
         """
@@ -1441,9 +1434,7 @@ def _analyze_inst(self, label, scope, equiv_set, inst, redefined):
         elif type(inst) in array_analysis_extensions:
             # let external calls handle stmt if type matches
             f = array_analysis_extensions[type(inst)]
-            # Use a copy to isolate state mutation
-            copied_array_analysis = self.copy()
-            pre, post = f(inst, equiv_set, self.typemap, copied_array_analysis)
+            pre, post = f(inst, equiv_set, self.typemap, self)
 
         return pre, post
 


Test patch:
diff --git a/numba/tests/test_array_analysis.py b/numba/tests/test_array_analysis.py
--- a/numba/tests/test_array_analysis.py
+++ b/numba/tests/test_array_analysis.py
@@ -1025,6 +1025,68 @@ def iface_stub(self, scope, equiv_set, loc, args, kws):
             with self.subTest(fname=k, sig=got):
                 self.assertEqual(got, expected)
 
+    @skip_unsupported
+    def test_array_analysis_extensions(self):
+        # Test that the `array_analysis` object in `array_analysis_extensions`
+        # can perform analysis on the scope using `equiv_sets`.
+        from numba.parfors.parfor import Parfor
+        from numba.parfors import array_analysis
+
+        orig_parfor = array_analysis.array_analysis_extensions[Parfor]
+
+        shared = {'counter': 0}
+
+        def testcode(array_analysis):
+            # Find call node corresponding to the ``A = empty(n)``
+            func_ir = array_analysis.func_ir
+            for call in func_ir.blocks[0].find_exprs('call'):
+                callee = func_ir.get_definition(call.func)
+                if getattr(callee, "value", None) is empty:
+                    if getattr(call.args[0], 'name', None) == 'n':
+                        break
+            else:
+                return
+
+            variable_A = func_ir.get_assignee(call)
+            # n must be equiv to
+            es = array_analysis.equiv_sets[0]
+            self.assertTrue(es.is_equiv('n', variable_A.name))
+            shared['counter'] += 1
+
+        def new_parfor(parfor, equiv_set, typemap, array_analysis):
+            """Recursive array analysis for parfor nodes.
+            """
+            testcode(array_analysis)
+            # Call original
+            return orig_parfor(
+                parfor, equiv_set, typemap, array_analysis,
+            )
+
+        try:
+            # Replace the array-analysis extension for Parfor node
+            array_analysis.array_analysis_extensions[Parfor] = new_parfor
+
+            empty = np.empty   # avoid scanning a getattr in the IR
+            def f(n):
+                A = empty(n)
+                for i in prange(n):
+                    S = np.arange(i)
+                    A[i] = S.sum()
+                return A
+
+            got = njit(parallel=True)(f)(10)
+            executed_count = shared['counter']
+            self.assertGreater(executed_count, 0)
+        finally:
+            # Re-install the original handler
+            array_analysis.array_analysis_extensions[Parfor] = orig_parfor
+
+        # Check normal execution
+        expected = njit(parallel=True)(f)(10)
+        self.assertPreciseEqual(got, expected)
+        # Make sure we have uninstalled the handler
+        self.assertEqual(executed_count, shared['counter'])
+
 
 if __name__ == '__main__':
     unittest.main()