Add pto.func helpers

ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md

@@ -2,10 +2,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
-the `@pto.jit` entry and module contracts, the two programming models, the two
-compilation backends, sub-kernel reference, parameter contracts, and boundary
+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.
 
 
@@ -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 with AST-rewritten control flow
 @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` 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. 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(returns=pto.i32)
+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
 

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,20 @@ 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. `@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
 
 By default, a native Python `if` becomes a device-side conditional:

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)
@@ -370,9 +370,49 @@ def _name(name: str, ctx=ast.Load()):
     return ast.Name(id=name, ctx=ctx)
 
 
+class _ControlFlowExitVisitor(ast.NodeVisitor):
+    def __init__(self, *, reject_bare_returns: bool):
+        self.exit_node = None
+        self._reject_bare_returns = reject_bare_returns
+
+    def visit_Return(self, 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
+
+    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, *, 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:
+            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):
+    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}"
@@ -468,6 +508,17 @@ def _rewrite_if(self, stmt, *, live_after, allow_loop_control=False):
             )
             return [stmt]
 
+        _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)
         else_info = _name_info(stmt.orelse)
@@ -630,6 +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_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; "

ptodsl/ptodsl/_cache_signature.py

@@ -0,0 +1,76 @@
+# 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 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,
+        getattr(code, "co_qualname", fn.__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:
+        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",
+    "function_cache_signature",
+]

ptodsl/ptodsl/_func.py

@@ -0,0 +1,118 @@
+# 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.
+"""``@pto.func`` decorator and reusable callable handle."""
+
+from __future__ import annotations
+
+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, function_cache_signature
+from ._tracing import current_runtime
+
+_RETURNS_UNSET = object()
+
+
+@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, returns=_RETURNS_UNSET):
+        self.spec = spec
+        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:
+            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):
+        """Emit this helper body into the currently active trace."""
+        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):
+        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,
+            function_cache_signature(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, returns=_RETURNS_UNSET):
+    """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,
+            returns=returns,
+        )
+
+    if fn is not None:
+        return decorator(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",
+    "func",
+]