Skip to content

[Feature] 将 with pto.simt(x, y, z) 改为 Inline SIMT 实现 #495

Description

@and0d0

Summary

当前simt(x,y,z)虽然在接口层面看上去是inline的,但在实现层面不是。

Motivation / use case

背景

PTODSL 当前同时支持两类 SIMT 写法。

第一类是显式 helper 形式:

  @pto.simt
  def body(...):
      ...

  body[dim_x, dim_y, dim_z](...)

或者:

pto.simt_launch(body, dims=(dim_x, dim_y, dim_z), ...)

这类语义天然是“定义一个 SIMT helper,再通过 pto.simt_launch 启动它”。

第二类是 inline context-manager 形式:

  with pto.simt(dim_x, dim_y, dim_z):
      ...

这个写法更像“在当前 kernel 里直接写一段 SIMT scope”,不需要显式定义 helper 函数。
但它只是语法上 inline,实际 lowering 时仍然会:

  1. 捕获 with pto.simt(...) 块里的 op;
  2. 收集外部 captured values;
  3. 创建一个新的 helper function;
  4. 给 helper function 加 pto.simt_entry;
  5. 在原位置创建:

pto.simt_launch @helper(...)

因此当前 IR 形态大致是:

  func.func @kernel(...) {
    pto.simt_launch @inline_simt_0(%dim_x, %dim_y, %dim_z, %captures...)
    ...
  }

  func.func @inline_simt_0(%captures...) attributes {pto.simt_entry} {
    ...
    return
  }

这不是用户直觉上的 inline SIMT scope。

当前实现带来的问题

当前实现把 inline context 写法和 helper launch 写法混在了一起,导致几个问题:

  1. with pto.simt(dim_x, dim_y, dim_z) 仍然生成 helper function 和 pto.simt_launch,不是直接在当前位置展开 SIMT section。
  2. inline scope 需要 capture analysis,把外部值变成 helper 参数,复杂度比真正 inline 高。
  3. 对后续 pto.alloc_buffer(...)、runtime scalar、临时 vector value、局部 SSA value 的作用域判断会更复杂。表面上是 inline scope,实际却跨了 helper ABI 边界。
  4. 用户读 MLIR 时会看到额外 helper symbol,不利于把 DSL 源码和 IR 对齐。
  5. 当前 _outline_inline_subkernel() 同时承担 simd/cube/simt 的 outline 逻辑。SIMT inline dims 分支实际是“inline syntax + launch helper semantics”,命名和行为不够一致。

目标

将:

  with pto.simt(dim_x, dim_y, dim_z):
      ...

实现为真正 inline 的 SIMT scope。

Proposed API / behavior

No response

Alternatives considered

No response

Additional context

No response

Metadata

Metadata

Assignees

Labels

enhancementNew feature or request

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions