SC22 Proceedings

The International Conference for High Performance Computing, Networking, Storage, and Analysis

Technical Papers Archive

Scalable Automatic Differentiation of Multiple Parallel Paradigms through Compiler Augmentation

Authors: William Moses (Massachusetts Institute of Technology (MIT)); Sri Hari Krishna Narayanan (Argonne National Laboratory (ANL)); Ludger Paehler (Technical University Munich); Valentin Churavy (Massachusetts Institute of Technology (MIT)); and Michel Schanen, Jan Hueckelheim, Johannes Doerfert, and Paul Hovland (Argonne National Laboratory (ANL))

Abstract: Derivatives are key to numerous science, engineering, and machine learning applications. While existing tools generate derivatives of programs in a single language, modern parallel applications combine a set of frameworks and languages to leverage available performance and function in an evolving hardware landscape.

We propose a scheme for differentiating arbitrary DAG-based parallelism that preserves scalability and efficiency, implemented into the LLVM-based Enzyme automatic differentiation framework. By integrating with a full-fledged compiler backend, Enzyme can differentiate numerous parallel frameworks and directly control code generation. This flexibility permits Enzyme to leverage parallel and differentiation-specific optimizations far beyond existing tools.

We differentiate nine distinct versions of the LULESH and miniBUDE applications, written in different programming languages (C++, Julia) and parallel frameworks (OpenMP, MPI, RAJA, Julia tasks, MPI.jl), demonstrating similar scalability to the original program and with a differentiation overhead of 1.0-11.7x on 64 threads/nodes.

Back to Technical Papers Archive Listing