Skip to content

[Feature] PTODSL 缺少用于显式 Mixed Kernel 的公共 Cube/Vector Section 接口 【阻塞】 #1076

Description

@erhsh

Summary / 问题背景

PTOAS 后端已经支持以下能力:

  • pto.section.cubepto.section.vector IR。
  • 将包含 Cube/Vector Section 的逻辑 Kernel 拆分为对应的物理 Module。
  • 生成 Cube/Vector Mixed Fatobj。
  • 为同一逻辑 Kernel 生成统一的 Host Stub。

TileLang 的 Mixed Kernel 已经在 TIR 中明确划分了 CUBEVECTOR Block,因此需要将这两个 Block 映射到 PTOAS 的 Cube/Vector Section。

但 PTODSL 当前没有公开且稳定的接口,用于在同一个显式 Kernel 中直接表达 Cube/Vector Section。

Motivation / use case / 当前问题

TileLang 目前只能通过以下方式实现 Mixed Kernel,但都不适合作为正式方案:

  1. 直接调用 MLIR PTO dialect 的私有接口,创建 SectionCubeOpSectionVectorOp
  2. 修改 _compiler._module_specModuleStyle 等 PTODSL 内部状态。
  3. 将 Section outline 成 helper,再要求 PTOAS 在 CV split 前进行额外内联。
  4. 由 TileLang 自行生成 Cube/Vector 双 Module,并负责公共代码复制、ABI 对齐和 Fatobj 组织。

其中前两种依赖 PTODSL 私有实现;后两种会把本应由 PTOAS/PTODSL 负责的 Mixed Kernel 结构处理转移到 TileLang。

Proposed API / behavior / 期望能力

希望 PTOAS/PTODSL 提供正式的公共接口,使外部前端能够表达类似以下结构:

@pto.jit(target="a5", mode="explicit")
def kernel(...):
    # Cube/Vector 共用的参数及标量前导计算

    with pto.cube_section():
        # Cube GEMM、L1/L0 搬运及跨核同步

    with pto.vector_section():
        # L0C -> UB -> GM 以及跨核同步

该接口需要明确并支持:

  • 一个逻辑 Kernel 中同时存在 Cube 和 Vector Section。
  • Section 可以使用函数参数以及 Section 外定义的 SSA 值。
  • Section 内定义的 SSA 值不得非法逃逸到另一个 Section。
  • 不要求调用方修改 kernel_kindModuleStyle 或 compiler 内部状态。
  • 生成结果能够直接被 VPTOSplitCVModule 和标准 VPTO Pipeline 处理。
  • 支持 Mixed Kernel 所需的跨核同步,包括动态 Event ID。
  • 明确 Section 采用直接 Region 语义还是 outlined helper 语义。

Alternatives considered / 验收场景

使用公共 PTODSL API 编写一个 Mixed GEMM Kernel:

  • Cube 侧执行 GM/L1/L0 搬运和 GEMM。
  • Vector 侧执行 L0C → UB → GM。
  • Cube 与 Vector 之间通过跨核 Flag 同步。
  • Event ID 可以由运行时标量计算得到。
  • Section 外的公共标量可以分别被两个 Section 使用。
  • 不依赖任何 PTODSL 私有 Python API。
  • 通过标准命令成功生成 Mixed Fatobj:
ptoas --pto-arch=a5 --pto-backend=vpto kernel.pto -o kernel.o

Additional context / 建议定性

这是一个 PTODSL 公共能力缺口。如果 PTOAS 已经承诺支持基于 Section 的 Mixed Kernel authoring,也可以视为“设计已有但前端实现尚未补齐”的功能性 Bug。

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions