PolyTOPS: Reconfigurable and Flexible Polyhedral Scheduler

TL;DR

PolyTOPS introduces a configurable, iterative polyhedral scheduler to address the rigidity of existing schedulers across architectures. It provides JSON and C++ interfaces to define per-dimension cost functions, constraints, and fusion/distribution strategies, enabling scenario- and kernel-specific optimizations. The approach integrates with MindSpore AKG and backends like isl/CLooG, achieving substantial speedups on Ascend NPU custom operators and competitive performance on PolyBench and PolyMage benchmarks. This configurability fills a gap in current polyhedral tooling, offering a practical path to architecture-aware optimization with potential for kernel-size and hardware-specific refinements. The work suggests that carefully designed fusion and strategy customization can surpass traditional schedulers in diverse settings, paving the way for broader adoption of configurable polyhedral optimization.

Abstract

Polyhedral techniques have been widely used for automatic code optimization in low-level compilers and higher-level processes. Loop optimization is central to this technique, and several polyhedral schedulers like Feautrier, Pluto, isl and Tensor Scheduler have been proposed, each of them targeting a different architecture, parallelism model, or application scenario. The need for scenario-specific optimization is growing due to the heterogeneity of architectures. One of the most critical cases is represented by NPUs (Neural Processing Units) used for AI, which may require loop optimization with different objectives. Another factor to be considered is the framework or compiler in which polyhedral optimization takes place. Different scenarios, depending on the target architecture, compilation environment, and application domain, may require different kinds of optimization to best exploit the architecture feature set. We introduce a new configurable polyhedral scheduler, PolyTOPS, that can be adjusted to various scenarios with straightforward, high-level configurations. This scheduler allows the creation of diverse scheduling strategies that can be both scenario-specific (like state-of-the-art schedulers) and kernel-specific, breaking the concept of a one-size-fits-all scheduler approach. PolyTOPS has been used with isl and CLooG as code generators and has been integrated in MindSpore AKG deep learning compiler. Experimental results in different scenarios show good performance: a geomean speedup of 7.66x on MindSpore (for the NPU Ascend architecture) hybrid custom operators over isl scheduling, a geomean speedup up to 1.80x on PolyBench on different multicore architectures over Pluto scheduling. Finally, some comparisons with different state-of-the-art tools are presented in the PolyMage scenario.

Figures (9)

• Figure 0: Left: original for-loop code. It is fully parallel and suitable for GPU architecture. The Right one is optimal for vectorization of the NPU (for both vectorized data loading and computing) thanks to the loop distribution and the interchange for the statement 0. Pluto or isl may be able to find the loop distribution (specifying the correct fusion heuristic), but no interchange would be found.
• Figure 1: PolyTOPS workflow representation, showing the major blocks, including the post-processing and the configuration. We support both Openscop and isl-representation as Polyhedral Representation
• Figure 2: JSON example showing most of the configurable features of PolyTOPS, including cost function control and definition, constraints definition, fusion control and directives
• Figure 3: An isl style configuration described using the C++ configuration. Feautrier is used as the fallback case when the Proximity fails to extract parallelism
• Figure 3: Custom Operator example.