NNSmith: Generating Diverse and Valid Test Cases for Deep Learning Compilers

To find bugs in deep-learning compilers, NNSmith grows random neural networks that are guaranteed valid by a constraint solver, then tunes their weights by gradient descent until no NaN or infinity can blur the verdict — and catches the compiler red-handed by comparison.

Why this problem is worth solving

A deep-learning compiler takes a trained model — represented as a computation graph whose nodes are tensor operators like convolution, matrix multiplication and reshape, and whose edges carry tensors of specific shapes and element types — and rewrites it into a faster executable. TVM alone applies up to 61 graph-level passes (constant folding, operator fusion) and 58 low-level ones (loop tiling, arithmetic simplification). These compilers now sit in front of interactive and safety-critical deployments.

Compiler bugs come in two flavors, and the second is the frightening one. A crash is loud and diagnosable. A miscompilation is silent: the compiled model runs happily and returns numbers that are simply not what the input model computes. Since neural network outputs are approximate by nature, nobody eyeballing predictions will notice a subtly wrong optimization pass — the error just becomes part of the application’s behavior.

The established weapon against compiler bugs is fuzzing plus differential testing: generate random programs, run them through the system under test and through a trusted reference, and compare. Decades of this (Csmith and its descendants) hardened C compilers. Doing the same for deep-learning compilers stumbles on two problems specific to the domain, and NNSmith’s contribution is solving both.

What people did before, and where it broke

The validity–diversity tension. A computation graph must “type check”: each operator constrains its input shapes and attributes (matrix multiplication needs inner dimensions to agree; reshape needs element counts to match). Random wiring violates these rules, and the compiler rejects the graph before any optimization pass runs — a wasted test. Prior fuzzers bought validity by sacrificing diversity. LEMON mutates existing pre-trained models but only with shape-preserving single-input operators — and running real models makes it up to 103× slower per test. GraphFuzzer connects operators freely but “fixes” shape mismatches by inserting slice and pad nodes. The paper shows what that costs with a real example: a Conv2d output of shape (1,2,1,48) added to a constant of shape (1,1,48) — a broadcast — triggers a layout-analysis crash in TVM. Insert a slice to align the shapes, and the bug vanishes. The patch that guarantees validity is precisely what hides the bug.

The NaN problem. Test verdicts come from comparing outputs. But floating-point arithmetic has exceptional values — NaN (“not a number”, from things like the square root of a negative) and infinity (from overflow) — and once one appears anywhere in a graph, comparisons stop meaning anything: semantically equivalent implementations may legitimately disagree on how exceptional values propagate. With PyTorch’s default random initialization, 56.8% of 20-node models hit NaN or infinity somewhere. Prior tools ignored this, and up to 41% of their bug reports could be false alarms — while real bugs also hide behind the noise: the paper shows a genuine TVM semantic bug masked because an infinity swallowed the evidence. Subtler still: the poison can be internal — an ArgMax can output normal values from NaN inputs — so checking only the final output for NaNs is not enough.

The idea

Specifications: teach the fuzzer the type rules, cheaply

For each operator, the user writes a small specification with three parts: the accepted input/output tensor types (data type and rank), a requires function returning validity constraints as logical predicates (kernel size positive, kernel fits within padded input, …), and a type_transfer function computing output shapes from input shapes symbolically. Shapes are symbolic integers — variables, not numbers — so the same specification describes every possible size at once. The engineering cost is deliberately tiny: 59 of NNSmith’s 73 specifications are at most four lines; even Conv2d, the worst, needs nine inequalities and seven lines. Written once, the specs serve every compiler that accepts ONNX models.

Generation: grow the graph, let a solver keep it legal

NNSmith builds each model incrementally, starting from a placeholder node and repeatedly inserting a randomly chosen operator, in one of two directions. Forward insertion attaches the new operator downstream of existing outputs; backward insertion replaces a placeholder, growing the graph upstream — together they produce multi-input, multi-output topologies like real multi-modal models. Each insertion point is first filtered by a cheap type-matching heuristic (don’t ask the solver about a boolean edge feeding a float operator), then the accumulated constraints — the new operator’s requirements joined with everything already in the graph — go to an SMT solver (Z3), a program that decides whether a set of logical constraints has a solution and produces one. Constraints are cached and solved incrementally, so each insertion pays only for what changed.

