无需重新训练的跨方言泛化：面向MLIR的基于模式约束解码的基准与评估

# 跨方言泛化无需重训练：基于模式导出约束解码的 MLIR 基准测试与评估  
来源：https://arxiv.org/html/2607.18254  
Plawan Kumar Rath Meta [email protected] 本研究为作者个人行为。本文所述观点仅代表作者，不代表 Meta 立场。  

###### 摘要  
多级中间表示（MLIR）是现代机器学习编译器基础设施的核心，支撑着 TensorFlow、通过 StableHLO 的 JAX、PyTorch、Inductor、IREE 等，但在代码语言模型预训练语料库中仅出现微量。MLIR 本身可扩展：每个应用领域都会推出新的方言，因此为每个方言维护一个微调模型难以规模化。我们探讨一个问题：从各方言的操作定义规范（ODS）中机械化推导出的推理先验，能否替代基于梯度的适应方法。我们做出两项贡献。第一，我们发布了四个跨三个方言的自然语言转 MLIR 基准测试集：MLIR-Spec-150、Linalg-Spec-30、StableHLO-Spec-30 和 StableHLO-Held-Out-200，共计 410 对 NL→MLIR 映射，外加一个 25 程序的 StableHLO 超语法压力集和一个手工编写的 n=30 功能参考集（总计 435 个实例）。所有资源在 Apache-2.0 许可下发布，并附带 Gebru 数据表与 Croissant 1.0 元数据。第二，在这些基准测试之上，我们构建了一个三层模式导出约束栈：基于操作签名的上下文无关文法（C1）、从 ODS 提取的类型格导出的类型域切分（C2），以及驱动五次重试拒绝采样的 SSA 作用域验证器（C3）。将这一栈从 arith+func+memref+linalg 移植到 StableHLO 无需任何新的约束层代码。实验表明，在验证器语义主要由结构约束主导的方言上，模式导出先验使得 SmolLM2-1.7B 能够匹配或超越 15B–34B 代码语言模型，同时每代生成速度快 8–25 倍：在 linalg 上，SmolLM2 达到 80.0% 的验证通过率（三种子均值，n=125，每种子均为 80.0%），比 CodeLlama-34B、Granite-Code-34B 和 StarCoder2-15B 高出 21 到 44 个百分点，且置信区间不重叠；在同系列 fp16 精度控制下依然保持优势。在 arith+func 以及模板参数化的 StableHLO-Held-Out-200（验证器语义主要依赖属性值而非结构）上，相同的基线模型匹配或超过了 SLM；我们明确将这些归为非获胜区域。我们发布了基准测试集、解码器、每个提示的生成结果以及可复现的 Docker 镜像。  

