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

arXiv: https://arxiv.org/abs/2607.18254 · HTML: https://mlsanity.com/paper/2607.18254/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 17 of its 21 references that are in the map.
- First posted on arXiv 14 May 2026.

## Builds on (most cited first, 17 of 17)

- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [Code Llama: Open Foundation Models for Code](https://mlsanity.com/paper/2308.12950/) (2023)
- [Competition-Level Code Generation with AlphaCode](https://mlsanity.com/paper/2203.07814/) (2022)
- [Datasheets for Datasets](https://mlsanity.com/paper/1803.09010/) (2018)
- [Measuring Coding Challenge Competence With APPS](https://mlsanity.com/paper/2105.09938/) (2021)
- [StarCoder 2 and The Stack v2: The Next Generation](https://mlsanity.com/paper/2402.19173/) (2024)
- [LeanDojo: Theorem Proving with Retrieval-Augmented Language Models](https://mlsanity.com/paper/2306.15626/) (2023)
- [SmolLM2: When Smol Goes Big -- Data-Centric Training of a Small Language Model](https://mlsanity.com/paper/2502.02737/) (2025)
- [VerilogEval: Evaluating Large Language Models for Verilog Code Generation](https://mlsanity.com/paper/2309.07544/) (2023)
- [Efficient Guided Generation for Large Language Models](https://mlsanity.com/paper/2307.09702/) (2023)
- [Synchromesh: Reliable code generation from pre-trained language models](https://mlsanity.com/paper/2201.11227/) (2022)
- [Granite Code Models: A Family of Open Foundation Models for Code Intelligence](https://mlsanity.com/paper/2405.04324/) (2024)
- [Croissant: A Metadata Format for ML-Ready Datasets](https://mlsanity.com/paper/2403.19546/) (2024)
- [MLIR: A Compiler Infrastructure for the End of Moore's Law](https://mlsanity.com/paper/2002.11054/) (2020)
- [Type-Constrained Code Generation with Language Models](https://mlsanity.com/paper/2504.09246/) (2025)
- [Correctness-Guaranteed Code Generation via Constrained Decoding](https://mlsanity.com/paper/2508.15866/) (2025)

## Abstract (arXiv)

Multi-Level Intermediate Representation (MLIR) underlies modern ML compiler infrastructure (TensorFlow, JAX/StableHLO, PyTorch Inductor, IREE), yet appears only in trace amounts in code-LM pretraining corpora. MLIR is also extensible by design: new dialects ship per application domain, so 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. 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-to-MLIR pairs, plus a 25-program out-of-grammar stress set and a hand-authored n=30 functional reference set, shipped under Apache-2.0 with Gebru datasheets and Croissant 1.0 metadata. Second, we build a three-layer schema-derived constraint stack: a CFG 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 from arith+func+memref+linalg to StableHLO required no new constraint-layer code. On dialects whose verifier semantics are dominated by structural constraints, schema-derived priors let SmolLM2-1.7B match or exceed 15B-34B code LMs at 8-25x the per-generation speed: on linalg, SmolLM2 reaches 80.0% verify-valid (three-seed mean, n=125), beating CodeLlama-34B, Granite-Code-34B, and StarCoder2-15B by 21-44 percentage points with non-overlapping CIs. 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 as non-win cells. We release benchmarks, decoder, all per-prompt generations, and a reproducibility Docker image.