Skip to content

[Architecture] TileLang DSL lowering backend 渐进迁移到 pybinding builder #237

@Zhendong404

Description

@Zhendong404

Background

当前 TileLang DSL 的 lowering 主路径是:

  • 前端语义分析得到 SemanticKernel
  • tilelang-dsl/python/tilelang_dsl/lowering.py 中的 AuthoringModule / _AuthoringRenderer 直接发射 authoring-form VPTO MLIR 文本
  • mlir_text() / emit() 直接消费这份文本

也就是说,当前 lowering backend 的核心职责是“文本 emitter”,而不是“构造 MLIR Python IR 对象再统一打印”。

这条路径在当前阶段有明显优点:

  • 实现直接、产物可读
  • authoring-form VPTO 输出稳定
  • 与现有 ptoas --pto-backend=vpto 验证路径天然兼容

但随着 TileLang DSL 功能继续扩张,当前字符串发射模式也在逐渐暴露维护边界:

  • SSA/result/type/region 结构需要由 lowering 自己以文本方式维持
  • helper function 合并、function attr 重写、注释元信息输出都与文本形态耦合
  • 如果后续希望在 Python 内直接做 IR 级别 verify / inspect / pass / rewrite,需要先 parse 回 MLIR module
  • 新增 op family、复杂控制流和中间态 rewrite 时,文本 emitter 的维护成本会继续上升

与此同时,仓内已经具备一部分 pybinding 基础设施:

  • python/pto/dialects/pto.py 暴露了 PTO dialect 的 Python types / attrs / ops
  • test/samples/Abs/abs.pytest/samples/VectorAddition/vadd_pto_ir.pytest/samples/MatMul/tmatmulk.py 等样例已经展示了 mlir.ir + mlir.dialects.{func,arith,scf,pto} 的 builder-style IR 构造方式

因此,值得讨论是否将 TileLang DSL lowering backend 从“纯文本 emitter”渐进迁移到“pybinding builder”。

Problem

当前文本 emitter 方案的主要问题不是“不能工作”,而是它对后续工程演进不够友好:

  1. Lowering 逻辑与打印格式耦合较深
  • lowering.py 当前承担了大量字符串层职责:SSA 命名、类型字符串渲染、helper 行拼接、属性首行重写等。
  • 这使得“IR 结构正确性”和“文本格式正确性”混在一起维护。
  1. Python 侧 IR 级操作链路不自然
  • 若希望在 DSL 侧直接做 verify()、pass pipeline、局部 IR inspect、结构化测试,目前往往需要先从文本再 parse。
  • 这会增加中间转换,并让部分调试手段停留在 grep / 文本对比层面。
  1. 新增复杂 lowering 时维护成本更高
  • 控制流、multi-result op、region、symbol/helper 管理等在 builder 模式下天然是结构化对象;在文本模式下则需要重复维护字符串模板。
  1. 测试容易被文本格式影响
  • 当前很多正确性验证天然更偏向 textual golden。
  • 对后续 printer、注释、格式化方式的变化不够稳健。

Suggestion

建议以“渐进迁移”而不是“一次性全量重写”的方式,引入 pybinding lowering backend。

Phase 0: 保持现有文本 backend 作为 reference

  • 保留当前 AuthoringModule.render() 路径,继续作为稳定参考实现。
  • 不在第一阶段就删除文本 emitter。

Phase 1: 先打通 parse / module 化路径

先做最小收益版本:

  • TileLang 继续输出当前 authoring-form VPTO 文本
  • 在 DSL 侧补齐统一的 Context + register_dialect + Module.parse(...) 路径
  • 让 Python 侧能直接拿到结构化 module 做 verify / inspect

这一阶段不要求重写 lowering,只是先把“文本产物”变成 DSL 侧可继续操作的 MLIR module。

Phase 2: 引入可选 pybinding backend

在不破坏现有文本路径的前提下,引入一个并行 backend,例如:

  • TextAuthoringBackend
  • PybindingAuthoringBackend

两者共享:

  • SemanticKernel
  • 大部分语义检查 / specialization 结果
  • 尽量统一的 lowering helper 抽象

但输出形式不同:

  • 文本 backend 负责现有 str 产物
  • pybinding backend 负责直接构造 mlir.ir.Module / Operation / Value

Phase 3: 先迁结构性最强的一批 lowering

建议优先迁移这些部分:

  • module / func / block 构造
  • arith.constant
  • scf.for / scf.if
  • TensorView / PartitionView / TileBuf 等 type 构造
  • AllocTile / MakeTensorView / PartitionView 这类结构性强的 PTO op

而将更复杂、文本模板较重的部分放到后续:

  • DMA 相关 lowering
  • 更复杂的 vector-family lowering
  • helper inline / 特殊 debug 注释导出

Phase 4: 双 backend 对拍并收敛

在迁移过程中,用双 backend 对拍降低风险:

  • pybinding backend 打印出的 module 文本与 text backend 对齐
  • 共享同一批 verify / unittest / sample case
  • 当覆盖面足够高时,再决定是否将 pybinding backend 转为默认实现

Expected outcome

  1. TileLang DSL 在 Python 侧具备更自然的 IR 级操作能力
  • 可直接 verify / inspect / pass / debug,而不必把文本作为唯一中间载体。
  1. Lowering backend 逐步从“文本模板维护”转向“结构化 IR 构造”
  • 更利于后续新增 op family、复杂控制流和中间态 rewrite。
  1. 测试方式可以从纯文本 golden 扩展到结构化断言
  • 提高对 printer/格式变化的鲁棒性。
  1. 迁移风险可控
  • 通过保留 text backend reference 和双 backend 对拍,避免一次性重写带来的大回归风险。

Possible acceptance criteria

  • Phase 1: TileLang DSL 能在 Python 内直接返回结构化 MLIR module,而不只是一份文本。
  • Phase 2: 存在一个可选 pybinding lowering backend,至少能覆盖一小批代表性 TileLang kernel。
  • Phase 3: 结构性最强的一批 op / control-flow lowering 能稳定通过现有 verify / unittest / sample 验证。
  • Phase 4: 文本 backend 与 pybinding backend 在共享 case 上能够完成稳定对拍。

Non-goals

这件事不应被理解为:

  • 立即删除现有文本 emitter
  • 立即在一个 patch 中全量重写 lowering.py
  • 只为了“换一种写法”而迁移

本提案的主要价值在于后续工程演进能力,而不是短期代码风格统一。

Related issues

  • #209 [Docs] 增加 DSL 框架对接协议文档:上层如何对接 TileLang/DSL 生成 VPTO IR
  • #210 [Architecture] VPTO IR cache 机制应迁移到 DSL 框架侧
  • #195 [TypeSystem] 明确 TileLang DSL 与 VPTO 的 scalar signedness 边界,并推动 scalar lowering 向 signless 收敛

其中:

  • #209 更关注框架对接协议
  • #210 更关注实例/cache 归属边界
  • #195 更关注 scalar type system / lowering contract

本 issue 聚焦的是 TileLang DSL lowering backend 的实现形态演进:从字符串 emitter 渐进引入 pybinding builder 能力。

Additional context

当前仓内可作为该提案直接参考的基础包括:

  • tilelang-dsl/python/tilelang_dsl/lowering.py
  • tilelang-dsl/python/tilelang_dsl/kernel.py
  • python/pto/dialects/pto.py
  • test/samples/Abs/abs.py
  • test/samples/VectorAddition/vadd_pto_ir.py
  • test/samples/MatMul/tmatmulk.py

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions