Numba: Array-oriented Python Compiler for NumPy

Numba: An array-oriented
Python compiler
SIAM Conference on Computational
Science and Engineering
Travis E. Oliphant

February 25, 2012

Big Picture
Empower domain experts with
high-level tools that exploit modern
hard-ware

?

Array Oriented Computing

Software Stack Future?
Plateaus of Code re-use + DSLs
SQL R
TDPL Matlab

Python

OBJC Julia C
FORTRAN C++

LLVM

Motivation
• Python is great for rapid development and
high-level thinking-in-code
• It is slow for interior loops because lack of
type information leads to a lot of indirection
and “extra” code.

Motivation
• NumPy users have had a lot of type
information for a long time --- but only
currently have one-size ﬁts all pre-
compiled, vectorized loops.
• Idea is to use this type information to
allow compilation of arbitrary
expressions involving NumPy arrays

Current approaches
• Cython
• Weave
• Write fast-code in C/C++/Fortran and “wrap” with
- SWIG or Cython
- f2py or fwrap
- hand-written

Numba philosophy : don’t wrap or rewrite
--- just decorate

Simple API

• jit --- provide type information (fastest to call)
• autojit --- detects input types, infers output, generates
code if needed, and dispatches (a little more
expensive to call)

@jit('void(double[:,:], double, double)')
#@autojit
def numba_update(u, dx2, dy2):
nx, ny = u.shape
for i in xrange(1,nx-1):
for j in xrange(1, ny-1):
u[i,j] = ((u[i+1,j] + u[i-1,j]) * dy2 +
(u[i,j+1] + u[i,j-1]) * dx2) / (2*(dx2+dy2))

~150x speed-up Real-time image
processing in Python (50
fps Mandelbrot)

Image Processing ~1500x speed-up

