How to use numba in Colaboratory - numba

Question

Anybody tried to use numba in google collaboratory? I just can not figure out how to set it up in this environment.
At the moment, I'm stuck with the error library nvvm not found.

Answer 1

Copy this code into cell. It works for me.
!apt-get install nvidia-cuda-toolkit
!pip3 install numba
import os
os.environ['NUMBAPRO_LIBDEVICE'] = "/usr/lib/nvidia-cuda-toolkit/libdevice"
os.environ['NUMBAPRO_NVVM'] = "/usr/lib/x86_64-linux-gnu/libnvvm.so"
from numba import cuda
import numpy as np
import time
#cuda.jit
def hello(data):
data[cuda.blockIdx.x, cuda.threadIdx.x] = cuda.blockIdx.x
numBlocks = 5
threadsPerBlock = 10
data = np.ones((numBlocks, threadsPerBlock), dtype=np.uint8)
hello[numBlocks, threadsPerBlock](data)
print(data)

Answer 2

I didn't have to install the packages #Algis suggested, but the paths to the drivers were different. So I had to do the following.
First determine the correct paths for the drivers
!find / -iname 'libdevice'
!find / -iname 'libnvvm.so'
# Output:
# /usr/local/cuda-9.2/nvvm/lib64/libnvvm.so
# /usr/local/cuda-9.2/nvvm/libdevice
Then set the paths as #Algis described
import os
os.environ['NUMBAPRO_LIBDEVICE'] = "/usr/local/cuda-9.2/nvvm/libdevice"
os.environ['NUMBAPRO_NVVM'] = "/usr/local/cuda-9.2/nvvm/lib64/libnvvm.so"

Answer 3

You can do #Stan's work in one simple sweep if you have this block at the beginning of your colab notebook (this also automatically updates as CUDA gets updated)
import os
dev_lib_path = !find / -iname 'libdevice'
nvvm_lib_path = !find / -iname 'libnvvm.so'
assert len(dev_lib_path)>0, "Device Lib Missing"
assert len(nvvm_lib_path)>0, "NVVM Missing"
os.environ['NUMBAPRO_LIBDEVICE'] = dev_lib_path[0]
os.environ['NUMBAPRO_NVVM'] = nvvm_lib_path[0]