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

Plawan Kumar Rath

cs.AIMay 14, 2026

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

Authors: Plawan Kumar Rath

Organizations: Meta

Abstract

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, StableHLO-Held-Out-200), totaling 410 in-scope NL-to-MLIR pairs, plus a 25-program out-of-grammar stress set, with functional references on 30 Spec prompts, under Apache-2.0 with Gebru datasheets and Croissant 1.1 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 it to StableHLO required no new constraint-layer code. On structurally dominated dialects, SmolLM2-1.7B with token-level constraints matches or exceeds 15B-34B code LMs running post-hoc rejection at a matched retry budget and 8-25x the per-generation speed: on linalg it reaches 80.0% verify-valid (three runs, n=125), 21 to 44 pp above CodeLlama-34B, Granite-Code-34B, and StarCoder2-15B with non-overlapping CIs. The comparison is between systems, not models: under the baselines' post-hoc protocol, SmolLM2 falls from 52.0% to 24.5% on arith+func. On arith+func and the templated StableHLO-Held-Out-200 the baselines match or beat the SLM; we scope these as non-win cells. We release benchmarks, decoder, every per-prompt generation, and a pinned reproducibility Dockerfile.

Figures & tables

Panel of Figure 1 — Figure 1: Roadmap. The released benchmarks (left) are generated against a schema-derived constraint stack (middle); each claim (right) points to the table or figure that supports it.

Benchmark	$n$	Dialect	Ops covered	E/M/H	Source	Verifier
MLIR-Spec-150	150	arith+func+memref	${\sim}50$ distinct	38/44/18%	hand	mlir-opt
Linalg-Spec-30	30	linalg (memref)	12 named	-	hand	mlir-opt
StableHLO-Spec-30	30	stablehlo	10 op families	40/40/20%	hand	iree-compile
StableHLO-Held-Out-200	200	stablehlo	7 fam $\times$ 6 dtypes $\times$ ranks	sweep	sweep	iree-compile
StableHLO-OutOfGrammar-25	25	stablehlo	23 fam, out of grammar	-	hand	iree-compile

Table 1: Benchmark summary: the four released benchmarks plus the out-of-grammar stress set. Total: 410 in-scope NL

\to

MLIR pairs across three dialects, with a 25-program out-of-scope supplement.

System	arith+func	linalg
SmolLM2-1.7B + $C_{1}$ + $C_{2}$ + $C_{3}$	53.2% ( $\pm$ 1.8pp)	80.0% ( $\pm$ 0.0pp)
CodeLlama-34B + $C_{1}$ + $C_{3}$	59.8% ( $\pm$ 0.8pp)	58.7% ( $\pm$ 2.4pp)
Granite-Code-34B + $C_{1}$ + $C_{3}$	51.5% ( $\pm$ 2.5pp)	35.7% ( $\pm$ 4.4pp)
StarCoder2-15B + $C_{1}$ + $C_{3}$	66.8% ( $\pm$ 1.5pp)	54.9% ( $\pm$ 1.2pp)
Granite-Code-8B-fp16 + $C_{1}$ + $C_{3}$ (fp16 control)	68.2% ( $\pm$ 3.0pp)	50.1% ( $\pm$ 1.6pp)

Table 2: Apples-to-apples comparison under a matched five-retry budget, mean of three runs (

n{=}200

per run arith+func;

n{=}125

per run linalg); across-run half-range in parentheses. Granite-Code-8B-fp16 is a same-family precision control on both dialects (§ 7 ). Reading guide: compare each baseline row with SmolLM2 per dialect; linalg is the win and arith+func the scoped non-win.

Benchmark	SmolLM2 1.7B	CodeLlama 34B	Granite-Code 34B	StarCoder2 15B
Spec-30 ( $n{=}30$ , hand-curated)	63.3%	36.7%	33.3%	60.0%
Held-Out-200 ( $n{=}200$ , parametric)	61.5%	100.0%	98.0%	100.0%
OutOfGrammar-25 ( $n{=}25$ , stress)	0% (intended)	—	—	—

Table 3: StableHLO cross-system comparison. SmolLM2-1.7B at full stack (

C_{1}+C_{3}

) vs the three Ollama baselines under matched

C_{1}+C_{3}

with 5-retry rejection, on the three released StableHLO benchmarks. SmolLM2 constraint-ladder on StableHLO:

26.7\to 63.3\to 63.3\%

(Spec-30) and

44.0\to 57.0\to 61.5\%

(Held-Out-200) for free

\to

C_{1}

\to

C_{1}+C_{3}

