Tapeflow: Streaming gradient tapes in automatic differentiation

Computing gradients is a fundamental task in a variety of fields, ranging from machine learning to physics simulations and scientific computing. Automatic differentiation (AD) is a widely used technique for computing gradients of arbitrary imperative code. In AD, the gradient tape is an auxiliary data structure that is used to transfer intermediary values required for gradient computation. However, managing the gradient tape in memory presents a significant challenge, as existing approaches suffer from limitations in spatial reuse, on-chip energy consumption, and cache size requirements. These limitations highlight the need for a new and efficient solution. We present Tapeflow, a novel compiler framework that efficiently orchestrates and manages the gradient tape. Our approach offers three key contributions. First, we introduce the concept of regions, which transform the tape layout into an array-of-structs format, enhancing spatial reuse. Second, we schedule the execution into layers and explicitly orchestrate the tape operands using a scratchpad, effectively reducing required cache size and on-chip energy. Third, we stream the tape from the DRAM by organizing it into a FIFO of tiles, with tape operands arriving just-in-time for each layer. To evaluate the effectiveness of our approach, we conducted experiments on a wide range of algorithms across various domains. Our results show that Tapeflow outperforms Enzyme, the state-of-the-art compiler, by 1.3-2.5×, reduces on-chip SRAM usage by 5-40×, and saves 8× on-chip energy. By providing a general solution that does not rely on domain-specific knowledge, Tapeflow has the potential to significantly improve gradient tape management in various applications.