Helper function call improvements

- Allow recursive function calls. Add a recursion limit instead.
- Include cuTile stack trace in error messages, rather than a single
  immediate location.
- Don't set Loc.call_site in ast2hir, delay until hir2ir. This way, a
  single HIR object could in theory be reused for multiple calls
  (handy for implementing closures/lambdas, as well as potentially
  caching the HIR for faster kernel specialization)

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

Author: gbonik

changelog.d/allow-recursion.md

@@ -0,0 +1,10 @@
+<!--- SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!--- SPDX-License-Identifier: Apache-2.0 -->
+
+- Lift the ban on recursive helper function calls. Instead, add a limit on recursion depth.
+- Add a new exception class `TileRecursionError`, thrown at compile time when the recursion limit
+  is reached during function call inlining.
+- Include a full cuTile traceback in error messages. Improve formatting of code locations:
+  include function names, remove unnecessary characters to reduce line lengths.
+- Expose the `TileError` base class in the public API.
+

src/cuda/tile/__init__.py

@@ -42,7 +42,9 @@
 from cuda.tile._exception import (
     TileCompilerExecutionError,
     TileCompilerTimeoutError,
+    TileError,
     TileInternalError,
+    TileRecursionError,
     TileSyntaxError,
     TileTypeError,
     TileValueError,
@@ -172,7 +174,9 @@
 
     "TileCompilerExecutionError",
     "TileCompilerTimeoutError",
+    "TileError",
     "TileInternalError",
+    "TileRecursionError",
     "TileSyntaxError",
     "TileTypeError",
     "TileValueError",

src/cuda/tile/_compile.py

@@ -75,7 +75,7 @@ def wrapper(*args, **kwargs):
 
 def _get_final_ir(pyfunc, args, tile_context) -> ir.Block:
     ir_ctx = ir.IRContext(tile_context)
-    func_hir: hir.Function = get_function_hir(pyfunc, ir_ctx, call_site=None)
+    func_hir: hir.Function = get_function_hir(pyfunc, ir_ctx, entry_point=True)
 
     ir_args = _bind_kernel_arguments(func_hir.param_names, args, get_constant_annotations(pyfunc))
     func_body = hir2ir(func_hir, ir_args, ir_ctx)

src/cuda/tile/_exception.py

@@ -1,24 +1,34 @@
 # SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
-
+import dataclasses
 import linecache
 import re
 from dataclasses import dataclass
-from typing import Optional, Callable
+from typing import Optional
 from unicodedata import east_asian_width
 
 
+@dataclass(eq=False, frozen=True)
+class FunctionDesc:
+    name: str
+    filename: str
+    line: int
+
+
 @dataclass(slots=True)
 class Loc:
     line: int
     col: int
     filename: Optional[str] = None
     last_line: Optional[int] = None
     end_col: Optional[int] = None
-    function: Optional[Callable] = None
+    function: Optional[FunctionDesc] = None
     call_site: Optional["Loc"] = None
 
+    def with_call_site(self, call_site) -> "Loc":
+        return dataclasses.replace(self, call_site=call_site)
+
     def __str__(self) -> str:
         if self.filename:
             return f"{self.filename}:{self.line}:{self.col}"
@@ -49,6 +59,14 @@ def _wcwidth(s: str) -> int:
 
 
 def format_location(loc: Loc):
+    frames = []
+    while loc is not None:
+        frames.append(loc)
+        loc = loc.call_site
+    return "".join(_format_location_frame(x) for x in reversed(frames))
+
+
+def _format_location_frame(loc: Loc) -> str:
     if loc.is_unknown():
         return "Unknown location"
 
@@ -74,12 +92,14 @@ def format_location(loc: Loc):
         cols_str = f"col {visual_col + 1}"
     else:
         end_visual_col = _wcwidth(line_bytes[:end_col].decode())
-        cols_str = f"col {visual_col + 1}--{end_visual_col}"
+        cols_str = f"col {visual_col + 1}-{end_visual_col}"
 
     spaces = " " * visual_col
     carets = "^" * (end_visual_col - visual_col)
 
-    return (f'  In file "{loc.filename}", {lines_str}, {cols_str}:\n'
+    func_str = "" if loc.function is None else f", in {loc.function.name}"
+
+    return (f'  "{loc.filename}", {lines_str}, {cols_str}{func_str}:\n'
             f"    {line_text}\n"
             f"    {spaces}{carets}\n")
 
@@ -103,6 +123,12 @@ class TileTypeError(TileError):
     pass
 
 
+class TileRecursionError(TileError):
+    """Thrown at compile time to indicate that the recursion limit has been reached
+    when inlining a function call.
+    """
+
+
 class TileValueError(TileError):
     """Exception when an unexpected python value is encountered."""
     pass

src/cuda/tile/_ir/hir.py

@@ -19,7 +19,7 @@
 from textwrap import indent
 from typing import Any, Set, Mapping
 
-from cuda.tile._exception import Loc
+from cuda.tile._exception import Loc, FunctionDesc
 from cuda.tile._ir.ir import Var
 
 
@@ -97,6 +97,7 @@ def jump_str(self):
 
 @dataclass
 class Function:
+    desc: FunctionDesc
     body: Block
     param_names: tuple[str, ...]
     param_locs: tuple[Loc, ...]

src/cuda/tile/_ir/ir.py

@@ -327,6 +327,7 @@ def nested_block(name: str, loc: Loc, params: Sequence[Var] = (),
     new_loop_info = loop_info or prev_builder.loop_info
     new_if_else_info = if_else_info or prev_builder.if_else_info
     scope = prev_builder.scope
+    loc = loc.with_call_site(scope.call_site)
     with Builder(prev_builder.ir_ctx, loc, scope, new_loop_info, new_if_else_info) as builder, \
             scope.local.enter_branch():
         yield block
@@ -410,6 +411,7 @@ def __setitem__(self, key, value):
 class Scope:
     local: LocalScope
     frozen_globals: Mapping[str, Any]
+    call_site: Loc | None
 
 
 class Builder:

src/cuda/tile/_ir2bytecode.py

@@ -3,7 +3,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import functools
-import inspect
 import os
 from contextlib import contextmanager
 from typing import Dict, Tuple, Sequence, Any, List, Optional, Iterator, Set
@@ -14,7 +13,8 @@
 import cuda.tile._bytecode as bc
 from cuda.tile._compiler_options import CompilerOptions
 from cuda.tile._debug import CUDA_TILE_TESTING_DISABLE_DIV
-from cuda.tile._exception import TileInternalError, TileError, ConstFoldNotImplementedError
+from cuda.tile._exception import TileInternalError, TileError, ConstFoldNotImplementedError, \
+    FunctionDesc
 from cuda.tile._ir.ir import Block, Loc, Var, IRContext
 from cuda.tile._ir.ops_utils import (
     padding_mode_to_bytecode, rounding_mode_to_bytecode,
@@ -472,27 +472,24 @@ def __init__(self, debug_attr_table: bc.DebugAttrTable, linkage_name: str, anony
         self._linkage_name = linkage_name
         self._anonymize = anonymize
 
-    def get_subprogram(self, pyfunc) -> bc.DebugAttrId:
+    def get_subprogram(self, func_desc: FunctionDesc) -> bc.DebugAttrId:
         try:
-            return self._subprogram_cache[pyfunc]
+            return self._subprogram_cache[func_desc]
         except KeyError:
             pass
 
-        func_name = pyfunc.__name__
-        func_filename = inspect.getfile(pyfunc)
-        _, func_line = inspect.findsource(pyfunc)
-        func_dirname, func_basename = os.path.split(func_filename)
+        func_dirname, func_basename = os.path.split(func_desc.filename)
         file_attr = self._debug_attr_table.file(func_basename, func_dirname)
         compile_unit_attr = self._debug_attr_table.compile_unit(file_attr)
         ret = self._debug_attr_table.subprogram(
             file=file_attr,
-            line=func_line,
-            name=func_name,
+            line=func_desc.line,
+            name=func_desc.name,
             linkage_name=self._linkage_name,
             compile_unit=compile_unit_attr,
-            scope_line=func_line,
+            scope_line=func_desc.line,
         )
-        self._subprogram_cache[pyfunc] = ret
+        self._subprogram_cache[func_desc] = ret
         return ret
 
     def get_debugattr(self, loc: Loc) -> bc.DebugAttrId:

src/cuda/tile/_passes/ast2hir.py

@@ -10,14 +10,14 @@
 from typing import List, Sequence, Optional, Any, Dict, Type, Callable
 
 from cuda.tile import _datatype as datatype
-from cuda.tile._exception import TileSyntaxError, Loc
+from cuda.tile._exception import TileSyntaxError, Loc, FunctionDesc
 from cuda.tile._ir.ir import IRContext, Var
 from cuda.tile._ir import hir
 
 
 def get_function_hir(pyfunc: Callable,
                      ir_ctx: IRContext,
-                     call_site: Optional[Loc]) -> hir.Function:
+                     entry_point: bool) -> hir.Function:
     # Get the original function from the decorated function if it exists.
     pyfunc = getattr(pyfunc, "__wrapped__", pyfunc)
 
@@ -56,14 +56,17 @@ def get_function_hir(pyfunc: Callable,
         for name, cell in zip(pyfunc.__code__.co_freevars, pyfunc.__closure__):
             func_globals[name] = cell.cell_contents
 
-    ctx = _Context(inspect.getfile(pyfunc), first_line, call_site, pyfunc, ir_ctx)
+    filename = inspect.getfile(pyfunc)
+    desc = FunctionDesc(func_def.name, filename, first_line)
+    ctx = _Context(filename, first_line, desc, entry_point, ir_ctx)
     assert isinstance(func_def, ast.FunctionDef)
     body = _ast2hir(func_def, ctx)
     all_ast_args = _get_all_parameters(func_def, ctx)
     param_names = tuple(p.arg for p in all_ast_args)
     param_locs = tuple(ctx.get_loc(p) for p in all_ast_args)
     body.stored_names.update(param_names)
-    return hir.Function(body, param_names, param_locs, func_globals)
+
+    return hir.Function(desc, body, param_names, param_locs, func_globals)
 
 
 # Translate the 1-based line number of the chunk we passed to the AST parser
@@ -81,13 +84,12 @@ class LoopKind(Enum):
 
 
 class _Context:
-    def __init__(self, filename: str, first_line: int, call_site: Optional[Loc],
-                 function: Callable, ir_ctx: IRContext):
+    def __init__(self, filename: str, first_line: int, function_desc: FunctionDesc,
+                 entry_point: bool, ir_ctx: IRContext):
         self.filename = filename
         self.first_line = first_line
-        self.entry_point = call_site is None
-        self.call_site = call_site
-        self.function = function
+        self.function_desc = function_desc
+        self.entry_point = entry_point
         self.parent_loops: List[LoopKind] = []
         self.current_loc = Loc.unknown()
         self.current_block: Optional[hir.Block] = None
@@ -145,8 +147,7 @@ def get_loc(self, node: ast.AST) -> Loc:
         # Subtract 1 from the column offset to correct for an extra level
         # of indentation we inserted for the dummy "if True" block.
         return Loc(line_no, node.col_offset - 1, self.filename,
-                   last_line_no, node.end_col_offset - 1, self.function,
-                   self.call_site)
+                   last_line_no, node.end_col_offset - 1, self.function_desc)
 
     def syntax_error(self, message: str, loc=None) -> TileSyntaxError:
         if loc is None:

src/cuda/tile/_passes/hir2ir.py

@@ -6,11 +6,11 @@
 import sys
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any, Callable, Sequence
+from typing import Any, Sequence
 
 from .ast2hir import get_function_hir
 from .. import TileTypeError
-from .._exception import Loc, TileSyntaxError, TileInternalError, TileError
+from .._exception import Loc, TileSyntaxError, TileInternalError, TileError, TileRecursionError
 from .._ir import hir, ir
 from .._ir.ir import Var, IRContext, Argument, Scope, LocalScope
 from .._ir.op_impl import op_implementations, impl
@@ -20,10 +20,13 @@
 from .._ir.typing_support import get_signature
 
 
+MAX_RECURSION_DEPTH = 50
+
+
 def hir2ir(func_hir: hir.Function,
            args: tuple[Argument, ...],
            ir_ctx: IRContext) -> ir.Block:
-    scope = _create_scope(func_hir, ir_ctx)
+    scope = _create_scope(func_hir, ir_ctx, call_site=None)
     ir_params = tuple(scope.local.redefine(name, loc)
                       for name, loc in zip(func_hir.param_names, func_hir.param_locs, strict=True))
     preamble = []
@@ -50,9 +53,9 @@ def hir2ir(func_hir: hir.Function,
     return ret
 
 
-def _create_scope(func_hir: hir.Function, ir_ctx: IRContext):
+def _create_scope(func_hir: hir.Function, ir_ctx: IRContext, call_site: Loc | None) -> Scope:
     local_scope = LocalScope(func_hir.body.stored_names, ir_ctx)
-    return Scope(local_scope, func_hir.frozen_globals)
+    return Scope(local_scope, func_hir.frozen_globals, call_site)
 
 
 def dispatch_hir_block(block: hir.Block):
@@ -83,7 +86,8 @@ def _dispatch_hir_block_inner(preamble: Sequence[hir.Call],
         if not _dispatch_hir_calls(state, builder):
             return ()
         result_vars = tuple(_resolve_operand(x) for x in block.results)
-        with _wrap_exceptions(block.jump_loc), builder.change_loc(block.jump_loc):
+        loc = _add_call_site(block.jump_loc, builder)
+        with _wrap_exceptions(loc), builder.change_loc(loc):
             _dispatch_hir_jump(block.jump, result_vars)
     except Exception:
         if 'CUTILEIR' in builder.ir_ctx.tile_ctx.config.log_keys:
@@ -120,7 +124,8 @@ def _dispatch_hir_jump(jump: hir.Jump,
 def _dispatch_hir_calls(state: _State, cur_builder: ir.Builder) -> bool:
     while len(state.todo_stack) > 0:
         with state.next_call() as call:
-            with _wrap_exceptions(call.loc), cur_builder.change_loc(call.loc):
+            loc = _add_call_site(call.loc, cur_builder)
+            with _wrap_exceptions(loc), cur_builder.change_loc(loc):
                 _dispatch_call(call, cur_builder, state.todo_stack)
             if cur_builder.is_terminated:
                 # The current block has been terminated, e.g. by flattening an if-else
@@ -130,6 +135,10 @@ def _dispatch_hir_calls(state: _State, cur_builder: ir.Builder) -> bool:
     return True
 
 
+def _add_call_site(loc: Loc, builder: ir.Builder) -> Loc:
+    return loc.with_call_site(builder.scope.call_site)
+
+
 @contextmanager
 def _wrap_exceptions(loc: Loc):
     with loc:
@@ -183,14 +192,14 @@ def _dispatch_call(call: hir.Call, builder: ir.Builder, todo_stack: list[hir.Cal
                 builder.ops[i] = builder.ops[i].clone(mapper)
     else:
         # Callee is a user-defined function.
-        _check_recursive_call(call.loc, callee)
+        _check_recursive_call(builder.loc)
         sig = get_signature(callee)
         for param_name, param in sig.parameters.items():
             if param.kind in (inspect.Parameter.VAR_POSITIONAL,
                               inspect.Parameter.VAR_KEYWORD):
                 raise TileSyntaxError("Variadic parameters in user-defined"
                                       " functions are not supported")
-        callee_hir = get_function_hir(callee, builder.ir_ctx, call_site=call.loc)
+        callee_hir = get_function_hir(callee, builder.ir_ctx, entry_point=False)
 
         # Since `todo_stack` is a stack, we push things backwards. First, we push identity()
         # calls to assign the temporary return values back to the original result variables.
@@ -202,7 +211,7 @@ def _dispatch_call(call: hir.Call, builder: ir.Builder, todo_stack: list[hir.Cal
         # For this purpose, we push a call to the special _set_scope stub.
         old_scope = builder.scope
         todo_stack.append(hir.Call((), _set_scope, (old_scope,), (), call.loc))
-        builder.scope = _create_scope(callee_hir, builder.ir_ctx)
+        builder.scope = _create_scope(callee_hir, builder.ir_ctx, call_site=builder.loc)
 
         # Now push the function body.
         todo_stack.extend(reversed(callee_hir.body.calls))
@@ -219,11 +228,14 @@ def _is_freshly_defined(var: Var, builder: ir.Builder, first_idx: int):
                for r in builder.ops[i].result_vars)
 
 
-def _check_recursive_call(call_loc: Loc, callee: Callable):
+def _check_recursive_call(call_loc: Loc):
+    depth = 1
     while call_loc is not None:
-        if call_loc.function is callee:
-            raise TileTypeError("Recursive function call detected")
+        depth += 1
         call_loc = call_loc.call_site
+    if depth > MAX_RECURSION_DEPTH:
+        raise TileRecursionError(f"Maximum recursion depth ({MAX_RECURSION_DEPTH}) reached"
+                                 f" while inlining a function call")
 
 
 def _get_callee_and_self(callee_var: Var) -> tuple[Any, tuple[()] | tuple[Var]]: