IIntroduction to AI Systems Performance EngineeringTechnical textbook
Chapter 12: ML Compiler Architecture and Optimization
12.4

Debugging the compiler stack

Locate the first compiler stage that causes a correctness or performance divergence

Before you begin

Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.

  • Graph capture
  • IR levels
  • Performance evidence levels

A model runs correctly in PyTorch eager mode but becomes 18 percent slower after torch.compile for sequence length 127. Compiler logs show one captured graph, so graph fragmentation is not the cause. The optimized graph reveals that a generated narrow matrix schedule replaced a cuBLASLt call. Nsight Compute confirms lower matrix throughput and more register pressure in the generated kernel. The first differing compiler decision is therefore library selection, not a vague compiler problem.

Compiler debugging follows one source program through named artifacts: capture logs, guards, high-level graph IR, optimized IR, generated GPU code, profiler timelines, counter reports, and numerical outputs. Each artifact is produced by a declared tool or compiler option. Comparing adjacent stages locates the earliest divergence and prevents a learner from opening a huge final assembly file without knowing which decision to investigate.

One source expression can trigger Python tracing, graph partitioning, symbolic-shape guards, tensor rewrites, fusion, autotuning, code generation, assembly, cache lookup, and runtime launch. A report that only states that the compiler is slow does not identify which phase requires investigation.

Debugging proceeds like performance profiling: start at the symptom, divide time and behavior by phase, then inspect the lowest representation necessary. Preserve a known eager or library baseline, isolate the smallest reproducer that retains the failure, and compare graphs and generated artifacts before changing compiler flags at random.

A useful investigation produces a first divergent stage. Before that stage, the reference and compiled path have the expected structure or result. At that stage, a guard, graph partition, rewrite, schedule, instruction selection, or runtime event differs. The first divergence narrows both a workaround and an upstream report.

Classify the observed failure

Compilation latency, repeated compilation, slow device execution, excessive memory, graph breaks, crashes, and wrong results are distinct failure classes. Measure cold and warm execution separately. Use a host-device timeline to identify whether time is spent tracing, compiling, waiting on the host, launching kernels, or executing kernels.

For memory growth, separate compiler caches, generated-module memory, persistent runtime allocations, and peak device buffers. For wrong results, determine whether the failure depends on optimization level, shape, datatype, or repeated execution. A precise symptom determines which artifacts should be inspected next.

Follow performance through the compiler stages

Begin with the source region and captured graph. Confirm graph boundaries, operators, shapes, and guards. Inspect the optimized graph for expected fusion, decomposition, layout conversion, and library selection. Then inspect generated kernel source or IR and the native profile for transactions, instructions, occupancy, and stalls.

At each compiler stage, compare one expectation with evidence. If launch count is high, find the graph or fusion split before reading assembly. If one generated kernel is long, inspect the schedule and native behavior. If the device is idle between kernels, inspect host orchestration or compilation. The diagnostic order avoids collecting low-level detail for the wrong phase.

Find wrong-code failures by controlled comparison

Establish a trusted eager or high-precision reference and record the smallest input that fails. Compare outputs after meaningful subregions when possible. Disable or isolate passes systematically to identify the first transformation after which results diverge. Record the exact setting for every comparison.

Preserve numerical expectations. A small floating-point difference can be valid under a changed reduction order, while a shape-dependent out-of-bounds value is not. Test determinism, non-finite inputs, boundary dimensions, and repeated launches. A device memory checker can distinguish an invalid access from an algebraic rewrite error.

Minimize without removing the cause

Remove unrelated model layers and inputs while keeping the decisive shape, stride, datatype, dynamic guard, and consumer structure. A transpose view can be the reason a layout conversion appears. A second consumer can be the reason fusion stops. Removing either can produce a small program that no longer reproduces the problem.

Automate the reproducer and run it in a clean process when cache or initialization state matters. State the software versions, target device, command, input construction, expected result, observed result, and measurement method. Include compact graph, IR, generated-code, and profiler excerpts that locate the first divergence.

Separate a local workaround from the compiler defect

A user can often restore performance by changing a graph boundary, forcing a library call, padding a shape, or disabling one optimization for a region. The local change can be a valid operational workaround. The workaround does not replace identifying the cost-model error, missing legality rule, or code-generation defect that produced the regression.

