Summary / 问题背景
PTOAS 后端已经支持以下能力:
pto.section.cube 和 pto.section.vector IR。
- 将包含 Cube/Vector Section 的逻辑 Kernel 拆分为对应的物理 Module。
- 生成 Cube/Vector Mixed Fatobj。
- 为同一逻辑 Kernel 生成统一的 Host Stub。
TileLang 的 Mixed Kernel 已经在 TIR 中明确划分了 CUBE 和 VECTOR Block,因此需要将这两个 Block 映射到 PTOAS 的 Cube/Vector Section。
但 PTODSL 当前没有公开且稳定的接口,用于在同一个显式 Kernel 中直接表达 Cube/Vector Section。
Motivation / use case / 当前问题
TileLang 目前只能通过以下方式实现 Mixed Kernel,但都不适合作为正式方案:
- 直接调用 MLIR PTO dialect 的私有接口,创建
SectionCubeOp 和 SectionVectorOp。
- 修改
_compiler._module_spec、ModuleStyle 等 PTODSL 内部状态。
- 将 Section outline 成 helper,再要求 PTOAS 在 CV split 前进行额外内联。
- 由 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_kind、ModuleStyle 或 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。
Summary / 问题背景
PTOAS 后端已经支持以下能力:
pto.section.cube和pto.section.vectorIR。TileLang 的 Mixed Kernel 已经在 TIR 中明确划分了
CUBE和VECTORBlock,因此需要将这两个 Block 映射到 PTOAS 的 Cube/Vector Section。但 PTODSL 当前没有公开且稳定的接口,用于在同一个显式 Kernel 中直接表达 Cube/Vector Section。
Motivation / use case / 当前问题
TileLang 目前只能通过以下方式实现 Mixed Kernel,但都不适合作为正式方案:
SectionCubeOp和SectionVectorOp。_compiler._module_spec、ModuleStyle等 PTODSL 内部状态。其中前两种依赖 PTODSL 私有实现;后两种会把本应由 PTOAS/PTODSL 负责的 Mixed Kernel 结构处理转移到 TileLang。
Proposed API / behavior / 期望能力
希望 PTOAS/PTODSL 提供正式的公共接口,使外部前端能够表达类似以下结构:
该接口需要明确并支持:
kernel_kind、ModuleStyle或 compiler 内部状态。VPTOSplitCVModule和标准 VPTO Pipeline 处理。Alternatives considered / 验收场景
使用公共 PTODSL API 编写一个 Mixed GEMM Kernel:
Additional context / 建议定性
这是一个 PTODSL 公共能力缺口。如果 PTOAS 已经承诺支持基于 Section 的 Mixed Kernel authoring,也可以视为“设计已有但前端实现尚未补齐”的功能性 Bug。