Package: drjit
About this projectDr.Jitis ajust-in-time(JIT) compiler for ordinary and differentiable
computation. It was originally created as the numerical foundation ofMitsuba
3, a differentiableMonte
Carlorenderer. However,Dr.Jitis a general-purpose tool that can also help with various other types
of embarrassingly parallel computation.Dr.Jitprincipally facilitates three steps:Vectorization and tracing: WhenDr.Jitencounters an operation (e.g. an
additiona + b) it does not execute it right away: instead, it remembers
that an addition will be needed at some later point by recording it into a
graph representation (this is calledtracing). Eventually, it willjust-in-time(JIT) compile the recorded operations into afusedkernel
using eitherLLVM(when targeting the
CPU) orCUDA(when targeting the
GPU). The valuesaandbwill typically be arrays with many elements,
and the system parallelizes their evaluation using multi-core parallelism and
vector instruction sets likeAVX512orARM Neon.Dr.Jitis ideal for Monte Carlo methods, where the same computation must be
repeated for millions of random samples.Dr.Jitdynamically generates
specialized parallel code for the target platform.
As a fallback, Dr.Jit can also be used without JIT-compilation, which turns
the project into a header-only vector library without external dependencies.Differentiation: If desired,Dr.Jitcan compute derivatives usingautomatic differentiation(AD), using eitherforward or reverse-mode
accumulation.
Differentiation and tracing go hand-in-hand to produce specialized derivative
evaluation code.Python:Dr.Jittypes are accessible within C++17 and Python. Code can be
developed in either language, or even both at once. Combinations of Python
and C++ code can be jointly traced and differentiated.Dr.Jithandles large programs with custom data structures, side effects, and
polymorphism. It includes a mathematical support library including
transcendental functions and types like vectors, matrices, complex numbers,
quaternions, etc.Difference to machine learning frameworksWhy did we createDr.Jit, when dynamic derivative compilation is already
possible using Python-based ML frameworks likeJAX,Tensorflow,
andPyTorchalong with backends likeXLAandTorchScript?The reason is related to the typical workloads: machine learning involves
small-ish computation graphs that are, however, made of arithmetically intense
operations like convolutions, matrix multiplications, etc. The application
motivatingDr.Jit(differentiable rendering) creates giant and messy
computation graphs consisting of 100K to millions of “trivial” nodes
(elementary arithmetic operations). In our experience, ML compilation backends
use internal representations and optimization passes that aretoo richfor
this type of input, causing them to crash or time out during compilation. If
you have encountered such issues, you may findDr.Jituseful.CloningDr.Jit recursively depends on two other repositories:pybind11for Python bindings, anddrjit-coreproviding core components of
the JIT-compiler.To fetch the entire project including these dependencies, clone the project
using the--recursiveflag as follows:$gitclone--recursivehttps://github.com/mitsuba-renderer/drjitDocumentationPlease see Dr.Jit’s page onreadthedocs.iofor example code and reference documentation.References, citationsPlease see the paperDr.Jit: A Just-In-Time Compiler for Differentiable
Renderingfor the
nitty-gritty details and details on the problem motivating this project. There
is also avideo presentationexplaining the design decisions at a higher level.If you useDr.Jitin your own research, please cite it using the following
BibTeX entry:@article{Jakob2022DrJit,author={Wenzel Jakob and S{\'e}bastien Speierer and Nicolas Roussel and Delio Vicini},title={Dr.Jit: A Just-In-Time Compiler for Differentiable Rendering},journal={Transactions on Graphics (Proceedings of SIGGRAPH)},volume={41},number={4},year={2022},month=jul,doi={10.1145/3528223.3530099}}Logo and historyTheDr.Jitlogo was generously created byOtto Jakob. The “Dr.” prefix simultaneously abbreviatesdifferentiable renderingwith the stylized partial derivativeD, while also
conveying a medical connotation that is emphasized by theRod of Asclepius. Differentiable rendering
algorithms are growing beyond our control in terms of conceptual and
implementation-level complexity. A doctor is a person, who can offer help in
such a time of great need.Dr.Jittries to fill this role to to improve the
well-being of differentiable rendering researchers.Dr.Jitis the successor of theEnokiproject, and its high-level API
still somewhat resembles that of Enoki. The system evolved towards a different
approach and has an all-new implementation, hence the decision to switch to a
different project name.