Repository: numba/numba

Issue:
Allow overloading of cuda device function
<!---
@huboard:{"order":5.75}
-->



Discussion:
I've just run into this issue (I think). Using the code:

``` python
from numba import cuda
import numpy as np

@cuda.jit(device=True)
def inner(arg):
    return arg + 1

@cuda.jit
def outer(argin, argout):
    argout[0] = inner(argin[0]) + inner(2)

a = np.zeros(1)
b = np.zeros(1)

stream = cuda.stream()
d_a = cuda.to_device(a, stream)
d_b = cuda.to_device(b, stream)

outer[1, 1, stream](d_a, d_b)

d_b.to_host(stream)

print(d_b[0])
```

I get this traceback:

```
$ python dfnd.py 
Traceback (most recent call last):
  File "/home/gmarkall/work/numba/numba/typeinfer.py", line 112, in propagate
    constrain(context, typevars)
  File "/home/gmarkall/work/numba/numba/typeinfer.py", line 229, in __call__
    self.resolve(context, typevars, fnty)
  File "/home/gmarkall/work/numba/numba/typeinfer.py", line 239, in resolve
    sig = context.resolve_function_type(fnty, args, ())
  File "/home/gmarkall/work/numba/numba/typing/context.py", line 53, in resolve_function_type
    return func.template(self).apply(args, kws)
  File "/home/gmarkall/work/numba/numba/typing/templates.py", line 198, in apply
    sig = generic(args, kws)
  File "/home/gmarkall/work/numba/numba/cuda/compiler.py", line 101, in generic
    return dft.compile(args)
  File "/home/gmarkall/work/numba/numba/cuda/compiler.py", line 82, in compile
    cres.target_context.insert_user_function(self, cres.fndesc, libs)
  File "/home/gmarkall/work/numba/numba/targets/base.py", line 170, in insert_user_function
    self.users[func] = type(name, baseclses, glbls)
  File "/home/gmarkall/work/numba/numba/utils.py", line 135, in __setitem__
    raise KeyError('Key %s already present in UniqueDict' % key)
KeyError: 'Key <numba.cuda.compiler.DeviceFunctionTemplate object at 0x7fea784fb7f0> already present in UniqueDict'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "dfnd.py", line 19, in <module>
    outer[1, 1, stream](d_a, d_b)
  File "/home/gmarkall/work/numba/numba/cuda/compiler.py", line 434, in __call__
    kernel = self.specialize(*args)
  File "/home/gmarkall/work/numba/numba/cuda/compiler.py", line 446, in specialize
    **self.targetoptions)
  File "/home/gmarkall/work/numba/numba/cuda/compiler.py", line 48, in compile_kernel
    cres = compile_cuda(pyfunc, types.void, args, debug=debug, inline=inline)
  File "/home/gmarkall/work/numba/numba/cuda/compiler.py", line 39, in compile_cuda
    locals={})
  File "/home/gmarkall/work/numba/numba/compiler.py", line 560, in compile_extra
    return pipeline.compile_extra(func)
  File "/home/gmarkall/work/numba/numba/compiler.py", line 266, in compile_extra
    return self.compile_bytecode(res.result, func_attr=self.func_attr)
  File "/home/gmarkall/work/numba/numba/compiler.py", line 278, in compile_bytecode
    return self._compile_bytecode()
  File "/home/gmarkall/work/numba/numba/compiler.py", line 509, in _compile_bytecode
    return self._run_pipeline(pipelines)
  File "/home/gmarkall/work/numba/numba/compiler.py", line 534, in _run_pipeline
    raise _raise_error(msg, res.exception)
  File "/home/gmarkall/work/numba/numba/compiler.py", line 166, in __call__
    res = self.func()
  File "/home/gmarkall/work/numba/numba/compiler.py", line 360, in stage_nopython_frontend
    self.locals)
  File "/home/gmarkall/work/numba/numba/compiler.py", line 678, in type_inference_stage
    infer.propagate()
  File "/home/gmarkall/work/numba/numba/typeinfer.py", line 386, in propagate
    self.constrains.propagate(self.context, self.typevars)
  File "/home/gmarkall/work/numba/numba/typeinfer.py", line 118, in propagate
    loc=constrain.loc)
numba.typeinfer.TypingError: Failed at nopython frontend
Internal error at <numba.typeinfer.CallConstrain object at 0x7fea607a8048>:
'Key <numba.cuda.compiler.DeviceFunctionTemplate object at 0x7fea784fb7f0> already present in UniqueDict'
File "dfnd.py", line 10
```

