Repository: numba/numba

Issue:
overload inlining fails for builtin operators
<!--

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:

-->

- [x] 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).
- [x] 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!

-->
`overload` with `inline='always'` fails for builtin operators:
```python
import operator
import numba
from numba.extending import overload, register_model, models, unbox, NativeValue, typeof_impl

class Foo:
    pass

class FooType(numba.types.Type):
    def __init__(self):
        super(FooType, self).__init__(name='Foo')


foo_type = FooType()
register_model(FooType)(models.OpaqueModel)


@typeof_impl.register(Foo)
def _typ_foo(val, c):
  return FooType()


@unbox(FooType)
def unbox_foo(typ, obj, c):
  return NativeValue(c.context.get_dummy_value())


@overload(operator.eq, inline='always')
def foo_eq(a, b):
    if a == FooType():
        return lambda a, b: True


@numba.njit
def f(a):
    return a == "A"


print(f(Foo()))

```
```
  File "/opt/anaconda3/envs/TEST3/lib/python3.7/site-packages/numba/typed_passes.py", line 349, in run_pass
    lower.lower()
  File "/opt/anaconda3/envs/TEST3/lib/python3.7/site-packages/numba/lowering.py", line 195, in lower
    self.lower_normal_function(self.fndesc)
  File "/opt/anaconda3/envs/TEST3/lib/python3.7/site-packages/numba/lowering.py", line 248, in lower_normal_function
    entry_block_tail = self.lower_function_body()
  File "/opt/anaconda3/envs/TEST3/lib/python3.7/site-packages/numba/lowering.py", line 273, in lower_function_body
    self.lower_block(block)
  File "/opt/anaconda3/envs/TEST3/lib/python3.7/site-packages/numba/lowering.py", line 288, in lower_block
    self.lower_inst(inst)
  File "/opt/anaconda3/envs/TEST3/lib/python3.7/contextlib.py", line 130, in __exit__
    self.gen.throw(type, value, traceback)
  File "/opt/anaconda3/envs/TEST3/lib/python3.7/site-packages/numba/errors.py", line 725, in new_error_context
    six.reraise(type(newerr), newerr, tb)
  File "/opt/anaconda3/envs/TEST3/lib/python3.7/site-packages/numba/six.py", line 669, in reraise
    raise value
numba.errors.LoweringError: Failed in nopython mode pipeline (step: nopython mode backend)
No definition for lowering _EmptyImplementationEntry('always inlined')(Foo, Literal[str](A)) -> bool

File "test_builtins_inline.py", line 35:
def f(a):
    return a == "A"
    ^

[1] During: lowering "$6compare_op.2 = a == $const4.1" at test_builtins_inline.py (35)
```

Discussion:

Patch:
diff --git a/numba/core/inline_closurecall.py b/numba/core/inline_closurecall.py
index be56fe051..a3ea57967 100644
--- a/numba/core/inline_closurecall.py
+++ b/numba/core/inline_closurecall.py
@@ -688,6 +688,8 @@ def _get_callee_args(call_expr, callee, loc, func_ir):
         args = list(call_expr.args)
     elif call_expr.op == 'getattr':
         args = [call_expr.value]
+    elif ir_utils.is_operator_or_getitem(call_expr):
+        args = call_expr.list_vars()
     else:
         raise TypeError("Unsupported ir.Expr.{}".format(call_expr.op))
 
@@ -1269,6 +1271,7 @@ def _fix_nested_array(func_ir):
         size_tuple_def.items += extra_dims
         # In-place modify rhs_def to be getitem
         rhs_def.op = 'getitem'
+        rhs_def.fn = operator.getitem
         rhs_def.value = get_definition(func_ir, lhs, lhs_only=True)
         rhs_def.index = stmt.index
         del rhs_def._kws['func']
diff --git a/numba/core/ir.py b/numba/core/ir.py
index d07430fe5..b4df49025 100644
--- a/numba/core/ir.py
+++ b/numba/core/ir.py
@@ -499,7 +499,8 @@ class Expr(Inst):
         assert isinstance(index, Var)
         assert isinstance(loc, Loc)
         op = 'getitem'
