Add pto.func helpers#952
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the @pto.func decorator to define reusable PTODSL helper functions with AST-rewritten control flow, along with documentation, cache-signature helpers, tracing/session lowering support, and comprehensive tests. The review feedback highlights critical improvements for robustness, including handling string annotations when from __future__ import annotations is enabled, correcting the parameter validation order to prevent premature type errors, and adding support for implicit conversion of float literals.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| def __init__(self, spec: FuncSpec, py_fn, *, ast_rewrite: bool = True, returns=_RETURNS_UNSET): | ||
| self.spec = spec | ||
| self.py_fn = py_fn | ||
| self._ast_rewrite = ast_rewrite | ||
| self.signature = inspect.signature(py_fn) | ||
| if returns is not _RETURNS_UNSET: | ||
| self.declared_returns = returns | ||
| elif self.signature.return_annotation is not inspect.Signature.empty: | ||
| self.declared_returns = self.signature.return_annotation | ||
| else: | ||
| raise TypeError( | ||
| "@pto.func helpers must explicitly declare return types with " | ||
| "@pto.func(returns=...) or a Python return annotation; use " | ||
| "returns=None or -> None for helpers that do not return values" | ||
| ) | ||
| update_wrapper(self, py_fn) |
There was a problem hiding this comment.
当用户在定义 @pto.func 的 Python 文件中启用了 from __future__ import annotations 时，inspect.signature(py_fn) 返回的参数和返回值注解将是字符串（例如 'pto.i32'），而不是实际的类型对象。这会导致后续的类型解析和缓存签名生成失败。\n\n建议使用标准库中的 typing.get_type_hints 来安全地解析和获取真实的类型注解，从而完美兼容 from __future__ import annotations。
def __init__(self, spec: FuncSpec, py_fn, *, ast_rewrite: bool = True, returns=_RETURNS_UNSET):
self.spec = spec
self.py_fn = py_fn
self._ast_rewrite = ast_rewrite
self.signature = inspect.signature(py_fn)
import typing
try:
self.type_hints = typing.get_type_hints(py_fn)
except Exception:
self.type_hints = {}
if returns is not _RETURNS_UNSET:
self.declared_returns = returns
elif "return" in self.type_hints:
self.declared_returns = self.type_hints["return"]
elif self.signature.return_annotation is not inspect.Signature.empty:
self.declared_returns = self.signature.return_annotation
else:
raise TypeError(
"@pto.func helpers must explicitly declare return types with "
"@pto.func(returns=...) or a Python return annotation; use "
"returns=None or -> None for helpers that do not return values"
)
update_wrapper(self, py_fn)| for name, param in func_template.signature.parameters.items(): | ||
| value = self._normalize_ptodsl_func_argument( | ||
| name, | ||
| param, | ||
| bound.arguments[name], | ||
| ) | ||
| normalized_values[name] = value | ||
| ordered_arg_values.append(value) | ||
| if param.kind not in { | ||
| inspect.Parameter.POSITIONAL_ONLY, | ||
| inspect.Parameter.POSITIONAL_OR_KEYWORD, | ||
| inspect.Parameter.KEYWORD_ONLY, | ||
| }: | ||
| raise TypeError("@pto.func helpers do not support var-positional or var-keyword parameters yet") |
There was a problem hiding this comment.
这里有两个可以改进的地方：\n1. 参数类型校验顺序问题：目前参数类型的校验（param.kind）是在尝试解析和转换参数之后进行的。如果用户传入了不支持的参数类型（例如 *args 或 **kwargs），_normalize_ptodsl_func_argument 会先抛出一个通用的 TypeError，导致更具体的错误信息永远无法被触发。建议将 param.kind 的校验移动到循环的最开始。\n2. 兼容 from __future__ import annotations：当用户启用了 from __future__ import annotations 时，param.annotation 将是字符串，导致类型解析失败。建议配合 FuncTemplate 中解析好的 type_hints，将解析后的真实类型注解传递给 _normalize_ptodsl_func_argument。
| for name, param in func_template.signature.parameters.items(): | |
| value = self._normalize_ptodsl_func_argument( | |
| name, | |
| param, | |
| bound.arguments[name], | |
| ) | |
| normalized_values[name] = value | |
| ordered_arg_values.append(value) | |
| if param.kind not in { | |
| inspect.Parameter.POSITIONAL_ONLY, | |
| inspect.Parameter.POSITIONAL_OR_KEYWORD, | |
| inspect.Parameter.KEYWORD_ONLY, | |
| }: | |
| raise TypeError("@pto.func helpers do not support var-positional or var-keyword parameters yet") | |
| for name, param in func_template.signature.parameters.items(): | |
| if param.kind not in { | |
| inspect.Parameter.POSITIONAL_ONLY, | |
| inspect.Parameter.POSITIONAL_OR_KEYWORD, | |
| inspect.Parameter.KEYWORD_ONLY, | |
| }: | |
| raise TypeError("@pto.func helpers do not support var-positional or var-keyword parameters yet") | |
| annotation = getattr(func_template, "type_hints", {}).get(name, param.annotation) | |
| value = self._normalize_ptodsl_func_argument( | |
| name, | |
| annotation, | |
| bound.arguments[name], | |
| ) | |
| normalized_values[name] = value | |
| ordered_arg_values.append(value) |
| def _normalize_ptodsl_func_argument(self, name: str, param, value): | ||
| raw_value = unwrap_surface_value(value) | ||
| if hasattr(raw_value, "type"): | ||
| return raw_value | ||
| if param.annotation is not inspect.Parameter.empty: | ||
| try: | ||
| target_type = _resolve(param.annotation) | ||
| except Exception: | ||
| target_type = param.annotation |
There was a problem hiding this comment.
配合 lower_ptodsl_func_call 中的改进，建议将 _normalize_ptodsl_func_argument 的第二个参数从 param 改为直接接收解析后的 annotation，以完美兼容 from __future__ import annotations 带来的字符串注解问题。
| def _normalize_ptodsl_func_argument(self, name: str, param, value): | |
| raw_value = unwrap_surface_value(value) | |
| if hasattr(raw_value, "type"): | |
| return raw_value | |
| if param.annotation is not inspect.Parameter.empty: | |
| try: | |
| target_type = _resolve(param.annotation) | |
| except Exception: | |
| target_type = param.annotation | |
| def _normalize_ptodsl_func_argument(self, name: str, annotation, value): | |
| raw_value = unwrap_surface_value(value) | |
| if hasattr(raw_value, "type"): | |
| return raw_value | |
| if annotation is not inspect.Parameter.empty: | |
| try: | |
| target_type = _resolve(annotation) | |
| except Exception: | |
| target_type = annotation |
| if isinstance(raw_value, bool): | ||
| return const(int(raw_value), dtype=int1) | ||
| if isinstance(raw_value, int): | ||
| return const(raw_value) |
There was a problem hiding this comment.
在 _normalize_ptodsl_func_argument 中，目前仅支持将 bool 和 int 类型的 Python 字面量隐式转换为 PTODSL 常量。如果用户在未标注类型的参数中传入了 float 字面量（例如 0.5），将会触发 TypeError。\n\n建议增加对 float 类型字面量的支持，允许其自动转换为 PTODSL 常量，从而提升易用性。
if isinstance(raw_value, bool):
return const(int(raw_value), dtype=int1)
if isinstance(raw_value, int):
return const(raw_value)
if isinstance(raw_value, float):
return const(raw_value)
Zhendong404
left a comment
Zhendong404
left a comment
There was a problem hiding this comment.
Review 完成。以下结论基于 PR head（ea90014a）的完整源码分析，其中最关键的一条我用 PR 里的 _ast_rewrite.py 做了独立运行验证。按严重程度排列：
严重：静默错误代码生成
1. return 出现在被重写的 if/for 内部时，生成语义完全错误的 IR，且无任何报错
rewriter 把动态 if 改写成 with pto.if_(cond) as br: with br.then_: ...，但对分支体里的 return 不做任何处理（break/continue 有显式报错，return 没有）。我把 PR head 的 _ast_rewrite.py（纯 stdlib 可独立运行）配上 fake pto 实测：
@pto.func(returns=pto.i32)
def f(x: pto.i32):
if x > 0:
return x # early return
return pto.const(0, dtype=pto.i32)trace 时 return x 在 with 块内直接从 Python 函数返回——else_ 分支和 if 之后的所有语句永远不会被 trace，随后 lower_ptodsl_func_call 用 x 发出 func.ReturnOp。结果：helper 无条件返回 x，IR 能通过 module.verify()，但语义全错。for 循环体内 return 同理（循环只 trace 一次就退出）。
这是该 PR 最大的风险：@pto.func 的全部价值就在于返回值，而"分支内提前 return"是写多分支 helper 最自然的写法。修复方向：rewriter 在动态分支/循环体内检测到 ast.Return（以及 Yield）时抛 PTODSLAstRewriteError，至少在支持前 fail-fast。
2. 非标量参数（TensorView / PartitionTensorView / TileValue）的 surface 元数据在 helper 边界丢失
_normalize_ptodsl_func_argument（session.py:824）返回的是 unwrap_surface_value 后的裸 MLIR value，随后 wrap_like_surface_value(normalized_values[name], entry_arg)（session.py:579）拿到的 template 是裸 value，退化为 wrap_surface_value：
TensorViewValue→shape=None, strides=NonePartitionTensorViewValue→root_tensor_view=None, offsets=None, sizes=None，.shape/.strides全丢TileValue→valid_shape等 metadata 全丢
对比 subkernel 路径（session.py:505-530）：arg_templates = tuple(args) 保留原始 surface wrapper，wrap_like_surface_value 显式逐字段复制元数据——说明这些元数据是 body 内 op 所依赖的。后果：传 partition/tile 参数的 helper，体内再做 partition 或读 shape 会失败或静默走动态路径。修复很简单：归一化时同时保留原始 surface value 作为 wrap template。返回值一侧同理（wrap_surface_value 默认包装），但目前 returns 基本只能声明标量，问题主要在参数侧。
高：正确性与兼容性
3. 递归/互递归无任何防护
PR body 承认"递归且有 return 的情况没考虑"，但代码层面是静默放行：get_or_create_helper_function 在 body emission 前就注册（session.py:944），自递归会得到 created=False 并机械地产出递归 func.call——MLIR 语法合法、能通过 verify，但 vpto/emitc 后端和设备执行几乎肯定不支持递归，用户只会在很下游收到莫名错误。若每层递归调用参数类型不同（触发新特化），还会 trace 期无限递归直到 RecursionError。建议在 session 维护"正在 emission 的 helper key"集合，重入时给出明确诊断。
4. PEP 563（from __future__ import annotations）破坏注解处理
用户文件写了 future annotations 后（ptodsl 自己的代码全是这个风格，用户极可能照抄），inspect.signature 拿到的注解全是字符串：
-> pto.i32→declared_returns = "pto.i32"；_resolve()对非_DType原样返回（_types.py:185），字符串直接进了func.FunctionType.get(..., ["pto.i32"])→ pybind 层崩溃，报错莫名。-> None变成"None"，绕过is None/is type(None)检查（session.py:852），同样崩。- 参数注解
"pto.i64"→ coerce 抛TypeError被except: pass吞掉（session.py:833-840）→ 落到 int 兜底 → 静默用默认 dtype 而非注解 dtype，类型错了但没有任何提示。
修复：用 typing.get_type_hints(py_fn) 解析注解，而不是直接读 inspect.signature 的字符串。
5. helper 符号名哈希跨进程不稳定，破坏可复现构建
FuncTemplate.__ptodsl_cache_signature__ 含 id(self.py_fn)；且 cache_signature_atom(returns=pto.i32) 对 _DType 原样返回对象，其 __repr__ 是 <pto.dtype <function ... at 0x7f...>>（_types.py:84），带内存地址。两者经 repr(cache_key) 进 sha1（session.py:66-69）→ 同一份源码每次编译生成不同的 __ptodsl_<hash> 后缀。PR 自己的 probe 注释也承认 "may differ across runs"，测试里也只能用 regex 匹配哈希。后果：AOT/落盘产物符号名不可复现；任何按模块文本做 key 的下游缓存/比对失效。建议：identity 改用 __module__ + __qualname__ + 源码/字节码哈希；给 _DType 实现稳定 repr 或 __ptodsl_cache_signature__（atom 钩子已支持）。
6. HelperFunctionSpec.cache_key 新增 identity 字段导致所有既有 helper 符号哈希全部变化
即使 identity 为空，repr 也多了 , () → sha1 变 → 所有 subkernel / kernel-module helper 的既有 __ptodsl_<hash> 全部改名。这是无谓的 churn（identity 为空时不纳入 key 即可），下游若有 golden 文件或工具依赖这些名字会全灭。
7. pto.const_expr 参数静默降级为运行时参数
const_expr 不是 _DType，_resolve 原样返回 → coerce 失败被吞 → int 兜底 const(...) → 用户声明的编译期常量静默变成运行时 i32 参数；函数体若在 trace-time 上下文使用它（static_range 等）会以莫名的 native_python_control_flow_error 失败。应显式支持 constexpr 参数或显式拒绝。
中：缺失的防护
8. @pto.func 与 subkernel/simt 域互相调用无护栏。 lower_ptodsl_func_call 用 suspend_subkernel_scope() 清栈，导致：(a) 在 @pto.simt 体内调 @pto.func → 普通 helper 被 simt 函数调用；(b) 在 @pto.func 体内调 @pto.simt → 因栈已空，绕过 session.py:646 的 simt 嵌套检查，pto.simt_entry helper 被普通 func 调用、StoreVfSimtInfoOp 出现在任意函数里。这些组合的合法性完全未验证、无测试。PR 目标声明覆盖 "pto.func/pto.jit/pto.simt 函数域"，但跨域组合要么显式拒绝要么补测试。
9. 闭包捕获 traced SSA 值 → 跨函数 SSA 引用。 helper 定义在 kernel 体内并捕获 trace 中的 runtime value 时，rewrite 会把闭包变量注入为 kwonly 默认值，helper body 直接引用 caller 函数的 SSA → dominance 违反，verify 报错难懂。应检测自由变量中的 surface value 并给清晰诊断（subkernel 路径同样存在，但 func 更容易踩到）。
10. 参数 kind 检查是死代码且报错顺序误导。 *args/**kwargs 会在归一化阶段先以 "expects a traced runtime value" 报错，而不是后面的 "do not support var-positional..."；同样的 raise 在 session.py 出现了两次。另外归一化可能已往 caller 体内发了 const op 才抛错（无害但脏）。
低：卫生与测试
tmp/ptodsl_func_chain_probe.py和.mlir是开发探针，且 .mlir 内含不稳定哈希、无回归价值，不应合入仓库。test/lit/vpto/mte_ub_gm_l2_cache_ctl.pto的改动（LABEL 加(）与本 PR 无关，是 drive-by 修复，建议拆分。- 测试只做到
compile().mlir_text()+ parse/verify，即只覆盖 tracing 层；多返回值 helper 经过 vpto/emitc 后端 lowering 的端到端路径完全没测。错误路径（返回值个数/类型不匹配）有实现但只有一条expect_raises。缺：early return、递归、PEP 563、const_expr、keyword-only/默认值、非标量参数、跨域调用的用例。 - 文档未写明上述限制：分支内 return 的行为、while 不支持、递归未支持、future-annotations 限制、非 trace 环境调用
@pto.func会 RuntimeError。 emit_body每次 compile 重新做 AST rewrite（subkernel 也一样，属既有模式，但 kernel 路径有tracing_callbackmemoize，这里可以顺手对齐）。
做得对的地方
返回值显式声明（returns=/注解）的契约设计、按参数类型特化 + 同签名复用同一 helper、source-less 时回退原函数并有测试覆盖、_cache_signature.py 抽取去重——这些都是合理的。
总结：建议合并前至少修掉 #1（early-return 静默错码，必须 fail-fast）和 #2（surface 元数据丢失），补上 #3 的递归防护和 #4 的 PEP 563 处理；#5/#6 的哈希稳定性涉及可复现构建，也值得在本 PR 内一并解决，否则这个不稳定后缀会随着 @pto.func 的推广沉到所有下游产物里。
|
2 participants
TODO：
暂无
#946