(note that the exception message is one I added whilst looking into what my problem was - Numba master has `assert key not in self` in `UniqueDict.__setitem__`, which gives a blank message.

The error disappears if I remove the call to `inner(2)`, which seems to suggest that the problem would go away if overloading device functions were supported.

Does specifying the type signatures for `inner()` explicitly workaround this issue?

It does for this example, if I specify the type of `arg` to be `float64`, and I think that's the case because `inner(2)` gets 2 cast to float64 - however, in the general case it's still a problem because I might want to call the same device function with two different structured types. 

Unless there's someway to specify multiple argument types for a device function that I've missed?

Using:

``` python
@cuda.jit(argtypes=['int32(int32,)', 'float64(float64,)'], device=True)
```

fails ("Too many argument types") as does

``` python
@cuda.jit(argtypes=['int32(int32,)', 'float64(float64,)'], device=True)
```

("TypeError: 'DeviceFunctionTemplate' object is not callable")


Resolving patch:
diff --git a/numba/cuda/compiler.py b/numba/cuda/compiler.py
--- a/numba/cuda/compiler.py
+++ b/numba/cuda/compiler.py
@@ -77,11 +77,20 @@ def compile(self, args):
         if args not in self._compileinfos:
             cres = compile_cuda(self.py_func, None, args, debug=self.debug,
                                 inline=self.inline)
+            first_definition = not self._compileinfos
             self._compileinfos[args] = cres
             libs = [cres.library]
-            cres.target_context.insert_user_function(self, cres.fndesc, libs)
+
+            if first_definition:
+                # First definition
+                cres.target_context.insert_user_function(self, cres.fndesc,
+                                                         libs)
+            else:
+                cres.target_context.add_user_function(self, cres.fndesc, libs)
+
         else:
             cres = self._compileinfos[args]
+
         return cres.signature
 
 
diff --git a/numba/targets/base.py b/numba/targets/base.py
--- a/numba/targets/base.py
+++ b/numba/targets/base.py
@@ -169,6 +169,13 @@ def insert_user_function(self, func, fndesc, libs=()):
         name = "CallTemplate(%s)" % fndesc.mangled_name
         self.users[func] = type(name, baseclses, glbls)
 
+    def add_user_function(self, func, fndesc, libs=()):
+        if func not in self.users:
+            msg = "{func} is not a registered user function"
+            raise KeyError(msg.format(func=func))
+        imp = user_function(func, fndesc, libs)
+        self.defns[func].append(imp)
+
     def insert_class(self, cls, attrs):
         clsty = types.Object(cls)
         for name, vtype in utils.iteritems(attrs):


Test patch:
diff --git a/numba/cuda/tests/cudapy/test_cuda_autojit.py b/numba/cuda/tests/cudapy/test_cuda_autojit.py
--- a/numba/cuda/tests/cudapy/test_cuda_autojit.py
+++ b/numba/cuda/tests/cudapy/test_cuda_autojit.py
@@ -52,6 +52,28 @@ def driver(A, B):
 
         np.testing.assert_allclose(Acopy + Acopy + Bcopy + Bcopy + 1, B)
 
+    def test_device_auto_jit_2(self):
+        @cuda.jit(device=True)
+        def inner(arg):
+            return arg + 1
+
+        @cuda.jit
+        def outer(argin, argout):
+            argout[0] = inner(argin[0]) + inner(2)
+
+        a = np.zeros(1)
+        b = np.zeros(1)
+
+        stream = cuda.stream()
+        d_a = cuda.to_device(a, stream)
+        d_b = cuda.to_device(b, stream)
+
+        outer[1, 1, stream](d_a, d_b)
+
+        d_b.copy_to_host(b, stream)
+
+        self.assertEqual(b[0], (a[0] + 1) + (2 + 1))
+
 
 if __name__ == '__main__':
     unittest.main()