Repository: numba/numba

Issue:
Support prefer_literal option for intrinsic decorator
- [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).
- [x] I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'.

Hi, 

I noticed that `prefer_literal` option that was added recently for overload decorators is not supported for `intrinsic`.
It seems that `make_intrinsic_template` does not have prefer_literal argument (unlike for example `make_overload_template`) and moreover objects of `_Intrinsic` class do not store any options at all. At the same time intrinsic templates are created from base `FunctionTemplate`, which has `prefer_literal` hardcoded to False:
https://github.com/numba/numba/blob/4f985a5d90c4202c5f77bbd05fcff894200752de/numba/core/typing/templates.py#L261

This results in multiple calls to `get_call_type` with non-literal arguments during type-inference and may cause subtle issues if intrinsic code does not have correct input type checks. For instance the output of the following script

```
import numba
from numba import types
from numba.extending import overload, intrinsic
from numba.core.typing.templates import signature
import numpy as np

@intrinsic
# @intrinsic(prefer_literal=True)
def tuple_take(typingctx, data, idxs):

    print("DEBUG: tuple_take: data", data)
    print("DEBUG: tuple_take: idxs", idxs)
    if not isinstance(data, types.Tuple):
        return None
    if not (isinstance(idxs, types.Tuple)
            and all(isinstance(x, types.IntegerLiteral) for x in idxs)):
        return None

    old_ncols = len(data)
    unlit_keep_idxs = [a.literal_value for a in idxs]
    keep_col_idxs, new_typs = zip(*((i, t) for i, t in enumerate(data) if i in unlit_keep_idxs))
    def codegen(context, builder, signature, args):
        old_data = args[0]
        new_data = [
            builder.extract_value(
                old_data,
                types.intp(i)
            ) for i in range(old_ncols) if i in keep_col_idxs
        ]

        res = context.make_tuple(builder, ret_typ, new_data)
        return res

    ret_typ = types.Tuple(new_typs)
    sig = signature(ret_typ, data, idxs)
    return sig, codegen

def foo(data, idxs):
    pass

@overload(foo, prefer_literal=True)
def foo_ovld(data, idxs):
    print(f"DEBUG: foo: data={data}")
    print(f"DEBUG: foo: idxs={idxs}")

    def foo_impl(data, idxs):
        return tuple_take(data, idxs)
    return foo_impl

@numba.njit
def test_func():
    data = ('a', [2, 3], np.nan, 600)
    keep_idxs = (1, 3)
    return foo(data, keep_idxs)

res = test_func()
print(f"DEBUG: res={res}")
```
shows that we entered `tuple_take` code twice with non-literal arguments
> DEBUG: foo: data=Tuple(Literal[str](a), list(int64)<iv=[2, 3]>, float64, Literal[int](600))
> DEBUG: foo: idxs=Tuple(Literal[int](1), Literal[int](3))
> DEBUG: tuple_take: data Tuple(unicode_type, list(int64)<iv=None>, float64, int64)
> DEBUG: tuple_take: idxs UniTuple(int64 x 2)
> DEBUG: tuple_take: data Tuple(Literal[str](a), list(int64)<iv=[2, 3]>, float64, Literal[int](600))
> DEBUG: tuple_take: idxs Tuple(Literal[int](1), Literal[int](3))
> DEBUG: tuple_take: data Tuple(unicode_type, list(int64)<iv=None>, float64, int64)
> DEBUG: tuple_take: idxs UniTuple(int64 x 2)
> DEBUG: res=([2, 3], 600)

and in the ideal case it can be this way:
> DEBUG: foo: data=Tuple(Literal[str](a), list(int64)<iv=[2, 3]>, float64, Literal[int](600))
> DEBUG: foo: idxs=Tuple(Literal[int](1), Literal[int](3))
> DEBUG: tuple_take: data Tuple(Literal[str](a), list(int64)<iv=[2, 3]>, float64, Literal[int](600))
> DEBUG: tuple_take: idxs Tuple(Literal[int](1), Literal[int](3))
> DEBUG: res=([2, 3], 600)

Can this be changed in the future and prefer_literal added to intrinsic too?

BR,
Alexey.


Discussion:
The requested changes makes sense and it should be very easy to accomplish.
Tips for implementation:

- Plumbing needed to set `prefer_literal` at https://github.com/numba/numba/blob/4f985a5d90c4202c5f77bbd05fcff894200752de/numba/core/typing/templates.py#L878-L879 from `@intrinsic(prefer_literal=True)`
- other changes needed at https://github.com/numba/numba/blob/4f985a5d90c4202c5f77bbd05fcff894200752de/numba/core/extending.py#L280-L282, and https://github.com/numba/numba/blob/4f985a5d90c4202c5f77bbd05fcff894200752de/numba/core/extending.py#L308
- a test can be created from the reproducer code in this issue.
I would love to contribute back to numba, can I work on this one? @sklam  

