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

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.

References

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).
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).
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).
4. Benedikt Bitterli. 2016. Rendering resources. https://benedikt-bitterli.me/resources/.
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/jax
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).
7. Brent Burley. 2012. Physically Based Shading at Disney. In ACM SIGGRAPH Talks.
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.
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.0759
10. Conal Elliott. 2018. The Simple Essence of Automatic Differentiation. Proc. ACM Program. Lang. 2, ICFP (July 2018).
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).
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.
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.
14. Google. 2017. XLA: Optimizing Compiler for Machine Learning. https://www.tensorflow.org/xla (Accessed Jan 16, 2022).
15. Andreas Griewank et al. 1989. On automatic differentiation. Mathematical Programming: recent developments and applications 6, 6 (1989).
16. Andreas Griewank and Andrea Walther. 2008. Evaluating derivatives: principles and techniques of algorithmic differentiation. Vol. 105. SIAM.
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.
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).
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).
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).
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).
22. Michael Innes. 2019. Don't Unroll Adjoint: Differentiating SSA-Form Programs. arXiv:1810.07951
23. Wenzel Jakob. 2019. Enoki: structured vectorization and differentiation on modern processor architectures. https://github.com/mitsuba-renderer/enoki. (Accessed: Jan 16, 2022).
24. James T. Kajiya. 1986. The Rendering Equation. SIGGRAPH Comput. Graph. 20, 4 (Aug. 1986).
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.
26. Chris Lattner and Vikram Adve. 2004. LLVM: A Compilation Framework for Lifelong Program Analysis and Transformation. San Jose, CA, USA.
27. João Rui Leal and Brad Bell. 2017. CppAD 2017.
28. Mark Lee, Brian Green, Feng Xie, and Eric Tabellion. 2017. Vectorized production path tracing. In Proceedings of High Performance Graphics. ACM.
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).
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).
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).
32. Seppo Linnainmaa. 1976. Taylor expansion of the accumulated rounding error. BIT Numerical Mathematics 16, 2 (1976).
33. Guillaume Loubet, Nicolas Holzschuch, and Wenzel Jakob. 2019. Reparameterizing Discontinuous Integrands for Differentiable Rendering. ACM Trans. Graph. (SIGGRAPH Asia) 38, 6 (Nov. 2019).
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).
35. Michael McCool, Stefanus Du Toit, Tiberiu Popa, Bryan Chan, and Kevin Moule. 2004. Shader Algebra. ACM Trans. Graph. (SIGGRAPH) 23, 3 (Aug. 2004).
36. Michael D McCool, Zheng Qin, and Tiberiu S Popa. 2002. Shader metaprogramming. In Proceedings of the SIGGRAPH/EUROGRAPHICS conference on Graphics hardware.
37. William S Moses and Valentin Churavy. 2020. Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients. (2020). arXiv:2010.01709
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.
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).
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).
41. John F Nolan. 1953. Analytical differentiation on a digital computer. Ph.D. Dissertation. Massachusetts Institute of Technology.
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.
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).
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).
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).
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).
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-v4
48. Matt Pharr and William R Mark. 2012. ispc: A SPMD compiler for high-performance CPU programming. In 2012 Innovative Parallel Computing (InPar). IEEE.
49. Lev Pontryagin. 1962. Mathematical theory of optimal processes. CRC Press.
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).
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).
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).
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).
54. Bert Speelpenning. 1980. Compiling fast partial derivatives of functions given by algorithms. Ph.D. Dissertation. University of Illinois at Urbana-Champaign.
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).
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).
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).
58. Robert Edwin Wengert. 1964. A simple automatic derivative evaluation program. Commun. ACM 7, 8 (1964).
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).
60. Cheng Zhang, Bailey Miller, Kan Yan, Ioannis Gkioulekas, and Shuang Zhao. 2020. Path-space differentiable rendering. ACM Trans. Graph. (SIGGRAPH) 39, 4 (2020).

Index Terms

• Computing methodologies

  • Computer graphics

    • Rendering

• Mathematics of computing

  • Mathematical analysis

    • Numerical analysis

      • Automatic differentiation

• Software and its engineering

  • Software notations and tools

    • Compilers

      • Just-in-time compilers