@jit('void(f8[:,:],f8[:,:],f8[:,:])')
def filter(image, filt, output):
M, N = image.shape
m, n = filt.shape
for i in range(m//2, M-m//2):
for j in range(n//2, N-n//2):
result = 0.0
for k in range(m):
for l in range(n):
result += image[i+k-m//2,j+l-n//2]*filt[k, l]
output[i,j] = result

NumPy + Mamba = Numba
Python Function Machine Code

LLVM-PY

LLVM Library
ISPC OpenCL OpenMP CUDA CLANG

Intel AMD Nvidia Apple ARM

Example
@jit(‘f8(f8)’)
def sinc(x):
if x==0.0:
return 1.0
else:
return sin(x*pi)/(pi*x)

Numba

Compiler Overview

C++
x86

C
LLVM IR ARM
Fortran

PTX
Python

Numba turns Python into a “compiled
language” (but much more ﬂexible)

Compile NumPy array expressions

import numbapro
from numba import autojit

@autojit
def formula(a, b, c):
a[1:,1:] = a[1:,1:] + b[1:,:-1] + c[1:,:-1]

@autojit
def express(m1, m2):
m2[1:-1:2,0,...,::2] = (m1[1:-1:2,...,::2] *
m1[-2:1:-2,...,::2])
return m2

Fast vectorize
NumPy’s ufuncs take “kernels” and
apply the kernel element-by-element
over entire arrays
Write kernels in
from numbapro import vectorize
from math import sin
Python!
@vectorize([‘f8(f8)’, ‘f4(f4)’])
def sinc(x):
if x==0.0:
return 1.0
else:

Updated Laplace Example
https://github.com/teoliphant/speed.git
Version Time Speed Up
NumPy 3.19 1.0
Numba 2.32 1.38
Vect. Numba 2.33 1.37
Cython 2.38 1.34
Weave 2.47 1.29
Numexpr 2.62 1.22
Fortran Loops 2.30 1.39
Vect. Fortran 1.50 2.13

Many Advanced Features
• Extension classes (jit a class)
• Struct support (NumPy arrays can be structs)
• SSA --- can refer to local variables as different types
• Typed lists and typed dictionaries coming
• pointer support
• calling ctypes and CFFI functions natively
• pycc (create stand-alone dynamic library and executable)
• pycc --python (create static extension module for Python)

Ufuncs

Generalized
UFuncs
Python
Function
Window
Kernel
Funcs

Function-
Uses of Numba

based
Indexing

Memory
Filters
Numba

NumPy Runtime
I/O Filters

Reduction
Filters

Computed
Columns
function pointer

Status
• Main team sponsored by Continuum Analytics and
composed of:
- Travis Oliphant (NumPy, SciPy)
- Jon Riehl (Mython, PyFront, Basil, ...)
- Mark Florrison (minivect, Cython)
- Siu Kwan Lam (pymothoa, llvmpy)
• Rapid progress this year
• Version 0.6 released ﬁrst of February
• Version 0.7 next week numba.pydata.org
• Version 0.8 ﬁrst of April
• Stable API (jit, autojit) easy to use
• Full Python support by 1.0 at end of summer
• Should be able to write equivalent of NumPy and
SciPy with Numba

Introducing NumbaPro

• Create parallel-for loops
• Parallel execution of
ufuncs
fast development and fast • Run ufuncs on the GPU
execution! • Write CUDA directly in
Python!
Python and NumPy compiled to
• Free for Academics
Parallel Architectures
(GPUs and multi-core
machines)

Create parallel-for loops

import numbapro # import first to make prange available
from numba import autojit, prange

@autojit
def parallel_sum2d(a):
sum = 0.0
for i in prange(a.shape[0]):
for j in range(a.shape[1]):
sum += a[i,j]

Ufuncs in parallel (multi-core or GPU)
from numbapro import vectorize
from math import sin

@vectorize([‘f8(f8)’, ‘f4(f4)’], target=‘gpu’)
def sinc(x):
if x==0.0:
return 1.0
else:

@vectorize([‘f8(f8)’, ‘f4(f4)’], target=‘parallel’)
def sinc2(x):
if x==0.0:
return 1.0
else:

Introducing CUDA-Python
from numbapro import cuda
from numba import autojit

@autojit(target=‘gpu’)
def array_scale(src, dst, scale):
tid = cuda.threadIdx.x
blkid = cuda.blockIdx.x
blkdim = cuda.blockDim.x CUDA Development
i = tid + blkid * blkdim

if i >= n: using Python syntax!
return

dst[i] = src[i] * scale

src = np.arange(N, dtype=np.float)
dst = np.empty_like(src)

array_scale[grid, block](src, dst, 5.0)

Example: Matrix multiply
@cuda.jit(argtypes=[f4[:,:], f4[:,:], f4[:,:]])
def cu_square_matrix_mul(A, B, C):
sA = cuda.shared.array(shape=(tpb, tpb),
dtype=f4)
sB = cuda.shared.array(shape=(tpb, tpb),
dtype=f4) bpg = 50
tpb = 32
tx = cuda.threadIdx.x n = bpg * tpb
ty = cuda.threadIdx.y
bx = cuda.blockIdx.x A = np.array(np.random.random((n, n)),
by = cuda.blockIdx.y dtype=np.float32)
bw = cuda.blockDim.x
bh = cuda.blockDim.y
B = np.array(np.random.random((n, n)),
dtype=np.float32)
x = tx + bx * bw C = np.empty_like(A)
y = ty + by * bh
stream = cuda.stream()
acc = 0. with stream.auto_synchronize():
for i in range(bpg): dA = cuda.to_device(A, stream)
if x < n and y < n: dB = cuda.to_device(B, stream)
sA[ty, tx] = A[y, tx + i * tpb]
sB[ty, tx] = B[ty + i * tpb, x]
dC = cuda.to_device(C, stream)
cu_square_matrix_mul[(bpg, bpg),
cuda.syncthreads() (tpb, tpb), stream](dA, dB, dC)
dC.to_host(stream)
if x < n and y < n:
for j in range(tpb):
acc += sA[ty, j] * sB[j, tx]

cuda.syncthreads()

if x < n and y < n:
C[y, x] = acc

Early Performance Results
core GeForce GTX
i7 560 Ti
Already about 6x
faster on the
GPU.

Example: Black-Scholes
@cuda.jit(argtypes=(double,), restype=double, device=True, inline=True)
def cnd_cuda(d):
@cuda.jit(argtypes=(double[:], double[:],
A1 = 0.31938153
double[:],
A2 = -0.356563782
double[:], double[:],
A3 = 1.781477937
double, double))
A4 = -1.821255978
def black_scholes_cuda(callResult, putResult,
A5 = 1.330274429
S, X, T, R, V):
RSQRT2PI = 0.39894228040143267793994605993438
# S = stockPrice
K = 1.0 / (1.0 + 0.2316419 * math.fabs(d))
# X = optionStrike
ret_val = (RSQRT2PI * math.exp(-0.5 * d * d) *
# T = optionYears
(K * (A1 + K * (A2 + K * (A3 + K * (A4 + K * A5))))))
# R = Riskfree
if d > 0:
# V = Volatility
ret_val = 1.0 - ret_val
i = cuda.threadIdx.x + cuda.blockIdx.x *
return ret_val
cuda.blockDim.x
blockdim = 1024, 1 if i >= S.shape[0]:
griddim = int(math.ceil(float(OPT_N)/blockdim[0])), 1 return
stream = cuda.stream() sqrtT = math.sqrt(T[i])
d_callResult = cuda.to_device(callResultNumbapro, d1 = (math.log(S[i] / X[i]) +
stream) (R + 0.5 * V * V) * T[i]) / (V *
d_putResult = cuda.to_device(putResultNumbapro, sqrtT)
stream) d2 = d1 - V * sqrtT
d_stockPrice = cuda.to_device(stockPrice, stream) cndd1 = cnd_cuda(d1)
d_optionStrike = cuda.to_device(optionStrike, stream) cndd2 = cnd_cuda(d2)
d_optionYears = cuda.to_device(optionYears, stream)
for i in range(iterations): expRT = math.exp((-1. * R) * T[i])
black_scholes_cuda[griddim, blockdim, stream]( callResult[i] = (S[i] * cndd1 - X[i] *
d_callResult, d_putResult, d_stockPrice, expRT * cndd2)
d_optionStrike, putResult[i] = (X[i] * expRT * (1.0 -
d_optionYears, RISKFREE, VOLATILITY) cndd2) -
d_callResult.to_host(stream) S[i] * (1.0 -
d_putResult.to_host(stream) cndd1))
stream.synchronize()

Black-Scholes: Results
core GeForce GTX
i7 560 Ti
Already
about 6x
faster on the
GPU

NumFOCUS
www.numfocus.org
501(c)3 Public Charity

NumFOCUS Mission

• Sponsor development of high-level languages and
libraries for science
• Foster teaching of array-oriented and higher-order
computational approaches and applied
computational science
• Promote the use of open code in science and
encourage reproducible and accessible research

NumFOCUS Activities

• Sponsor sprints and conferences
• Provide scholarships and grants
• Provide bounties and prizes for code development
• Pay for freely-available documentation and basic
course development
• Equipment grants
• Sponsor BootCamps
• Raise funds from industries using open source
high-level languages

Numba: Array-oriented Python Compiler for NumPy

Recommended

More Related Content

What's hot (20)

Viewers also liked (11)

Similar to Numba: Array-oriented Python Compiler for NumPy (20)

More from Travis Oliphant (18)

Recently uploaded (20)

Numba: Array-oriented Python Compiler for NumPy