Stateful Testing for Deep Learning Compilers

Data flow underlies deep learning (DL) framework execution, including automatic differentiation and operator-fusion optimization. Bugs involving data flow can cause incorrect results or unexpected behavior. Existing testing methods represent models primarily by operator topology, omitting runtime tensor values and properties and operator input requirements. They therefore cannot distinguish data flows in models with similar topologies, even when these differences affect framework behavior. We introduce the Tensor Flow Graph (TFG) to capture runtime tensor states and operator input requirements. Its nodes represent operators and their input requirements, while its edges record tensor values and properties. Building on the TFG, we propose FlowMuT, a mutation-based testing method grounded in data flow. Mutations are represented as constrained subgraph replacements, with mutation operators derived from prior work and an analysis of mutable data flow objects. In each round, FlowMuT builds a TFG for the seed model, pairs mutation operators with TFG sites, and filters illegal candidates using data flow constraints. It uses LinUCB to select a candidate based on coverage and impact, executes the mutant, updates selection feedback, and uses each accepted mutant as the next seed. We evaluate FlowMuT on PyTorch and MindSpore using representative models from diverse application scenarios. FlowMuT submitted 55 bug reports for cases missed by all baseline methods; developers confirmed 41 reports and fixed the bugs described in 7 reports. FlowMuT achieves 43.4% higher data flow diversity than the strongest baseline and a legal mutant rate of 96.7%, compared with 86.4% without constraints. These results show that testing from a data flow perspective is practical and effective.

Performance optimization of AI infrastructure is key to the fast adoption of large language models (LLMs). The PyTorch compiler (torch.compile), a core optimization tool for deep learning (DL) models (including LLMs), has received due attention. However, torch.compile is prone to correctness bugs, which cause incorrect outputs of compiled DL models without triggering exceptions, crashes, or warnings. These bugs pose a serious threat to the reliability of downstream LLM applications. Data from the PyTorch community shows that 19.2% of high-priority issues are incorrect outputs of compiled DL models induced by torch.compile bugs, the second-most-common bug category (only behind program crashes at 19.57%). However, no systematic study has been conducted to specifically characterize and thereby detect these bugs. In this paper, we present the first empirical study of the correctness bugs in torch.compile, examine their characteristics, and assess the effectiveness of existing fuzzers in detecting them. Based on our findings, we propose a proof-of-concept testing technique named AlignGuard, tailored specifically for detecting correctness bugs in torch.compile. AlignGuard incorporates bug characteristics distilled from our empirical study, applying LLM-based test mutation to existing test cases for correctness bug detection. At the time of writing, AlignGuard has successfully detected 23 new correctness bugs in recent torch.compile. All these bugs have been confirmed or fixed by the PyTorch development team, and over half (14/23) of them are even marked as high-priority bugs, underscoring the usefulness of our technique.

Tensor programs, as used in deep learning models, are a prime target for optimization, as small performance improvements can have a large impact across training or inference workloads. However, such optimizations are complicated and can produce subtle bugs. Traditionally, correctness is assumed when differential testing against a reference on random inputs fails to reveal bugs. However, the inputs to these programs are massive tensors, and finding bugs can require generating extremely low likelihood inputs with precise relationships among their values. We propose a novel way to find bugs more consistently by flipping the quantifiers. Rather than generating a single input and checking all output tensor locations for equivalence, what if you could check a single output tensor location's equivalence for all inputs? We implement this idea in a system, \dirigo, by using a novel symbolic execution strategy. We demonstrate that \dirigo can find bugs effectively in a public dataset of 6,988 AI-written CUDA kernels that are all marked correct by differential testing. Of these, \dirigo finds 600 kernels that are actually buggy, and finds 97.3% of those bugs within two minutes.