Repository: numba/numba

Issue:
Fix up the removed refcount pruner for arrays, make it more strict but still work!
Do as title.

Fixes #5746 #5743 

Discussion:

Patch:
diff --git a/docs/source/developer/repomap.rst b/docs/source/developer/repomap.rst
index 47a414a2e..b27e4acfe 100644
--- a/docs/source/developer/repomap.rst
+++ b/docs/source/developer/repomap.rst
@@ -393,6 +393,8 @@ typing and implementation to be specified together.
   needed to link generated code
 - :ghfile:`numba/core/fastmathpass.py` - Rewrite pass to add fastmath
   attributes to function call sites and binary operations
+- :ghfile:`numba/core/removerefctpass.py` - Rewrite pass to remove
+  unnecessary incref/decref pairs
 - :ghfile:`numba/core/descriptors.py` - empty base class for all target
   descriptors (is this needed?)
 - :ghfile:`numba/cpython/builtins.py` - Python builtin functions and
diff --git a/numba/core/callconv.py b/numba/core/callconv.py
index faf1455f1..ec1541288 100644
--- a/numba/core/callconv.py
+++ b/numba/core/callconv.py
@@ -503,6 +503,27 @@ class CPUCallConv(BaseCallConv):
                 if isinstance(a.type, ir.PointerType):
                     a.add_attribute("nocapture")
                     a.add_attribute("noalias")
+
+        # Add metadata to mark functions that may need NRT
+        # thus disabling aggressive refct pruning in removerefctpass.py
+        def type_may_always_need_nrt(ty):
+            # Returns True if it's a non-Array type that is contains MemInfo
+            if not isinstance(ty, types.Array):
+                dmm = self.context.data_model_manager
+                if dmm[ty].contains_nrt_meminfo():
+                    return True
+            return False
+
+        args_may_always_need_nrt = any(
+            map(type_may_always_need_nrt, fe_argtypes)
+        )
+
+        if args_may_always_need_nrt:
+            nmd = fn.module.add_named_metadata(
+                'numba_args_may_always_need_nrt',
+            )
+            nmd.add(fn.module.add_metadata([fn]))
+
         return fn
 
     def get_arguments(self, func):
diff --git a/numba/core/lowering.py b/numba/core/lowering.py
index 04e1ed68e..a8da51e71 100644
--- a/numba/core/lowering.py
+++ b/numba/core/lowering.py
@@ -7,7 +7,7 @@ from llvmlite.llvmpy.core import Constant, Type, Builder
 
 from numba import _dynfunc
 from numba.core import (typing, utils, types, ir, debuginfo, funcdesc,
-                        generators, config, ir_utils, cgutils)
+                        generators, config, ir_utils, cgutils, removerefctpass)
 from numba.core.errors import (LoweringError, new_error_context, TypingError,
                                LiteralTypingError, UnsupportedError)
 from numba.core.funcdesc import default_mangler
@@ -204,6 +204,12 @@ class BaseLower(object):
                 print(self.module)
             print('=' * 80)
 
+        # Special optimization to remove NRT on functions that do not need it.
+        if self.context.enable_nrt and self.generator_info is None:
+            removerefctpass.remove_unnecessary_nrt_usage(self.function,
+                                                         context=self.context,
+                                                         fndesc=self.fndesc)
+
         # Run target specific post lowering transformation
         self.context.post_lowering(self.module, self.library)
 
diff --git a/numba/core/removerefctpass.py b/numba/core/removerefctpass.py
new file mode 100644
index 000000000..e22343aa1
--- /dev/null
+++ b/numba/core/removerefctpass.py
@@ -0,0 +1,121 @@
+"""
+Implement a rewrite pass on a LLVM module to remove unnecessary 
+refcount operations.
+"""
+
+from llvmlite.ir.transforms import CallVisitor
+
+from numba.core import types
+
+
+class _MarkNrtCallVisitor(CallVisitor):
+    """
+    A pass to mark all NRT_incref and NRT_decref.
+    """
+    def __init__(self):
+        self.marked = set()
+
+    def visit_Call(self, instr):
+        if getattr(instr.callee, 'name', '') in _accepted_nrtfns:
+            self.marked.add(instr)
+
+
+def _rewrite_function(function):
+    # Mark NRT usage
+    markpass = _MarkNrtCallVisitor()
+    markpass.visit_Function(function)
+    # Remove NRT usage
+    for bb in function.basic_blocks:
+        for inst in list(bb.instructions):
+            if inst in markpass.marked:
+                bb.instructions.remove(inst)
+
+
+_accepted_nrtfns = 'NRT_incref', 'NRT_decref'
+
+
+def _legalize(module, dmm, fndesc):
+    """
+    Legalize the code in the module.
+    Returns True if the module is legal for the rewrite pass that removes
+    unnecessary refcounts.
+    """
+
+    def valid_output(ty):
+        """
+        Valid output are any type that does not need refcount
+        """
+        model = dmm[ty]
+        return not model.contains_nrt_meminfo()
+
+    def valid_input(ty):
+        """
+        Valid input are any type that does not need refcount except Array.
+        """
+        return valid_output(ty) or isinstance(ty, types.Array)
+
+
+    # Ensure no reference to function marked as
+    # "numba_args_may_always_need_nrt"
+    try:
+        nmd = module.get_named_metadata("numba_args_may_always_need_nrt")
+    except KeyError:
+        # Nothing marked
+        pass
+    else:
+        # Has functions marked as "numba_args_may_always_need_nrt"
+        if len(nmd.operands) > 0:
+            # The pass is illegal for this compilation unit.
+            return False
+
+    # More legalization base on function type
+    argtypes = fndesc.argtypes
+    restype = fndesc.restype
+    calltypes = fndesc.calltypes
+
+    # Legalize function arguments
+    for argty in argtypes:
+        if not valid_input(argty):
+            return False
+
+    # Legalize function return
+    if not valid_output(restype):
+        return False
+
+    # Legalize all called functions
+    for callty in calltypes.values():
+        if callty is not None and not valid_output(callty.return_type):
+            return False
+
+    # Ensure no allocation
+    for fn in module.functions:
+        if fn.name.startswith("NRT_"):
+            if fn.name not in _accepted_nrtfns:
+                return False
+
+    return True
+
+
+def remove_unnecessary_nrt_usage(function, context, fndesc):
+    """
+    Remove unnecessary NRT incref/decref in the given LLVM function.
+    It uses highlevel type info to determine if the function does not need NRT.
+    Such a function does not:
+
+    - return array object(s);
+    - take arguments that need refcounting except array;
+    - call function(s) that return refcounted object.
+
+    In effect, the function will not capture or create references that extend
+    the lifetime of any refcounted objects beyound the lifetime of the
+    function.
+
+    The rewrite is performed in place.
+    If rewrite has happened, this function returns True, otherwise, it returns False.
+    """
+    dmm = context.data_model_manager
+    if _legalize(function.module, dmm, fndesc):
+        _rewrite_function(function)
+        return True
+    else:
+        return False


Test patch:
diff --git a/numba/tests/test_array_reductions.py b/numba/tests/test_array_reductions.py
index 05822eaaf..4c648a6d9 100644
--- a/numba/tests/test_array_reductions.py
+++ b/numba/tests/test_array_reductions.py
@@ -992,7 +992,6 @@ class TestArrayReductionsExceptions(MemoryLeakMixin, TestCase):
         with self.assertRaises(ValueError) as e:
             cfunc(self.zero_size)
         self.assertIn(msg, str(e.exception))
-        self.disable_leak_check()
 
     @classmethod
     def install(cls):