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

In Transactions on Graphics (Proceedings of SIGGRAPH 2022)

Dr.Jit is a do­main-spe­cif­ic com­piler for phys­ic­ally-based (dif­fer­en­ti­able) ren­der­ing. When Dr.Jit ex­ecutes a ren­der­ing al­gorithm, it gen­er­ates a trace: a large graph com­prised of arith­met­ic, loops, ray tra­cing op­er­a­tions, and poly­morph­ic calls that ex­change in­form­a­tion between the ren­der­ing al­gorithm and scene ob­jects (shapes, BSD­Fs, tex­tures, emit­ters, etc.). Dr.Jit spe­cial­izes this graph to the provided scene and com­piles it in­to a large data-par­al­lel ker­nel (“mega­ker­nel') via LLVM or Op­tiX backends, achiev­ing geo­met­ric mean GPU spee­dups of 3.70 × (vs. Mit­suba 2) and 2.14 × (vs. PBRT 4). While help­ful for or­din­ary ren­der­ing, the main pur­pose of Dr.Jit is to dy­nam­ic­ally com­pile dif­fer­en­tial sim­u­la­tions. Re­cent meth­ods in this area de­com­pose a lar­ger dif­fer­en­ti­ation task in­to a series of in­cre­ment­al steps, which re­quires an un­usu­ally fine-grained ap­proach to auto­mat­ic dif­fer­en­ti­ation. Dr.Jit sup­ports such trans­form­a­tions in for­ward and re­verse modes: the former com­putes a per­turb­a­tion in im­age space, which is help­ful for de­bug­ging and visu­al­iz­a­tion. The lat­ter provides de­riv­at­ives in para­met­er space (e.g. texels of the wall­pa­per) for sim­ul­tan­eous op­tim­iz­a­tion of large num­bers of un­knowns.


Dr.Jit is a new just-in-time com­piler for phys­ic­ally based ren­der­ing and its de­riv­at­ive. Dr.Jit ex­ped­ites re­search on these top­ics in two ways: first, it traces high-level sim­u­la­tion code (e.g., writ­ten in Py­thon) and ag­gress­ively sim­pli­fies and spe­cial­izes the res­ult­ing pro­gram rep­res­ent­a­tion, pro­du­cing data-par­al­lel ker­nels with state-of-the-art per­form­ance on CPUs and GPUs.

Second, it sim­pli­fies the de­vel­op­ment of dif­fer­en­ti­able ren­der­ing al­gorithms. Ef­fi­cient meth­ods in this area turn the de­riv­at­ive of a sim­u­la­tion in­to a sim­u­la­tion of the de­riv­at­ive. Dr.Jit provides fine-grained con­trol over the pro­cess of auto­mat­ic dif­fer­en­ti­ation to help with this trans­form­a­tion.

Spe­cial­iz­a­tion is par­tic­u­larly help­ful in the con­text of dif­fer­en­ti­ation, since large parts of the sim­u­la­tion ul­ti­mately do not in­flu­ence the com­puted gradi­ents. Dr.Jit tracks data de­pend­en­cies glob­ally to find and re­move re­dund­ant com­pu­ta­tion.


Text citation

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

    author = {Wenzel Jakob and Sé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}