An upstream report should state both. Show the workaround and its scope, then provide the smaller compiler-level case and evidence. Separating workaround from root cause lets maintainers repair the general decision while application owners deploy a bounded change. Re-test the end-to-end workload after a compiler fix because a local microbenchmark improvement may move another bottleneck.

Worked example: Diagnosing a compiled regression

Compiled execution is 20% slower than eager execution for one model block even after warm-up.

  1. Confirm equivalent outputs and stable warmed timing. Use a timeline to see whether the loss comes from more launches, longer kernels, or host gaps.
  2. Record cold compilation separately and run enough warmed iterations to show that the 20% loss belongs to generated execution.
  3. Compare graphs. The compiler fused an elementwise chain but replaced a tuned library GEMM with a generated schedule for an awkward narrow shape.
  4. Profile the generated GEMM and library baseline. Show lower matrix throughput and resource pressure, while saved elementwise traffic is insufficient to compensate.
  5. Adjust the partition or force the library path, then verify both local and full-model improvement. Preserve a minimal reproducer for the cost-model gap.
  6. Attach the narrow matrix dimensions, generated schedule, library timing, compiler timing, and surrounding fusion. The recorded artifacts let a cost-model owner reproduce the incorrect choice.

Conclusion: The evidence identifies a profitability and library-selection problem, not a generic failure of compilation. The precise classification guides both workaround and upstream fix.

Common errors

  • Disabling optimization passes until performance changes without recording mechanisms. Pass toggling finds correlations but rarely produces a maintainable fix.
  • Comparing cold compiled time with warm eager time. Compilation and execution are separate operating modes.
  • Minimizing away the decisive shape or consumer. Compiler behavior depends on context and guards, not source text alone.
  • Starting with the final assembly for a graph-fragmentation problem. The missing optimization opportunity disappeared before target code generation.

Reference summary

  • Minimize the program while preserving the miscompile or regression.
  • Dump graphs and IR before and after relevant passes.
  • Compare eager/reference outputs and use higher precision or interpreters when available.
  • Bisect passes, flags, commits, shapes, or fusions to find the first divergence.
  • Separate capture overhead, compile time, generated-kernel time, and runtime orchestration.
  • Attach a standalone reproducer and environment to every report.

XLA tooling, PyTorch compiler logs, MLIR pass infrastructure, and generated Triton/PTX dumps all support this isolation style. The essential skill is controlling the experiment across levels rather than reading an enormous final IR and guessing.

Knowledge check

What is the smallest set of artifacts you would attach to a compiler performance regression report?

Show the answer guide
  • Attach a minimal executable reproducer with fixed inputs or an input generator, together with the exact shapes, dtypes, strides, devices, and compiler options required to trigger the regression.
  • Record the software and hardware environment and provide before-and-after timing distributions with warm-up, synchronization, and the comparison revision identified. A single unsynchronized timing is not sufficient evidence.
  • Include capture and compiler diagnostics that locate the first changed stage: graph-break and guard logs plus the smallest relevant IR or generated-code diff.
  • Attach one runtime trace or profiler report that connects the compiler change to launches, memory traffic, occupancy, spills, collectives, or host gaps. State correctness results so a fast but wrong path is excluded.
  • The smallest useful report is therefore reproducer, environment, benchmark, first divergent compiler artifact, runtime evidence, and correctness status; larger dumps without this chain are not a substitute.
Practice

Assemble a minimal compiler regression report

Create a controlled regression in a compiled PyTorch function by changing one shape assumption, disabling one fusion, or selecting a deliberately poor schedule for a declared input such as [1024,4096]. Reduce the program until removing any remaining operation eliminates the slowdown, then compare the good and bad configurations over 100 synchronized iterations.

Evidence to keep: Submit the minimal reproducer, environment manifest, correctness test, before-and-after latency distribution, relevant compiler logs and IR diff, and one profiler trace that explains the runtime consequence. Finish with a causal statement that names the first divergent stage and the resulting measured cost.

Part IV extends the work, byte, latency, and overlap calculations from one accelerator to distributed execution.

References

Sources and further reading

The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.

  1. torch.compile Programming Model
  2. MLIR Language Reference
  3. MLIR GPU Dialect
  4. XLA Architecture
  5. XLA:GPU Architecture Overview
  6. Nsight Compute Profiling Guide