; the

+C_{2}

layer is silent on StableHLO because operand types are constrained structurally rather than by named domain split. Reading guide: Spec-30 discriminates between systems, while Held-Out-200 saturates the baselines, so the lead holds on Spec-30 only.

Dialect	Verify-valid (three runs)	Fixed-input output-match (10 references)	Gold-differential (spec prompts)
arith+func	53.2% ( $n{=}200$ )	80.0% (8/10)	63.2% (43/68)
linalg	80.0% ( $n{=}125$ )	20.0% (2/10)	36.4% (8/22)
StableHLO (Spec-30)	63.3% ( $n{=}30$ )	50.0% (5/10)	n/a

Table 4: Verify-valid vs functional correctness for SmolLM2 +

C_{1}

C_{2}

C_{3}

. Verify-valid is the three-run mean on the evaluation pool. Fixed-input output-match is on the hand-authored functional references (10 per dialect). Gold-differential is the share of verify-valid seed-0 generations on the spec prompts that match the gold module on all five random draws; non-conforming interfaces count as failures. The StableHLO golds are not part of the differential run. Reading guide: read across a row to see how much of the structural rate survives execution.

Figure 2: SmolLM2-1.7B verify-valid vs constraint level on the frozen released pools. Different layers dominate on different dialects:

C_{3}

drives the arith+func lift (

+6.5

pp,

p=0.002

);

C_{2}

drives the linalg lift (

+9.6

pp,

p<0.0001

). Error bars are 95% bootstrap CIs. Reading guide: the slope between adjacent points is the paired effect of adding one layer.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Generations	arith+func	linalg
SmolLM2-1.7B ( $C_{1}$ rung)	3 / 89	1 / 80
SmolLM2-1.7B ( $C_{1}{+}C_{2}$ rung)	5 / 91	1 / 92
CodeLlama-34B	0 / 121	0 / 71
StarCoder2-15B	0 / 134	0 / 59
Granite-Code-34B	0 / 94	0 / 42

Appendix

Table 5: c3_scope false rejects against mlir-opt on parse-valid generations from the frozen released pools (arith+func

n{=}200

, linalg

n{=}125

). Each cell is false rejects / verify-valid programs. Every baseline cell is zero, so the shared

C_{3}

gate never discards a program the verifier accepts; the few SmolLM2 false rejects come from multi-line operations, as the text explains.

Figure 3: Headline: mlir-opt / iree-compile verify-valid rates on the two dialects where we report a win (linalg, StableHLO). SmolLM2-1.7B +

C_{1}

C_{2}

C_{3}

attains a higher rate than each 15B/34B baseline under the same 5-retry rejection budget. Linalg bars are means of three runs at

n{=}125

with across-run range whiskers; the SmolLM2 cell is 80.0% (range 0pp). The StableHLO panel is corpus-conditional and reports the seed-0 Spec-30 corpus (

n{=}30

); the larger Held-Out-200 corpus saturates the 15B/34B baselines (verify rate

\approx

1.0) and is therefore reported in a separate table rather than in the headline. arith+func is intentionally omitted here and reported in Table 10 as a non-win cell (mean of three runs 53.2%

\pm

1.8pp for SmolLM2 +

C_{1}

C_{2}

C_{3}

); “verify-valid” is a structural property, not functional correctness (§ 7 ).

Figure 4: Efficiency frontier: wall-clock per generation (log scale) vs verify-valid rate (mean of three runs). Five systems

\times

two dialects, including a Granite-Code-8B-fp16 same-family precision control. The SmolLM2 +

C_{1}

C_{2}

C_{3}

linalg point sits at 1.86 s / 80.0% (top-left, ours) and the arith+func point at 1.65 s / 53.2%; on linalg the SmolLM2 cell is Pareto-dominant by both axes against every other cell. The 15B/34B 4-bit baselines run 6–24

\times

slower per generation; the Granite-Code-8B-fp16 control runs 7

\times

slower on arith+func and 25

\times

slower on linalg. Color = system, marker = dialect (circles = arith+func, squares = linalg).

Figure 5: Apples-to-apples comparison, mean of three runs with across-run half-range whiskers. Five systems

\times

two dialects, including a Granite-Code-8B-fp16 precision control on both dialects. linalg: SmolLM2-1.7B +

C_{1}

C_{2}

C_{3}

(red border, “ours”) = 80.0%, leading CodeLlama-34B (58.7%) by

+21.3

pp, Granite-34B (35.7%) by

+44.3

pp, StarCoder2-15B (54.9%) by

+25.1