One subtlety earns its own mechanism: solvers are lazy and return boundary values — ask for shapes ≥ 1 and every dimension comes back exactly 1, trivial models that miss bugs (the paper’s Conv2d-broadcast crash needs a non-trivial constant shape to fire). Attribute binning adds randomized range constraints, choosing for each attribute an exponentially sized bucket (like AFL’s hit-count buckets) and sampling a sub-range from it — with a retry-and-drop scheme if the extra constraints make the system unsatisfiable. Binning alone yields 2.07× more unique operator instances and roughly doubles unique branch coverage.

Numerical validity: descend until nothing explodes

Rather than rewriting vulnerable operators to be safe (which changes the graphs being tested — the diversity mistake again), NNSmith fixes the numbers. Each vulnerable operator’s safe domain is written as one or two tensor inequalities — Sqrt needs X ≥ 0; Pow(X, Y) is under-approximated as X > 0 and Y·log(X) ≤ 40 — and each inequality converts mechanically into a loss function that is positive exactly when violated. The search runs the model, finds the first operator (in topological order) producing NaN or infinity, and takes a gradient step on that operator’s loss with respect to all weights and inputs, using the Adam optimizer; where gradients are zero or undefined (ReLU’s flat side, Floor’s steps), hand-chosen proxy derivatives keep the descent moving. Iterate until the whole graph runs clean.

The economics are decisive: the search succeeds for 98% of generated models in 3.5 milliseconds on average — 4.2% of the 83 milliseconds it takes to generate a 10-node model — and beats random sampling by growing margins (1.16–1.34×) as models get bigger.

The verdict: differential testing with a careful oracle

The compiled model runs on the searched inputs, and outputs are compared against PyTorch — whose results the search already computed for free, and whose interpreter is the best-tested implementation available. When PyTorch and the compiler disagree, one more check assigns blame: recompile at optimization level zero. If the unoptimized compiler agrees with PyTorch, the optimization passes are guilty.

graph: [placeholder p0]   constraints: {}

Why it works

Because validity is encoded as solver constraints rather than achieved by restricting operators or patching shapes, the generator explores the full combinatorial space of operator connections — including the broadcasts, multi-input joins and odd attribute values where transformation bugs live. → Because specifications are symbolic, one four-line spec covers every shape an operator can take, and the solver — not hand-written enumeration — finds concrete instantiations. → Because binning forces the solver off its boundary-value habits, attribute diversity comes for free on top of topological diversity. → Because gradient search makes 98% of models numerically clean before any comparison, disagreements between compiler and reference are actual evidence — not floating-point noise — cutting both false alarms and masked bugs. → Therefore each generated test is simultaneously acceptable to the compiler, unusual enough to reach rarely exercised code, and judgeable against a reference — and 72 real bugs followed.

What it costs

• Someone writes the specifications. The framework and meta-types make each one cheap, but coverage of an operator library is still a human effort, and new compilers with exotic operators need new specs.
• Constraint solving isn’t free. Generation runs slower per test than shape-patching fuzzers; NNSmith wins on coverage per test, not tests per second.
• The oracle is imperfect. PyTorch is trusted but not infallible — one of the 72 bugs surfaced ambiguity in PyTorch’s own converter, which is why disagreements get re-checked against the compiler’s unoptimized mode.
• Vulnerable-operator domains are approximations. Safe domains are under-approximated by simple inequalities; models with genuinely no valid inputs make the search fail (the main failure mode within its 2%).
• Small models by design. Default graphs of 10–20 nodes keep fuzzing throughput high; bugs that need very deep or very large models to manifest are outside this net.

The evaluation

Four systems under test — TVM, ONNXRuntime, TensorRT, and PyTorch (via its ONNX path) — over roughly seven months of fuzzing, plus controlled 4-hour coverage runs against the state-of-the-art generators LEMON and GraphFuzzer and the TVM-specific fuzzer Tzer.