Numba
lastmod 2024-09-11created 2019-04-0421 edits
Numba is an open source JIT compiler that translates a subset of Python and NumPy code into fast machine code.
Install
jit
| Option | Description | Caveats | Example |
| nopython | Compiles instead of running Python through the interpreter | Using libraries that exist only in Python (pandas, etc.) raises an Error | |
| nogil | Does not use the GIL (Global Interpreter Lock). | May cause thread-safety issues; memory usage may increase | |
| cache | Caches compile results on disk so they can be reused later | Cache files may become meaningless when the Numba version or the code changes. | |
| parallel | Used for parallel processing; parallelizes loops and array operations to make use of CPU cores | May add memory usage and overhead; parallelization does not always guarantee better performance. | |
njit
@njit == @jit(nopython=True).
Parallel loop optimization
See #prange.
Caching
Overrides the location of the cache directory. If defined, this must be a valid directory path.
If not defined, Numba picks the cache directory in the following order:
- In-tree cache. Puts the cache next to the corresponding source file under a
__pycache__directory, following how.pycfiles are stored. - User-wide cache. Puts the cache in the user's application directory using
appdirs.user_cache_dirfrom the Appdirs package. - IPython cache. Puts the cache in an IPython-specific application directory. Stores are made under
numba_cachein the directory returned byIPython.paths.get_ipython_cache_dir().
See also the docs on cache sharing and on cache clearing.
NumPy support
-
numpy.nditer - Only the first Argument may be passed. If you need a nested loop, use #prange.
Automatic parallelization with @jit
prange
from numba import njit, prange
import numpy as np
@njit(parallel=True)
def test(x):
n = x.shape[0]
a = np.sin(x)
b = np.cos(a * a)
acc = 0
for i in prange(n - 2):
for j in prange(n - 1):
acc += b[i] + b[j + 1]
return acc
test(np.arange(10))
test.parallel_diagnostics(level=4)
Multithreading without the GIL
Releasing the GIL lets multithreading deliver its full performance on CPU-bound work.
#!/usr/bin/env python
from __future__ import print_function, division, absolute_import
import math
import threading
from timeit import repeat
import numpy as np
from numba import jit
nthreads = 4
size = 10**6
def func_np(a, b):
"""
Control function using Numpy.
"""
return np.exp(2.1 * a + 3.2 * b)
@jit('void(double[:], double[:], double[:])', nopython=True, nogil=True)
def inner_func_nb(result, a, b):
"""
Function under test.
"""
for i in range(len(result)):
result[i] = math.exp(2.1 * a[i] + 3.2 * b[i])
def timefunc(correct, s, func, *args, **kwargs):
"""
Benchmark *func* and print out its runtime.
"""
print(s.ljust(20), end=" ")
# Make sure the function is compiled before we start the benchmark
res = func(*args, **kwargs)
if correct is not None:
assert np.allclose(res, correct), (res, correct)
# time it
print('{:>5.0f} ms'.format(min(repeat(lambda: func(*args, **kwargs),
number=5, repeat=2)) * 1000))
return res
def make_singlethread(inner_func):
"""
Run the given function inside a single thread.
"""
def func(*args):
length = len(args[0])
result = np.empty(length, dtype=np.float64)
inner_func(result, *args)
return result
return func
def make_multithread(inner_func, numthreads):
"""
Run the given function inside *numthreads* threads, splitting its
arguments into equal-sized chunks.
"""
def func_mt(*args):
length = len(args[0])
result = np.empty(length, dtype=np.float64)
args = (result,) + args
chunklen = (length + numthreads - 1) // numthreads
# Create argument tuples for each input chunk
chunks = [[arg[i * chunklen:(i + 1) * chunklen] for arg in args]
for i in range(numthreads)]
# Spawn one thread per chunk
threads = [threading.Thread(target=inner_func, args=chunk)
for chunk in chunks]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
return result
return func_mt
func_nb = make_singlethread(inner_func_nb)
func_nb_mt = make_multithread(inner_func_nb, nthreads)
a = np.random.rand(size)
b = np.random.rand(size)
correct = timefunc(None, "numpy (1 thread)", func_np, a, b)
timefunc(correct, "numba (1 thread)", func_nb, a, b)
timefunc(correct, "numba (%d threads)" % nthreads, func_nb_mt, a, b)
How to use CUDA devices
- Numba for CUDA GPUs — Numba 0.50.1 documentation
- Computing the Jaccard score with Numba's GPU processing | by Chiwon Son | Standigm | Medium
- ok97465 | Example CUDA program using Numba
Freeing CUDA memory
Sharing CUDA Memory
Inter-process sharing is implemented using the legacy CUDA IPC API (functions whose names start with cuIpc) and is supported only on Linux.
Exporting a device array to another process
A device array can be shared with another process on the same machine using the CUDA IPC API.
To do so, use the .get_ipc_handle() method on the device array to get an #IpcArrayHandle object, which can be transferred to another process.
-
DeviceNDArray.get_ipc_handle() - Returns an #IpcArrayHandle object that is safe to serialize and transfer to another process to share the local allocation.
IpcArrayHandle
The class signature is class numba.cuda.cudadrv.devicearray.IpcArrayHandle(ipc_handle, array_desc).
An IPC array handle can be serialized and passed to another process on the same machine, and is used to share a GPU allocation.
In the destination process (the receiving side), use the .open() method to create a new DeviceNDArray object that shares the allocation of the original process.
To release the resource, call the .close() method. After that, the destination process (the receiving side) can no longer use the shared array object. At this point, the weak reference to the resource is removed.
This object supports the context manager (with syntax), which calls the .open() and .close() methods automatically.
with the_ipc_array_handle as ipc_array:
# use ipc_array here as a normal gpu array object
some_code(ipc_array)
# ipc_array is dead at this point
Importing IPC memory from another process
The cuda.open_ipc_array(shape, dtype, strides=None, offset=0) function is used to open an IPC handle from another process as a device array.
It is a context manager that opens an IPC handle (CUipcMemHandle) represented as a sequence of bytes (e.g. bytes, a tuple of int) and represents it as an array with shape, dtype, and strides.
strides can be omitted, in which case the array is assumed to be a 1D C1 contiguous array.
Creates a device array.
When the context manager exits, the IPC handle is closed automatically.
Example
Host-side (sharing) process:
import numpy as np
from numba import cuda
from time import sleep
from pickle import dumps
# CUDA 장치에서 배열을 생성합니다.
data = np.arange(10, dtype=np.float32)
d_data = cuda.to_device(data)
# IPC 핸들을 생성합니다.
ipc_handle = d_data.get_ipc_handle()
# IPC 핸들을 직렬화하여 다른 프로세스에 전달할 수 있습니다.
serialized_handle = dumps(ipc_handle)
# serialized_handle을 파일이나 다른 방법으로 다른 프로세스에 전달합니다.
with open('ipc_handle.bin', 'wb') as f:
f.write(serialized_handle)
while True:
sleep(1)
print(".", end="", flush=True)
Guest-side (receiving) process:
from numba import cuda
from pickle import loads
# serialized_handle을 파일이나 다른 방법으로 수신합니다.
with open('ipc_handle.bin', 'rb') as f:
serialized_handle = f.read()
# IPC 핸들을 역직렬화합니다.
ipc_handle = loads(serialized_handle)
# 핸들을 통해 배열에 접근합니다.
d_data = ipc_handle.open()
data = d_data.copy_to_host()
print(data)
See also
Favorite site
- Numba web site
- Numba User Manual
- numba - Performance up!
- Numba vs Cython: How to Choose – Stephan Hoyer
- Boosting Python with Cython and Numba | by Rodrigo Arancibia | Rodrigo Arancibia tech site | Medium