pp, and the Granite-Code-8B-fp16 control (50.1%

\pm

1.6pp) by

+29.9

pp; CIs non-overlapping across all five systems. arith+func: SmolLM2 (53.2%

\pm

1.8pp) trails CodeLlama-34B (59.8%) by

-6.6

pp (within CI), narrowly leads Granite-34B (51.5%) by

+1.7

pp (within CI), trails StarCoder2-15B (66.8%) by

-13.6

pp, and trails the Granite-Code-8B-fp16 control (68.2%

\pm

3.0pp) by

-15.0

pp.

n{=}200

per run for arith+func,

n{=}125

per run for linalg, three runs each.

Figure 6: Error-category progression across constraint levels, SmolLM2-1.7B, frozen released pools. Each layer shrinks its target bucket. arith+func:

C_{3}

cuts type_ssa errors

61\to 37

(

-39\%

). linalg:

C_{2}

eliminates syntax errors

19\to 0

; type_ssa stays small (at most 7) because linalg’s ins(...) outs(...) re-declares types at each use site.

Figure 7: Paired bootstrap

\Delta

verify per constraint layer with 95% CIs, frozen released pools.

C_{1}

C_{2}

C_{3}-C_{1}

C_{2}

+6.5

pp on arith+func (CI

[+2.0,+11.0]

) and

+4.0

pp on linalg (CI

[+0.0,+8.0]

);

C_{1}

C_{2}-C_{1}

+9.6

pp on linalg (CI

[+4.8,+15.2]

) and

+1.0

