DOCSEARCHbyAlgolia
Submit Docs

Dr.Jit Documentation

repository·master·Indexed 21 days ago

https://github.com/mitsuba-renderer/drjit

Dr.Jit is a Just-In-Time compiler for differentiable and parallel computation, optimized for the complex computation graphs found in differentiable rendering. It supports forward and reverse mode automatic differentiation, custom differentiable operations via dr.CustomOp, and specialized differentiation strategies for loops, including optimized sum loops and general loops with trajectory storage.

Tokens
99.1K
Snippets
295
Records
434
Agent score
73%

What's inside Dr.Jit

  1. What is Dr.Jit and how does it work?

    master

    Dr.Jit is a just-in-time (JIT) compiler designed for ordinary and differentiable computation. It is optimized for workloads with extremely large and complex computation graphs (e.g., millions of elementary arithmetic nodes) that might cause standard machine learning frameworks like JAX, TensorFlow, or PyTorch to crash or time out.

    Dr.Jit operates through three primary mechanisms:

    1. Tracing and Vectorization: Instead of executing arithmetic operations immediately, Dr.Jit records them into a computation graph. This graph is then JIT-compiled into efficient fused kernels. Supported backends include:
      • GPU: Metal (macOS) or CUDA (other platforms).
      • CPU: Host CPU using vector instruction sets like AVX512 or NEON via LLVM.
      • Header-only mode: Can be used without JIT-compilation as a standard header-only vector library.
    2. Differentiation: Supports automatic differentiation (AD) using both forward and reverse-mode accumulation. Tracing and differentiation are integrated to produce specialized derivative evaluation code.
    3. Multi-language Support: Dr.Jit types are accessible in both C++17 and Python. You can develop in either language or mix them, as combinations of Python and C++ code can be jointly traced and differentiated.

  3. Ensure thread safety with new scopes

    master

    When executing Dr.Jit computation on different threads, you must ensure that basic blocks are separated between threads. If a thread $T_2$ references Dr.Jit arrays created by another thread $T_1$, $T_2$ must create a new scope to guarantee that dependencies are correctly tracked and ordered during compilation.

    Failure to do this will result in an exception when attempting to evaluate incorrectly ordered expressions. Use drjit.detail.new_scope() to set a unique, new scope identifier (a 32-bit integer) to separate the current computation from the previous basic block.

  4. Asynchronous execution in Dr.Jit

    master

    Dr.Jit employs two layers of asynchrony:

    1. Tracing: Operations are recorded as a graph rather than executed immediately.
    2. Evaluation: The dr.eval() call is asynchronous. It appends a work item to the device (GPU/CPU) command queue and returns control to the host immediately. This allows the host to continue tracing the next block of code while the device is busy executing the previous kernel.

    This behavior is transparent; Dr.Jit automatically handles synchronization and waiting when necessary.

    dr.eval(x)  # Returns immediately; work is queued on the device

  5. Unsupported operations in frozen functions: Compress and Pointers

    master

    Two additional technical limitations exist:

    1. drjit.compress: Using dr.compress inside a frozen function often renders the feature useless. Because compress produces output sizes that depend on input content, functions that require known array sizes in advance (like dr.block_reduce or dr.scatter_reduce on LLVM) will be forced to re-trace on every call.
    2. Pointers with offsets (C++): In custom C++ code, pointers that point inside a memory region (e.g., UInt32::load_(x.data() + 4)) are not supported because Dr.Jit identifies variables via their base data pointers.

  6. Enable Debug mode for error detection

    master

    Dr.Jit provides a debug mode (JitFlag_Debug) to uncover errors in application code.

    What it does:

    • Enables assertion checks for drjit.assert_true, drjit.assert_false, and drjit.assert_equal.
    • Intercepts out-of-bounds reads/writes in operations like drjit.scatter, drjit.gather, drjit.scatter_reduce, and drjit.scatter_inc.
    • Detects invalid callables in drjit.switch and drjit.dispatch.

    Warning: Debug mode comes at a significant cost. It interferes with kernel caching, reduces tracing performance, and produces slower kernels. Use it only for debugging or periodically before a release.

    # Example of an out-of-bounds warning in debug mode
>>> dr.gather(dtype=UInt, source=UInt(1, 2, 3), index=UInt(0, 1, 100))
RuntimeWarning: drjit.gather(): out-of-bounds read from position 100 in an array of size 3.

  7. Unsupported operations in frozen functions: Array access

    master

    Frozen functions cannot perform operations that depend on the specific values of array elements to determine control flow. While frozen functions can accept PyTrees (containing scalars or Dr.Jit arrays), you cannot extract scalar elements from a Dr.Jit array to influence logic (e.g., if x[1] > 0:). Doing so would 'bake' the observed constant into the kernel, making it impossible to replay with different data. Dr.Jit will raise an exception if such an operation is attempted.

    @dr.freeze
