Repository: numba/numba

Issue:
refcount op in branching code affects performance
The refcount ops emitted by this code is splitted across basicblocks.  Numba refcount pruning pass is not able to prune them.

```python
import numba
import numpy as np
from timeit import default_timer as timer


def bench(nrt):
    print("set _nrt=%s" % nrt)

    jit = numba.jit(nopython=True, nogil=True, _nrt=nrt)

    @jit
    def _append_non_na(x, y, agg, field):
        if not np.isnan(field):
            agg[y, x] += 1

    @jit
    def _append(x, y, agg, field):
        if not np.isnan(field):
            if np.isnan(agg[y, x]):
                agg[y, x] = field
            else:
                agg[y, x] += field

    @jit
    def append(x, y, agg, field):
        _append_non_na(x, y, agg, field)
        _append(x, y, agg, field)


    @jit
    def extend(arr, field):
        for i in range(arr.shape[0]):
            for j in range(arr.shape[1]):
                append(j, i, arr, field)


    arr = np.ones((10, 100))

    for i in range(10):
        ts = timer()
        extend(arr, 10)
        te = timer()
        print("%.5fms" % ((te - ts) * 1000))

    # extend.inspect_cfg(extend.signatures[0]).display(view=True)


bench(nrt=True)
bench(nrt=False)
```

Output
```
set _nrt=True
224.90708ms
0.02676ms
0.02478ms
0.02456ms
0.02460ms
0.02464ms
0.02724ms
0.02477ms
0.02457ms
0.02448ms
set _nrt=False
157.52096ms
0.00259ms
0.00072ms
0.00065ms
0.00062ms
0.00062ms
0.00065ms
0.00062ms
0.00065ms
0.00066ms
```



Discussion:

Resolving patch:
diff --git a/numba/datamodel/models.py b/numba/datamodel/models.py
--- a/numba/datamodel/models.py
+++ b/numba/datamodel/models.py
@@ -94,6 +94,12 @@ def traverse(self, builder):
         """
         return []
 
+    def traverse_models(self):
+        """
+        Recursively list all models involved in this model.
+        """
+        return [self._dmm[t] for t in self.traverse_types()]
+
     def traverse_types(self):
         """
         Recursively list all frontend types involved in this model.
@@ -116,6 +122,12 @@ def get_nrt_meminfo(self, builder, value):
     def has_nrt_meminfo(self):
         return False
 
+    def contains_nrt_meminfo(self):
+        """
+        Recursively check all contained types for need for NRT meminfo.
+        """
+        return any(model.has_nrt_meminfo() for model in self.traverse_models())
+
     def _compared_fields(self):
         return (type(self), self._fe_type)
 
diff --git a/numba/lowering.py b/numba/lowering.py
--- a/numba/lowering.py
+++ b/numba/lowering.py
@@ -1,16 +1,14 @@
 from __future__ import print_function, division, absolute_import
 
 from collections import namedtuple
-import sys
 from functools import partial
 
-from llvmlite.ir import Value
 from llvmlite.llvmpy.core import Constant, Type, Builder
 
 from . import (_dynfunc, cgutils, config, funcdesc, generators, ir, types,
                typing, utils)
 from .errors import LoweringError, new_error_context
-from .targets import imputils
+from .targets import removerefctpass
 from .funcdesc import default_mangler
 from . import debuginfo
 
@@ -147,6 +145,12 @@ def lower(self):
             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/targets/removerefctpass.py b/numba/targets/removerefctpass.py
new file mode 100644
--- /dev/null
+++ b/numba/targets/removerefctpass.py
@@ -0,0 +1,109 @@
+"""
+Implement a rewrite pass on LLVM module to remove unnecessary refcount
+operation.
+"""
+from __future__ import absolute_import, print_function
+
+from llvmlite.ir.transforms import CallVisitor
+
+from numba 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 instr.callee.name in ('NRT_incref', 'NRT_decref'):
+            self.marked.add(instr)
+
+
+def _rewrite_function(function):
+    # Mark NRT usage
+    markpass = _MarkNrtCallVisitor()
+    markpass.visit_Function(function)
+    marked = markpass.marked
+    # Remove NRT usage
+    for bb in function.basic_blocks:
+        for inst in list(bb.instructions):
+            if inst in 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 remove
+    unnecessary refcount.
+    """
+
+    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)
+
+    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;
+    - take arguments that need refcount except array;
+    - call function 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 performs inplace.
+    If rewrite has happen, this function return True. Otherwise, return 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_datamodel.py b/numba/tests/test_datamodel.py
--- a/numba/tests/test_datamodel.py
+++ b/numba/tests/test_datamodel.py
@@ -176,5 +176,30 @@ def test_nested_empty_tuples(self):
         self._test_as_arguments(fe_args)
 
 
+class TestMemInfo(unittest.TestCase):
+    def setUp(self):
+        self.dmm = datamodel.default_manager
+
+    def test_number(self):
+        ty = types.int32
+        dm = self.dmm[ty]
+        self.assertFalse(dm.contains_nrt_meminfo())
+
+    def test_array(self):
+        ty = types.int32[:]
+        dm = self.dmm[ty]
+        self.assertTrue(dm.contains_nrt_meminfo())
+
+    def test_tuple_of_number(self):
+        ty = types.UniTuple(dtype=types.int32, count=2)
+        dm = self.dmm[ty]
+        self.assertFalse(dm.contains_nrt_meminfo())
+
+    def test_tuple_of_array(self):
+        ty = types.UniTuple(dtype=types.int32[:], count=2)
+        dm = self.dmm[ty]
+        self.assertTrue(dm.contains_nrt_meminfo())
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/numba/tests/test_nrt.py b/numba/tests/test_nrt.py
--- a/numba/tests/test_nrt.py
+++ b/numba/tests/test_nrt.py
@@ -3,6 +3,7 @@
 import math
 import os
 import sys
+import re
 
 import numpy as np
 
@@ -435,6 +436,42 @@ def test_refct_pruning_op_recognize(self):
         # no other lines
         self.assertEqual(len(list(pruned_lines.splitlines())), len(combined))
 
+    def test_refct_pruning_with_branches(self):
+        '''testcase from #2350'''
+        @njit
+        def _append_non_na(x, y, agg, field):
+            if not np.isnan(field):
+                agg[y, x] += 1
+
+        @njit
+        def _append(x, y, agg, field):
+            if not np.isnan(field):
+                if np.isnan(agg[y, x]):
+                    agg[y, x] = field
+                else:
+                    agg[y, x] += field
+
+        @njit
+        def append(x, y, agg, field):
+            _append_non_na(x, y, agg, field)
+            _append(x, y, agg, field)
+
+        # Disable python wrapper to avoid detecting necessary
+        # refcount inside it
+        @njit(no_cpython_wrapper=True)
+        def extend(arr, field):
+            for i in range(arr.shape[0]):
+                for j in range(arr.shape[1]):
+                    append(j, i, arr, field)
+
+        # Compile
+        extend.compile("(f4[:,::1], f4)")
+
+        # Test there are no reference count operations
+        llvmir = str(extend.inspect_llvm(extend.signatures[0]))
+        refops = list(re.finditer(r'(NRT_incref|NRT_decref)\([^\)]+\)', llvmir))
+        self.assertEqual(len(refops), 0)
+
 
 if __name__ == '__main__':
     unittest.main()