Skip to content
Search

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

pip install numba

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

@jit(nopython=True)

nogil

Does not use the GIL (Global Interpreter Lock).

May cause thread-safety issues; memory usage may increase

@jit(nogil=True)

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.

@jit(cache=True)

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.

@jit(parallel=True)

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:

  1. In-tree cache. Puts the cache next to the corresponding source file under a __pycache__ directory, following how .pyc files are stored.
  2. User-wide cache. Puts the cache in the user's application directory using appdirs.user_cache_dir from the Appdirs package.
  3. IPython cache. Puts the cache in an IPython-specific application directory. Stores are made under numba_cache in the directory returned by IPython.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

Freeing CUDA memory

from numba import cuda

# 현재 사용 중인 디바이스 해제
cuda.select_device(0)
cuda.close()

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

References


  1. A one-dimensional C-style memory layout where array.flags['C_CONTIGUOUS'] is True. ↩

author Zer0 sitemap