From d35be5504eb0f6b99062098c43a06a62b1658b0f Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 16 Jul 2026 16:32:18 +0800 Subject: [PATCH 01/15] Add reusable pto.func helpers --- .../03-kernel-entry-and-subkernels.md | 32 +++- ptodsl/docs/user_guide/05-control-flow.md | 16 +- ptodsl/ptodsl/_cache_signature.py | 58 ++++++ ptodsl/ptodsl/_func.py | 80 +++++++++ ptodsl/ptodsl/_kernel_compilation.py | 51 +----- ptodsl/ptodsl/_tracing/runtime.py | 5 + ptodsl/ptodsl/_tracing/session.py | 165 +++++++++++++++++- ptodsl/ptodsl/pto.py | 1 + ptodsl/tests/test_jit_compile.py | 95 ++++++++++ 9 files changed, 446 insertions(+), 57 deletions(-) create mode 100644 ptodsl/ptodsl/_cache_signature.py create mode 100644 ptodsl/ptodsl/_func.py diff --git a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md index 08987e0ea1..b636f379f3 100644 --- a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md +++ b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md @@ -3,10 +3,11 @@ PTODSL provides one kernel decorator (`@pto.jit`) with two roles (`entry=True` / `entry=False`), two compilation backends (`vpto` / `emitc`), and two reusable compute helper decorators (`@pto.tileop` and `@pto.simt`), -plus inline unit-specific context managers. This chapter covers +plus reusable PTODSL helper functions (`@pto.func`) and inline unit-specific +context managers. This chapter covers the `@pto.jit` entry and module contracts, the two programming models, the two -compilation backends, sub-kernel reference, parameter contracts, and boundary -constraints. +compilation backends, helper functions, sub-kernel reference, parameter +contracts, and boundary constraints. ## 3.1 `@pto.jit` — roles, backends, and modes @@ -22,6 +23,7 @@ Decorator overview: mode="explicit" micro-instruction authoring, user-managed staging @pto.tileop Single-core Tile/scalar compute helper with inferred Vector/Cube kind +@pto.func Reusable PTODSL helper, no host/module ABI boundary @pto.simt Explicitly launched SIMT helper with pointer/scalar ABI ``` @@ -59,6 +61,30 @@ manual-address, user-managed staging contract of explicit kernels. (`@pto.tileop` and `@pto.simt`) define sub-kernels that are called from within `@pto.jit` bodies. +`@pto.func` is the lightweight helper boundary for reusable PTODSL code. It +does not create a host-launchable entry, a kernel-module ABI, or a hardware-unit +sub-kernel section. When traced PTODSL code calls a `@pto.func` helper, PTODSL +materializes one helper `func.func` in the caller's active compilation context +and emits `func.call` at call sites. Supported native Python `if` and +`for range(...)` in the helper body use the same AST rewrite path as `@pto.jit` +and named sub-kernels. The helper can return PTODSL runtime values, including +multiple values via a tuple. + +```python +@pto.func +def add_rows(total: pto.i32, rows: pto.i32): + one = pto.const(1, dtype=pto.i32) + for _ in range(rows): + total = total + one + return total + + +@pto.jit(target="a5") +def kernel(rows: pto.i32): + total = add_rows(pto.const(0, dtype=pto.i32), rows) + _ = total +``` + ## 3.2 `entry=True` — host-launchable kernel entry diff --git a/ptodsl/docs/user_guide/05-control-flow.md b/ptodsl/docs/user_guide/05-control-flow.md index 0afca4f409..049d6b498f 100644 --- a/ptodsl/docs/user_guide/05-control-flow.md +++ b/ptodsl/docs/user_guide/05-control-flow.md @@ -4,7 +4,7 @@ PTODSL uses a **tracing** compilation model. When you call `kernel.compile(...)` This has one critical implication for how you write loops and branches: -- **Python native `for`/`if`** is rewritten to device-side control flow by default in `@pto.jit` bodies and named `@pto.tileop` / `@pto.simt` helpers. A `for i in range(rows)` loop records a device loop, and a runtime `if` records both branches. +- **Python native `for`/`if`** is rewritten to device-side control flow by default in `@pto.jit` bodies, `@pto.func` helpers, and named `@pto.tileop` / `@pto.simt` helpers. A `for i in range(rows)` loop records a device loop, and a runtime `if` records both branches. - **`pto.const_expr` / `pto.static_range`** keep compile-time Python behavior when you want trace-time specialization or unrolling. - **`pto.for_` / `pto.if_`** produce device-side control flow. The loop bound or branch condition can be a runtime value, and the hardware will execute the loop or take the branch dynamically. @@ -249,11 +249,19 @@ This lets you write a single kernel that specializes into different strategies b ## 5.5 Native Python control-flow rewrite -`@pto.jit` rewrites supported native Python control flow before tracing. In the -default mode, plain Python `if` and `for range(...)` in the rewritten scope -become device-side control flow. Use `pto.const_expr(...)` and +`@pto.jit`, `@pto.func`, and named `@pto.cube` / `@pto.simd` / `@pto.simt` +callables rewrite supported native Python control flow before tracing their +bodies. In the default mode, plain Python `if` and `for range(...)` in the +rewritten scope become device-side control flow. Use `pto.const_expr(...)` and `pto.static_range(...)` when you want trace-time behavior. +PTODSL does not recursively rewrite arbitrary undecorated Python callees. If an +external helper should contain runtime native control flow, decorate it with +`@pto.func` or one of the other PTODSL callable decorators. A plain Python +helper is still executed during tracing; static `range(...)` loops in such a +helper unroll at trace time, and runtime loop bounds or branch conditions are +not converted into `scf.for` / `scf.if`. + ### Runtime branches By default, a native Python `if` becomes a device-side conditional: diff --git a/ptodsl/ptodsl/_cache_signature.py b/ptodsl/ptodsl/_cache_signature.py new file mode 100644 index 0000000000..ebfee276ec --- /dev/null +++ b/ptodsl/ptodsl/_cache_signature.py @@ -0,0 +1,58 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +"""Shared cache-signature helpers for PTODSL tracing frontends.""" + +from __future__ import annotations + + +def closure_cache_signature(fn): + """Return one stable cache signature for the closure state captured by *fn*.""" + try: + import inspect + + closure_vars = inspect.getclosurevars(fn) + except TypeError: + return () + return tuple( + (name, cache_signature_atom(value)) + for name, value in sorted(closure_vars.nonlocals.items()) + ) + + +def cache_signature_atom(value): + """Return one hashable cache-signature atom for arbitrary captured values.""" + cache_signature = getattr(value, "__ptodsl_cache_signature__", None) + if callable(cache_signature): + return ("ptodsl-cache-signature", cache_signature_atom(cache_signature())) + try: + hash(value) + except TypeError: + if isinstance(value, dict): + items = ( + (cache_signature_atom(key), cache_signature_atom(item)) + for key, item in value.items() + ) + return ("dict", tuple(sorted(items, key=repr))) + if isinstance(value, (list, tuple)): + return ( + type(value).__name__, + tuple(cache_signature_atom(item) for item in value), + ) + if isinstance(value, set): + return ( + "set", + tuple(sorted((cache_signature_atom(item) for item in value), key=repr)), + ) + return (type(value).__name__, repr(value)) + return value + + +__all__ = [ + "cache_signature_atom", + "closure_cache_signature", +] diff --git a/ptodsl/ptodsl/_func.py b/ptodsl/ptodsl/_func.py new file mode 100644 index 0000000000..62fa9989b4 --- /dev/null +++ b/ptodsl/ptodsl/_func.py @@ -0,0 +1,80 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +"""``@pto.func`` decorator and reusable callable handle.""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import update_wrapper +import inspect + +from ._ast_rewrite import rewrite_jit_function +from ._cache_signature import closure_cache_signature +from ._tracing import current_runtime + + +@dataclass(frozen=True) +class FuncSpec: + """Declarative metadata for a rewrite-capable reusable PTODSL callable.""" + + symbol_name: str + + +class FuncTemplate: + """Callable decorated PTODSL helper surface.""" + + def __init__(self, spec: FuncSpec, py_fn, *, ast_rewrite: bool = True): + self.spec = spec + self.py_fn = py_fn + self._ast_rewrite = ast_rewrite + self.signature = inspect.signature(py_fn) + update_wrapper(self, py_fn) + + def emit_body(self, *args, **kwargs): + """Emit this helper body into the currently active trace.""" + py_fn = rewrite_jit_function(self.py_fn) if self._ast_rewrite else self.py_fn + return py_fn(*args, **kwargs) + + def __call__(self, *args, **kwargs): + runtime = current_runtime() + if runtime is None: + raise RuntimeError( + "@pto.func helpers may only be called while tracing a compatible PTODSL kernel" + ) + return runtime.dispatch_ptodsl_func_call(self, *args, **kwargs) + + def __ptodsl_cache_signature__(self): + return ( + type(self).__name__, + self.spec.symbol_name, + id(self.py_fn), + self._ast_rewrite, + closure_cache_signature(self.py_fn), + ) + + +def func(fn=None, *, name: str | None = None, ast_rewrite: bool = True): + """Decorate a Python function as a reusable PTODSL callable helper.""" + + def decorator(py_fn): + return FuncTemplate( + FuncSpec(symbol_name=name or py_fn.__name__), + py_fn, + ast_rewrite=ast_rewrite, + ) + + if fn is not None: + return decorator(fn) + return decorator + + +__all__ = [ + "FuncSpec", + "FuncTemplate", + "func", +] diff --git a/ptodsl/ptodsl/_kernel_compilation.py b/ptodsl/ptodsl/_kernel_compilation.py index 3d2217ee44..dc3561f020 100644 --- a/ptodsl/ptodsl/_kernel_compilation.py +++ b/ptodsl/ptodsl/_kernel_compilation.py @@ -9,9 +9,8 @@ from __future__ import annotations -import inspect - from ._ast_rewrite import rewrite_jit_function +from ._cache_signature import closure_cache_signature from ._diagnostics import ( jit_source_compile_constexpr_error, kernel_module_compile_error, @@ -101,7 +100,7 @@ def compile(self, **constexpr_bindings): if self._ast_rewrite: kernel_identity = ( kernel_identity, - _closure_cache_signature(self._callback), + closure_cache_signature(self._callback), ) specialization_key = self._kernel_signature.specialization_key( kernel_identity, @@ -165,52 +164,6 @@ def cached_specializations(self): return tuple(self._compiled_cache.values()) -def _closure_cache_signature(fn): - try: - closure_vars = inspect.getclosurevars(fn) - except TypeError: - return () - return tuple( - (name, _cache_signature_atom(value)) - for name, value in sorted(closure_vars.nonlocals.items()) - ) - - -def _cache_signature_atom(value): - cache_signature = getattr(value, "__ptodsl_cache_signature__", None) - if callable(cache_signature): - return ("ptodsl-cache-signature", _cache_signature_atom(cache_signature())) - try: - hash(value) - except TypeError: - if isinstance(value, dict): - items = ( - (_cache_signature_atom(key), _cache_signature_atom(item)) - for key, item in value.items() - ) - return ( - "dict", - tuple(sorted(items, key=repr)), - ) - if isinstance(value, (list, tuple)): - return ( - type(value).__name__, - tuple(_cache_signature_atom(item) for item in value), - ) - if isinstance(value, set): - return ( - "set", - tuple( - sorted( - (_cache_signature_atom(item) for item in value), - key=repr, - ) - ), - ) - return (type(value).__name__, repr(value)) - return value - - __all__ = [ "CompiledKernelHandle", "KernelCompiler", diff --git a/ptodsl/ptodsl/_tracing/runtime.py b/ptodsl/ptodsl/_tracing/runtime.py index e876c25f05..34ba416ef4 100644 --- a/ptodsl/ptodsl/_tracing/runtime.py +++ b/ptodsl/ptodsl/_tracing/runtime.py @@ -67,6 +67,11 @@ def dispatch_subkernel_call(self, subkernel, *args, **kwargs): return session.lower_helper_subkernel(subkernel, *args, **kwargs) return subkernel.emit_body(*args, **kwargs) + def dispatch_ptodsl_func_call(self, func_template, *args, **kwargs): + """Dispatch one ``@pto.func`` helper call in the active trace.""" + session = require_active_session("@pto.func") + return session.lower_ptodsl_func_call(func_template, *args, **kwargs) + def dispatch_kernel_module_call(self, kernel_handle, *args, **kwargs): """Dispatch one ``@pto.jit(entry=False)`` kernel-module call in the active trace.""" session = require_active_session("@pto.jit(entry=False)") diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index 46adcb0f1b..16a295ee8c 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -12,12 +12,14 @@ from contextlib import contextmanager from dataclasses import dataclass import hashlib +import inspect from .._diagnostics import ( inline_subkernel_value_escape_error, inline_tileop_capture_type_error, subkernel_kernel_kind_mismatch_error, ) +from .._scalar_coercion import coerce_scalar_to_type from .._kernel_signature import RuntimeScalarParameterSpec from .._ops import const from .._surface_values import ( @@ -25,13 +27,14 @@ is_tile_ir_type, unwrap_surface_value, wrap_like_surface_value, + wrap_surface_value, ) from .control_flow import ( build_carry_loop_frame, finish_carry_loop_frame, yield_carry_loop_state, ) -from .._types import _strip_integer_signedness +from .._types import _resolve, _strip_integer_signedness, int1 from .module_builder import create_container_child_module from mlir.dialects import arith, func @@ -45,6 +48,7 @@ IntegerType, Operation, StringAttr, + TypeAttr, UnitAttr, ) @@ -57,6 +61,7 @@ class HelperFunctionSpec: arg_types: tuple result_types: tuple = () attributes: tuple[tuple[str, object], ...] = () + identity: tuple = () def cache_key(self) -> tuple: """Return one stable ABI-sensitive cache key for this helper signature.""" @@ -65,6 +70,7 @@ def cache_key(self) -> tuple: tuple(str(arg_type) for arg_type in self.arg_types), tuple(str(result_type) for result_type in self.result_types), tuple((attr_name, str(attr_value)) for attr_name, attr_value in self.attributes), + self.identity, ) def specialized_symbol_name(self) -> str: @@ -536,6 +542,81 @@ def lower_helper_subkernel(self, subkernel, *args, **kwargs): func.CallOp(helper_fn, [unwrap_surface_value(arg) for arg in arg_templates]) return None + def lower_ptodsl_func_call(self, func_template, *args, **kwargs): + """Lower one ``@pto.func`` helper call in the active trace.""" + bound = func_template.signature.bind(*args, **kwargs) + bound.apply_defaults() + ordered_arg_values = [] + normalized_values = {} + 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") + + arg_templates = tuple(ordered_arg_values) + owner_symbol_name = self.current_function_owner_symbol_name + helper_spec = HelperFunctionSpec( + symbol_name=func_template.spec.symbol_name, + arg_types=tuple(unwrap_surface_value(arg).type for arg in arg_templates), + attributes=(("pto.ptodsl.callable_kind", StringAttr.get("func")),), + identity=func_template.__ptodsl_cache_signature__(), + ) + helper_fn, created = self.get_or_create_helper_function( + helper_spec, + owner_symbol_name=owner_symbol_name, + ) + + if created: + entry_block = helper_fn.add_entry_block() + entry_args = tuple(entry_block.arguments) + wrapped_args = [] + wrapped_kwargs = {} + entry_arg_index = 0 + for name, param in func_template.signature.parameters.items(): + entry_arg = entry_args[entry_arg_index] + entry_arg_index += 1 + wrapped_value = wrap_like_surface_value(normalized_values[name], entry_arg) + if param.kind in {inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD}: + wrapped_args.append(wrapped_value) + elif param.kind == inspect.Parameter.KEYWORD_ONLY: + wrapped_kwargs[name] = wrapped_value + else: + raise TypeError("@pto.func helpers do not support var-positional or var-keyword parameters yet") + with ( + self.enter_function(helper_fn, owner_symbol_name=owner_symbol_name), + self.suspend_subkernel_scope(), + InsertionPoint(entry_block), + ): + result = func_template.emit_body(*wrapped_args, **wrapped_kwargs) + return_values = self._normalize_ptodsl_func_return_values( + result, + func_template=func_template, + ) + result_types = tuple(value.type for value in return_values) + helper_fn.attributes["function_type"] = TypeAttr.get( + func.FunctionType.get( + list(helper_spec.arg_types), + list(result_types), + ) + ) + if return_values: + func.ReturnOp([unwrap_surface_value(value) for value in return_values]) + else: + func.ReturnOp([]) + + call_op = func.CallOp(helper_fn, [unwrap_surface_value(arg) for arg in arg_templates]) + return self._wrap_ptodsl_func_call_results(call_op.results) + def begin_carry_loop(self, start, stop, step, state_items): """Materialize one authored ``pto.for_(...).carry(...)`` loop body.""" frame = build_carry_loop_frame(start, stop, step, state_items) @@ -755,6 +836,88 @@ def lookup_helper(self, symbol_name: str): return helper return None + 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 + try: + return coerce_scalar_to_type( + raw_value, + target_type, + context=f"@pto.func parameter {name!r}", + ) + except TypeError: + pass + if isinstance(raw_value, bool): + return const(int(raw_value), dtype=int1) + if isinstance(raw_value, int): + return const(raw_value) + raise TypeError( + f"@pto.func parameter {name!r} expects a traced runtime value or a supported literal, " + f"got {raw_value!r}" + ) + + def _normalize_ptodsl_func_return_values(self, result, *, func_template): + if result is None: + return () + if isinstance(result, tuple): + values = result + elif isinstance(result, list): + values = tuple(result) + else: + values = (result,) + + normalized = [] + return_annotation = func_template.signature.return_annotation + target_type = None + if return_annotation is not inspect.Signature.empty: + try: + target_type = _resolve(return_annotation) + except Exception: + target_type = None + + for index, value in enumerate(values): + raw_value = unwrap_surface_value(value) + if hasattr(raw_value, "type"): + normalized.append(raw_value) + continue + if target_type is not None: + try: + normalized.append( + coerce_scalar_to_type( + raw_value, + target_type, + context=f"@pto.func return value {index}", + ) + ) + continue + except TypeError: + pass + if isinstance(raw_value, bool): + normalized.append(const(int(raw_value), dtype=int1)) + continue + if isinstance(raw_value, int): + normalized.append(const(raw_value)) + continue + raise TypeError( + f"@pto.func return value {index} must be a traced runtime value or supported literal, " + f"got {raw_value!r}" + ) + return tuple(normalized) + + def _wrap_ptodsl_func_call_results(self, results): + if not results: + return None + wrapped = tuple(wrap_surface_value(result) for result in results) + if len(wrapped) == 1: + return wrapped[0] + return wrapped + def _attach_ptodsl_logical_name_attr(self, func_op, logical_name: str) -> None: """Mark one ABI-specialized PTODSL symbol with its authored logical name.""" func_op.attributes["pto.ptodsl.logical_name"] = StringAttr.get(logical_name) diff --git a/ptodsl/ptodsl/pto.py b/ptodsl/ptodsl/pto.py index 669a7b2ac4..fb845f4f81 100644 --- a/ptodsl/ptodsl/pto.py +++ b/ptodsl/ptodsl/pto.py @@ -159,6 +159,7 @@ # ── Decorator ───────────────────────────────────────────────────────────────── from ._jit import jit, KernelHandle, merge_jit_modules # noqa: F401 +from ._func import func # noqa: F401 from ._subkernels import cube, simd, simt, tileop # noqa: F401 from ._pipe_namespace import pipe # noqa: F401 diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 9cf79da38d..155803d917 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -1268,6 +1268,46 @@ def helper(limit, enabled): _ = value +@pto.func +def func_runtime_for_return_helper(limit: pto.i32, initial: pto.i32): + one = pto.const(1, dtype=pto.i32) + total = initial + for _ in range(limit): + total = total + one + return total + + +@pto.func +def func_runtime_if_return_helper(lhs: pto.i32, rhs: pto.i32): + if lhs > rhs: + total = lhs + rhs + else: + total = rhs + lhs + return total + + +@pto.func +def func_multi_return_helper(value: pto.i32): + one = pto.const(1, dtype=pto.i32) + return value, value + one + + +@pto.func +def func_void_helper(): + pto.pipe_barrier(pto.Pipe.ALL) + + +@pto.jit(target="a5") +def ptodsl_func_call_probe(rows: pto.i32): + init = pto.const(0, dtype=pto.i32) + total = func_runtime_for_return_helper(rows, init) + merged = func_runtime_if_return_helper(total, init) + first, second = func_multi_return_helper(merged) + _ = first + second + func_void_helper() + func_void_helper() + + @pto.jit(target="a5") def ast_nested_helper_freevar_if_merge_probe(): lhs = pto.const(4, dtype=pto.i32) @@ -1462,6 +1502,29 @@ def sourceless_subkernel_entry_probe(*, TRACE_TOKEN: pto.const_expr = 0): sourceless_subkernel_entry_probe = make_sourceless_subkernel_entry() +def make_sourceless_ptodsl_func_probe(): + namespace = {"pto": pto} + exec( + """ +@pto.func +def sourceless_ptodsl_func_helper(): + if True: + pto.pipe_barrier(pto.Pipe.ALL) +""", + namespace, + ) + helper = namespace["sourceless_ptodsl_func_helper"] + + @pto.jit(target="a5") + def sourceless_ptodsl_func_probe(*, TRACE_TOKEN: pto.const_expr = 0): + helper() + + return sourceless_ptodsl_func_probe + + +sourceless_ptodsl_func_probe = make_sourceless_ptodsl_func_probe() + + def make_entry_closure_kernel_module_probe(): @pto.jit(target="a5", entry=False) def closure_helper(): @@ -4994,6 +5057,28 @@ def _enter_inline_simt_with_resource_attr(): "rewritten nested helpers should preserve loop-carried and branch live-out values", ) + ptodsl_func_call_text = ptodsl_func_call_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify(ptodsl_func_call_text, "@pto.func helper call specialization") + expect( + re.search(r"func\.func @func_runtime_for_return_helper__ptodsl_[0-9a-f]+\(.*\) -> i32", ptodsl_func_call_text) + is not None, + "@pto.func helpers that return one runtime value should materialize a typed helper result", + ) + expect( + re.search(r"func\.func @func_multi_return_helper__ptodsl_[0-9a-f]+\(.*\) -> \(i32, i32\)", ptodsl_func_call_text) + is not None, + "@pto.func helpers should support multiple returned runtime values", + ) + expect( + ptodsl_func_call_text.count("scf.for") >= 1 and ptodsl_func_call_text.count("scf.if") >= 1, + "@pto.func helper bodies should use native control-flow AST rewrite", + ) + expect( + len(re.findall(r"func\.func @func_void_helper__ptodsl_[0-9a-f]+", ptodsl_func_call_text)) == 1 + and len(re.findall(r"call @func_void_helper__ptodsl_[0-9a-f]+", ptodsl_func_call_text)) == 2, + "repeated @pto.func calls should reuse one materialized helper artifact", + ) + ast_nested_helper_freevar_if_merge_text = ast_nested_helper_freevar_if_merge_probe.compile().mlir_text() expect_parse_roundtrip_and_verify( ast_nested_helper_freevar_if_merge_text, @@ -5117,6 +5202,16 @@ def _enter_inline_simt_with_resource_attr(): "source-less subkernels should fall back to original trace-time Python execution", ) + sourceless_ptodsl_func_text = sourceless_ptodsl_func_probe.compile(TRACE_TOKEN=1).mlir_text() + expect_parse_roundtrip_and_verify( + sourceless_ptodsl_func_text, + "source-less @pto.func AST rewrite fallback specialization", + ) + expect( + sourceless_ptodsl_func_text.count("pto.barrier ") == 1, + "source-less @pto.func helpers should fall back to original trace-time Python execution", + ) + ast_python_bool_guard_enabled_text = ast_python_bool_guard_probe.compile().mlir_text() expect_parse_roundtrip_and_verify( ast_python_bool_guard_enabled_text, From 2e224b1a3a89b096948ed76abccd592279e44b02 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 16 Jul 2026 17:20:00 +0800 Subject: [PATCH 02/15] Add pto.func chain probe artifacts --- tmp/ptodsl_func_chain_probe.mlir | 78 ++++++++++++++++++++++++++++ tmp/ptodsl_func_chain_probe.py | 89 ++++++++++++++++++++++++++++++++ 2 files changed, 167 insertions(+) create mode 100644 tmp/ptodsl_func_chain_probe.mlir create mode 100644 tmp/ptodsl_func_chain_probe.py diff --git a/tmp/ptodsl_func_chain_probe.mlir b/tmp/ptodsl_func_chain_probe.mlir new file mode 100644 index 0000000000..1ff5babfa9 --- /dev/null +++ b/tmp/ptodsl_func_chain_probe.mlir @@ -0,0 +1,78 @@ +// Generated from tmp/ptodsl_func_chain_probe.py while developing issue #946. +// The __ptodsl_ suffixes are specialization hashes and may differ across runs. +// +// Observed behavior: +// 1. Plain undecorated helper: the internal `if True + range(2)` executes during tracing. +// The caller contains two expanded `pto.barrier ` ops and no helper func. +// 2. @pto.func dynamic loop: `for _ in range(limit)` is AST-rewritten and the helper body contains `scf.for`. +// 3. @pto.func dynamic if: `if lhs > rhs` is AST-rewritten and the helper body contains `scf.if`. +// 4. @pto.func(ast_rewrite=False): static `if True + range(2)` does not generate scf. +// It trace-time expands inside the helper body into two `arith.addi` ops. +// 5. Chained calls: `multi_return_helper -> chain_mid -> dyn_loop_helper/dyn_if_helper/no_rewrite_static_helper`. +// 6. Multiple returns: `multi_return_helper` returns `(i32, i32)`. +// 7. Reuse: `dyn_if_helper` is defined once and called twice. +// +// Counts: +// scf.for: 1 +// scf.if: 1 +// func.func @dyn_loop_helper__ptodsl_: 1 +// func.func @dyn_if_helper__ptodsl_: 1 +// func.func @no_rewrite_static_helper__ptodsl_: 1 +// func.func @chain_mid__ptodsl_: 1 +// func.func @multi_return_helper__ptodsl_: 1 +// call @dyn_if_helper__ptodsl_: 2 +// pto.barrier : 2 + +module attributes {pto.target_arch = "a5"} { + module attributes {pto.backend = "vpto", pto.kernel_kind = #pto.kernel_kind, pto.target_arch = "a5"} { + func.func @func_chain_probe(%arg0: i32) attributes {pto.entry} { + %c0_i32 = arith.constant 0 : i32 + %0:2 = call @multi_return_helper__ptodsl_406ab2a269(%arg0, %c0_i32) : (i32, i32) -> (i32, i32) + %1 = call @dyn_if_helper__ptodsl_82f7abb427(%0#0, %0#1) : (i32, i32) -> i32 + pto.barrier + pto.barrier + return + } + func.func @multi_return_helper__ptodsl_406ab2a269(%arg0: i32, %arg1: i32) -> (i32, i32) attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "multi_return_helper"} { + %0 = call @chain_mid__ptodsl_be1027dd8d(%arg0, %arg1) : (i32, i32) -> i32 + %c1_i32 = arith.constant 1 : i32 + %1 = arith.addi %0, %c1_i32 : i32 + return %0, %1 : i32, i32 + } + func.func @chain_mid__ptodsl_be1027dd8d(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "chain_mid"} { + %0 = call @dyn_loop_helper__ptodsl_795afe8017(%arg0, %arg1) : (i32, i32) -> i32 + %1 = call @dyn_if_helper__ptodsl_82f7abb427(%0, %arg1) : (i32, i32) -> i32 + %2 = call @no_rewrite_static_helper__ptodsl_eccc8a4ce9(%1) : (i32) -> i32 + return %2 : i32 + } + func.func @dyn_loop_helper__ptodsl_795afe8017(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "dyn_loop_helper"} { + %c1_i32 = arith.constant 1 : i32 + %c0 = arith.constant 0 : index + %0 = arith.index_cast %arg0 : i32 to index + %c1 = arith.constant 1 : index + %1 = scf.for %arg2 = %c0 to %0 step %c1 iter_args(%arg3 = %arg1) -> (i32) { + %2 = arith.addi %arg3, %c1_i32 : i32 + scf.yield %2 : i32 + } + return %1 : i32 + } + func.func @dyn_if_helper__ptodsl_82f7abb427(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "dyn_if_helper"} { + %0 = arith.cmpi sgt, %arg0, %arg1 : i32 + %1 = scf.if %0 -> (i32) { + %2 = arith.subi %arg0, %arg1 : i32 + scf.yield %2 : i32 + } else { + %2 = arith.subi %arg1, %arg0 : i32 + scf.yield %2 : i32 + } + return %1 : i32 + } + func.func @no_rewrite_static_helper__ptodsl_eccc8a4ce9(%arg0: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "no_rewrite_static_helper"} { + %c1_i32 = arith.constant 1 : i32 + %0 = arith.addi %arg0, %c1_i32 : i32 + %c1_i32_0 = arith.constant 1 : i32 + %1 = arith.addi %0, %c1_i32_0 : i32 + return %1 : i32 + } + } +} diff --git a/tmp/ptodsl_func_chain_probe.py b/tmp/ptodsl_func_chain_probe.py new file mode 100644 index 0000000000..74aec3803e --- /dev/null +++ b/tmp/ptodsl_func_chain_probe.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +"""Probe for mixed ``@pto.func`` AST rewrite and trace-time expansion behavior.""" + +from ptodsl import pto + + +def plain_trace_helper(): + if True: + for _ in range(2): + pto.pipe_barrier(pto.Pipe.ALL) + + +@pto.func +def dyn_loop_helper(limit: pto.i32, value: pto.i32): + one = pto.const(1, dtype=pto.i32) + total = value + for _ in range(limit): + total = total + one + return total + + +@pto.func +def dyn_if_helper(lhs: pto.i32, rhs: pto.i32): + if lhs > rhs: + chosen = lhs - rhs + else: + chosen = rhs - lhs + return chosen + + +@pto.func(ast_rewrite=False) +def no_rewrite_static_helper(value: pto.i32): + total = value + if True: + for _ in range(2): + total = total + pto.const(1, dtype=pto.i32) + return total + + +@pto.func +def chain_mid(limit: pto.i32, seed: pto.i32): + looped = dyn_loop_helper(limit, seed) + branched = dyn_if_helper(looped, seed) + static_expanded = no_rewrite_static_helper(branched) + return static_expanded + + +@pto.func +def multi_return_helper(limit: pto.i32, seed: pto.i32): + value = chain_mid(limit, seed) + return value, value + pto.const(1, dtype=pto.i32) + + +@pto.jit(target="a5") +def func_chain_probe(limit: pto.i32): + zero = pto.const(0, dtype=pto.i32) + first, second = multi_return_helper(limit, zero) + merged = dyn_if_helper(first, second) + _ = merged + plain_trace_helper() + + +def main(): + text = func_chain_probe.compile().mlir_text() + print(text) + print("\n=== COUNTS ===") + for needle in [ + "scf.for", + "scf.if", + "func.func @dyn_loop_helper__ptodsl_", + "func.func @dyn_if_helper__ptodsl_", + "func.func @no_rewrite_static_helper__ptodsl_", + "func.func @chain_mid__ptodsl_", + "func.func @multi_return_helper__ptodsl_", + "call @dyn_if_helper__ptodsl_", + "pto.barrier ", + ]: + print(f"{needle}: {text.count(needle)}") + + +if __name__ == "__main__": + main() From b5afdbc781aeadff912d597eacb3839c15e6ebb7 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 16 Jul 2026 18:51:06 +0800 Subject: [PATCH 03/15] Refine pto.func docs and probe header --- .../03-kernel-entry-and-subkernels.md | 28 +++++++++---------- tmp/ptodsl_func_chain_probe.mlir | 7 +++++ 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md index b636f379f3..ecf484030e 100644 --- a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md +++ b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md @@ -2,12 +2,12 @@ PTODSL provides one kernel decorator (`@pto.jit`) with two roles (`entry=True` / `entry=False`), two compilation backends (`vpto` / `emitc`), -and two reusable compute helper decorators (`@pto.tileop` and `@pto.simt`), -plus reusable PTODSL helper functions (`@pto.func`) and inline unit-specific -context managers. This chapter covers -the `@pto.jit` entry and module contracts, the two programming models, the two -compilation backends, helper functions, sub-kernel reference, parameter -contracts, and boundary constraints. +two reusable compute helper decorators (`@pto.tileop` and `@pto.simt`), and +reusable PTODSL helper functions (`@pto.func`) plus inline unit-specific +context managers. This chapter covers the `@pto.jit` entry and module +contracts, the two programming models, the two compilation backends, helper +functions, sub-kernel reference, parameter contracts, and boundary +constraints. ## 3.1 `@pto.jit` — roles, backends, and modes @@ -23,7 +23,7 @@ Decorator overview: mode="explicit" micro-instruction authoring, user-managed staging @pto.tileop Single-core Tile/scalar compute helper with inferred Vector/Cube kind -@pto.func Reusable PTODSL helper, no host/module ABI boundary +@pto.func Reusable PTODSL helper with AST-rewritten control flow @pto.simt Explicitly launched SIMT helper with pointer/scalar ABI ``` @@ -61,14 +61,12 @@ manual-address, user-managed staging contract of explicit kernels. (`@pto.tileop` and `@pto.simt`) define sub-kernels that are called from within `@pto.jit` bodies. -`@pto.func` is the lightweight helper boundary for reusable PTODSL code. It -does not create a host-launchable entry, a kernel-module ABI, or a hardware-unit -sub-kernel section. When traced PTODSL code calls a `@pto.func` helper, PTODSL -materializes one helper `func.func` in the caller's active compilation context -and emits `func.call` at call sites. Supported native Python `if` and -`for range(...)` in the helper body use the same AST rewrite path as `@pto.jit` -and named sub-kernels. The helper can return PTODSL runtime values, including -multiple values via a tuple. +`@pto.func` defines reusable PTODSL helper functions. Use it when a helper +contains PTODSL operations or native Python `if` / `for range(...)` that should +compile as device-side control flow. By default, supported native control flow +in a `@pto.func` body is AST-rewritten just like in `@pto.jit` and named +sub-kernels. The helper can return PTODSL runtime values, including multiple +values via a tuple. ```python @pto.func diff --git a/tmp/ptodsl_func_chain_probe.mlir b/tmp/ptodsl_func_chain_probe.mlir index 1ff5babfa9..0915e839b3 100644 --- a/tmp/ptodsl_func_chain_probe.mlir +++ b/tmp/ptodsl_func_chain_probe.mlir @@ -1,3 +1,10 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. // Generated from tmp/ptodsl_func_chain_probe.py while developing issue #946. // The __ptodsl_ suffixes are specialization hashes and may differ across runs. // From 4b984ca0d770d8c12f1400f9de87944b65cb0a5e Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 16 Jul 2026 18:54:16 +0800 Subject: [PATCH 04/15] Fix pto.func license header --- ptodsl/ptodsl/_func.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ptodsl/ptodsl/_func.py b/ptodsl/ptodsl/_func.py index 62fa9989b4..5e07a911a5 100644 --- a/ptodsl/ptodsl/_func.py +++ b/ptodsl/ptodsl/_func.py @@ -2,7 +2,7 @@ # This program is free software, you can redistribute it and/or modify it under the terms and conditions of # CANN Open Software License Agreement Version 2.0 (the "License"). # Please refer to the License for details. You may not use this file except in compliance with the License. -# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. """``@pto.func`` decorator and reusable callable handle.""" From a2c4b95ca1769331efac2b52687ffb61a0ff8cb4 Mon Sep 17 00:00:00 2001 From: andodo Date: Fri, 17 Jul 2026 09:30:02 +0800 Subject: [PATCH 05/15] Require explicit pto.func return types --- .../03-kernel-entry-and-subkernels.md | 6 +- ptodsl/docs/user_guide/05-control-flow.md | 9 ++- ptodsl/ptodsl/_func.py | 20 ++++- ptodsl/ptodsl/_tracing/session.py | 74 ++++++++++--------- ptodsl/tests/test_jit_compile.py | 15 ++-- tmp/ptodsl_func_chain_probe.mlir | 22 +++--- tmp/ptodsl_func_chain_probe.py | 10 +-- 7 files changed, 90 insertions(+), 66 deletions(-) diff --git a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md index ecf484030e..bf42891b0b 100644 --- a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md +++ b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md @@ -66,10 +66,12 @@ contains PTODSL operations or native Python `if` / `for range(...)` that should compile as device-side control flow. By default, supported native control flow in a `@pto.func` body is AST-rewritten just like in `@pto.jit` and named sub-kernels. The helper can return PTODSL runtime values, including multiple -values via a tuple. +values via a tuple. Every `@pto.func` helper must declare its return type with +`returns=...` or a Python return annotation; use `returns=None` or `-> None` for +helpers that do not return values. ```python -@pto.func +@pto.func(returns=pto.i32) def add_rows(total: pto.i32, rows: pto.i32): one = pto.const(1, dtype=pto.i32) for _ in range(rows): diff --git a/ptodsl/docs/user_guide/05-control-flow.md b/ptodsl/docs/user_guide/05-control-flow.md index 049d6b498f..137d2493e3 100644 --- a/ptodsl/docs/user_guide/05-control-flow.md +++ b/ptodsl/docs/user_guide/05-control-flow.md @@ -257,10 +257,11 @@ rewritten scope become device-side control flow. Use `pto.const_expr(...)` and PTODSL does not recursively rewrite arbitrary undecorated Python callees. If an external helper should contain runtime native control flow, decorate it with -`@pto.func` or one of the other PTODSL callable decorators. A plain Python -helper is still executed during tracing; static `range(...)` loops in such a -helper unroll at trace time, and runtime loop bounds or branch conditions are -not converted into `scf.for` / `scf.if`. +`@pto.func` or one of the other PTODSL callable decorators. `@pto.func` helpers +must explicitly declare their return type with `returns=...` or a Python return +annotation. A plain Python helper is still executed during tracing; static +`range(...)` loops in such a helper unroll at trace time, and runtime loop +bounds or branch conditions are not converted into `scf.for` / `scf.if`. ### Runtime branches diff --git a/ptodsl/ptodsl/_func.py b/ptodsl/ptodsl/_func.py index 5e07a911a5..e834efe735 100644 --- a/ptodsl/ptodsl/_func.py +++ b/ptodsl/ptodsl/_func.py @@ -14,9 +14,11 @@ import inspect from ._ast_rewrite import rewrite_jit_function -from ._cache_signature import closure_cache_signature +from ._cache_signature import cache_signature_atom, closure_cache_signature from ._tracing import current_runtime +_RETURNS_UNSET = object() + @dataclass(frozen=True) class FuncSpec: @@ -28,11 +30,21 @@ class FuncSpec: class FuncTemplate: """Callable decorated PTODSL helper surface.""" - def __init__(self, spec: FuncSpec, py_fn, *, ast_rewrite: bool = True): + 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) def emit_body(self, *args, **kwargs): @@ -54,11 +66,12 @@ def __ptodsl_cache_signature__(self): self.spec.symbol_name, id(self.py_fn), self._ast_rewrite, + cache_signature_atom(self.declared_returns), closure_cache_signature(self.py_fn), ) -def func(fn=None, *, name: str | None = None, ast_rewrite: bool = True): +def func(fn=None, *, name: str | None = None, ast_rewrite: bool = True, returns=_RETURNS_UNSET): """Decorate a Python function as a reusable PTODSL callable helper.""" def decorator(py_fn): @@ -66,6 +79,7 @@ def decorator(py_fn): FuncSpec(symbol_name=name or py_fn.__name__), py_fn, ast_rewrite=ast_rewrite, + returns=returns, ) if fn is not None: diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index 16a295ee8c..bb2922dc5f 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -48,7 +48,6 @@ IntegerType, Operation, StringAttr, - TypeAttr, UnitAttr, ) @@ -568,6 +567,7 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): helper_spec = HelperFunctionSpec( symbol_name=func_template.spec.symbol_name, arg_types=tuple(unwrap_surface_value(arg).type for arg in arg_templates), + result_types=self._declared_ptodsl_func_result_types(func_template), attributes=(("pto.ptodsl.callable_kind", StringAttr.get("func")),), identity=func_template.__ptodsl_cache_signature__(), ) @@ -601,13 +601,7 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): return_values = self._normalize_ptodsl_func_return_values( result, func_template=func_template, - ) - result_types = tuple(value.type for value in return_values) - helper_fn.attributes["function_type"] = TypeAttr.get( - func.FunctionType.get( - list(helper_spec.arg_types), - list(result_types), - ) + result_types=helper_spec.result_types, ) if return_values: func.ReturnOp([unwrap_surface_value(value) for value in return_values]) @@ -862,8 +856,21 @@ def _normalize_ptodsl_func_argument(self, name: str, param, value): f"got {raw_value!r}" ) - def _normalize_ptodsl_func_return_values(self, result, *, func_template): + def _declared_ptodsl_func_result_types(self, func_template): + declared_returns = func_template.declared_returns + if declared_returns is None or declared_returns is type(None): + return () + if isinstance(declared_returns, (tuple, list)): + return tuple(_resolve(return_type) for return_type in declared_returns) + return (_resolve(declared_returns),) + + def _normalize_ptodsl_func_return_values(self, result, *, func_template, result_types): if result is None: + if result_types: + raise TypeError( + f"@pto.func {func_template.spec.symbol_name!r} must return " + f"{len(result_types)} value(s) matching its declared return type" + ) return () if isinstance(result, tuple): values = result @@ -872,41 +879,36 @@ def _normalize_ptodsl_func_return_values(self, result, *, func_template): else: values = (result,) - normalized = [] - return_annotation = func_template.signature.return_annotation - target_type = None - if return_annotation is not inspect.Signature.empty: - try: - target_type = _resolve(return_annotation) - except Exception: - target_type = None + if len(values) != len(result_types): + raise TypeError( + f"@pto.func {func_template.spec.symbol_name!r} returned {len(values)} value(s), " + f"but its declared return type expects {len(result_types)}" + ) - for index, value in enumerate(values): + normalized = [] + for index, (value, target_type) in enumerate(zip(values, result_types)): raw_value = unwrap_surface_value(value) if hasattr(raw_value, "type"): + if str(raw_value.type) != str(target_type): + raise TypeError( + f"@pto.func return value {index} has type {raw_value.type}, " + f"but the declared return type is {target_type}" + ) normalized.append(raw_value) continue - if target_type is not None: - try: - normalized.append( - coerce_scalar_to_type( - raw_value, - target_type, - context=f"@pto.func return value {index}", - ) + try: + normalized.append( + coerce_scalar_to_type( + raw_value, + target_type, + context=f"@pto.func return value {index}", ) - continue - except TypeError: - pass - if isinstance(raw_value, bool): - normalized.append(const(int(raw_value), dtype=int1)) - continue - if isinstance(raw_value, int): - normalized.append(const(raw_value)) + ) continue + except TypeError: + pass raise TypeError( - f"@pto.func return value {index} must be a traced runtime value or supported literal, " - f"got {raw_value!r}" + f"@pto.func return value {index} must match declared type {target_type}, got {raw_value!r}" ) return tuple(normalized) diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 155803d917..972d0930ef 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -1268,7 +1268,7 @@ def helper(limit, enabled): _ = value -@pto.func +@pto.func(returns=pto.i32) def func_runtime_for_return_helper(limit: pto.i32, initial: pto.i32): one = pto.const(1, dtype=pto.i32) total = initial @@ -1277,7 +1277,7 @@ def func_runtime_for_return_helper(limit: pto.i32, initial: pto.i32): return total -@pto.func +@pto.func(returns=pto.i32) def func_runtime_if_return_helper(lhs: pto.i32, rhs: pto.i32): if lhs > rhs: total = lhs + rhs @@ -1286,13 +1286,13 @@ def func_runtime_if_return_helper(lhs: pto.i32, rhs: pto.i32): return total -@pto.func +@pto.func(returns=(pto.i32, pto.i32)) def func_multi_return_helper(value: pto.i32): one = pto.const(1, dtype=pto.i32) return value, value + one -@pto.func +@pto.func(returns=None) def func_void_helper(): pto.pipe_barrier(pto.Pipe.ALL) @@ -1506,7 +1506,7 @@ def make_sourceless_ptodsl_func_probe(): namespace = {"pto": pto} exec( """ -@pto.func +@pto.func(returns=None) def sourceless_ptodsl_func_helper(): if True: pto.pipe_barrier(pto.Pipe.ALL) @@ -5078,6 +5078,11 @@ def _enter_inline_simt_with_resource_attr(): and len(re.findall(r"call @func_void_helper__ptodsl_[0-9a-f]+", ptodsl_func_call_text)) == 2, "repeated @pto.func calls should reuse one materialized helper artifact", ) + expect_raises( + TypeError, + lambda: pto.func(lambda value: value), + "must explicitly declare return types", + ) ast_nested_helper_freevar_if_merge_text = ast_nested_helper_freevar_if_merge_probe.compile().mlir_text() expect_parse_roundtrip_and_verify( diff --git a/tmp/ptodsl_func_chain_probe.mlir b/tmp/ptodsl_func_chain_probe.mlir index 0915e839b3..4bcc3d85ed 100644 --- a/tmp/ptodsl_func_chain_probe.mlir +++ b/tmp/ptodsl_func_chain_probe.mlir @@ -34,25 +34,25 @@ module attributes {pto.target_arch = "a5"} { module attributes {pto.backend = "vpto", pto.kernel_kind = #pto.kernel_kind, pto.target_arch = "a5"} { func.func @func_chain_probe(%arg0: i32) attributes {pto.entry} { %c0_i32 = arith.constant 0 : i32 - %0:2 = call @multi_return_helper__ptodsl_406ab2a269(%arg0, %c0_i32) : (i32, i32) -> (i32, i32) - %1 = call @dyn_if_helper__ptodsl_82f7abb427(%0#0, %0#1) : (i32, i32) -> i32 + %0:2 = call @multi_return_helper__ptodsl_c1d36bffde(%arg0, %c0_i32) : (i32, i32) -> (i32, i32) + %1 = call @dyn_if_helper__ptodsl_e04f1ff14e(%0#0, %0#1) : (i32, i32) -> i32 pto.barrier pto.barrier return } - func.func @multi_return_helper__ptodsl_406ab2a269(%arg0: i32, %arg1: i32) -> (i32, i32) attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "multi_return_helper"} { - %0 = call @chain_mid__ptodsl_be1027dd8d(%arg0, %arg1) : (i32, i32) -> i32 + func.func @multi_return_helper__ptodsl_c1d36bffde(%arg0: i32, %arg1: i32) -> (i32, i32) attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "multi_return_helper"} { + %0 = call @chain_mid__ptodsl_f1c09ab46a(%arg0, %arg1) : (i32, i32) -> i32 %c1_i32 = arith.constant 1 : i32 %1 = arith.addi %0, %c1_i32 : i32 return %0, %1 : i32, i32 } - func.func @chain_mid__ptodsl_be1027dd8d(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "chain_mid"} { - %0 = call @dyn_loop_helper__ptodsl_795afe8017(%arg0, %arg1) : (i32, i32) -> i32 - %1 = call @dyn_if_helper__ptodsl_82f7abb427(%0, %arg1) : (i32, i32) -> i32 - %2 = call @no_rewrite_static_helper__ptodsl_eccc8a4ce9(%1) : (i32) -> i32 + func.func @chain_mid__ptodsl_f1c09ab46a(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "chain_mid"} { + %0 = call @dyn_loop_helper__ptodsl_41121a518e(%arg0, %arg1) : (i32, i32) -> i32 + %1 = call @dyn_if_helper__ptodsl_e04f1ff14e(%0, %arg1) : (i32, i32) -> i32 + %2 = call @no_rewrite_static_helper__ptodsl_3485521ccf(%1) : (i32) -> i32 return %2 : i32 } - func.func @dyn_loop_helper__ptodsl_795afe8017(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "dyn_loop_helper"} { + func.func @dyn_loop_helper__ptodsl_41121a518e(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "dyn_loop_helper"} { %c1_i32 = arith.constant 1 : i32 %c0 = arith.constant 0 : index %0 = arith.index_cast %arg0 : i32 to index @@ -63,7 +63,7 @@ module attributes {pto.target_arch = "a5"} { } return %1 : i32 } - func.func @dyn_if_helper__ptodsl_82f7abb427(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "dyn_if_helper"} { + func.func @dyn_if_helper__ptodsl_e04f1ff14e(%arg0: i32, %arg1: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "dyn_if_helper"} { %0 = arith.cmpi sgt, %arg0, %arg1 : i32 %1 = scf.if %0 -> (i32) { %2 = arith.subi %arg0, %arg1 : i32 @@ -74,7 +74,7 @@ module attributes {pto.target_arch = "a5"} { } return %1 : i32 } - func.func @no_rewrite_static_helper__ptodsl_eccc8a4ce9(%arg0: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "no_rewrite_static_helper"} { + func.func @no_rewrite_static_helper__ptodsl_3485521ccf(%arg0: i32) -> i32 attributes {pto.ptodsl.callable_kind = "func", pto.ptodsl.logical_name = "no_rewrite_static_helper"} { %c1_i32 = arith.constant 1 : i32 %0 = arith.addi %arg0, %c1_i32 : i32 %c1_i32_0 = arith.constant 1 : i32 diff --git a/tmp/ptodsl_func_chain_probe.py b/tmp/ptodsl_func_chain_probe.py index 74aec3803e..d47af6e29e 100644 --- a/tmp/ptodsl_func_chain_probe.py +++ b/tmp/ptodsl_func_chain_probe.py @@ -17,7 +17,7 @@ def plain_trace_helper(): pto.pipe_barrier(pto.Pipe.ALL) -@pto.func +@pto.func(returns=pto.i32) def dyn_loop_helper(limit: pto.i32, value: pto.i32): one = pto.const(1, dtype=pto.i32) total = value @@ -26,7 +26,7 @@ def dyn_loop_helper(limit: pto.i32, value: pto.i32): return total -@pto.func +@pto.func(returns=pto.i32) def dyn_if_helper(lhs: pto.i32, rhs: pto.i32): if lhs > rhs: chosen = lhs - rhs @@ -35,7 +35,7 @@ def dyn_if_helper(lhs: pto.i32, rhs: pto.i32): return chosen -@pto.func(ast_rewrite=False) +@pto.func(ast_rewrite=False, returns=pto.i32) def no_rewrite_static_helper(value: pto.i32): total = value if True: @@ -44,7 +44,7 @@ def no_rewrite_static_helper(value: pto.i32): return total -@pto.func +@pto.func(returns=pto.i32) def chain_mid(limit: pto.i32, seed: pto.i32): looped = dyn_loop_helper(limit, seed) branched = dyn_if_helper(looped, seed) @@ -52,7 +52,7 @@ def chain_mid(limit: pto.i32, seed: pto.i32): return static_expanded -@pto.func +@pto.func(returns=(pto.i32, pto.i32)) def multi_return_helper(limit: pto.i32, seed: pto.i32): value = chain_mid(limit, seed) return value, value + pto.const(1, dtype=pto.i32) From 822ecea158042f274d68c45eca6c79bf0e91d534 Mon Sep 17 00:00:00 2001 From: andodo Date: Fri, 17 Jul 2026 15:45:25 +0800 Subject: [PATCH 06/15] Fix mte cache ctl FileCheck labels --- test/lit/vpto/mte_ub_gm_l2_cache_ctl.pto | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/lit/vpto/mte_ub_gm_l2_cache_ctl.pto b/test/lit/vpto/mte_ub_gm_l2_cache_ctl.pto index a4dd4911ab..0331130532 100644 --- a/test/lit/vpto/mte_ub_gm_l2_cache_ctl.pto +++ b/test/lit/vpto/mte_ub_gm_l2_cache_ctl.pto @@ -48,11 +48,11 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, !pto.ptr, i64, i64, i64, i64{{$}} -// EXPAND-EXPLICIT-LABEL: func.func @mte_ub_gm_l2_cache_ctl( -// EXPAND-EXPLICIT: %[[L2:.*]] = arith.constant 5 : i64 -// EXPAND-EXPLICIT: pto.copy_ubuf_to_gm %{{[^,]+}}, %arg0, %{{[^,]+}}, %{{[^,]+}}, %{{[^,]+}}, %[[L2]], %{{[^,]+}}, %{{[^:]+}} -// EXPAND-EXPLICIT-SAME: : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 // EXPAND-DEFAULT-LABEL: func.func @mte_ub_gm_l2_cache_ctl_default( // EXPAND-DEFAULT: %[[ZERO:.*]] = arith.constant 0 : i64 // EXPAND-DEFAULT: pto.copy_ubuf_to_gm %{{[^,]+}}, %arg0, %[[ZERO]], %{{[^,]+}}, %{{[^,]+}}, %[[ZERO]], %{{[^,]+}}, %{{[^:]+}} // EXPAND-DEFAULT-SAME: : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 +// EXPAND-EXPLICIT-LABEL: func.func @mte_ub_gm_l2_cache_ctl( +// EXPAND-EXPLICIT: %[[L2:.*]] = arith.constant 5 : i64 +// EXPAND-EXPLICIT: pto.copy_ubuf_to_gm %{{[^,]+}}, %arg0, %{{[^,]+}}, %{{[^,]+}}, %{{[^,]+}}, %[[L2]], %{{[^,]+}}, %{{[^:]+}} +// EXPAND-EXPLICIT-SAME: : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 From 0f0698a53dee1455f0772ceca16d441862030c04 Mon Sep 17 00:00:00 2001 From: andodo Date: Tue, 21 Jul 2026 20:43:47 +0800 Subject: [PATCH 07/15] Reject early returns in rewritten control flow --- ptodsl/ptodsl/_ast_rewrite.py | 41 ++++++++++++++++++++++++ ptodsl/tests/test_jit_compile.py | 54 ++++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/ptodsl/ptodsl/_ast_rewrite.py b/ptodsl/ptodsl/_ast_rewrite.py index be766b58f5..1343517db0 100644 --- a/ptodsl/ptodsl/_ast_rewrite.py +++ b/ptodsl/ptodsl/_ast_rewrite.py @@ -370,6 +370,43 @@ def _name(name: str, ctx=ast.Load()): return ast.Name(id=name, ctx=ctx) +class _ControlFlowExitVisitor(ast.NodeVisitor): + def __init__(self): + self.exit_node = None + + def visit_Return(self, node): + self.exit_node = node + + def visit_Yield(self, node): + self.exit_node = node + + def visit_YieldFrom(self, node): + self.exit_node = node + + def visit_FunctionDef(self, node): + return + + def visit_AsyncFunctionDef(self, node): + return + + def visit_Lambda(self, node): + return + + def visit_ClassDef(self, node): + return + + +def _reject_control_flow_exits(stmts, context: str): + visitor = _ControlFlowExitVisitor() + for stmt in stmts: + visitor.visit(stmt) + if visitor.exit_node is not None: + raise PTODSLAstRewriteError( + f"ast_rewrite=True does not support return/yield inside rewritten {context}; " + "assign values to locals and return after the rewritten control flow" + ) + + class _ControlFlowRewriter: def __init__(self): self._counter = 0 @@ -468,6 +505,9 @@ def _rewrite_if(self, stmt, *, live_after, allow_loop_control=False): ) return [stmt] + _reject_control_flow_exits(stmt.body, "if branches") + _reject_control_flow_exits(stmt.orelse, "if branches") + cond_name = self._fresh("cond") then_info = _name_info(stmt.body) else_info = _name_info(stmt.orelse) @@ -630,6 +670,7 @@ def _rewrite_for(self, stmt, *, live_after, allow_loop_control=False): raise PTODSLAstRewriteError("ast_rewrite=True does not support for-else on runtime loops") if not isinstance(stmt.target, ast.Name): raise PTODSLAstRewriteError("ast_rewrite=True runtime for-loops require a simple name target") + _reject_control_flow_exits(stmt.body, "for-loop bodies") if stmt.target.id in live_after: raise PTODSLAstRewriteError( "ast_rewrite=True runtime for-loops cannot expose the loop induction variable outside the loop yet; " diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 972d0930ef..c80c2ce0c4 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -20,6 +20,7 @@ from ptodsl import pto, scalar from ptodsl import _types as pto_types +from ptodsl._ast_rewrite import PTODSLAstRewriteError from ptodsl._bootstrap import make_context from ptodsl._kernel_signature import DeviceParameterSpec, HelperMarkerParameterSpec, RuntimeScalarParameterSpec from ptodsl._tracing.runtime import SignatureTracingRuntime @@ -1286,6 +1287,26 @@ def func_runtime_if_return_helper(lhs: pto.i32, rhs: pto.i32): return total +@pto.func(returns=pto.i32) +def func_runtime_if_early_return_helper(lhs: pto.i32, rhs: pto.i32): + if lhs > rhs: + return lhs + return rhs + + +@pto.func(returns=pto.i32) +def func_runtime_for_early_return_helper(limit: pto.i32, initial: pto.i32): + for _ in range(limit): + return initial + return initial + + +@pto.func(returns=None) +def func_runtime_if_yield_helper(lhs: pto.i32, rhs: pto.i32): + if lhs > rhs: + yield lhs + + @pto.func(returns=(pto.i32, pto.i32)) def func_multi_return_helper(value: pto.i32): one = pto.const(1, dtype=pto.i32) @@ -1308,6 +1329,24 @@ def ptodsl_func_call_probe(rows: pto.i32): func_void_helper() +@pto.jit(target="a5") +def ptodsl_func_if_early_return_probe(rows: pto.i32): + init = pto.const(0, dtype=pto.i32) + _ = func_runtime_if_early_return_helper(rows, init) + + +@pto.jit(target="a5") +def ptodsl_func_for_early_return_probe(rows: pto.i32): + init = pto.const(0, dtype=pto.i32) + _ = func_runtime_for_early_return_helper(rows, init) + + +@pto.jit(target="a5") +def ptodsl_func_if_yield_probe(rows: pto.i32): + init = pto.const(0, dtype=pto.i32) + func_runtime_if_yield_helper(rows, init) + + @pto.jit(target="a5") def ast_nested_helper_freevar_if_merge_probe(): lhs = pto.const(4, dtype=pto.i32) @@ -5078,6 +5117,21 @@ def _enter_inline_simt_with_resource_attr(): and len(re.findall(r"call @func_void_helper__ptodsl_[0-9a-f]+", ptodsl_func_call_text)) == 2, "repeated @pto.func calls should reuse one materialized helper artifact", ) + expect_raises( + PTODSLAstRewriteError, + lambda: ptodsl_func_if_early_return_probe.compile().mlir_text(), + "return/yield inside rewritten if branches", + ) + expect_raises( + PTODSLAstRewriteError, + lambda: ptodsl_func_for_early_return_probe.compile().mlir_text(), + "return/yield inside rewritten for-loop bodies", + ) + expect_raises( + PTODSLAstRewriteError, + lambda: ptodsl_func_if_yield_probe.compile().mlir_text(), + "return/yield inside rewritten if branches", + ) expect_raises( TypeError, lambda: pto.func(lambda value: value), From 25b8bf6549b9f46cdb15790dd657bfe71c9d63ed Mon Sep 17 00:00:00 2001 From: andodo Date: Wed, 22 Jul 2026 09:04:54 +0800 Subject: [PATCH 08/15] Preserve func helper surface templates --- ptodsl/ptodsl/_tracing/session.py | 17 ++++++++++------- ptodsl/tests/test_jit_compile.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 7 deletions(-) diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index bb2922dc5f..011a36c554 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -546,15 +546,16 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): bound = func_template.signature.bind(*args, **kwargs) bound.apply_defaults() ordered_arg_values = [] - normalized_values = {} + arg_templates = [] for name, param in func_template.signature.parameters.items(): + original_value = bound.arguments[name] value = self._normalize_ptodsl_func_argument( name, param, - bound.arguments[name], + original_value, ) - normalized_values[name] = value ordered_arg_values.append(value) + arg_templates.append(original_value) if param.kind not in { inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD, @@ -562,11 +563,12 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): }: raise TypeError("@pto.func helpers do not support var-positional or var-keyword parameters yet") - arg_templates = tuple(ordered_arg_values) + arg_values = tuple(ordered_arg_values) + arg_templates = tuple(arg_templates) owner_symbol_name = self.current_function_owner_symbol_name helper_spec = HelperFunctionSpec( symbol_name=func_template.spec.symbol_name, - arg_types=tuple(unwrap_surface_value(arg).type for arg in arg_templates), + arg_types=tuple(unwrap_surface_value(arg).type for arg in arg_values), result_types=self._declared_ptodsl_func_result_types(func_template), attributes=(("pto.ptodsl.callable_kind", StringAttr.get("func")),), identity=func_template.__ptodsl_cache_signature__(), @@ -584,8 +586,9 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): entry_arg_index = 0 for name, param in func_template.signature.parameters.items(): entry_arg = entry_args[entry_arg_index] + arg_template = arg_templates[entry_arg_index] entry_arg_index += 1 - wrapped_value = wrap_like_surface_value(normalized_values[name], entry_arg) + wrapped_value = wrap_like_surface_value(arg_template, entry_arg) if param.kind in {inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD}: wrapped_args.append(wrapped_value) elif param.kind == inspect.Parameter.KEYWORD_ONLY: @@ -608,7 +611,7 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): else: func.ReturnOp([]) - call_op = func.CallOp(helper_fn, [unwrap_surface_value(arg) for arg in arg_templates]) + call_op = func.CallOp(helper_fn, [unwrap_surface_value(arg) for arg in arg_values]) return self._wrap_ptodsl_func_call_results(call_op.results) def begin_carry_loop(self, start, stop, step, state_items): diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index c80c2ce0c4..63a1fdda70 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -1318,6 +1318,11 @@ def func_void_helper(): pto.pipe_barrier(pto.Pipe.ALL) +@pto.func(returns=pto.i32) +def func_partition_metadata_helper(part: pto.PartitionTensorView, cols: pto.i32): + return part.sizes[0] + cols + + @pto.jit(target="a5") def ptodsl_func_call_probe(rows: pto.i32): init = pto.const(0, dtype=pto.i32) @@ -1329,6 +1334,16 @@ def ptodsl_func_call_probe(rows: pto.i32): func_void_helper() +@pto.jit(target="a5") +def ptodsl_func_partition_metadata_probe( + A_ptr: pto.ptr(pto.f32, "gm"), + cols: pto.i32, +): + a_view = pto.make_tensor_view(A_ptr, shape=[1, 16], strides=[16, 1]) + part = pto.partition_view(a_view, offsets=[0, 0], sizes=[1, 16]) + _ = func_partition_metadata_helper(part, cols) + + @pto.jit(target="a5") def ptodsl_func_if_early_return_probe(rows: pto.i32): init = pto.const(0, dtype=pto.i32) @@ -5117,6 +5132,19 @@ def _enter_inline_simt_with_resource_attr(): and len(re.findall(r"call @func_void_helper__ptodsl_[0-9a-f]+", ptodsl_func_call_text)) == 2, "repeated @pto.func calls should reuse one materialized helper artifact", ) + ptodsl_func_partition_metadata_text = ptodsl_func_partition_metadata_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + ptodsl_func_partition_metadata_text, + "@pto.func partition metadata specialization", + ) + expect( + re.search(r"func\.func @func_partition_metadata_helper__ptodsl_[0-9a-f]+", ptodsl_func_partition_metadata_text) + is not None + and re.search(r"func\.func @func_partition_metadata_helper__ptodsl_[0-9a-f]+\(.*\) -> i32", ptodsl_func_partition_metadata_text) + is not None + and ptodsl_func_partition_metadata_text.count("pto.partition_view") >= 2, + "@pto.func should preserve partition metadata across the helper boundary", + ) expect_raises( PTODSLAstRewriteError, lambda: ptodsl_func_if_early_return_probe.compile().mlir_text(), From 41a33a096f50e197cf844c35ccafd9c433f6c76c Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 23 Jul 2026 10:21:42 +0800 Subject: [PATCH 09/15] Support postponed annotations in pto.func --- ptodsl/ptodsl/_func.py | 20 +++++++++ ptodsl/ptodsl/_tracing/session.py | 24 ++++++----- ptodsl/tests/test_jit_compile.py | 72 +++++++++++++++++++++++++++++++ 3 files changed, 105 insertions(+), 11 deletions(-) diff --git a/ptodsl/ptodsl/_func.py b/ptodsl/ptodsl/_func.py index e834efe735..8cf2a64602 100644 --- a/ptodsl/ptodsl/_func.py +++ b/ptodsl/ptodsl/_func.py @@ -12,6 +12,7 @@ from dataclasses import dataclass from functools import update_wrapper import inspect +import typing from ._ast_rewrite import rewrite_jit_function from ._cache_signature import cache_signature_atom, closure_cache_signature @@ -35,8 +36,18 @@ def __init__(self, spec: FuncSpec, py_fn, *, ast_rewrite: bool = True, returns=_ self.py_fn = py_fn self._ast_rewrite = ast_rewrite self.signature = inspect.signature(py_fn) + try: + self.type_hints = typing.get_type_hints(py_fn) + except Exception as exc: + if _has_annotations(self.signature): + raise TypeError( + f"failed to resolve @pto.func annotations for {py_fn.__qualname__!r}" + ) from exc + 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: @@ -87,6 +98,15 @@ def decorator(py_fn): return decorator +def _has_annotations(signature: inspect.Signature) -> bool: + if signature.return_annotation is not inspect.Signature.empty: + return True + return any( + param.annotation is not inspect.Parameter.empty + for param in signature.parameters.values() + ) + + __all__ = [ "FuncSpec", "FuncTemplate", diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index 011a36c554..30d20feec4 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -547,21 +547,23 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): bound.apply_defaults() ordered_arg_values = [] arg_templates = [] + type_hints = getattr(func_template, "type_hints", {}) 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") original_value = bound.arguments[name] + annotation = type_hints.get(name, param.annotation) value = self._normalize_ptodsl_func_argument( name, - param, + annotation, original_value, ) ordered_arg_values.append(value) arg_templates.append(original_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") arg_values = tuple(ordered_arg_values) arg_templates = tuple(arg_templates) @@ -833,15 +835,15 @@ def lookup_helper(self, symbol_name: str): return helper return None - def _normalize_ptodsl_func_argument(self, name: str, param, value): + 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 param.annotation is not inspect.Parameter.empty: + if annotation is not inspect.Parameter.empty: try: - target_type = _resolve(param.annotation) + target_type = _resolve(annotation) except Exception: - target_type = param.annotation + target_type = annotation try: return coerce_scalar_to_type( raw_value, diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 63a1fdda70..6fbc3e1fd9 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -1323,6 +1323,41 @@ def func_partition_metadata_helper(part: pto.PartitionTensorView, cols: pto.i32) return part.sizes[0] + cols +def _make_future_annotations_ptodsl_func_helpers(): + namespace = {"pto": pto} + exec( + """ +from __future__ import annotations + +@pto.func +def func_future_i32_return_helper(x: pto.i32) -> pto.i32: + return x + pto.const(1, dtype=pto.i32) + +@pto.func +def func_future_i64_literal_helper(x: pto.i64) -> pto.i64: + return x + pto.const(1, dtype=pto.i64) + +@pto.func +def func_future_void_helper(x: pto.i32) -> None: + _ = x + pto.pipe_barrier(pto.Pipe.ALL) +""", + namespace, + ) + return ( + namespace["func_future_i32_return_helper"], + namespace["func_future_i64_literal_helper"], + namespace["func_future_void_helper"], + ) + + +( + func_future_i32_return_helper, + func_future_i64_literal_helper, + func_future_void_helper, +) = _make_future_annotations_ptodsl_func_helpers() + + @pto.jit(target="a5") def ptodsl_func_call_probe(rows: pto.i32): init = pto.const(0, dtype=pto.i32) @@ -1344,6 +1379,14 @@ def ptodsl_func_partition_metadata_probe( _ = func_partition_metadata_helper(part, cols) +@pto.jit(target="a5") +def ptodsl_func_future_annotations_probe(rows: pto.i32): + i32_value = func_future_i32_return_helper(rows) + i64_value = func_future_i64_literal_helper(1) + func_future_void_helper(i32_value) + _ = i64_value + + @pto.jit(target="a5") def ptodsl_func_if_early_return_probe(rows: pto.i32): init = pto.const(0, dtype=pto.i32) @@ -5145,6 +5188,35 @@ def _enter_inline_simt_with_resource_attr(): and ptodsl_func_partition_metadata_text.count("pto.partition_view") >= 2, "@pto.func should preserve partition metadata across the helper boundary", ) + ptodsl_func_future_annotations_text = ptodsl_func_future_annotations_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + ptodsl_func_future_annotations_text, + "@pto.func future annotations specialization", + ) + expect( + re.search( + r"func\.func @func_future_i32_return_helper__ptodsl_[0-9a-f]+\(.*i32.*\) -> i32", + ptodsl_func_future_annotations_text, + ) + is not None, + "@pto.func should resolve PEP 563 string return annotations to PTO result types", + ) + expect( + re.search( + r"func\.func @func_future_i64_literal_helper__ptodsl_[0-9a-f]+\(.*i64.*\) -> i64", + ptodsl_func_future_annotations_text, + ) + is not None, + "@pto.func should resolve PEP 563 string parameter annotations before materializing literals", + ) + expect( + re.search( + r"func\.func @func_future_void_helper__ptodsl_[0-9a-f]+\(.*i32.*\) attributes", + ptodsl_func_future_annotations_text, + ) + is not None, + "@pto.func should resolve PEP 563 -> None annotations as void helpers", + ) expect_raises( PTODSLAstRewriteError, lambda: ptodsl_func_if_early_return_probe.compile().mlir_text(), From 4a329e18980d6c44668c185a64c27a31cc2340ad Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 23 Jul 2026 10:32:14 +0800 Subject: [PATCH 10/15] Relax pto.func partition metadata test --- ptodsl/tests/test_jit_compile.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 6fbc3e1fd9..4fd295a245 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -5185,7 +5185,7 @@ def _enter_inline_simt_with_resource_attr(): is not None and re.search(r"func\.func @func_partition_metadata_helper__ptodsl_[0-9a-f]+\(.*\) -> i32", ptodsl_func_partition_metadata_text) is not None - and ptodsl_func_partition_metadata_text.count("pto.partition_view") >= 2, + and re.search(r"%c1_i32 = arith\.constant 1 : i32", ptodsl_func_partition_metadata_text) is not None, "@pto.func should preserve partition metadata across the helper boundary", ) ptodsl_func_future_annotations_text = ptodsl_func_future_annotations_probe.compile().mlir_text() From 12ffb4b166bf22e875585e499ffbd8eaa077691c Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 23 Jul 2026 11:10:31 +0800 Subject: [PATCH 11/15] Stabilize pto.func helper identity --- ptodsl/ptodsl/_cache_signature.py | 18 ++++++++++++++++++ ptodsl/ptodsl/_func.py | 4 ++-- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/ptodsl/ptodsl/_cache_signature.py b/ptodsl/ptodsl/_cache_signature.py index ebfee276ec..c401977786 100644 --- a/ptodsl/ptodsl/_cache_signature.py +++ b/ptodsl/ptodsl/_cache_signature.py @@ -10,6 +10,23 @@ from __future__ import annotations +def function_cache_signature(fn): + """Return one stable cache signature for a Python function body.""" + code = fn.__code__ + return ( + "python-function", + code.co_filename, + code.co_firstlineno, + code.co_qualname, + code.co_argcount, + code.co_posonlyargcount, + code.co_kwonlyargcount, + code.co_names, + code.co_consts, + code.co_code, + ) + + def closure_cache_signature(fn): """Return one stable cache signature for the closure state captured by *fn*.""" try: @@ -55,4 +72,5 @@ def cache_signature_atom(value): __all__ = [ "cache_signature_atom", "closure_cache_signature", + "function_cache_signature", ] diff --git a/ptodsl/ptodsl/_func.py b/ptodsl/ptodsl/_func.py index 8cf2a64602..be317087bc 100644 --- a/ptodsl/ptodsl/_func.py +++ b/ptodsl/ptodsl/_func.py @@ -15,7 +15,7 @@ import typing from ._ast_rewrite import rewrite_jit_function -from ._cache_signature import cache_signature_atom, closure_cache_signature +from ._cache_signature import cache_signature_atom, closure_cache_signature, function_cache_signature from ._tracing import current_runtime _RETURNS_UNSET = object() @@ -75,7 +75,7 @@ def __ptodsl_cache_signature__(self): return ( type(self).__name__, self.spec.symbol_name, - id(self.py_fn), + function_cache_signature(self.py_fn), self._ast_rewrite, cache_signature_atom(self.declared_returns), closure_cache_signature(self.py_fn), From e65c0ad208cdd80d2f03d1e98ead1645a5f35c76 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 23 Jul 2026 14:16:12 +0800 Subject: [PATCH 12/15] Preserve helper ABI cache key --- ptodsl/ptodsl/_tracing/session.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index 30d20feec4..5ad7d6f481 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -69,12 +69,17 @@ def cache_key(self) -> tuple: tuple(str(arg_type) for arg_type in self.arg_types), tuple(str(result_type) for result_type in self.result_types), tuple((attr_name, str(attr_value)) for attr_name, attr_value in self.attributes), - self.identity, ) + def specialization_key(self) -> tuple: + """Return one stable semantic cache key for this helper specialization.""" + if not self.identity: + return self.cache_key() + return (*self.cache_key(), self.identity) + def specialized_symbol_name(self) -> str: - """Return one stable symbol name that is unique for this helper ABI.""" - digest = hashlib.sha1(repr(self.cache_key()).encode("utf-8")).hexdigest()[:10] + """Return one stable symbol name that is unique for this helper specialization.""" + digest = hashlib.sha1(repr(self.specialization_key()).encode("utf-8")).hexdigest()[:10] return f"{self.symbol_name}__ptodsl_{digest}" @@ -939,7 +944,7 @@ def get_or_create_helper_function(self, spec: HelperFunctionSpec, *, owner_symbo owner_symbol_name = ( self.current_function_owner_symbol_name if owner_symbol_name is None else owner_symbol_name ) - cache_key = (owner_symbol_name, spec.cache_key()) + cache_key = (owner_symbol_name, spec.specialization_key()) helper = self._helpers.get(cache_key) if helper is not None: return helper, False @@ -962,7 +967,7 @@ def get_or_create_helper_function(self, spec: HelperFunctionSpec, *, owner_symbo def get_or_create_kernel_module_primary_function(self, spec: HelperFunctionSpec, module_spec): """Look up or create the primary definition for one kernel-module callee.""" - cache_key = spec.cache_key() + cache_key = spec.specialization_key() helper = self._kernel_module_primary_functions.get(cache_key) if helper is not None: return helper, False From 5d1674c54e4036ac2bc7da0dffb1e3ed04034332 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 23 Jul 2026 14:19:31 +0800 Subject: [PATCH 13/15] Support older code object metadata --- ptodsl/ptodsl/_cache_signature.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ptodsl/ptodsl/_cache_signature.py b/ptodsl/ptodsl/_cache_signature.py index c401977786..92a111fae0 100644 --- a/ptodsl/ptodsl/_cache_signature.py +++ b/ptodsl/ptodsl/_cache_signature.py @@ -17,7 +17,7 @@ def function_cache_signature(fn): "python-function", code.co_filename, code.co_firstlineno, - code.co_qualname, + getattr(code, "co_qualname", fn.__qualname__), code.co_argcount, code.co_posonlyargcount, code.co_kwonlyargcount, From 35d9040d71cb8477aaf975ebe8e45dfb70e06513 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 23 Jul 2026 15:48:22 +0800 Subject: [PATCH 14/15] Support const_expr pto.func parameters --- ptodsl/ptodsl/_tracing/session.py | 58 +++++++++++++++++++++++-------- ptodsl/tests/test_jit_compile.py | 44 +++++++++++++++++++++++ 2 files changed, 87 insertions(+), 15 deletions(-) diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index 5ad7d6f481..0ffc5a361b 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -19,9 +19,11 @@ inline_tileop_capture_type_error, subkernel_kernel_kind_mismatch_error, ) +from .._cache_signature import cache_signature_atom from .._scalar_coercion import coerce_scalar_to_type from .._kernel_signature import RuntimeScalarParameterSpec from .._ops import const +from .._surface_types import const_expr as _const_expr_marker from .._surface_values import ( is_runtime_scalar_ir_type, is_tile_ir_type, @@ -550,8 +552,10 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): """Lower one ``@pto.func`` helper call in the active trace.""" bound = func_template.signature.bind(*args, **kwargs) bound.apply_defaults() - ordered_arg_values = [] - arg_templates = [] + runtime_arg_values = [] + runtime_arg_templates = [] + param_bindings = [] + constexpr_bindings = [] type_hints = getattr(func_template, "type_hints", {}) for name, param in func_template.signature.parameters.items(): if param.kind not in { @@ -562,23 +566,35 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): raise TypeError("@pto.func helpers do not support var-positional or var-keyword parameters yet") original_value = bound.arguments[name] annotation = type_hints.get(name, param.annotation) + if annotation is _const_expr_marker: + value = self._normalize_ptodsl_func_constexpr_argument(name, original_value) + param_bindings.append(("constexpr", name, param, value)) + constexpr_bindings.append((name, cache_signature_atom(value))) + continue value = self._normalize_ptodsl_func_argument( name, annotation, original_value, ) - ordered_arg_values.append(value) - arg_templates.append(original_value) - - arg_values = tuple(ordered_arg_values) - arg_templates = tuple(arg_templates) + param_bindings.append(("runtime", name, param, original_value)) + runtime_arg_values.append(value) + runtime_arg_templates.append(original_value) + + runtime_arg_values = tuple(runtime_arg_values) + runtime_arg_templates = tuple(runtime_arg_templates) + identity = func_template.__ptodsl_cache_signature__() + if constexpr_bindings: + identity = ( + identity, + ("constexprs", tuple(constexpr_bindings)), + ) owner_symbol_name = self.current_function_owner_symbol_name helper_spec = HelperFunctionSpec( symbol_name=func_template.spec.symbol_name, - arg_types=tuple(unwrap_surface_value(arg).type for arg in arg_values), + arg_types=tuple(unwrap_surface_value(arg).type for arg in runtime_arg_values), result_types=self._declared_ptodsl_func_result_types(func_template), attributes=(("pto.ptodsl.callable_kind", StringAttr.get("func")),), - identity=func_template.__ptodsl_cache_signature__(), + identity=identity, ) helper_fn, created = self.get_or_create_helper_function( helper_spec, @@ -591,11 +607,14 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): wrapped_args = [] wrapped_kwargs = {} entry_arg_index = 0 - for name, param in func_template.signature.parameters.items(): - entry_arg = entry_args[entry_arg_index] - arg_template = arg_templates[entry_arg_index] - entry_arg_index += 1 - wrapped_value = wrap_like_surface_value(arg_template, entry_arg) + for binding_kind, name, param, value in param_bindings: + if binding_kind == "constexpr": + wrapped_value = value + else: + entry_arg = entry_args[entry_arg_index] + arg_template = runtime_arg_templates[entry_arg_index] + entry_arg_index += 1 + wrapped_value = wrap_like_surface_value(arg_template, entry_arg) if param.kind in {inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD}: wrapped_args.append(wrapped_value) elif param.kind == inspect.Parameter.KEYWORD_ONLY: @@ -618,7 +637,7 @@ def lower_ptodsl_func_call(self, func_template, *args, **kwargs): else: func.ReturnOp([]) - call_op = func.CallOp(helper_fn, [unwrap_surface_value(arg) for arg in arg_values]) + call_op = func.CallOp(helper_fn, [unwrap_surface_value(arg) for arg in runtime_arg_values]) return self._wrap_ptodsl_func_call_results(call_op.results) def begin_carry_loop(self, start, stop, step, state_items): @@ -866,6 +885,15 @@ def _normalize_ptodsl_func_argument(self, name: str, annotation, value): f"got {raw_value!r}" ) + def _normalize_ptodsl_func_constexpr_argument(self, name: str, value): + raw_value = unwrap_surface_value(value) + if hasattr(raw_value, "type"): + raise TypeError( + f"@pto.func const_expr parameter {name!r} expects a compile-time Python value, " + f"got traced runtime value of type {raw_value.type}" + ) + return raw_value + def _declared_ptodsl_func_result_types(self, func_template): declared_returns = func_template.declared_returns if declared_returns is None or declared_returns is type(None): diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 4fd295a245..3b351b2f32 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -1323,6 +1323,15 @@ def func_partition_metadata_helper(part: pto.PartitionTensorView, cols: pto.i32) return part.sizes[0] + cols +@pto.func(returns=pto.i32) +def func_constexpr_static_helper(value: pto.i32, *, BLOCK: pto.const_expr = 2): + total = value + one = pto.const(1, dtype=pto.i32) + for _ in pto.static_range(BLOCK): + total = total + one + return total + + def _make_future_annotations_ptodsl_func_helpers(): namespace = {"pto": pto} exec( @@ -1387,6 +1396,18 @@ def ptodsl_func_future_annotations_probe(rows: pto.i32): _ = i64_value +@pto.jit(target="a5") +def ptodsl_func_constexpr_probe(rows: pto.i32): + first = func_constexpr_static_helper(rows, BLOCK=2) + second = func_constexpr_static_helper(rows, BLOCK=4) + _ = first + second + + +@pto.jit(target="a5") +def ptodsl_func_constexpr_runtime_value_probe(rows: pto.i32): + _ = func_constexpr_static_helper(rows, BLOCK=rows) + + @pto.jit(target="a5") def ptodsl_func_if_early_return_probe(rows: pto.i32): init = pto.const(0, dtype=pto.i32) @@ -5217,6 +5238,29 @@ def _enter_inline_simt_with_resource_attr(): is not None, "@pto.func should resolve PEP 563 -> None annotations as void helpers", ) + ptodsl_func_constexpr_text = ptodsl_func_constexpr_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify( + ptodsl_func_constexpr_text, + "@pto.func const_expr specialization", + ) + constexpr_helper_names = re.findall( + r"func\.func @(func_constexpr_static_helper__ptodsl_[0-9a-f]+)\(%arg0: i32\) -> i32", + ptodsl_func_constexpr_text, + ) + expect( + len(set(constexpr_helper_names)) == 2, + "@pto.func const_expr parameters should specialize helpers without entering the runtime ABI", + ) + expect( + ptodsl_func_constexpr_text.count("call @func_constexpr_static_helper__ptodsl_") == 2 + and ptodsl_func_constexpr_text.count("arith.addi") >= 6, + "@pto.func const_expr values should remain available for static_range unrolling", + ) + expect_raises( + TypeError, + lambda: ptodsl_func_constexpr_runtime_value_probe.compile().mlir_text(), + "const_expr parameter 'BLOCK' expects a compile-time Python value", + ) expect_raises( PTODSLAstRewriteError, lambda: ptodsl_func_if_early_return_probe.compile().mlir_text(), From 856a819ff719ca71a2b7e6bc4d950f226830dba4 Mon Sep 17 00:00:00 2001 From: andodo Date: Thu, 23 Jul 2026 16:25:27 +0800 Subject: [PATCH 15/15] Keep bare returns compatible in AST rewrite --- ptodsl/ptodsl/_ast_rewrite.py | 35 +++++++++++++++++++++++++---------- ptodsl/ptodsl/_func.py | 6 +++++- 2 files changed, 30 insertions(+), 11 deletions(-) diff --git a/ptodsl/ptodsl/_ast_rewrite.py b/ptodsl/ptodsl/_ast_rewrite.py index 1343517db0..fff4ef6493 100644 --- a/ptodsl/ptodsl/_ast_rewrite.py +++ b/ptodsl/ptodsl/_ast_rewrite.py @@ -20,7 +20,7 @@ class PTODSLAstRewriteError(SyntaxError): """Raised when AST rewrite sees unsupported Python control flow.""" -def rewrite_jit_function(fn): +def rewrite_jit_function(fn, *, reject_bare_returns: bool = False): """Return a function whose Python if/for control flow lowers to PTODSL APIs.""" try: source = inspect.getsource(fn) @@ -42,7 +42,7 @@ def rewrite_jit_function(fn): closure_vars = inspect.getclosurevars(fn) _inject_closure_defaults(function_def, closure_vars.nonlocals) _sanitize_signature_for_exec(function_def) - rewriter = _ControlFlowRewriter() + rewriter = _ControlFlowRewriter(reject_bare_returns=reject_bare_returns) function_def.body = rewriter.rewrite_block(function_def.body, live_after=set()) tree = ast.Module(body=[function_def], type_ignores=[]) ast.fix_missing_locations(tree) @@ -371,11 +371,13 @@ def _name(name: str, ctx=ast.Load()): class _ControlFlowExitVisitor(ast.NodeVisitor): - def __init__(self): + def __init__(self, *, reject_bare_returns: bool): self.exit_node = None + self._reject_bare_returns = reject_bare_returns def visit_Return(self, node): - self.exit_node = node + if self._reject_bare_returns or node.value is not None: + self.exit_node = node def visit_Yield(self, node): self.exit_node = node @@ -396,8 +398,8 @@ def visit_ClassDef(self, node): return -def _reject_control_flow_exits(stmts, context: str): - visitor = _ControlFlowExitVisitor() +def _reject_control_flow_exits(stmts, context: str, *, reject_bare_returns: bool): + visitor = _ControlFlowExitVisitor(reject_bare_returns=reject_bare_returns) for stmt in stmts: visitor.visit(stmt) if visitor.exit_node is not None: @@ -408,8 +410,9 @@ def _reject_control_flow_exits(stmts, context: str): class _ControlFlowRewriter: - def __init__(self): + def __init__(self, *, reject_bare_returns: bool = False): self._counter = 0 + self._reject_bare_returns = reject_bare_returns def _fresh(self, prefix: str) -> str: value = f"__pto_ast_{prefix}_{self._counter}" @@ -505,8 +508,16 @@ def _rewrite_if(self, stmt, *, live_after, allow_loop_control=False): ) return [stmt] - _reject_control_flow_exits(stmt.body, "if branches") - _reject_control_flow_exits(stmt.orelse, "if branches") + _reject_control_flow_exits( + stmt.body, + "if branches", + reject_bare_returns=self._reject_bare_returns, + ) + _reject_control_flow_exits( + stmt.orelse, + "if branches", + reject_bare_returns=self._reject_bare_returns, + ) cond_name = self._fresh("cond") then_info = _name_info(stmt.body) @@ -670,7 +681,11 @@ def _rewrite_for(self, stmt, *, live_after, allow_loop_control=False): raise PTODSLAstRewriteError("ast_rewrite=True does not support for-else on runtime loops") if not isinstance(stmt.target, ast.Name): raise PTODSLAstRewriteError("ast_rewrite=True runtime for-loops require a simple name target") - _reject_control_flow_exits(stmt.body, "for-loop bodies") + _reject_control_flow_exits( + stmt.body, + "for-loop bodies", + reject_bare_returns=self._reject_bare_returns, + ) if stmt.target.id in live_after: raise PTODSLAstRewriteError( "ast_rewrite=True runtime for-loops cannot expose the loop induction variable outside the loop yet; " diff --git a/ptodsl/ptodsl/_func.py b/ptodsl/ptodsl/_func.py index be317087bc..f79c72ff02 100644 --- a/ptodsl/ptodsl/_func.py +++ b/ptodsl/ptodsl/_func.py @@ -60,7 +60,11 @@ def __init__(self, spec: FuncSpec, py_fn, *, ast_rewrite: bool = True, returns=_ def emit_body(self, *args, **kwargs): """Emit this helper body into the currently active trace.""" - py_fn = rewrite_jit_function(self.py_fn) if self._ast_rewrite else self.py_fn + py_fn = ( + rewrite_jit_function(self.py_fn, reject_bare_returns=True) + if self._ast_rewrite + else self.py_fn + ) return py_fn(*args, **kwargs) def __call__(self, *args, **kwargs):