-        return cls(op=op, loc=loc, value=value, index=index)
+        fn = operator.getitem
+        return cls(op=op, loc=loc, value=value, index=index, fn=fn)
 
     @classmethod
     def typed_getitem(cls, value, dtype, index, loc):
@@ -515,8 +516,9 @@ class Expr(Inst):
         assert index_var is None or isinstance(index_var, Var)
         assert isinstance(loc, Loc)
         op = 'static_getitem'
+        fn = operator.getitem
         return cls(op=op, loc=loc, value=value, index=index,
-                   index_var=index_var)
+                   index_var=index_var, fn=fn)
 
     @classmethod
     def cast(cls, value, loc):
diff --git a/numba/core/ir_utils.py b/numba/core/ir_utils.py
index 1d58c5c8b..38b4d3fd5 100644
--- a/numba/core/ir_utils.py
+++ b/numba/core/ir_utils.py
@@ -1737,6 +1737,12 @@ def dump_blocks(blocks):
         for stmt in block.body:
             print("    ", stmt)
 
+def is_operator_or_getitem(expr):
+    """true if expr is unary or binary operator or getitem"""
+    return (isinstance(expr, ir.Expr)
+            and getattr(expr, 'op', False)
+            and expr.op in ['unary', 'binop', 'inplace_binop', 'getitem', 'static_getitem'])
+
 def is_get_setitem(stmt):
     """stmt is getitem assignment or setitem (and static cases)"""
     return is_getitem(stmt) or is_setitem(stmt)
diff --git a/numba/core/typed_passes.py b/numba/core/typed_passes.py
index 3756a3e7c..e5b50497f 100644
--- a/numba/core/typed_passes.py
+++ b/numba/core/typed_passes.py
@@ -17,7 +17,8 @@ from numba.core.ir_utils import (raise_on_unsupported_feature, warn_deprecated,
                                  check_and_legalize_ir, guard,
                                  dead_code_elimination, simplify_CFG,
                                  get_definition, remove_dels,
-                                 build_definitions, compute_cfg_from_blocks)
+                                 build_definitions, compute_cfg_from_blocks,
+                                 is_operator_or_getitem)
 from numba.core import postproc
 
 
@@ -508,15 +509,11 @@ class InlineOverloads(FunctionPass):
         while work_list:
             label, block = work_list.pop()
             for i, instr in enumerate(block.body):
+                # TO-DO: other statements (setitem)
                 if isinstance(instr, ir.Assign):
                     expr = instr.value
                     if isinstance(expr, ir.Expr):
-                        if expr.op == 'call':
-                            workfn = self._do_work_call
-                        elif expr.op == 'getattr':
-                            workfn = self._do_work_getattr
-                        else:
-                            continue
+                        workfn = self._do_work_expr
 
                         if guard(workfn, state, work_list, block, i, expr,
                                  inline_worker):
@@ -548,69 +545,73 @@ class InlineOverloads(FunctionPass):
             print(''.center(80, '-'))
         return True
 
-    def _do_work_getattr(self, state, work_list, block, i, expr, inline_worker):
+    def _get_attr_info(self, state, expr):
         recv_type = state.type_annotation.typemap[expr.value.name]
         recv_type = types.unliteral(recv_type)
         matched = state.typingctx.find_matching_getattr_template(
             recv_type, expr.attr,
         )
         if not matched:
-            return False
+            return None
+
         template = matched['template']
         if getattr(template, 'is_method', False):
             # The attribute template is representing a method.
             # Don't inline the getattr.
-            return False
+            return None
 
-        inline_type = getattr(template, '_inline', None)
-        if inline_type is None:
-            # inline not defined
-            return False
+        templates = [template]
         sig = typing.signature(matched['return_type'], recv_type)
         arg_typs = sig.args
+        is_method = False
 
-        if not inline_type.is_never_inline:
-            try:
-                impl = template._overload_func(recv_type)
-                if impl is None:
-                    raise Exception  # abort for this template
-            except Exception:
-                return False
-        else:
-            return False
+        return templates, sig, arg_typs, is_method
 
-        is_method = False
-        return self._run_inliner(
-            state, inline_type, sig, template, arg_typs, expr, i, impl, block,
-            work_list, is_method, inline_worker,
-        )
+    def _get_callable_info(self, state, expr):
 
