Skip to content

[VPTO] PTODSL 默认 kernel_kind 被序列化为显式 kind,导致 C/V split 冲突 #1004

Description

@jimmychou0

背景与目标

TileLang 正在将手写 AIC/AIV GEMM 示例 examples/ascend/example_gemm_mix_manual.py 迁移到 VPTO 路径的 PTODSL。迁移入口需要支持未显式指定 kernel_kind@pto.jit(..., backend="vpto"):同一 entry 中的 Vector TileOp 和 Cube TileOp 应被 VPTO 分别编译为 Vector 与 Cube child。

本 issue 不建议引入 kernel_kind="mixed"。PTODSL 现有默认值 module_spec.kernel_kind="vector" 及其 kernel_kind_explicit=False 需要继续保留,供 native build 决定 host 编译架构;问题在于该默认值不应被当作用户写入 PTO IR 的显式 kernel kind。

复现 POC

下面的 POC 使用真实 pto.tileop(),故 PTODSL 会根据 TileOp 自动形成 pto.section.vectorpto.section.cube。它有意没有传入 kernel_kind=

from ptodsl import pto

@pto.jit(name="default_kind_cv_tileop_probe", target="a5")
def default_kind_cv_tileop_probe():
    with pto.tileop():
        zero = pto.const(0, dtype=pto.i64)
        source = pto.castptr(zero, pto.ptr(pto.f32, "ub"))
        destination = pto.castptr(zero, pto.ptr(pto.f32, "ub"))
        value = pto.vlds(source, 0, pto.vreg_type(64, pto.f32))
        pto.vsts(value, destination, 0, pto.pset_b32(pto.MaskPattern.ALL))

    with pto.tileop():
        zero = pto.const(0, dtype=pto.i64)
        lhs = pto.castptr(zero, pto.ptr(pto.f16, "left"))
        rhs = pto.castptr(zero, pto.ptr(pto.f16, "right"))
        acc = pto.castptr(zero, pto.ptr(pto.f32, "acc"))
        extent = pto.const(16, dtype=pto.i64)
        pto.mad(lhs, rhs, acc, extent, extent, extent)

该 POC 用于隔离 PTODSL trace 与 C/V split,不是完整可执行的 VPTO kernel:其中 pto.vlds 尚未放入合法的 pto.vecscope。这不会影响对 split 的观察,因为默认路径在到达 vector verifier 前已于 split pass 失败。

已验证的实际行为

在上游 maind852dd2 上:

  1. pto.jit 在省略 kernel_kind 时正确保留 module_spec.kernel_kind="vector",并标记 kernel_kind_explicit=False

  2. ptodsl/ptodsl/_tracing/module_builder.py::_apply_child_module_attrs 没有检查 kernel_kind_explicit。因此它仍将默认值写入 backend child:

    module attributes {pto.backend = "vpto",
                       pto.kernel_kind = #pto.kernel_kind<vector>} {
      func.func @default_kind_cv_tileop_probe() attributes {pto.entry} {
        pto.section.vector { ... }
        pto.section.cube { ... }
        return
      }
    }
  3. VPTOSplitCVModule 将这个 pto.kernel_kind<vector> 当作显式 authored intent。在同一 module 内发现 pto.section.cube 后,verifyExplicitKernelKindMatchesSections 直接报错并终止 emission:

    error: conflicts with explicit pto.kernel_kind on its module
    Error: VPTO emission pipeline failed.
    

因此此前“只生成 vector child、未生成 cube child”的结论不准确。实际是 C/V split 在冲突校验阶段失败,不会产生可继续使用的单 vector 输出。

单变量对照

仅手工删除真实生成 .pto 的 backend child pto.kernel_kind<vector> 属性,不改动任何 operation 或 section:

  • VPTOSplitCVModule 成功拆出一个 vector child 与一个 cube child;
  • 两个 child 分别保留对应的 vector/cube 指令;
  • split 之后才因 POC 中 pto.vlds 缺少 pto.vecscope 触发独立的 vector verifier 错误。

这证明现有 split pass 已能处理“无 kind 的 backend-partitioned child 含 C/V section”这一形态。本 issue 的失败由前端把默认值错误物化为 IR 属性引起。

期望语义与修复边界

PTODSL tracing 应区分运行时默认值和用户显式语义:

  1. 默认路径继续保留 module_spec.kernel_kind="vector"kernel_kind_explicit=False,不改变 native build 的选择逻辑。
  2. 仅当用户显式传入 kernel_kind=,即 kernel_kind_explicit=True 时,_apply_child_module_attrs 才向 PTO IR backend child 写入 pto.kernel_kind
  3. 默认且同时有 Cube/Vector section 的 backend child 不带 pto.kernel_kind;现有 VPTOSplitCVModule 据此 materialize 出 Cube 和 Vector 两个 child。
  4. 显式 kernel_kind="cube"kernel_kind="vector" 继续保持现有的单-kind 行为。

针对这个默认-kind 冲突,不需要先改动 PTOAS 后端 split pass;首要修改点是 PTODSL module_builder.py 的 child 属性写入条件。

影响与回归

该问题阻塞手写 Cube+Vector GEMM 的 VPTO 迁移:默认 entry 无法通过 C/V split,因此不能进入后续 L0C 到 UB、AIC/AIV 同步和代码生成验证。纯 Cube GEMM 与 PTOGemmL1Template 不受此问题影响。

回归应使用一个补齐 pto.vecscope 的合法 TileOp POC,至少覆盖:

  1. 未显式 kernel_kind 且同时含 C/V section:trace child 无 pto.kernel_kind--emit-vpto 产生一个 Cube child 和一个 Vector child;
  2. 显式 kernel_kind="vector":仍只 materialize Vector;
  3. 显式 kernel_kind="cube":仍只 materialize Cube。

环境

复核基线:PTOAS upstream main commit d852dd2,VPTO A5 emission。上述结论基于生成 PTO IR、split pass IR dump 及单变量去属性对照。

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions