Repository: numba/numba

Issue:
Reassign to an IR Variable Produces a Memory Leak
<!--

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).

-->

## Reporting a bug

<!--

Before submitting a bug report please ensure that you can check off these boxes:

-->

- [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).
- [ ] I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'.

<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->
Starting with Numba 0.54, some of Bodo's tests are reporting memory leaks. This appears to be because we are failing to generate DEL nodes if an IR variable is assigned to twice in the same basic block. The IR generated doesn't appear to change, so it seems likely that Numba 0.54 is just the first version where we accurately detect these memory leaks. 
 
For context, we define a variable twice in the same basic block to achieve the behavior of `inplace=True` with a sort operation. We allocate a new array somewhere and then as a last step we assign the old IR variable to this new array.

I cannot produce a reproducer in just Numba, so I've included a segment of the IR before and after del nodes are added (omitting a bunch of unnecessary lines that add clutter). It should be possible to reproduce this with a pipeline extension, but if that's not necessary I want to avoid generating that reproducer.

Before adding del nodes:
```python
label 0:
    $2load_global.0.1090 = global(_get_rank: CPUDispatcher(<function get_rank at 0x10d7e2790>)) ['$2load_global.0.1090']
    rank_var.1093 = call $2load_global.0.1090(func=$2load_global.0.1090, args=[], kws=(), vararg=None, target=None) ['$2load_global.0.1090', 'rank_var.1093']
    $2load_global.0.1094 = global(_get_size: CPUDispatcher(<function get_size at 0x123627d30>)) ['$2load_global.0.1094']
    n_pes_var.1097 = call $2load_global.0.1094(func=$2load_global.0.1094, args=[], kws=(), vararg=None, target=None) ['$2load_global.0.1094', 'n_pes_var.1097']
    df = arg(0, name=df)                     ['df']
    col_ind.842 = const(int, 0)              ['col_ind.842']
    $2load_global.0.845 = global(bodo: <module 'bodo' from '/Users/nicholasriasanovsky/Documents/bodo/Bodo/bodo/__init__.py'>) ['$2load_global.0.845']
    $4load_attr.1.846 = getattr(value=$2load_global.0.845, attr=hiframes) ['$2load_global.0.845', '$4load_attr.1.846']
    $6load_attr.2.847 = getattr(value=$4load_attr.1.846, attr=pd_dataframe_ext) ['$4load_attr.1.846', '$6load_attr.2.847']
    $8load_method.3.848 = getattr(value=$6load_attr.2.847, attr=get_dataframe_data) ['$6load_attr.2.847', '$8load_method.3.848']
    $data0.894.1026 = call $8load_method.3.848(df, col_ind.842, func=$8load_method.3.848, args=[Var(df, test_sort_values.py:278), Var(col_ind.842, test_sort_values.py:278)], kws=(), vararg=None, target=None) ['$8load_method.3.848', '$data0.894.1026', 'col_ind.842', 'df']
    col_ind.851 = const(int, 1)              ['col_ind.851']

    ...(more IR omited)

    $data0.894.1026 = static_getitem(value=key_data.1168, index=0, index_var=None, fn=<built-in function getitem>) ['$data0.894.1026', 'key_data.1168']
    del key_data.1168                        []
    $data1.895.1027 = static_getitem(value=sort_data.1169, index=0, index_var=None, fn=<built-in function getitem>) ['$data1.895.1027', 'sort_data.1169']
    $arr.886.1024 = static_getitem(value=sort_data.1169, index=1, index_var=None, fn=<built-in function getitem>) ['$arr.886.1024', 'sort_data.1169']
...(more IR omited)
```

After adding del nodes:
```python
label 0:
    $data0.894.1026 = call $8load_method.3.848(df, col_ind.842, func=$8load_method.3.848, args=[Var(df, test_sort_values.py:278), Var(col_ind.842, test_sort_values.py:278)], kws=(), vararg=None, target=None) ['$8load_method.3.848', '$data0.894.1026', 'col_ind.842', 'df']
    del col_ind.842                          []
    del $8load_method.3.848                  []
    col_ind.851 = const(int, 1)              ['col_ind.851']

    ...(more IR omited but no del for $data0.894.1026 )

    $data0.894.1026 = static_getitem(value=key_data.1168, index=0, index_var=None, fn=<built-in function getitem>) ['$data0.894.1026', 'key_data.1168']
    del key_data.1168                        []
    $data1.895.1027 = static_getitem(value=sort_data.1169, index=0, index_var=None, fn=<built-in function getitem>) ['$data1.895.1027', 'sort_data.1169']
    $arr.886.1024 = static_getitem(value=sort_data.1169, index=1, index_var=None, fn=<built-in function getitem>) ['$arr.886.1024', 'sort_data.1169']
    del sort_data.1169                       []
...(more IR omited but no del for $data0.894.1026 )
    del $data1.895.1027                      []
    del $data0.894.1026                      []
    
```