def func(x: Float, y: Float):
    # PROHIBITED: Accessing elements of x to influence control flow
    if x[1] > 0:
       return y + 1
    else:
       return y - 1

  8. Integrate Hash Grid Encodings into an nn.Module

    master

    Hash grid weights cannot be included in the standard nn.pack() process because they use a different memory layout and potentially incompatible types. To use a hash grid within an nn.Module (like nn.Sequential), you must wrap it in an nn.HashEncodingLayer.

    When using nn.HashEncodingLayer, you must:

    1. Provide a prefix to the nn.HashEncodingLayer (or the containing module) to prevent parameter name collisions in the optimizer.
    2. Optimize the encoding parameters independently from the packed module weights.
    3. Manually write the optimized encoding parameters back to the encoding object during the training loop.
    # 1. Create the encoding
enc = nn.HashGridEncoding(Float16, 2, rng=rng)

# 2. Wrap it in a HashEncodingLayer within a Sequential module
net = nn.Sequential(
    nn.HashEncodingLayer(enc),
    nn.Cast(Float16),
    nn.Linear(-1, -1, bias=False),
    # ... other layers
    prefix='mlp'
)

# 3. Optimize separately
opt = Adam(lr=1e-3)
opt.update(net)
opt['enc.params'] = Float32(enc.params)

# 4. In the training loop, manually update the encoding
enc.params[:] = Float16(opt['enc.params'])

  9. How Quaternion types behave in Dr.Jit

    master

    Quaternion types (e.g., drjit.scalar.Quaternion4f, drjit.cuda.ad.Quaternion4f64) represent quaternion-valued scalars and arrays.

    Key behaviors:

    • Broadcasting: Constructing a quaternion from non-quaternionic values broadcasts to the identity element.
    • Multiplication: The * operator performs a quaternion product.
    • Division: True division (arg0 / arg1) with a quaternion denominator involves a quaternion inverse.
    • Mathematical Operations: Generalizations exist for drjit.fma, drjit.rcp, drjit.abs, drjit.sqrt, drjit.rsqrt, drjit.log2, drjit.exp2, and drjit.power.

    Note: Trigonometric functions (sin, cos, etc.) are currently undefined for quaternions.

    >>> dr.scalar.Quaternion4f(1, 2, 3, 4) + 10
1i+2j+3k+14

  10. Differentiate PyTrees using `dr.grad()`

    master
    Dr.Jit functions can operate on PyTrees (arbitrarily nested tuples, lists, or dictionaries). When working with nested data structures, do not use the .grad attribute of individual arrays, as that only exists on Dr.Jit arrays. Instead, use the dr.grad() function to access the gradients of the entire nested structure.

  11. Understand Dr.Jit's approach to tensor operations

    master

    Dr.Jit is not a general-purpose tensor framework like NumPy or PyTorch. It is designed for high-performance parallel evaluations where programs are built from flat and nested array operations. These operations are fused into large, self-contained kernels.

    Using a tensor-heavy development style (treating Dr.Jit like a standard tensor library) can interfere with this kernel fusion process and degrade performance. For optimal results, focus on building programs that allow the system to fuse operations effectively.

  12. Perform atomic increments with scatter_inc()

    master

    The drjit.scatter_inc() function atomically increments values in an unsigned 32-bit integer array and returns the value prior to the update. This is a key building block for stream compaction.

    Stream Compaction Recipe

    To compact data, use scatter_inc to request unique indices for active elements, then use those indices to scatter the data into a new buffer.

    1. Create a counter ctr initialized to 0.
    2. Call idx = dr.scatter_inc(target=ctr, index=UInt32(0), mask=active).
    3. Use idx to dr.scatter your data into output buffers.
    4. Use dr.reshape with shrink=True to resize the output buffers to the actual number of active elements.

    Important: The 'Consumed' Index Pattern

    The return value of scatter_inc is an instantaneous state that is not reproducible. If you attempt to reuse the index in a subsequent kernel without materializing it, Dr.Jit will raise an exception.

    To fix reuse errors: You must evaluate the index at the same time as the first kernel to materialize it into a stored representation:

    dr.eval(data_compact_1, my_index)

    Comparison with drjit.compress()

    • drjit.compress(): Simpler (prefix sum), but requires evaluating the variables being reduced, which can be memory-intensive.
    • drjit.scatter_inc(): Can operate on symbolic arrays that exceed available device memory, making it more scalable for large datasets.
    data_1 = ...
data_2 = ...
active = drjit.ones(Bool, len(data_1))

# This will hold the counter
ctr = UInt32(0)

# Allocate output buffers
max_size = 1024
data_compact_1 = dr.empty(Float, max_size)
data_compact_2 = dr.empty(Float, max_size)

idx = dr.scatter_inc(target=ctr, index=UInt32(0), mask=active)

# Disable dr.scatter() operations below in case of a buffer overflow
active &= idx < max_size

dr.scatter(
    target=data_compact_1,
    value=data_1,
    index=idx, # Note: using the returned index
    mask=active
)

dr.scatter(
    target=data_compact_2,
    value=data_2,
    index=idx,
    mask=active
)