How to run numba.jit decorated function on GPU?

Question

I am trying to get into numba gpu processing. I have this MWE:

import numpy as np
import numba

@numba.njit
def function():
    ar = np.zeros((3, 3))
    for i in range(3):
        ar[i] = (1, 2, 3)
    return ar

ar = function()
print(ar)

Output:

[[1. 2. 3.]
 [1. 2. 3.]
 [1. 2. 3.]]

Now I want to run it on my gpu. I tried using following decorators:

@numba.njit(target='cuda')
@numba.njit(target='gpu')
@numba.cuda.jit

none of which work. Here are the error messages of above decorators:

Traceback (most recent call last):
  File "/home/amu/Desktop/RL_framework/help_functions/test.py", line 4, in <module>
    @numba.jit(target='cuda')
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/decorators.py", line 171, in jit
    targetoptions=options, **dispatcher_args)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/decorators.py", line 179, in _jit
    dispatcher = registry.dispatcher_registry[target]
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/registry.py", line 96, in __getitem__
    return super(TargetRegistry, self).__getitem__(item)
KeyError: 'cuda'

Traceback (most recent call last):
  File "/home/amu/Desktop/RL_framework/help_functions/test.py", line 4, in <module>
    @numba.njit(target='gpu')
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/decorators.py", line 236, in njit
    return jit(*args, **kws)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/decorators.py", line 171, in jit
    targetoptions=options, **dispatcher_args)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/decorators.py", line 179, in _jit
    dispatcher = registry.dispatcher_registry[target]
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/registry.py", line 96, in __getitem__
    return super(TargetRegistry, self).__getitem__(item)
KeyError: 'gpu'

Traceback (most recent call last):
  File "/home/amu/Desktop/RL_framework/help_functions/test.py", line 4, in <module>
    @numba.cuda.jit()
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/__init__.py", line 140, in __getattr__
    ) from None
AttributeError: module 'numba' has no attribute 'cuda'

numba.cuda is not even recognized as module function. I have numba 49.1 and cudatoolkit 9.0 installed.

Do I have to change the function for this to work? I have a huge numba.njit function that I need to run on a gpu.

I was trying to learn more about numba's cuda implementation from: http://numba.pydata.org/numba-doc/0.16.0/modules/numba.cuda.html

Thank you in advance.

Edit:

As @talonmies proposed I imported cuda explicitly from the numba module and outsourced the array creation:

import numpy as np
import numba
from numba import cuda

@numba.njit(target='cuda')
def function(ar=None):
    for i in range(3):
        ar[i] = (1, 2, 3)
    return ar

ar = np.zeros((3, 3))
ar_result = function(ar=ar)
print(ar_result)

Output:

Traceback (most recent call last):
  File "/home/amu/Desktop/RL_framework/help_functions/test.py", line 12, in <module>
    ar_result = function(ar=ar)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/cuda/dispatcher.py", line 40, in __call__
    return self.compiled(*args, **kws)
TypeError: __call__() got an unexpected keyword argument 'ar'

This error occurs with every of the aforementioned decorators except @numba.njit.

Edit_2:

When I try to run:

import numpy as np
import numba
from numba import cuda

@numba.jit(target='cuda')
def function(ar):
    for i in range(3):
        ar[i] = (1,2,3)

ar = np.zeros((3, 3))
function(ar)
print(ar)

the output is:

Traceback (most recent call last):
  File "/home/amu/Desktop/RL_framework/help_functions/test.py", line 11, in <module>
    function(ar)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/cuda/dispatcher.py", line 40, in __call__
    return self.compiled(*args, **kws)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/cuda/compiler.py", line 758, in __call__
    kernel = self.specialize(*args)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/cuda/compiler.py", line 769, in specialize
    kernel = self.compile(argtypes)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/cuda/compiler.py", line 785, in compile
    **self.targetoptions)
  File "/home/amu/anaconda3/lib/python3.7/site-packages/numba/core/compiler_lock.py", line 32, in _acquire_compile_lock
    return func(*args, **kwargs)
TypeError: compile_kernel() got an unexpected keyword argument 'boundscheck'

Answer 1

Three points:

1. You have to explicitly import the cuda module from numba to use it (this isn't specific to numba, all python libraries work like this)
2. The nopython mode (njit) doesn't support the CUDA target
3. Array creation, return values, keyword arguments are not supported in Numba for CUDA code

I can fix all that like this:

...: import numpy as np 
...: import numba 
...: from numba import cuda 
...:  
...: @numba.jit(target='cuda') 
...: def function(ar): 
...:     for i in range(3): 
...:         ar[i] = (1,2,3)   
...:  
...: ar = np.zeros((3, 3)) 
...: function(ar) 
...: print(ar) 

[[1. 2. 3.]
 [1. 2. 3.]
 [1. 2. 3.]]