“DR.JIT: a just-in-time compiler for differentiable rendering” by Jakob, Speierer, Roussel and Vicini

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.

1. Sai Praveen Bangaru, Tzu-Mao Li, and Frédo Durand. 2020. Unbiased Warped-Area Sampling for Differentiable Rendering. ACM Trans. Graph. (SIGGRAPH Asia) 39, 6 (Nov. 2020).Google Scholar
2. Atilim Gunes Baydin, Barak A Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. 2018. Automatic differentiation in machine learning: a survey. Journal of machine learning research 18 (2018).Google Scholar
3. Christian Bischof, Alan Carle, George Corliss, Andreas Griewank, and Paul Hovland. 1992. ADIFOR-generating derivative codes from Fortran programs. Scientific Programming 1, 1 (1992).Google Scholar
4. Benedikt Bitterli. 2016. Rendering resources. https://benedikt-bitterli.me/resources/.Google Scholar
5. James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. 2018. JAX: composable transformations of Python+NumPy programs. http://github.com/google/jaxGoogle Scholar
6. Ian Buck, T. Foley, Daniel Horn, Jeremy Sugerman, Kayvon Fatahalian, Mike Houston, and Pat Hanrahan. 2004. Brook for GPUs: stream computing on graphics hardware. ACM Trans. Graph. (SIGGRAPH) 23, 3 (2004).Google ScholarDigital Library
7. Brent Burley. 2012. Physically Based Shading at Disney. In ACM SIGGRAPH Talks.Google Scholar
8. Brent Burley. 2015. Physically Based Shading in Theory and Practice: Extending the Disney BRDF to a BSDF with Integrated Subsurface Scattering. In SIGGRAPH Courses.Google Scholar
9. Sharan Chetlur, Cliff Woolley, Philippe Vandermersch, Jonathan Cohen, John Tran, Bryan Catanzaro, and Evan Shelhamer. 2014. cuDNN: Efficient Primitives for Deep Learning. (2014). arXiv:1410.0759Google Scholar
10. Conal Elliott. 2018. The Simple Essence of Automatic Differentiation. Proc. ACM Program. Lang. 2, ICFP (July 2018).Google ScholarDigital Library
11. Luca Fascione, Johannes Hanika, Mark Leone, Marc Droske, Jorge Schwarzhaupt, Tomáš Davidovič, Andrea Weidlich, and Johannes Meng. 2018. Manuka: A batch-shading architecture for spectral path tracing in movie production. ACM Trans. Graph. 37, 3 (2018).Google ScholarDigital Library
12. Tess Foley and Pat Hanrahan. 2011. Spark: Modular, Composable Shaders for Graphics Hardware. ACM Trans. Graph. (SIGGRAPH) 30, 4, Article 107 (July 2011), 12 pages.Google ScholarDigital Library
13. Ioannis Gkioulekas, Anat Levin, and Todd Zickler. 2016. An evaluation of computational imaging techniques for heterogeneous inverse scattering. In European Conference on Computer Vision (ECCV). Springer.Google ScholarCross Ref
14. Google. 2017. XLA: Optimizing Compiler for Machine Learning. https://www.tensorflow.org/xla (Accessed Jan 16, 2022).Google Scholar
15. Andreas Griewank et al. 1989. On automatic differentiation. Mathematical Programming: recent developments and applications 6, 6 (1989).Google Scholar
16. Andreas Griewank and Andrea Walther. 2008. Evaluating derivatives: principles and techniques of algorithmic differentiation. Vol. 105. SIAM.Google Scholar
17. Pat Hanrahan and Jim Lawson. 1990. A language for shading and lighting calculations. In Proceedings of the 17th annual conference on Computer graphics and interactive techniques.Google ScholarDigital Library
18. Laurent Hascoet and Valérie Pascual. 2013. The Tapenade automatic differentiation tool: Principles, model, and specification. ACM Transactions on Mathematical Software (TOMS) 39, 3 (2013).Google ScholarDigital Library
19. Yong He, Kayvon Fatahalian, and T. Foley. 2018. Slang: language mechanisms for extensible real-time shading systems. ACM Trans. Graph. (SIGGRAPH) 37, 4 (2018).Google ScholarDigital Library
20. Yuanming Hu, Luke Anderson, Tzu-Mao Li, Qi Sun, Nathan Carr, Jonathan Ragan-Kelley, and Frédo Durand. 2020. DiffTaichi: Differentiable Programming for Physical Simulation. ICLR (2020).Google Scholar
21. Yuanming Hu, Tzu-Mao Li, Luke Anderson, Jonathan Ragan-Kelley, and Frédo Durand. 2019. Taichi: A Language for High-Performance Computation on Spatially Sparse Data Structures. ACM Trans. Graph. (SIGGRAPH Asia) 38, 6 (Nov. 2019).Google ScholarDigital Library
22. Michael Innes. 2019. Don’t Unroll Adjoint: Differentiating SSA-Form Programs. arXiv:1810.07951Google Scholar
23. Wenzel Jakob. 2019. Enoki: structured vectorization and differentiation on modern processor architectures. https://github.com/mitsuba-renderer/enoki. (Accessed: Jan 16, 2022).Google Scholar
24. James T. Kajiya. 1986. The Rendering Equation. SIGGRAPH Comput. Graph. 20, 4 (Aug. 1986).Google ScholarDigital Library
25. Samuli Laine, Tero Karras, and Timo Aila. 2013. Megakernels considered harmful: Wavefront path tracing on GPUs. In Proceedings of the 5th High-Performance Graphics Conference.Google ScholarDigital Library
26. Chris Lattner and Vikram Adve. 2004. LLVM: A Compilation Framework for Lifelong Program Analysis and Transformation. San Jose, CA, USA.Google Scholar
27. João Rui Leal and Brad Bell. 2017. CppAD 2017. Google ScholarCross Ref
28. Mark Lee, Brian Green, Feng Xie, and Eric Tabellion. 2017. Vectorized production path tracing. In Proceedings of High Performance Graphics. ACM.Google ScholarDigital Library
29. Roland Leißa, Klaas Boesche, Sebastian Hack, Arsène Pérard-Gayot, Richard Membarth, Philipp Slusallek, André Müller, and Bertil Schmidt. 2018. AnyDSL: A Partial Evaluation Framework for Programming High-Performance Libraries. Proceedings of the ACM on Programming Languages (PACMPL) 2, OOPSLA (Nov. 2018).Google Scholar
30. Tzu-Mao Li, Miika Aittala, Frédo Durand, and Jaakko Lehtinen. 2018a. Differentiable Monte Carlo Ray Tracing Through Edge Sampling. ACM Trans. Graph. (SIGGRAPH Asia) 37, 6 (Dec. 2018).Google Scholar
31. Tzu-Mao Li, Michaël Gharbi, Andrew Adams, Frédo Durand, and Jonathan Ragan-Kelley. 2018b. Differentiable programming for image processing and deep learning in Halide. ACM Trans. Graph. (SIGGRAPH) 37, 4 (2018).Google ScholarDigital Library
32. Seppo Linnainmaa. 1976. Taylor expansion of the accumulated rounding error. BIT Numerical Mathematics 16, 2 (1976).Google Scholar
33. Guillaume Loubet, Nicolas Holzschuch, and Wenzel Jakob. 2019. Reparameterizing Discontinuous Integrands for Differentiable Rendering. ACM Trans. Graph. (SIGGRAPH Asia) 38, 6 (Nov. 2019).Google ScholarDigital Library
34. William R. Mark, R. Steven Glanville, Kurt Akeley, and Mark J. Kilgard. 2003. Cg: A System for Programming Graphics Hardware in a C-like Language. ACM Trans. Graph. (SIGGRAPH) 22, 3 (July 2003).Google ScholarDigital Library
35. Michael McCool, Stefanus Du Toit, Tiberiu Popa, Bryan Chan, and Kevin Moule. 2004. Shader Algebra. ACM Trans. Graph. (SIGGRAPH) 23, 3 (Aug. 2004).Google ScholarDigital Library
36. Michael D McCool, Zheng Qin, and Tiberiu S Popa. 2002. Shader metaprogramming. In Proceedings of the SIGGRAPH/EUROGRAPHICS conference on Graphics hardware.Google ScholarDigital Library
37. William S Moses and Valentin Churavy. 2020. Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients. (2020). arXiv:2010.01709Google Scholar
38. William S Moses, Valentin Churavy, Ludger Paehler, Jan Hückelheim, Sri Hari Krishna Narayanan, Michel Schanen, and Johannes Doerfert. 2021. Reverse-mode automatic differentiation and optimization of GPU kernels via enzyme. In Proceedings of SC21.Google ScholarDigital Library
39. Merlin Nimier-David, Sébastien Speierer, Benoît Ruiz, and Wenzel Jakob. 2020. Radiative Backpropagation: An Adjoint Method for Lightning-Fast Differentiable Rendering. ACM Trans. Graph. (SIGGRAPH) 39, 4 (July 2020).Google ScholarDigital Library
40. Merlin Nimier-David, Delio Vicini, Tizian Zeltner, and Wenzel Jakob. 2019. Mitsuba 2: A Retargetable Forward and Inverse Renderer. ACM Trans. Graph. (SIGGRAPH Asia) 38, 6 (Nov. 2019).Google ScholarDigital Library
41. John F Nolan. 1953. Analytical differentiation on a digital computer. Ph.D. Dissertation. Massachusetts Institute of Technology.Google Scholar
42. Melissa E. O’Neill. 2014. PCG: A Family of Simple Fast Space-Efficient Statistically Good Algorithms for Random Number Generation. Technical Report HMC-CS-2014-0905. Harvey Mudd College, Claremont, CA.Google Scholar
43. Steven G. Parker, James Bigler, Andreas Dietrich, Heiko Friedrich, Jared Hoberock, David Luebke, David McAllister, Morgan McGuire, Keith Morley, Austin Robison, and Martin Stich. 2010. OptiX: A General Purpose Ray Tracing Engine. ACM Trans. Graph. (SIGGRAPH) 29, 4 (July 2010).Google ScholarDigital Library
44. Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. 2017. Automatic differentiation in PyTorch. (2017).Google Scholar
45. Adam Paszke, Daniel D. Johnson, David Duvenaud, Dimitrios Vytiniotis, Alexey Radul, Matthew J. Johnson, Jonathan Ragan-Kelley, and Dougal Maclaurin. 2021. Getting to the Point: Index Sets and Parallelism-Preserving Autodiff for Pointful Array Programming. Proc. ACM Program. Lang. 5, ICFP (Aug. 2021).Google ScholarDigital Library
46. Barak A Pearlmutter and Jeffrey Mark Siskind. 2008. Reverse-mode AD in a functional framework: Lambda the ultimate backpropagator. ACM Transactions on Programming Languages and Systems (TOPLAS) 30, 2 (2008).Google ScholarDigital Library
47. Matt Pharr, Wenzel Jakob, and Greg Humphreys. 2020. Implementation of the forthcoming 4th edition of Physically Based Rendering: From Theory to Implementation. https://github.com/mmp/pbrt-v4Google Scholar
48. Matt Pharr and William R Mark. 2012. ispc: A SPMD compiler for high-performance CPU programming. In 2012 Innovative Parallel Computing (InPar). IEEE.Google Scholar
49. Lev Pontryagin. 1962. Mathematical theory of optimal processes. CRC Press.Google Scholar
50. Arsène Pérard-Gayot, Richard Membarth, Roland Leißa, Sebastian Hack, and Philipp Slusallek. 2019. Rodent: Generating Renderers without Writing a Generator. ACM Trans. Graph. (SIGGRAPH) 38, 4 (July 2019).Google ScholarDigital Library
51. Jonathan Ragan-Kelley, Andrew Adams, Sylvain Paris, Marc Levoy, Saman Amarasinghe, and Frédo Durand. 2012. Decoupling Algorithms from Schedules for Easy Optimization of Image Processing Pipelines. ACM Trans. Graph. (SIGGRAPH) 31, 4, Article 32 (July 2012).Google ScholarDigital Library
52. Amir Shaikhha, Andrew Fitzgibbon, Dimitrios Vytiniotis, and Simon Peyton Jones. 2019. Efficient Differentiable Programming in a Functional Array-Processing Language. Proc. ACM Program. Lang. 3, ICFP (July 2019).Google ScholarDigital Library
53. Jeffrey Mark Siskind and Barak A Pearlmutter. 2018. Divide-and-conquer checkpointing for arbitrary programs with no user annotation. Optimization Methods and Software 33, 4–6 (2018).Google ScholarCross Ref
54. Bert Speelpenning. 1980. Compiling fast partial derivatives of functions given by algorithms. Ph.D. Dissertation. University of Illinois at Urbana-Champaign.Google ScholarDigital Library
55. Delio Vicini, Sébastien Speierer, and Wenzel Jakob. 2021. Path Replay Backpropagation: Differentiating Light Paths Using Constant Memory and Linear Time. ACM Trans. Graph. (SIGGRAPH) 40, 4 (Aug. 2021).Google ScholarDigital Library
56. Yu M Volin and GM Ostrovskii. 1985. Automatic computation of derivatives with the use of the multilevel differentiating technique-1. algorithmic basis. Computers & mathematics with applications 11, 11 (1985).Google Scholar
57. Ingo Wald, Sven Woop, Carsten Benthin, Gregory S. Johnson, and Manfred Ernst. 2014. Embree: A Kernel Framework for Efficient CPU Ray Tracing. ACM Trans. Graph. (SIGGRAPH) 33, 4 (July 2014).Google ScholarDigital Library
58. Robert Edwin Wengert. 1964. A simple automatic derivative evaluation program. Commun. ACM 7, 8 (1964).Google Scholar
59. Tizian Zeltner, Sébastien Speierer, Iliyan Georgiev, and Wenzel Jakob. 2021. Monte Carlo Estimators for Differential Light Transport. ACM Trans. Graph. (SIGGRAPH) 40, 4 (Aug. 2021).Google ScholarDigital Library
60. Cheng Zhang, Bailey Miller, Kan Yan, Ioannis Gkioulekas, and Shuang Zhao. 2020. Path-space differentiable rendering. ACM Trans. Graph. (SIGGRAPH) 39, 4 (2020).Google ScholarDigital Library