## 1 引言  
MLIR（多级中间表示）(Lattner et al., 2020 (https://arxiv.org/html/2607.18254#bib.bib6)) 已成为现代机器学习系统的首选编译器基础设施，支撑着 TensorFlow 的编译器栈、通过 StableHLO 的 JAX、PyTorch 的 Inductor 后端以及 IREE 等开源运行时。然而，在用于预训练开源权重语言模型的公开代码语料库中，MLIR 程序仅出现微量；当模型被要求从自然语言生成 MLIR 时，通常会输出语法错误的代码，引用目标方言中不存在的操作，或者构造出验证器当场拒绝的类型。微调在原则上可以解决这一问题，但 MLIR 是可扩展的：社区会为每个应用领域引入新的方言，为每个方言维护一个微调模型难以规模化。本文将生成问题视为结构先验问题而非数据问题。MLIR 的操作定义规范（ODS）——其 TableGen 记录声明了每个操作的操作数、结果、类型约束和结构不变量——提供了一个跨方言统一的、机器可读的模式。我们研究：在那些验证器语义主要由结构约束主导的方言上（操作数秩对等、类型域切分、SSA 作用域），模式导出先验能否使小型语言模型在匹配的采样预算下匹配或超越 15B–34B 代码语言模型？并且这一构造能否机械地跨方言移植？答案因情况而异：推导过程可跨所有三个方言移植，无需新的约束层代码；但 SLM 与基线之间的差距随方言和语料库条件变化——在结构主导的 linalg 上表现稳健，在 arith+func 以及属性值约束主导的模板参数化 StableHLO 上则持平或落后。SmolLM2-1.7B-Instruct (Allal et al., 2025 (https://arxiv.org/html/2607.18254#bib.bib17)) 在完整 C1+C2+C3 栈下，与 CodeLlama-34B (Rozière et al., 2023 (https://arxiv.org/html/2607.18254#bib.bib18))、Granite-Code-34B (Mishra et al., 2024 (https://arxiv.org/html/2607.18254#bib.bib19)) 和 StarCoder2-15B (Lozhkov et al., 2024 (https://arxiv.org/html/2607.18254#bib.bib20)) 在相同的五次重试拒绝采样预算下进行评估，在 linalg 上稳健胜出（80.0% 验证通过，三种子均值，n=125，半范围±0.0 个百分点，比基线高出 21 到 44 个百分点，置信区间不重叠）；在手工程序集 StableHLO-Spec-30 上胜过 34B 基线；在模板参数化的 StableHLO-Held-Out-200（基线达到 98–100%）上被超越；在 arith+func 上为非获胜区域。在单台 Apple M4 Max 笔记本上进行推理时，每代生成速度比 34B 基线快 8–25 倍，无需微调、强化学习或蒸馏：整个适应过程均在推理时完成。我们将此定位为数据与评价贡献，同时附带一个使用这些基准测试的方法贡献。  
(i) **数据集**：四个 NL→MLIR 基准测试集：MLIR-Spec-150、Linalg-Spec-30、StableHLO-Spec-30 和 StableHLO-Held-Out-200（后者为 200 程序参数化操作签名扫描，用于检查 Spec-30 中作者筛选偏差；跨三个方言共 410 对映射），外加一个 25 程序的 StableHLO 超语法压力集和一个手写 n=30 功能参考集（每个生成降级为可执行 IR 以进行输出匹配测试），总计 435 个实例，在发布时全部通过验证器检查，附带难度标签、Gebru 数据表 (Gebru et al., 2021 (https://arxiv.org/html/2607.18254#bib.bib10)) 记录和 Croissant 1.0 (MLCommons Croissant Working Group, 2024 (https://arxiv.org/html/2607.18254#bib.bib11)) 元数据，采用 Apache-2.0 许可；据我们所知，这是首批公开的 NL→MLIR 基准测试集。  
(ii) **方法**：一个从 ODS 机械提取的三层模式导出约束流水线（操作签名上的 CFG、从 ODS 格导出的类型域切分、带五次重试拒绝的 SSA 作用域验证器）；从 arith+func+memref+linalg 移植到 StableHLO——一个语法不同、作用域语义不同、下游验证器不同（使用 iree-compiler 而非 mlir-opt）的方言——无需任何新的约束层代码。C3 验证器既支持事后拒绝采样形式（用于我们的测量），也支持内联耦合解码器形式，附录中证明了接受的字符串集等价。  
(iii) **协议**：对所有模型应用相同的五次重试拒绝采样预算，从而消除混淆“带语法的小模型 vs. 自由解码的大模型”的约束不对称差距。  
(iv) **发现**：linalg 上的稳健胜利、条件相关的 StableHLO 结果，以及 arith+func 上诚实的非获胜结果。  

## 2 相关工作  
#### CFG 指导的解码  
基于上下文无关文法（CFG）指导的 token 级掩码是代码模型推理中一个成熟的子领域。Willard 和 Louf (Willard and Louf, 2023 (https://arxiv.org/html/2607.18254#bib.bib1)) 形式化地描述了通过正则表达式和上下文无关文法的有限状态转换实现高效掩码构造的方法；Outlines 库实现了这一方法，我们在其 llguidance 后端上使用它。Grammar-Aligned Decoding (Park et al., 2024 (https://arxiv.org/html/2607.18254#bib.bib2)) 指出，朴素的 token 级掩码会偏离它本该采样的 CFG 条件分布，并提出了一种原则性的重新归一化方法。Schall 和 de Melo (Schall and de Melo, 2025 (https://arxiv.org/html/2607.18254#bib.bib3)) 记录了“结构的隐藏成本”：在几个小模型上，约束解码的通过率低于自由解码，且该效果对提示格式敏感。我们的 C1 层是一个标准的 CFG 掩码，覆盖 MLIR 操作签名；其区别在于，语法是从 ODS 机械生成而非手工编写，因此相同的推导过程无需修改即可跨方言应用。我们在三样本启始提示下使用限标识符语法复制了 Schall 和 de Melo 的设置，并未观察到通过率反转 (§6 (https://arxiv.org/html/2607.18254#S6))；我们将我们的消融实验解读为对提示格式敏感性主张的佐证数据点，而非反驳。  

#### 上下文敏感的约束解码  
与我们 C3 层更接近的是针对形式化目标进行上下文敏感约束解码的较小研究领域。Synchromesh (Poesia et al., 2022 (https://arxiv.org/html/2607.18254#bib.bib7)) 引入了约束语义解码（CSD），一种每步掩码过程，利用目标语言的静态分析器（类型检查器、作用域分析器）在采样期间限制 token 词汇，并在 SQL 和 Python 合成上进行了评估。Type-Constrained Code Generation (Mündler et al., 2025 (https://arxiv.org/html/2607.18254#bib.bib8)) 将静态类型检查器集成到每步解码中，适用于静态类型语言，在 1B 规模模型上大幅降低了类型错误率。Correctness-Guaranteed Code Generation (Li et al., 2025 (https://arxiv.org/html/2607.18254#bib.bib9)) 进一步扩展了这一方向，采用上下文敏感的解析器树驱动 token 级解码，确保生成的程序可解析。我们的工作与它们有三点不同：(i) 目标是一个可扩展的中间表示族而非通用编程语言，且我们的 C1/C2/C3 层是从机器可读模式（ODS）机械推导的，而非针对每种语言的类型系统手工实现，因此一个单一的推导过程即可覆盖所有提供 ODS 记录的方言（MLIR 核心、StableHLO 以及任何遵循相同约定的外部方言）。(ii) 我们证明了内联耦合解码器与事后拒绝采样在接受的字符串集上等价，并附带 BPE 边界论证处理跨多个语法终结符的子词 token (定理 A (https://arxiv.org/html/2607.18254#A1.SSx4)，附录 A (https://arxiv.org/html/2607.18254#A1))。(iii) 我们发布了四个 NL→MLIR 基准测试集；Synchromesh 和类型约束解码工作是在现有代码基准测试集上评估的，但未发布特定目标的数据集。  

#### 面向形式化目标的小语言模型  
针对形式化目标的语言模型方法主要集中在证明助手（如 Lean 策略生成 (Yang et al., 2023 (https://arxiv.org/html/2607.18254#bib.bib21)) 和 Isabelle 完整证明合成 (First et al., 2023 (https://arxiv.org/html/2607.18254#bib.bib22))），以及硬件描述语言（如 Verilog (Liu et al., 2023 (https://arxiv.org/html/2607.18254#bib.bib23))）。尽管 MLIR 在编译器和机器学习系统基础设施中扮演核心角色，但作为直接生成目标得到的关注相对较少；据我们所知，这是首项在方言无关层面将 MLIR 视为无需训练的语言模型生成目标的工作。我们的贡献无需训练；微调是一个正交方向，我们在此未进行探索，但可以与模式导出栈组合使用。  

## 3 方法  
#### 流水线概述  
我们的方法接受自然语言描述 x，并通过从一个小型语言模型中进行采样，在一个从目标方言的操作定义规范（ODS）机械推导出的三层约束栈 C=(C1,C2,C3) 下生成 MLIR 代码 y。三层覆盖了 MLIR 合法性的不同方面：C1 是针对操作签名的上下文无关文法，C2 按操作数和结果类型域对该文法进行切分，C3 是强制 SSA 作用域正确性的动态验证器。C1 和 C2 在解码时作为 token 级掩码应用；C3 驱动一个围绕掩码解码器的五次重试拒绝采样器。相同的推导过程无需修改即可跨方言应用：将栈从 arith+func+memref+linalg 移植到 StableHLO 无需任何新的约束层代码。  

#### 模式提取  
我们通过 `llvm-tblgen --dump-json` 从 LLVM、IREE 和 StableHLO 的 TableGen 源码树中提取 ODS 记录，得到目标方言中 276 条操作记录。每条记录命名了操作的操作数、结果、属性和结构特征（例如，SameOperandsAndResultType、ElementwiseMappable）。我们展开多重继承类层次结构，并生成两个派生结构：(i) 每个方言的操作格，命名操作数和结果类型的约束；(ii) 每个方言的属性格，命名属性种类和值域。两者均被 C1 语法生成器和 C2 切分器使用。（TableGen 的 JSON 转储正确解析了多重继承层次；早期基于正则表达式的解析器仅能解析我们所需的 98 个 linalg 操作中的 2 个。）  

#### C1：操作签名的上下文无关文法  
C1 是一个经过 LALR 编译的上下文无关文法，覆盖 MLIR 表层语法，并限制在目标方言的操作集内。产生式从展开的 ODS 生成：每个操作贡献一个产生式，以具体 MLIR 语法命名其助记符、操作数数量、属性槽和结果类型约束（例如，`arith.addi %a, %b : i32` 而非 ODS 声明形式）。该文法被编译成 token 级自动机，并通过 Outlines (Willard and Louf, 2023 (https://arxiv.org/html/2607.18254#bib.bib1)) 及其 llguidance 后端作为 logits 级掩码提供。两个工程选择至关重要：显式空白产生式（替代 LARK 默认的 `%ignore WS`，后者会悄无声息地耗尽 `max_tokens` 而文法无进展）和有界的 SSA 标识符正则表达式（防止名称失控）。使用这些 C1 后，跨方言的解析有效输出率达到 96–99%。  

#### C2：类型域文法切分  
C2 通过从 ODS 类型格读取的类型域对 C1 进行细化切分。例如，`arith.addi` 被限制为整数操作数和结果类型，`arith.addf` 被限制为浮点类型；`arith.cmpi` 采用整数比较谓词语汇（eq, ne, slt, ...），而 `arith.cmpf` 采用浮点比较谓词语汇（oeq, olt, une, ...）。在 linalg 中，切分编码了更丰富的约束：`linalg.matmul on memref by memref` 要求操作数内部维度一致；`linalg.transpose` 约束置换属性必须匹配源 memref 的秩。每个切分由 ODS 记录的特征集机械确定（例如，`SameOperandsAndResultType` 意指所有位置共用一个类型域），我们不针对每个方言手工编写切分规则。  

#### C3：SSA 作用域验证器  
C3 是一个针对 SSA 使用的动态验证器。给定一个生成的 MLIR 片段，它逐行解析，维护一个函数的符号表（初始化为函数参数列表），在每个 SSA 定义点更新符号表，并针对每个使用点验证（名称在作用域内 ∧ 类型匹配声明）。验证失败时，该片段被拒绝，解码器重试。