The key takeaways are that we assigns `$data0.894.1026` twice, but the del is only generated once at the end.
Update postproc to general multiple del nodes if a variable is assigned multiple times in the same basic block
<!--

Thanks for wanting to contribute to Numba :)

First, if you need some help or want to chat to the core developers, please
visit https://gitter.im/numba/numba for real time chat or post to the Numba
forum https://numba.discourse.group/.

Here's some guidelines to help the review process go smoothly.

0. Please write a description in this text box of the changes that are being
   made.

1. Please ensure that you have written units tests for the changes made/features
   added.

2. If you are closing an issue please use one of the automatic closing words as
   noted here: https://help.github.com/articles/closing-issues-using-keywords/

3. If your pull request is not ready for review but you want to make use of the
   continuous integration testing facilities please label it with `[WIP]` and
   then remove the label when you'd like it to be reviewed.

4. Once review has taken place please do not add features or make changes out of
   the scope of those requested by the reviewer (doing this just add delays as
   already reviewed code ends up having to be re-reviewed/it is hard to tell
   what is new etc!). Further, please do not rebase your branch on master/force
   push/rewrite history, doing any of these causes the context of any comments
   made by reviewers to be lost. If conflicts occur against master they should
   be resolved by merging master into the branch used for making the pull
   request.

Many thanks in advance for your cooperation!

-->

In Bodo we assign to a variable multiple times to support "inplace" updates. For example a sort with `inplace=True` is accomplished by reassigning to the same IR variable after a sort is finished. This change fixes a memory leak that was deteced in Numba 0.54 because only 1 del is generated.

I'll add a reproducer with a Numba pipeline extension shortly.

<!-- You can link to an existing issue using the github syntax #<issue number>  -->



Discussion:
@sklam Is reassignment legal? Should the IR be in SSA form?
Is this to fix #7225?
No this fixes #7258, which I previously described on discourse but hadn't filed the issue. This seems like an easy fix and doesn't seem to interfere with any tests, so could we include it in the numba 0.54 release?
> I'll add a reproducer with a Numba pipeline extension shortly.

Is this still upcoming?

> This seems like an easy fix and doesn't seem to interfere with any tests, so could we include it in the numba 0.54 release?

I'm not 100% sure that reassignment ought to be legal, as mentioned in #7258 - @sklam's input may be needed here.
> > I'll add a reproducer with a Numba pipeline extension shortly.
> 
> Is this still upcoming?
> 

Yes. I'll attach it in the next comment.

> > This seems like an easy fix and doesn't seem to interfere with any tests, so could we include it in the numba 0.54 release?
> 
> I'm not 100% sure that reassignment ought to be legal, as mentioned in #7258 - @sklam's input may be needed here.

So from our perspective this is incredibly useful because we have some difficulties with inplace updates. For example, we may need to update an immutable array inplace in a way that changes its length (I can elaborate if necessary). We handle this  by updating the reference to where the data is stored, which requires reassignment.

I totally understand if this is not behavior that is needed when just using Numba as a compiler, but since Numba is also a tool for building compilers we believe we have a relevant usecase for why Numba shouldn't complete disallow this behavior. Perhaps there could be a flag used that would allow us to enable this behavior?
Here is the Python file that can be used to generate IR that produces reassign in a way that vaguely resembles our usecase. I tried to simplify it as much as possible so this isn't a general/complete solution.

