DR.JIT: a just-in-time compiler for differentiable rendering

derivative of a conceptually simple algorithm like path tracing [Kajiya 1986] with precautions for linear time complexity [Vicini et al. 2021] and unbiased visibility handling [Bangaru et al. 2020] turns into an enormously complicated function. Correct implementation of such a large and intricate program is near-impossible even for experts in the field. Mere correctness is also unsatisfactory: optimizations tend to run for thousands of iterations, hence the resulting program needs to be fast. It is evident that better tools are need bridge this conspicuous gap between PBDR theory and practice.

The design of Dr.Jit was guided by a single unifying objective: it should provide a practical and efficient foundation for work in this area. Most architectural decisions are direct consequences of this overarching goal. For example, consider the differentiation step that is implicit in differentiable rendering. Manual differentiation is tedious and error-prone, hence it is logical that the system should build on automatic differentiation (AD) to simplify development.

However, the needs of PBDR are more specific: standard use of AD to differentiate a rendering algorithm produces an inefficient and biased derivative that precludes many applications. Recent work addresses inefficiencies using physical reciprocity [Nimier-David et al. 2020] and arithmetic invertibility [Vicini et al. 2021] to turn the derivative of a simulation into a simulation of the derivative, while re-parameterizing the integration domain to remove bias [Loubet et al. 2019]. These steps move the differentiation operation into the random walk, where it introduces partial derivative terms at each scattering event. This has implications on the design of the system: the derivatives must somehow be (pre-)compiled, since the machinery of AD is too slow to be used dynamically at such high rates.

Differentiating a simulation changes the underlying computation, but the details of this change depend on the scene, simulation algorithm, and optimization task. When the optimization only targets a subset of the scene's parameters, it is desirable that the system uses this information to remove steps that cannot influence the computed gradient. The dynamic nature of this problem calls for a similarly dynamic approach to compilation, which is why we pursue an approach centered around just-in-time (JIT) compilation.

Effective use of modern computing hardware requires that the program is organized into data-parallel phases known as kernels. Several kernels are generally needed to handle data dependencies, which must then exchange information through device memory. This interkernel communication comes at a cost in terms of storage and memory bandwidth, hence the specific manner in which a computation is partitioned into kernels can have a pronounced impact on performance. In the case of PBDR, the simulation parallelizes over millions of Monte Carlo samples that represent a large amount of program state. In our experiments, we find that it is almost always preferable that the Monte Carlo integration occurs within a megakernel, i.e., a large kernel containing all program instructions needed to evaluate the integrand. Most sample state can then be stored in registers, reducing memory usage and inter-kernel communication.

Finally, physically-based rendering algorithms are commonly expressed using subtype polymorphism to dynamically dispatch method calls from abstract component interfaces (e.g., a material encountered by a ray) to concrete implementations (e.g., a woven fabric or a rough metallic surface). The ability to represent, differentiate, and optimize such polymorphic constructions benefits performance.