Cross-Dialect Generalization Without Retraining: Benchmarks and Evaluation of Schema-Derived Constrained Decoding for MLIR

Multi-Level Intermediate Representation (MLIR) underlies modern ML compiler infrastructure, including TensorFlow, JAX via StableHLO, PyTorch's Inductor backend, and IREE. Yet it appears only in trace amounts in code-LM pretraining corpora. MLIR is also extensible by design: new dialects ship per application domain, so maintaining a fine-tuned model per dialect does not scale. We ask whether inference-time priors derived mechanically from each dialect's Operation Definition Specification (ODS) can substitute for gradient-based adaptation. We make two contributions. First, we release four natural-language-to-MLIR benchmarks across three dialects (MLIR-Spec-150, Linalg-Spec-30, StableHLO-Spec-30, and StableHLO-Held-Out-200), totaling 410 in-scope NL→MLIR pairs, plus a 25-program StableHLO-OutOfGrammar stress set (435 benchmark instances). Thirty of the Spec prompts, ten per dialect, also carry hand-authored input and expected-output references for functional testing (MLIR-Functional-Reference-30). All artifacts ship under Apache-2.0 with Gebru datasheets and Croissant 1.1 metadata. Second, on top of these benchmarks we build a three-layer schema-derived constraint stack: a context-free grammar over op signatures (C1), type-domain splits from an ODS-extracted type lattice (C2), and an SSA-scope validator driving five-retry rejection sampling (C3). Porting the stack from arith+func+memref+linalg to StableHLO required no new constraint-layer code. Empirically, on dialects whose verifier semantics are dominated by structural constraints, SmolLM2-1.7B with token-level schema-derived constraints matches or exceeds 15B–34B code LMs that run the strongest constraint enforcement their serving stack admits (post-hoc rejection), at a matched retry budget and 8–25× the per-generation speed: on linalg it reaches 80.0% verify-valid (three runs, each 80.0%, n=125), 21 to 44 pp above CodeLlama-34B, Granite-Code-34B, and StarCoder2-15B with non-overlapping CIs, and the lead survives a same-family fp16 precision control. The comparison is between systems, not models: under the baselines' post-hoc protocol, SmolLM2 itself falls from 52.0% to 24.5% on arith+func, a +27.5 pp masking contribution. On arith+func and on the templated parametric StableHLO-Held-Out-200, where verifier semantics turn on attribute values rather than structure, the same baselines match or beat the SLM; we scope these explicitly as non-win cells. We release benchmarks, decoder, every per-prompt generation, and a pinned reproducibility Dockerfile.