Cross-Dialect Generalization Without Retraining: Benchmarks and Evaluation of Schema-Derived Constrained Decoding for MLIR
Not on the map. Why?
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.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 17 of its 21 references that are in the map. First posted on arXiv 14 May 2026.
Builds on
17 of 21 references.
References · 17
- 2021Evaluating Large Language Models Trained on Codecited by 8,053
- 2021Program Synthesis with Large Language Modelscited by 3,509
- 2023Code Llama: Open Foundation Models for Codecited by 2,105
- 2022Competition-Level Code Generation with AlphaCodecited by 1,427
- 2018Datasheets for Datasetscited by 1,371
- 2021Measuring Coding Challenge Competence With APPScited by 879
- 2024StarCoder 2 and The Stack v2: The Next Generationcited by 484
- 2023LeanDojo: Theorem Proving with Retrieval-Augmented Language Modelscited by 414
- 2025SmolLM2: When Smol Goes Big -- Data-Centric Training of a Small Language Modelcited by 307
- 2023VerilogEval: Evaluating Large Language Models for Verilog Code Generationcited by 246
- 2023Efficient Guided Generation for Large Language Modelscited by 192
- 2022Synchromesh: Reliable code generation from pre-trained language modelscited by 156
5 more
- 2024Granite Code Models: A Family of Open Foundation Models for Code Intelligencecited by 116
- 2024Croissant: A Metadata Format for ML-Ready Datasetscited by 79
- 2020MLIR: A Compiler Infrastructure for the End of Moore's Lawcited by 64
- 2025Type-Constrained Code Generation with Language Modelscited by 32
- 2025Correctness-Guaranteed Code Generation via Constrained Decodingcited by 7
Built on by
Not cited yet.