mitsuba-renderer/drjit

Dr.Jit — A Just-In-Time-Compiler for Differentiable Rendering

What it solves

Dr.Jit addresses the challenge of handling massive, complex computation graphs—often containing millions of trivial arithmetic nodes—which typically cause standard machine learning frameworks (like PyTorch or JAX) to crash or time out during compilation. It is specifically designed for workloads like differentiable rendering and other embarrassingly parallel computations.

How it works

Dr.Jit functions as a just-in-time (JIT) compiler that records arithmetic operations into a computation graph instead of executing them immediately. It then compiles these graphs into highly efficient fused kernels for GPUs (via CUDA or Metal) or CPUs (via LLVM using AVX512 or NEON). Additionally, it supports automatic differentiation (AD) using both forward and reverse-mode accumulation to produce specialized derivative evaluation code.

Who it’s for

It is intended for researchers and developers working on differentiable rendering, high-performance numerical computation, and those who need to perform automatic differentiation on giant, messy computation graphs that exceed the capabilities of traditional ML backends.

Highlights

  • Multi-Backend Support: Targets CUDA, Metal, and CPU (AVX512/NEON).
  • Hybrid Language Support: Accessible via C++17 and Python, allowing joint tracing and differentiation across both.
  • Automatic Differentiation: Supports both forward and reverse-mode AD.
  • Mathematical Library: Includes built-in support for vectors, matrices, complex numbers, quaternions, and transcendental functions.

Related

  • Project
    m4rs-mt/ILGPU

    A JIT compiler for high-performance GPU programming written in C#, enabling .NET developers to write GPU kernels without native dependencies.

  • Project
    Jittor/jittor

    Jittor is a high‑performance deep‑learning framework that JIT‑compiles the entire computation graph and uses meta‑operators to generate specialised CUDA/C++ kernels. It offers a PyTorch‑like Python API, supports CPU and GPU back‑ends (CUDA, ROCm, Hygon), and includes a model zoo for vision, rendering, geometric learning, and reinforcement learning.

  • Project
    brainpy/BrainPy

    A JIT-compiled framework for computational neuroscience and brain-inspired computation that enables efficient simulation, training, and analysis of brain dynamics models.

  • Project
    NVIDIA/warp

    A Python framework for GPU-accelerated, differentiable simulations used in robotics, physics, and machine learning.

  • Project
    jax-ml/jax

    A Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning.