Automated Batching and Differentiation of Scalar Code in Enzyme

Abstract

Derivatives are the key to many important problems in computing, such as machine learning, and optimization. Building on the Enzyme compiler plugin for automatic differentiation, we add forward mode automatic differentiation, batching and the emission of vectorization-ready IR for arbitrary scalar code to unlock significant performance boosts.