-    def _do_work_call(self, state, work_list, block, i, expr, inline_worker):
-        # try and get a definition for the call, this isn't always possible as
-        # it might be a eval(str)/part generated awaiting update etc. (parfors)
-        to_inline = None
-        try:
-            to_inline = state.func_ir.get_definition(expr.func)
-        except Exception:
-            return False
+        def get_func_type(state, expr):
+            func_ty = None
+            if expr.op == 'call':
+                # check this is a known and typed function
+                try:
+                    func_ty = state.type_annotation.typemap[expr.func.name]
+                except KeyError:
+                    # e.g. Calls to CUDA Intrinsic have no mapped type
+                    # so KeyError
+                    return None
+                if not hasattr(func_ty, 'get_call_type'):
+                    return None
+
+            elif is_operator_or_getitem(expr):
+                func_ty = state.typingctx.resolve_value_type(expr.fn)
+            else:
+                return None
 
-        # do not handle closure inlining here, another pass deals with that.
-        if getattr(to_inline, 'op', False) == 'make_function':
-            return False
+            return func_ty
 
-        # check this is a known and typed function
-        try:
-            func_ty = state.type_annotation.typemap[expr.func.name]
-        except KeyError:
-            # e.g. Calls to CUDA Intrinsic have no mapped type so KeyError
-            return False
-        if not hasattr(func_ty, 'get_call_type'):
-            return False
+        if expr.op == 'call':
+            # try and get a definition for the call, this isn't always
+            # possible as it might be a eval(str)/part generated
+            # awaiting update etc. (parfors)
+            to_inline = None
+            try:
+                to_inline = state.func_ir.get_definition(expr.func)
+            except Exception:
+                return None
+
+            # do not handle closure inlining here, another pass deals with that
+            if getattr(to_inline, 'op', False) == 'make_function':
+                return None
+
+        func_ty = get_func_type(state, expr)
+        if func_ty is None:
+            return None
 
         sig = state.type_annotation.calltypes[expr]
-        is_method = False
+        if not sig:
+            return None
 
-        # search the templates for this overload looking for "inline"
+        templates, arg_typs, is_method = None, None, False
         if getattr(func_ty, 'template', None) is not None:
             # @overload_method
             is_method = True
@@ -621,30 +622,53 @@ class InlineOverloads(FunctionPass):
             templates = getattr(func_ty, 'templates', None)
             arg_typs = sig.args
 
-        if templates is None:
-            return False
+        return templates, sig, arg_typs, is_method
 
-        impl = None
-        for template in templates:
-            inline_type = getattr(template, '_inline', None)
-            if inline_type is None:
-                # inline not defined
-                continue
-            if not inline_type.is_never_inline:
-                try:
-                    impl = template._overload_func(*arg_typs)
-                    if impl is None:
-                        raise Exception  # abort for this template
-                    break
-                except Exception:
+    def _do_work_expr(self, state, work_list, block, i, expr, inline_worker):
+
+        def select_template(templates, args):
+            if templates is None:
+                return None
+
+            impl = None
+            for template in templates:
+                inline_type = getattr(template, '_inline', None)
+                if inline_type is None:
+                    # inline not defined
+                    continue
+                if args not in template._inline_overloads:
+                    # skip overloads not matching signature
                     continue
+                if not inline_type.is_never_inline:
+                    try:
+                        impl = template._overload_func(*args)
+                        if impl is None:
+                            raise Exception  # abort for this template
+                        break
+                    except Exception:
+                        continue
+            else:
+                return None
+
+            return template, inline_type, impl
+
+        inlinee_info = None
+        if expr.op == 'getattr':
+            inlinee_info = self._get_attr_info(state, expr)
         else:
+            inlinee_info = self._get_callable_info(state, expr)
+
+        if not inlinee_info:
+            return False
+
+        templates, sig, arg_typs, is_method = inlinee_info
+        inlinee = select_template(templates, arg_typs)
+        if inlinee is None:
             return False
+        template, inlinee_type, impl = inlinee
 
-        # at this point we know we maybe want to inline something and there's
-        # definitely something that could be inlined.
         return self._run_inliner(
-            state, inline_type, sig, template, arg_typs, expr, i, impl, block,
+            state, inlinee_type, sig, template, arg_typs, expr, i, impl, block,
             work_list, is_method, inline_worker,
         )
 