Github didn't accept just the .py file so I placed it in a zip file. Sorry if that's annoying.
[dummy_reassign_ext.zip](https://github.com/numba/numba/files/6900861/dummy_reassign_ext.zip)

Currently, I believe Numba's IR is in SSA form only during typing. This change shouldn't affect SSA IR anyways.
~@njriasan, RE: the reproducer `dummy_reassign_ext.zip`. Where should I be looking to see the effect of the patch?~

*UPDATE* i can observe it via `get_allocation_stats()`


Resolving patch:
diff --git a/numba/core/lowering.py b/numba/core/lowering.py
--- a/numba/core/lowering.py
+++ b/numba/core/lowering.py
@@ -314,7 +314,14 @@ def _find_singly_assigned_variable(self):
                 if len(var_assign_map[var]) == 1:
                     # Usemap does not keep locally defined variables.
                     if len(var_use_map[var]) == 0:
-                        sav.add(var)
+                        # Ensure that the variable is not defined multiple times
+                        # the the block
+                        [defblk] = var_assign_map[var]
+                        assign_stmts = self.blocks[defblk].find_insts(ir.Assign)
+                        assigns = [stmt for stmt in assign_stmts
+                                   if stmt.target.name == var]
+                        if len(assigns) == 1:
+                            sav.add(var)
 
         self._singly_assigned_vars = sav
         self._blk_local_varmap = {}


Test patch:
diff --git a/numba/tests/test_ssa.py b/numba/tests/test_ssa.py
--- a/numba/tests/test_ssa.py
+++ b/numba/tests/test_ssa.py
@@ -8,9 +8,13 @@
 import numpy as np
 
 from numba import njit, jit, types
-from numba.core import errors
+from numba.core import errors, ir
+from numba.core.compiler_machinery import FunctionPass, register_pass
+from numba.core.compiler import DefaultPassBuilder, CompilerBase
+from numba.core.untyped_passes import ReconstructSSA, PreserveIR
+from numba.core.typed_passes import NativeLowering
 from numba.extending import overload
-from numba.tests.support import TestCase, override_config
+from numba.tests.support import MemoryLeakMixin, TestCase, override_config
 
 
 _DEBUG = False
@@ -493,3 +497,81 @@ def while_for(n, max_iter=1):
         self.assertPreciseEqual(while_for(10), while_for.py_func(10))
         # One phi?
         self.assertEqual(phi_counter, [1])
+
+
+class TestSROAIssues(MemoryLeakMixin, TestCase):
+    # This tests issues related to the SROA optimization done in lowering, which
+    # reduces time spent in the LLVM SROA pass. The optimization is related to
+    # SSA and tries to reduce the number of alloca statements for variables with
+    # only a single assignemnt.
+    def test_issue7258_multiple_assignment_post_SSA(self):
+        # This test adds a pass that will duplicate assignment statements to
+        # variables named "foobar".
+        # In the reported issue, the bug will cause a memory leak.
+        cloned = []
+
+        @register_pass(analysis_only=False, mutates_CFG=True)
+        class CloneFoobarAssignments(FunctionPass):
+            # A pass that clones variable assignments into "foobar"
+            _name = "clone_foobar_assignments_pass"
+
+            def __init__(self):
+                FunctionPass.__init__(self)
+
+            def run_pass(self, state):
+                mutated = False
+                for blk in state.func_ir.blocks.values():
+                    to_clone = []
+                    # find assignments to "foobar"
+                    for assign in blk.find_insts(ir.Assign):
+                        if assign.target.name == "foobar":
+                            to_clone.append(assign)
+                    # clone
+                    for assign in to_clone:
+                        clone = copy.deepcopy(assign)
+                        blk.insert_after(clone, assign)
+                        mutated = True
+                        # keep track of cloned statements
+                        cloned.append(clone)
+                return mutated
+
+        class CustomCompiler(CompilerBase):
+            def define_pipelines(self):
+                pm = DefaultPassBuilder.define_nopython_pipeline(
+                    self.state, "custom_pipeline",
+                )
+                pm._finalized = False
+                # Insert the cloning pass after SSA
+                pm.add_pass_after(CloneFoobarAssignments, ReconstructSSA)
+                # Capture IR post lowering
+                pm.add_pass_after(PreserveIR, NativeLowering)
+                pm.finalize()
+                return [pm]
+
+        @njit(pipeline_class=CustomCompiler)
+        def udt(arr):
+            foobar = arr + 1  # this assignment will be cloned
+            return foobar
+
+        arr = np.arange(10)
+        # Verify that the function works as expected
+        self.assertPreciseEqual(udt(arr), arr + 1)
+        # Verify that the expected statement is cloned
+        self.assertEqual(len(cloned), 1)
+        self.assertEqual(cloned[0].target.name, "foobar")
+        # Verify in the Numba IR that the expected statement is cloned
+        nir = udt.overloads[udt.signatures[0]].metadata['preserved_ir']
+        self.assertEqual(len(nir.blocks), 1,
+                         "only one block")
+        [blk] = nir.blocks.values()
+        assigns = blk.find_insts(ir.Assign)
+        foobar_assigns = [stmt for stmt in assigns
+                          if stmt.target.name == "foobar"]
+        self.assertEqual(
+            len(foobar_assigns), 2,
+            "expected two assignment statements into 'foobar'",
+        )
+        self.assertEqual(
+            foobar_assigns[0], foobar_assigns[1],
+            "expected the two assignment statements to be the same",
+        )