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..bf42891b0b 100644 --- a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md +++ b/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 diff --git a/ptodsl/docs/user_guide/05-control-flow.md b/ptodsl/docs/user_guide/05-control-flow.md index 0afca4f409..137d2493e3 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,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: diff --git a/ptodsl/ptodsl/_ast_rewrite.py b/ptodsl/ptodsl/_ast_rewrite.py index be766b58f5..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) @@ -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; " diff --git a/ptodsl/ptodsl/_cache_signature.py b/ptodsl/ptodsl/_cache_signature.py new file mode 100644 index 0000000000..92a111fae0 --- /dev/null +++ b/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", +] diff --git a/ptodsl/ptodsl/_func.py b/ptodsl/ptodsl/_func.py new file mode 100644 index 0000000000..f79c72ff02 --- /dev/null +++ b/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", +] 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..0ffc5a361b 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -12,26 +12,31 @@ 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 .._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, 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 @@ -57,6 +62,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.""" @@ -67,9 +73,15 @@ def cache_key(self) -> tuple: tuple((attr_name, str(attr_value)) for attr_name, attr_value in self.attributes), ) + 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}" @@ -536,6 +548,98 @@ 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() + 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 { + 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) + 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, + ) + 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 runtime_arg_values), + result_types=self._declared_ptodsl_func_result_types(func_template), + attributes=(("pto.ptodsl.callable_kind", StringAttr.get("func")),), + identity=identity, + ) + 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 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: + 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=helper_spec.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 runtime_arg_values]) + 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 +859,105 @@ def lookup_helper(self, symbol_name: str): return helper return None + def _normalize_ptodsl_func_argument(self, name: str, annotation, value): + raw_value = unwrap_surface_value(value) + if hasattr(raw_value, "type"): + return raw_value + if annotation is not inspect.Parameter.empty: + try: + target_type = _resolve(annotation) + except Exception: + target_type = annotation + 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_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): + 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 + elif isinstance(result, list): + values = tuple(result) + else: + values = (result,) + + 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)}" + ) + + 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 + try: + normalized.append( + coerce_scalar_to_type( + raw_value, + target_type, + context=f"@pto.func return value {index}", + ) + ) + continue + except TypeError: + pass + raise TypeError( + f"@pto.func return value {index} must match declared type {target_type}, 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) @@ -769,7 +972,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 @@ -792,7 +995,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 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..3b351b2f32 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 @@ -1268,6 +1269,163 @@ def helper(limit, enabled): _ = value +@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 + for _ in range(limit): + total = total + one + return total + + +@pto.func(returns=pto.i32) +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(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) + return value, value + one + + +@pto.func(returns=None) +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.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( + """ +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) + 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 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_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_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) + _ = 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) @@ -1462,6 +1620,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(returns=None) +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 +5175,113 @@ 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", + ) + 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 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() + 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", + ) + 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(), + "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), + "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( ast_nested_helper_freevar_if_merge_text, @@ -5117,6 +5405,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, 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 diff --git a/tmp/ptodsl_func_chain_probe.mlir b/tmp/ptodsl_func_chain_probe.mlir new file mode 100644 index 0000000000..4bcc3d85ed --- /dev/null +++ b/tmp/ptodsl_func_chain_probe.mlir @@ -0,0 +1,85 @@ +// 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. +// +// 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_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_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_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_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 + %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_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 + 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_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 + %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..d47af6e29e --- /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(returns=pto.i32) +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(returns=pto.i32) +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, returns=pto.i32) +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(returns=pto.i32) +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(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) + + +@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()