Test patch:
diff --git a/numba/tests/test_ir_inlining.py b/numba/tests/test_ir_inlining.py
index cf5f3f620..25a6c5db2 100644
--- a/numba/tests/test_ir_inlining.py
+++ b/numba/tests/test_ir_inlining.py
@@ -3,6 +3,7 @@ This tests the inline kwarg to @jit and @overload etc, it has nothing to do with
 LLVM or low level inlining.
 """
 
+import operator
 from itertools import product
 import numpy as np
 
@@ -125,10 +126,34 @@ class InliningBase(TestCase):
                 if getattr(expr, 'op', False) == 'call':
                     func_defn = fir.get_definition(expr.func)
                     found |= func_defn.name == k
+                elif ir_utils.is_operator_or_getitem(expr):
+                    found |= expr.fn.__name__ == k
             self.assertFalse(found == v)
 
         return fir  # for use in further analysis
 
+    def make_dummy_type(self):
+        # Use test_id to make sure no collision is possible.
+        test_id = self.id()
+        DummyType = type('DummyTypeFor{}'.format(test_id), (types.Opaque,), {})
+
+        dummy_type = DummyType("my_dummy")
+        register_model(DummyType)(OpaqueModel)
+
+        class Dummy(object):
+            pass
+
+        @typeof_impl.register(Dummy)
+        def typeof_dummy(val, c):
+            return dummy_type
+
+        @unbox(DummyType)
+        def unbox_dummy(typ, obj, c):
+            return NativeValue(c.context.get_dummy_value())
+
+        self.Dummy = Dummy
+        self.DummyType = DummyType
+
 
 # used in _gen_involved
 _GLOBAL = 1234
@@ -481,6 +506,12 @@ class TestRegisterJitableInlining(MemoryLeakMixin, InliningBase):
 
 class TestOverloadInlining(MemoryLeakMixin, InliningBase):
 
+    def setUp(self):
+        # only tests for operators need DummyType
+        if 'operator' in self._testMethodName:
+            self.make_dummy_type()
+        super(TestOverloadInlining, self).setUp()
+
     def test_basic_inline_never(self):
         def foo():
             pass
@@ -526,6 +557,139 @@ class TestOverloadInlining(MemoryLeakMixin, InliningBase):
 
         self.check(impl, inline_expect={'foo': True})
 
+    def test_inline_operators_unary(self):
+
+        def impl_inline(x):
+            return -x
+
+        def impl_noinline(x):
+            return +x
+
+        dummy_unary_impl = lambda x: True
+        setattr(self.Dummy, '__neg__', dummy_unary_impl)
+        setattr(self.Dummy, '__pos__', dummy_unary_impl)
+
+        @overload(operator.neg, inline='always')
+        def overload_dummy_neg(x):
+            if isinstance(x, self.DummyType):
+                return dummy_unary_impl
+
+        @overload(operator.pos, inline='never')
+        def overload_dummy_pos(x):
+            if isinstance(x, self.DummyType):
+                return dummy_unary_impl
+
+        self.check(impl_inline, self.Dummy(), inline_expect={'neg': True})
+        self.check(impl_noinline, self.Dummy(), inline_expect={'pos': False})
+
+    def test_inline_operators_binop(self):
+
+        def impl_inline(x):
+            return x == 1
+
+        def impl_noinline(x):
+            return x != 1
+
+        dummy_binop_impl = lambda a, b: True
+        setattr(self.Dummy, '__eq__', dummy_binop_impl)
+        setattr(self.Dummy, '__ne__', dummy_binop_impl)
+
+        @overload(operator.eq, inline='always')
+        def overload_dummy_eq(a, b):
+            if isinstance(a, self.DummyType):
+                return dummy_binop_impl
+
+        @overload(operator.ne, inline='never')
+        def overload_dummy_ne(a, b):
+            if isinstance(a, self.DummyType):
+                return dummy_binop_impl
+
+        self.check(impl_inline, self.Dummy(), inline_expect={'eq': True})
+        self.check(impl_noinline, self.Dummy(), inline_expect={'ne': False})
+
+    def test_inline_operators_inplace_binop(self):
+
+        def impl_inline(x):
+            x += 1
+
+        def impl_noinline(x):
+            x -= 1
+
+        dummy_inplace_binop_impl = lambda a, b: True
+        setattr(self.Dummy, '__iadd__', dummy_inplace_binop_impl)
+        setattr(self.Dummy, '__isub__', dummy_inplace_binop_impl)
+
+        @overload(operator.iadd, inline='always')
+        def overload_dummy_iadd(a, b):
+            if isinstance(a, self.DummyType):
+                return dummy_inplace_binop_impl
+
+        @overload(operator.isub, inline='never')
+        def overload_dummy_isub(a, b):
+            if isinstance(a, self.DummyType):
+                return dummy_inplace_binop_impl
+
+        # DummyType is not mutable, so lowering 'inplace_binop' Expr
+        # re-uses (requires) copying function definition
+        @overload(operator.add, inline='always')
+        def overload_dummy_add(a, b):
+            if isinstance(a, self.DummyType):
+                return dummy_inplace_binop_impl
+
+        @overload(operator.sub, inline='never')
+        def overload_dummy_sub(a, b):
+            if isinstance(a, self.DummyType):
+                return dummy_inplace_binop_impl
+
+        self.check(impl_inline, self.Dummy(), inline_expect={'iadd': True})
+        self.check(impl_noinline, self.Dummy(), inline_expect={'isub': False})
+
+    def test_inline_always_operators_getitem(self):
+
+        def impl(x, idx):
+            return x[idx]
+
+        def impl_static_getitem(x):
+            return x[1]
+
+        dummy_getitem_impl = lambda obj, idx: None
+        setattr(self.Dummy, '__getitem__', dummy_getitem_impl)
+
+        @overload(operator.getitem, inline='always')
+        def overload_dummy_getitem(obj, idx):
+            if isinstance(obj, self.DummyType):
+                return dummy_getitem_impl
+
+        # note getitem and static_getitem Exprs refer to operator.getitem
+        # hence they are checked using the same expected key
+        self.check(impl, self.Dummy(), 1, inline_expect={'getitem': True})
+        self.check(impl_static_getitem,
+                   self.Dummy(),
+                   inline_expect={'getitem': True})
+
+    def test_inline_never_operators_getitem(self):
+
+        def impl(x, idx):
+            return x[idx]
+
+        def impl_static_getitem(x):
+            return x[1]
+
+        dummy_getitem_impl = lambda obj, idx: None
+        setattr(self.Dummy, '__getitem__', dummy_getitem_impl)
+
+        @overload(operator.getitem, inline='never')
+        def overload_dummy_getitem(obj, idx):
+            if isinstance(obj, self.DummyType):
+                return dummy_getitem_impl
+
+        # noth getitem and static_getitem Exprs refer to opertor.getitem
+        # hence they are checked using the same expect key
+        self.check(impl, self.Dummy(), 1, inline_expect={'getitem': False})
+        self.check(impl_static_getitem,
+                   self.Dummy(),
+                   inline_expect={'getitem': False})
+
     def test_inline_stararg_error(self):
         def foo(a, *b):
             return a + b[0]
@@ -861,26 +1025,8 @@ class TestOverloadInlining(MemoryLeakMixin, InliningBase):
 
 class TestOverloadMethsAttrsInlining(InliningBase):
     def setUp(self):
-        # Use test_id to makesure no collision is possible.
-        test_id = self.id()
-        DummyType = type('DummyTypeFor{}'.format(test_id), (types.Opaque,), {})
-
-        dummy_type = DummyType("my_dummy")
-        register_model(DummyType)(OpaqueModel)
-
-        class Dummy(object):
-            pass
-
-        @typeof_impl.register(Dummy)
-        def typeof_Dummy(val, c):
-            return dummy_type
-
-        @unbox(DummyType)
-        def unbox_index(typ, obj, c):
-            return NativeValue(c.context.get_dummy_value())
-
-        self.Dummy = Dummy
-        self.DummyType = DummyType
+        self.make_dummy_type()
+        super(TestOverloadMethsAttrsInlining, self).setUp()
 
     def check_method(self, test_impl, args, expected, block_count,
                      expects_inlined=True):