TopoAxis: A Multi-Stage Compilation Optimization Framework for Neural Processing Units

Abstract

With the rapid growth of deep learning models, compiler optimization for heterogeneous accelerators has become critical to improving execution efficiency. Compared with GPUs, Neural Processing Units (NPUs) offer stronger domain-specific computing capability but impose more complex on-chip resource constraints. Existing compiler optimizations for general-purpose backends often face limited fusion-rule coverage, enlarged configuration search spaces, and high compilation overhead on NPUs. This paper proposes TopoAxis, an optimization framework for the TorchInductor NPU backend. TopoAxis first introduces a topology-weighted benefit analysis method to improve operator fusion by considering computation-graph topology and memory-access characteristics. It then presents an NPU-aware operator-generation optimization method that combines axis selection, axis splitting, hardware-resource constraint modeling, and priority-queue-based pruning to reduce invalid candidates and accelerate configuration search. Candidate reuse and multi-level configuration search are further incorporated to improve compilation efficiency. Experiments on BERT-Base, ResNet-18, and other models show that TopoAxis improves overall performance by 13.7% in training and 10.8% in inference on average. It also reduces the candidate search space by 42.82% on average and achieves model-level speedups of $2.82\times $ in inference and up to $3.79\times $ in training.