Don't ignore decorators of user-defined functions

Currently, ast2hir unwraps any decorators before getting the source
of the function. This is incorrect. For example:

    def decorate(func):
        @functools.wraps(func)
        def wrapper(x):
            return func(x + 3)
        return wrapper

    @decorate
    def decorated_helper(x):
        return x * 10

Calling decorated_helper(5) should return (5 + 3) * 10 == 80, but
currently it returns 5 * 10 == 50 because the wrapper is discarded.

This patch fixes this by only unwrapping the @ct.function decorator,
and respecting all other wrappers.

Signed-off-by: Greg Bonik <gbonik@nvidia.com>

Author: gbonik

changelog.d/fix-decorator.md

@@ -0,0 +1,2 @@
+- Fixed a bug where wrappers installed on a user-defined function
+  (e.g., by a decorator using `functools.wraps()`) were previously ignored.

src/cuda/tile/_execution.py

@@ -49,6 +49,7 @@ def decorator(func):
             def wrapped(*args, **kwargs):
                 return DispatchMode.get_current().call_tile_function_from_host(
                         wrapped, args, kwargs)
+            wrapped._cutile_function_wrapper = True
             return wrapped
 
     if func is None:
@@ -188,3 +189,7 @@ def is_stub(func) -> bool:
         func = getattr(func, "__wrapped__", None)
         if func is None:
             return False
+
+
+def is_function_wrapper(func) -> bool:
+    return getattr(func, "_cutile_function_wrapper", False)

src/cuda/tile/_ir/op_impl.py

@@ -171,7 +171,7 @@ def decorate(func):
             if len(fixed_args) > 0:
                 func = functools.partial(orig_func, *fixed_args)
 
-            func_sig = get_signature(func)
+            func_sig = inspect.signature(func)
             _verify_params_match(stub_sig, func_sig)
             is_coroutine = inspect.iscoroutinefunction(func)
             if is_coroutine:

src/cuda/tile/_ir/typing_support.py

@@ -16,7 +16,7 @@
 
 from .type import Type, DTypeConstructor, DTypeSpec, NONE, StringTy, \
     ELLIPSIS, SLICE, ModuleTy, FunctionTy, EnumTy, TypeTy, LooselyTypedScalar
-
+from .._execution import is_function_wrapper
 
 # Store mapping from 3rd party dtype objects
 # e.g. np.float32 -> float32, torch.bfloat16 -> bfloat16
@@ -131,7 +131,13 @@ def get_signature(f) -> inspect.Signature:
     elif is_dtype_constructor(f):
         # Data type constructors
         f = lambda x=0, /: None  # noqa: E731
-    return inspect.signature(f)
+
+    if isinstance(f, type):
+        return inspect.signature(f)
+
+    while is_function_wrapper(f):
+        f = f.__wrapped__
+    return inspect.signature(f, follow_wrapped=False)
 
 
 def is_supported_builtin_func(x: Any) -> bool:

src/cuda/tile/_passes/ast2hir.py

@@ -13,6 +13,7 @@
 
 from cuda.tile import _datatype as datatype
 from cuda.tile._exception import TileSyntaxError, Loc, FunctionDesc
+from cuda.tile._execution import is_function_wrapper
 from cuda.tile._ir.hir import make_value, ResolvedName, UNKNOWN_NAME
 from cuda.tile._ir import hir, hir_stubs
 from cuda.tile._ir.type import ClosureDefaultPlaceholder, FormattedPiece, StringFormat
@@ -22,10 +23,16 @@
 
 @lru_cache
 def get_function_hir(pyfunc: Callable, entry_point: bool) -> hir.Function:
-    # Get the original function from the decorated function if it exists.
-    pyfunc = getattr(pyfunc, "__wrapped__", pyfunc)
+    # Unwrap the @function decorator
+    while is_function_wrapper(pyfunc):
+        pyfunc = pyfunc.__wrapped__
+
+    # Use findsource() instead of getsourcelines() because the latter unwraps any decorators,
+    # which we don't want.
+    file_source_lines, first_line = inspect.findsource(pyfunc)
+    source_lines = inspect.getblock(file_source_lines[first_line:])
+    first_line += 1
 
-    source_lines, first_line = inspect.getsourcelines(pyfunc)
     # The source code of our function could be inside a class, an if-else block etc.
     # This means it can have extra indentation on the left. If we try to give it
     # to ast.parse() as is, we will get a parse error. The common workaround

test/test_helper_function.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
+import functools
 import inspect
 
 import pytest
@@ -364,3 +365,47 @@ def kernel(x):  # Line +8
     )
     with pytest.raises(TileTypeError, match=msg_regex):
         ct.launch(torch.cuda.current_stream(), (1,), kernel, (x,))
+
+
+def decorate(func):
+    @functools.wraps(func)
+    def wrapper(x):
+        return func(x + 3)
+    return wrapper
+
+
+@decorate
+def decorated_helper(x):
+    return x * 10
+
+
+def test_decorated_helper_function():
+    @ct.kernel
+    def kernel(y):
+        t = decorated_helper(5)
+        ct.scatter(y, (), t)
+    y = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (y,))
+    assert y.item() == 80
+
+
+def forward(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        return func(*args, **kwargs)
+    return wrapper
+
+
+@forward
+def forward_helper(x):
+    return x * 10
+
+
+def test_decorated_helper_function_forward():
+    @ct.kernel
+    def kernel(y):
+        t = forward_helper(5)
+        ct.scatter(y, (), t)
+    y = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (y,))
+    assert y.item() == 50