A Planner for Scalable Tensor Programs
Tanvir Ahmed Khan, Leonidas Fegaras · 2024
Current machine learning systems, such as TensorFlow and PyTorch, rely on high-performance linear algebra libraries for efficient tensor computations. Although they provide numerous fine-tuned array algorithms based on well-studied data placement and communication patterns, these libraries are hard to customize to capture irregular array programs and unconventional array storages. We present a framework for constructing distributed task workflows from ad-hoc tensor programs by partially evaluating these programs against the block coordinates of the tensors. In addition, we present a novel task scheduler based on pattern matching that assigns processes to tasks by recognizing certain patterns inside the task workflow. Although each such pattern applies to a small fixed number of tasks, when applied collectively, these patterns generate communication schemes that resemble optimal block-based algorithms, such as SUMMA. The tiling of the task workflow is based on pattern-matching and is done bottom-up, guided by cost.