Poster
in
Workshop: CODEML: Championing Open-source DEvelopment in Machine Learning
Control Flow Operators in PyTorch
Yidi Wu · Thomas Ortner · Edward Yang · Horace He · Yanan Cao
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 cond and reduce compilation time and peak memory usage when rewriting loops with map, scan, associativescan and whileloop.