Control Flow Operators in PyTorch

Yidi Wu, Thomas Ortner, Richard Zou, Edward Z. Yang, Adnan Akhundov, Horace He, Yanan Cao

Keywords: controlflow operators, PyTorch, machine learning, triton, compile

TL;DR: We introduce a controflow operator library that allows to holistically compile models with complex control flow structures

Abstract:

Representing control flows in machine learning (ML) compilers and intermediate representations (IRs) has been a long-standing problem, with many ML compilers opting to avoid supporting them altogether. A common practice is to trace the model and get a straight-line computational graph, for example, by specializing if predicates or by unrolling loops. Although this strategy may work well in some cases, it can result in performance issues and long compilation times. Furthermore, it is problematic in cases where the control flow of the program depends on data or a dynamic shape.

In this paper, we present the PyTorch control flow operator library, which addresses the challenge through the introduction of five control flow operators to PyTorch. We will explain their usage, present real world use cases and demonstrate their benefits such as the ability to capture dynamic control flows in the PyTorch 2 IR with $\texttt{cond}$ and reduce compilation time and peak memory usage when rewriting loops with $\texttt{map}$, $\texttt{scan}$, $\texttt{associative\\_scan}$ and $\texttt{while\\_loop}$.

Submission Number: 50