DyPARS: Dynamic-Shape DNN Optimization via Pareto-Aware MCTS for Graph Variants

Abstract

Dynamic-shape DNNs are widely used in applications such as variable-resolution image processing and language modeling with variable-length sequences. Existing DL (Deep-Learning) compilers apply rule-based rewriting to either transform a subgraph into a fixed variant at compile time (leading to suboptimal performance) or generate multiple variants at runtime, incurring significant overhead. The challenge is discovering and applying shape-dependent subgraph variants that maintain high efficiency across diverse inputs with minimal runtime cost.We propose DyPARS, a dynamic-shape DL compiler approach that discovers high-performance subgraph variants at compile time and applies the best ones at runtime. Leveraging Pareto-aware MCTS, DyPARS identifies shape-aware variants, incorporating shape-dependent kernel adaptations. These variants are integrated into a prediction-enhanced computational graph, enabling efficient variant selection based on input shapes with minimal overhead. DyPARS achieves average speedups of 1.31× and 1.80× over TorchInductor (JIT) and BladeDISC (non-JIT), respectively, across five DNN models, demonstrating robust efficiency across diverse inputs.