pp on arith+func (CI

[+0.0,+2.5]

Figure 8: Hidden Cost of Structure replication on the frozen arith+func pool (

n{=}200

), with null result.

C_{1}

is monotonically beneficial on both SLMs (

+24.5

pp SmolLM2,

+20.0

pp Phi-3.5-mini, both

p<0.0001

). Under 3-shot priming with a type-domain-aware grammar and bounded identifiers, the RANLP 2025 HCS reversal does not reproduce.

System	arith+func (n=200)	linalg (n=125)
SmolLM2-1.7B, free	20.0 [14.5, 25.5]	52.0 [43.2, 60.8]
SmolLM2-1.7B + C1	44.5 [37.5, 51.5]	64.0 [55.2, 72.0]
SmolLM2-1.7B + C1+C2	45.5 [38.5, 52.5]	73.6 [65.6, 80.8]
SmolLM2-1.7B + C1+C2+C3 (ours)	52.0 [45.0, 59.0]	77.6 [70.4, 84.8]
Granite-Code-34B, free	39.0 [32.5, 46.0]	27.2 [20.0, 35.2]
Granite-Code-34B + C1	47.0 [40.5, 54.0]	33.6 [25.6, 41.6]

Appendix

Table 6: Constraint-layer progression on the frozen released pools (seed-0 run;

n{=}200

arith+func,

n{=}125

linalg; verify-valid rate with 95% bootstrap CI in brackets). SmolLM2 rows form the constraint ladder; baseline rows use five-retry

C_{1}

rejection without

C_{3}

, and “free” is a single unconstrained attempt. Reading guide: compare adjacent SmolLM2 rows for each layer’s contribution (paired deltas in § 6 ); cross-system claims use the three-run means in Table 2 .

Figure 9: $C_{3}$ scope-validator confusion matrix vs mlir-opt on 200 SmolLM2 + $C_{1}$ arith+func generations (frozen released pool). The boxed (scope FAIL, verify PASS) cell counts false rejects: 3 of the 89 verify-valid programs. The same count is 0 for the $C_{1}$ outputs of every 15B/34B baseline on both dialects (Table 5 ).

System	arith C1-only	arith C1+C3	linalg C1-only	linalg C1+C3
CodeLlama-34B	60.5	59.8	56.8	58.7
StarCoder2-15B	67.0	66.8	47.2	54.9
Granite-Code-34B	47.0	51.5	33.6	35.7

Appendix

Table 7: Best-configuration baselines on the frozen released pools: verify-valid rate without

C_{3}

(five-retry

C_{1}

rejection, seed-0 run) and with

C_{3}

(mean of three runs).

C_{3}

never lowers a baseline beyond noise, so no baseline is represented by a

C_{3}

-penalized score; against each baseline’s best configuration SmolLM2’s linalg lead (80.0%) is at least

+21.3

pp.

Cell	arith+func all	arith+func unique	linalg all	linalg unique
SmolLM2 free	20.0	18.7	52.0	38.7
SmolLM2 + C1	44.5	43.5	64.0	56.0
SmolLM2 + C1+C2	45.5	44.6	73.6	69.3
SmolLM2 + C1+C2+C3 (seed-0 ladder)	52.0	50.8	77.6	75.0
SmolLM2 + C1+C2+C3 (3 runs)	53.2	52.2	80.0	75.7
CodeLlama-34B + C1+C3	59.8	59.4	58.7	63.6

Appendix

Table 8: Both aggregations for every pooled cell (verify-valid %). “All” counts every prompt row; “unique” gives each unique prompt string one value (the mean over its duplicates), per run, then averages runs. The pools hold 193 (arith+func) and 75 (linalg) unique strings; StableHLO pools contain no duplicates, so both aggregations coincide there. The linalg ranking is unchanged (smallest SmolLM2 margin

+12.1

pp) and arith+func remains a non-win.

Dialect	Constraint	passed	type	type_ssa	syntax	other
arith+func	none	40	1	3	0	156
arith+func	c1	89	13	51	0	47
arith+func	c1_c2	91	1	61	0	47
arith+func	c1_c2_c3	104	9	37	0	50
linalg	none	65	0	1	4	55
linalg	c1	80	0	5	19	21

Appendix

Table 9: Error categorization across constraint levels, SmolLM2-1.7B seed-0 ladder on the frozen released pools. Arity and dialect-misuse buckets are 0 in every cell and are omitted.

Figure 10: Per-op verify rate under SmolLM2 +

C_{1}

C_{2}

C_{3}

on linalg (seed-0 run of the three-run frozen-pool cell,

n{=}125

). Most named ops reach 100%; shape-reasoning ops ( transpose , broadcast ) and the small fill and exp cells trail, the residual gap motivating a future

C_{4}

shape-inference layer.

Cell	s0	s1	s2	mean	$\pm$ hr
SmolLM2 + $C_{1}$ + $C_{2}$ + $C_{3}$ , arith+func	52.0%	52.0%	55.5%	53.2%	$\pm 1.8$ pp
SmolLM2 + $C_{1}$ + $C_{2}$ + $C_{3}$ , linalg	80.0%	80.0%	80.0%	80.0%	$\pm 0.0$ pp
Granite-Code-34B-Q4 + $C_{1}$ + $C_{3}$ , arith+func	49.5%	54.5%	50.5%	51.5%	$\pm 2.5$ pp
Granite-Code-34B-Q4 + $C_{1}$ + $C_{3}$ , linalg	34.4%	40.8%	32.0%	35.7%	$\pm 4.4$ pp
CodeLlama-34B-Q4 + $C_{1}$ + $C_{3}$ , arith+func	60.5%	59.0%	60.0%	59.8%	$\pm 0.8$ pp
CodeLlama-34B-Q4 + $C_{1}$ + $C_{3}$ , linalg	59.2%	60.8%	56.0%	58.7%	$\pm 2.4$ pp

Appendix

Table 10: Per-seed verify-valid rates across the apples-to-apples cells. Four systems

\times

two dialects, plus a Granite-Code-8B-fp16 same-family precision control on both dialects, at

n{=}200

(arith+func) or

n{=}125

(linalg) per run, three runs (seeds 0, 1, 2). Baseline decoding is fully seeded; SmolLM2’s first attempt is greedy and its temperature-0.8 retries are not seeded, so its three runs are repeats with measured variance. Mean is the across-run average;

\pm

half-range is

(\max-\min)/2

. SmolLM2’s seed-0 arith+func cell equals Table 6 ’s top rung (52.0%); on linalg the ladder run’s 77.6% lies within the measured run-to-run spread (at most 2.4pp).

Claim	Status	Supporting evidence
linalg win (system-level)	Robust	Mean of three runs 80.0% (each run 80.0%); $+21.3$ to $+44.3$ pp over the baselines with non-overlapping CIs; at least $+21.3$ pp over each baseline’s best configuration; the fp16 control trails by 29.9pp
Enforcement asymmetry	Measured	SmolLM2 under the baselines’ post-hoc protocol: 24.5% vs 52.0% with masking on arith+func ( $+27.5$ pp, CI $[+20.0,+35.0]$ )
8–25 $\times$ per-generation speedup	Robust	Wall-clock on a single Apple M4 Max (Fig. 4 ); 1.65–1.86 s vs 16–40 s
Mechanical schema port across dialects	Robust	Zero new constraint-layer code from arith+func+memref+linalg $\to$ StableHLO
StableHLO win on hand-curated Spec-30	Corpus-conditional	$+26.6$ to $+30.0$ pp over the 34B baselines on Spec-30 only; baselines saturate at 98–100% on Held-Out-200
arith+func win	Not supported	53.2% vs CodeLlama-34B 59.8% (within CI), StarCoder2-15B 66.8% (non-overlapping), Granite-Code-8B-fp16 68.2% (paired $+15.0$ pp)

Appendix

Table 11: Summary of empirical claims and supporting evidence.

Benchmark	$n$	Dialect	Source	Type
MLIR-Spec-150	150	arith+func+memref	author-curated	in-grammar
Linalg-Spec-30	30	linalg (12 ops)	author-curated	in-grammar
StableHLO-Spec-30	30	stablehlo (10 fam.)	author-curated	in-grammar
StableHLO-Held-Out-200	200	stablehlo	parametric sweep	in-grammar (held-out)
StableHLO-OutOfGrammar-25	25	stablehlo	author-curated	out-of-grammar (stress)
MLIR-Functional-Reference-30	30	arith+func+memref / linalg+memref / stablehlo (10 ea.)	author-curated	evaluation artifact (NL+MLIR+inputs+expected output)

Appendix

Table 12: Per-benchmark composition. Spec-* benchmarks were author-written against public ODS specs and verifier-filtered; Held-Out-200 was generated by a parametric sweep and verifier-filtered; OutOfGrammar-25 was hand-authored against StableHLO ops the grammar does not cover; MLIR-Functional-Reference-30 adds hand-authored execution references to 30 Spec prompts and is used to bound the verify-valid

\to

output-match gap.

Constrained Decoding

Constrained decoding steers language model outputs to satisfy formal or task-specific rules while balancing validity, quality, and faithful probabilities.

See all 67 papers

An MLIR-Based Compilation Method for Large Language Models
Jul 17, 2026Pengchao Hu, Zhibin Xin, Yifan Chen +3LLM Serving LLM Inference Scheduling
Sophgo Inc.
DConstrain: Compiling Context-Parameterized Transitions for Diffusion Language Model Constrained Decoding
Oct 22, 2026Shidong Cao, Jing Ma, Wenhao Chai +3Constrained Decoding Grammar-Constrained Decoding
Hong Kong Baptist University · Princeton University · Ant Group
Selective Left-Shift: Turning Test-Time Compute and Difficulty-based Curation into Training Data for Low-Resource Code Generation
Jul 8, 2026Didula Samaraweera, Anjana Supun, Srinath PereraRL for Code Generation Code Generation
WSO2 Santa Clara, CA, USA

Jul 17, 2026J/K move · Enter open · S save

An MLIR-Based Compilation Method for Large Language Models

Pengchao Hu, Zhibin Xin, Yifan Chen, Yangyang Zhou+2

Sophgo Inc.

LLM Serving LLM Inference Scheduling LLM Inference Acceleration

Open paper AskPDF arXiv

Large Language Models (LLMs) have become the dominant workload on modern AI accelerators, yet deploying them on specialized hardware still faces two core challenges: how to import a trained model into a compiler-friendly intermediate representation, and how to efficiently schedule the autoregressive inference loop under limited on-chip memory. This paper presents an MLIR (Multi-Level Intermediate Representation) based compilation method for large language models, illustrated using two dialects of operators, TopOp and TpuOp. TopOp serves as a high-level graph dialect that is independent of both the source framework and the target chip, and is responsible for expressing model semantics; TpuOp serves as the target hardware dialect, carrying chip-related decisions such as quantization, layer groups, and memory layout. A model is first represented as TopOp, then lowered layer by layer to TpuOp, and finally a deployable binary is generated. In addition, each Transformer layer is split into three stages for static compilation: prefill, prefill_kv (prefill with historical key-value cache), and decode, so as to accommodate the different computational characteristics of prompt-parallel processing and per-token generation. The method has been implemented in the TPU-MLIR compiler {https://github.com/sophgo/tpu-mlir} and the LLM-TPU deployment project {https://github.com/sophgo/LLM-TPU}, supporting a variety of generative models including the Qwen, Llama, InternVL, and MiniCPM-V series, as well as multiple quantization and deployment forms such as GPTQ, AWQ, and AutoRound.

Abstract

Constrained Decoding

Explore similar work

An MLIR-Based Compilation Method for Large Language Models

DConstrain: Compiling Context-Parameterized Transitions for Diffusion Language Model Constrained Decoding

Selective Left-Shift: Turning Test-Time Compute and Difficulty-based Curation into Training Data for Low-Resource Code Generation