Resolving patch:
diff --git a/numba/core/extending.py b/numba/core/extending.py
--- a/numba/core/extending.py
+++ b/numba/core/extending.py
@@ -322,10 +322,11 @@ class _Intrinsic(ReduceMixin):
 
     __uuid = None
 
-    def __init__(self, name, defn, **kwargs):
+    def __init__(self, name, defn, prefer_literal=False, **kwargs):
         self._ctor_kwargs = kwargs
         self._name = name
         self._defn = defn
+        self._prefer_literal = prefer_literal
         functools.update_wrapper(self, defn)
 
     @property
@@ -354,7 +355,8 @@ def _register(self):
                                                  infer_global)
 
         template = make_intrinsic_template(self, self._defn, self._name,
-                                           self._ctor_kwargs)
+                                           prefer_literal=self._prefer_literal,
+                                           kwargs=self._ctor_kwargs)
         infer(template)
         infer_global(self, types.Function(template))
 
diff --git a/numba/core/typing/templates.py b/numba/core/typing/templates.py
--- a/numba/core/typing/templates.py
+++ b/numba/core/typing/templates.py
@@ -9,7 +9,7 @@
 import os.path
 from collections import namedtuple
 from collections.abc import Sequence
-from types import MethodType, FunctionType
+from types import MethodType, FunctionType, MappingProxyType
 
 import numba
 from numba.core import types, utils, targetconfig
@@ -999,15 +999,18 @@ def get_template_info(self):
         return info
 
 
-def make_intrinsic_template(handle, defn, name, kwargs):
+def make_intrinsic_template(handle, defn, name, *, prefer_literal=False,
+                            kwargs=None):
     """
     Make a template class for a intrinsic handle *handle* defined by the
     function *defn*.  The *name* is used for naming the new template class.
     """
+    kwargs = MappingProxyType({} if kwargs is None else kwargs)
     base = _IntrinsicTemplate
     name = "_IntrinsicTemplate_%s" % (name)
     dct = dict(key=handle, _definition_func=staticmethod(defn),
-               _impl_cache={}, _overload_cache={}, metadata=kwargs)
+               _impl_cache={}, _overload_cache={},
+               prefer_literal=prefer_literal, metadata=kwargs)
     return type(base)(name, (base,), dct)
 
 


Test patch:
diff --git a/numba/tests/test_extending.py b/numba/tests/test_extending.py
--- a/numba/tests/test_extending.py
+++ b/numba/tests/test_extending.py
@@ -8,6 +8,7 @@
 import re
 
 import numpy as np
+from llvmlite import ir
 
 from numba import njit, jit, vectorize, guvectorize, objmode
 from numba.core import types, errors, typing, compiler, cgutils
@@ -2067,5 +2068,50 @@ def check_non_lit(dummy, x):
         self.assertEqual(c, 300)
 
 
+class TestIntrinsicPreferLiteral(TestCase):
+    def test_intrinsic(self):
+        def intrin(context, x):
+            # This intrinsic will return 0xcafe if `x` is a literal `1`.
+            sig = signature(types.intp, x)
+            if isinstance(x, types.IntegerLiteral):
+                # With prefer_literal=False, this branch will not be reached
+                if x.literal_value == 1:
+                    def codegen(context, builder, signature, args):
+                        atype = signature.args[0]
+                        llrtype = context.get_value_type(atype)
+                        return ir.Constant(llrtype, 0xcafe)
+                    return sig, codegen
+                else:
+                    raise errors.TypingError('literal value')
+            else:
+                def codegen(context, builder, signature, args):
+                    atype = signature.return_type
+                    llrtype = context.get_value_type(atype)
+                    int_100 = ir.Constant(llrtype, 100)
+                    return builder.mul(args[0], int_100)
+                return sig, codegen
+
+        prefer_lit = intrinsic(prefer_literal=True)(intrin)
+        non_lit = intrinsic(prefer_literal=False)(intrin)
+
+        @njit
+        def check_prefer_lit(x):
+            return prefer_lit(1), prefer_lit(2), prefer_lit(x)
+
+        a, b, c = check_prefer_lit(3)
+        self.assertEqual(a, 0xcafe)
+        self.assertEqual(b, 200)
+        self.assertEqual(c, 300)
+
+        @njit
+        def check_non_lit(x):
+            return non_lit(1), non_lit(2), non_lit(x)
+
+        a, b, c = check_non_lit(3)
+        self.assertEqual(a, 100)
+        self.assertEqual(b, 200)
+        self.assertEqual(c, 300)
+
+
 if __name__ == "__main__":
     unittest.main()