Dr.Jit quickstart

Overview

This short tutorial recaps the basic functionalities and routines of the Dr.Jit library. You can also find more information on the Dr.Jit documentation.

Similarity with NumPy

On the Python side, the Dr.Jit syntax is very similar to NumPy. Moreover, as we will see later, both frameworks are interoperable.

Let’s first import both NumPy and Dr.Jit using the alias np and dr respectively

[1]:
import numpy as np
import drjit as dr

Unlike NumPy, Dr.Jit can perform array arithmetic on both CPU and GPU through various template variants which are exposed in top-level packages:

Variant

Description

drjit.scalar

Arrays built on top of scalars (float, int, etc.)

drjit.llvm

Arrays built on top of LLVMArray

drjit.cuda

Arrays built on top of CUDAArray

drjit.llvm.ad

Similar to drjit.llvm but with automatic differentiation support

drjit.cuda.ad

Similar to drjit.cuda but with automatic differentiation support

These packages all contains various types like: Bool, Float, Int, UInt, Array2f, Array2i, Matrix2f Matrix3f, ...

Let’s create some arrays using the drjit.llvm variants and play around with the NumPy interoperability:

[2]:
from drjit.llvm import Float, UInt32

# Create some floating-point arrays
a = Float([1.0, 2.0, 3.0, 4.0])
b = Float([4.0, 3.0, 2.0, 1.0])

# Perform simple arithmetic
c = a + 2.0 * b

print(f'c -> ({type(c)}) = {c}')

# Convert to NumPy array
d = np.array(c)

print(f'd -> ({type(d)}) = {d}')
c -> (<class 'drjit.llvm.Float'>) = [9.0, 8.0, 7.0, 6.0]
d -> (<class 'numpy.ndarray'>) = [9. 8. 7. 6.]

Array construction routines

This section provides an overview of various Dr.Jit routines (and their NumPy correspondence) to construct arrays.

[3]:
# Initialize floating-point array of size 5 with zeros
a = dr.zeros(Float, 5) # np.zeros(5)
print(f'dr.zeros: {a}')

# Initialize floating-point array of size 5 with a constant value
a = dr.full(Float, 0.1, 5) # np.ones(5, 0.4)
print(f'dr.full: {a}')

a = dr.arange(UInt32, 5) # np.arange(5)
print(f'dr.arange: {a}')

# Return evenly spaced numbers over a specified interval
a = dr.linspace(Float, 0.0, 2.0, 5) # np.linspace(0.0, 2.0, 5)
print(f'dr.linespace: {a}')
dr.zeros: [0.0, 0.0, 0.0, 0.0, 0.0]
dr.full: [0.10000000149011612, 0.10000000149011612, 0.10000000149011612, 0.10000000149011612, 0.10000000149011612]
dr.arange: [0, 1, 2, 3, 4]
dr.linespace: [0.0, 0.5, 1.0, 1.5, 2.0]

Masking

Writing codes using Dr.Jit often means working with large arrays at once. Therefore it is not possible to use regular if .. else .. statements based on concret values, as different elements in the array might branch differently. This is where masking comes to the rescue!

A mask (or Bool) is an array of boolean values that can be used to disable arithmetic operations on part of an array. It is possible to create such masks with any regular boolean arithmetic (e.g. >, <, >=, <=).

Often time, we combine masks with the dr.select(mask, a, b) statement which correspond to the ternary statement mask ? a : b. This is similar to the np.where function in NumPy.

[4]:
x = dr.arange(Float, 5)
m = x > 2.0 # True for all values of a that are greater than 2.0
y = dr.select(m, 4.0, 1.0) # Set the values greater than 2.0 to 4.0 otherwise to 1.0
print(f'x -> ({type(x)}) {x}')
print(f'm -> ({type(m)}) {m}')
print(f'y -> ({type(y)}) {y}')
x -> (<class 'drjit.llvm.Float'>) [0.0, 1.0, 2.0, 3.0, 4.0]
m -> (<class 'drjit.llvm.Bool'>) [False, False, False, True, True]
y -> (<class 'drjit.llvm.Float'>) [1.0, 1.0, 1.0, 4.0, 4.0]

Basic math arithmetic

All common math operators like +, -, /, *, *=, +=, %, //, ... are supported with Dr.Jit arrays.

Similarly to NumPy, Dr.Jit provides all kinds of math arithmetic that can be performed on the entire array in a single call. Here is a non-exaustive list of those math functions: abs, minimum, maximum, sqrt, pow, sin, cos, tan, atan2, sincos, sec, cot, asin, acos, atan, exp, exp2, log, log2, sinh, cosh, tanh, asinh, acosh, atanh, ...

Those routines are present in the root drjit package, hence can be used as follow:

[5]:
s, c = dr.sincos(a)
m = dr.minimum(s, c)
print(f'm: {m}')
m: [0.0, 0.4794255495071411, 0.5403022766113281, 0.07073719799518585, -0.41614681482315063]

Horizontal operations

Dr.Jit also provides operations that require a pass over the entire array and return a single scalar value. Those operations are expensive as they will trigger a syncronization point, hence it is better to avoid them if possible.

The following snippet of code explores a few of those:

[6]:
a = dr.arange(Float, 5) + 1
print(f'a: {a}')

# Horizontal sum
b = dr.sum(a) # np.sum(a)
print(f'dr.sum(a): {b}')

# Horizontal product
b = dr.prod(a) # np.prod(a)
print(f'dr.prod(a): {b}')

# Mean value over the entire array
b = dr.mean(a) # np.mean(a)
print(f'dr.mean(a): {b}')

m = a > 2
print(f'm: {m}')

# True if all value of the mask array are True
b = dr.all(m) # np.all(m)
print(f'dr.all(m): {b}')

# True if any value of the mask array are True
b = dr.any(m) # np.any(m)
print(f'dr.any(m): {b}')

# True if no value of the mask array are True
b = dr.none(m) # ~np.any(m)
print(f'dr.none(m): {b}')
a: [1.0, 2.0, 3.0, 4.0, 5.0]
dr.sum(a): [15.0]
dr.prod(a): [120.0]
dr.mean(a): [3.0]
m: [False, False, True, True, True]
dr.all(m): False
dr.any(m): True
dr.none(m): False

gather and scatter routines

In programming languages like C++ or Python, it is possible to access the i-th element of an array using the array[i] syntax. This can both be used to read or write values in an array. Similarly, Dr.Jit provides such read/write functionalities through the dr.gather and dr.scatter functions. Those are much more powerful than the regular array accessors as the index i can be an array itself! In which case the read operation (e.g. dr.gather) would return a array as well, not just a single value.

Here is how one should use the dr.gather routine to read entries from an Dr.Jit array:

[7]:
source = dr.linspace(Float, 0, 1, 5)
indices = UInt32([1, 2]) # Only read the 2nd and 3rd elements of the source array
result = dr.gather(Float, source, indices)
print(f'source: {source}')
print(f'indices: {indices}')
print(f'result: {result}')
source: [0.0, 0.25, 0.5, 0.75, 1.0]
indices: [1, 2]
result: [0.25, 0.5]

And here is how one can write entries at specific indices into a Dr.Jit array

[8]:
target = dr.zeros(Float, 5)
indices = UInt32([0, 3, 4]) # Write to the first and last two elements of the target array
source = Float([1.0, 2.0, 3.0])
dr.scatter(target, source, indices)
print(f'indices: {indices}')
print(f'source: {source}')
print(f'target: {target}')
indices: [0, 3, 4]
source: [1.0, 2.0, 3.0]
target: [1.0, 0.0, 0.0, 2.0, 3.0]