PolyMorphous: An MLIR-Based Polyhedral Compiler with Loop Transformation Primitives

Jinman Zhao, Seyed Aryan Vahabpour, Xingyu Yue, Kai-Ting Amy Wang, Tarek S. Abdelrahman

Abstract: We present PolyMorphous, an MLIR-based polyhedral compiler that exposes a set of loop-based scheduling primitives, providing users with ample control over optimizations for input code. The primitives are expressed in a new Schedule dialect that is based on MLIR's Transform dialect. PolyMorphous' polyhedral engine collapses the primitives into a polyhedral schedule and checks for its legality. If necessary, and possible, the engine corrects an illegal schedule into a legal one. PolyMorphous is evaluated using the PolyBench suite. The evaluation validates PolyMorphous' approach in two ways. First, it shows that PolyMorphous enables users to explore the optimization space, and that it results in well-optimized code. Second, it shows that PolyMorphous' correction of illegal schedules enables users to optimize code with fewer primitives and improves their productivity by alleviating the need for manual correction. Specifically, the evaluation shows that PolyMorphous optimized code performs as well as code optimized by Pluto+ with its optimization flags tuned to achieve the best performance for each benchmark. Further, for several benchmarks, PolyMorphous optimized code performs better, confirming the value of providing users with control over optimizations. Averaged over all the benchmarks, PolyMorphous optimized code has a speedup of $1.24 \mathrm{x} / 1.25 \mathrm{x}$ over that optimized by Pluto+ on Arm/X86 systems. PolyMorphous brings the approach of empowering users of polyhedral compilers with control over optimizations to a community-developed infrastructure, promoting the approach's adoptability. It does so while delivering performant code.