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

Wenzel Jakob1     Sébastien Speierer     Nicolas Roussel     Delio Vicini1    
École Polytechnique Fédérale de Lausanne (EPFL)
ACM Transactions on Graphics (Proceedings of SIGGRAPH), 2022
Teaser

Dr.Jit is a domain-specific compiler for physically-based (differentiable) rendering. When Dr.Jit executes a rendering algorithm, it generates a trace: a large graph comprised of arithmetic, loops, ray tracing operations, and polymorphic calls that exchange information between the rendering algorithm and scene objects (shapes, BSDFs, textures, emitters, etc.). Dr.Jit specializes this graph to the provided scene and compiles it into a large data-parallel kernel (“megakernel”) via LLVM or OptiX backends, achieving geometric mean GPU speedups of 3.70 × (vs. Mitsuba 2) and 2.14 × (vs. PBRT 4). While helpful for ordinary rendering, the main purpose of Dr.Jit is to dynamically compile differential simulations. Recent methods in this area decompose a larger differentiation task into a series of incremental steps, which requires an unusually fine-grained approach to automatic differentiation. Dr.Jit supports such transformations in forward and reverse modes: the former computes a perturbation in image space, which is helpful for debugging and visualization. The latter provides derivatives in parameter space (e.g. texels of the wallpaper) for simultaneous optimization of large numbers of unknowns

Abstract

Dr.Jit is a new just-in-time compiler for physically based rendering and its derivative. Dr.Jit expedites research on these topics in two ways: first, it traces high-level simulation code (e.g., written in Python) and aggressively simplifies and specializes the resulting program representation, producing data-parallel kernels with state-of-the-art performance on CPUs and GPUs. Second, it simplifies the development of differentiable rendering algorithms. Efficient methods in this area turn the derivative of a simulation into a simulation of the derivative. Dr.Jit provides fine-grained control over the process of automatic differentiation to help with this transformation. Specialization is particularly helpful in the context of differentiation, since large parts of the simulation ultimately do not influence the computed gradients. Dr.Jit tracks data dependencies globally to find and remove redundant computation

Video

Downloads & resources

Text Reference Copy to clipboard

Wenzel Jakob, Sébastien Speierer, Nicolas Roussel, Delio Vicini. Dr.Jit: A Just-In-Time Compiler for Differentiable Rendering. ACM Transactions on Graphics (Proceedings of SIGGRAPH), 41(4), July 2022.

BibTex Reference Copy to clipboard

@article{Jakob2022drjit,
  author  = {Jakob, Wenzel and Speierer, Sébastien and Roussel, Nicolas and Vicini, Delio},
  title   = {Dr.Jit: A Just-In-Time Compiler for Differentiable Rendering},
  journal = {ACM Transactions on Graphics (Proceedings of SIGGRAPH)},
  volume  = {41},
  number  = {4},
  year    = {2022},
  month   = jul,
  pages   = {124:1--124:19}
}