Use coroutines in hir2ir to simplify the logic and increase recursion limit

Currently, we use a `todo_stack` to store the pending HIR calls. When a
user-defined function is called, we push the `hir.Calls` of its body to the
`todo_stack` in reverse order.

This helps reduce the Python call stack depth but makes it harder to track the
state (i.e., which scope does a pending `hir.Call` belong to?). We also need to
do some awkward things like pushing special calls to `_set_scope` to reset the
scope after the user-defined function call returns, etc. The code would be much
more straightforward if we just used recursion.

Moreover, this doesn't actually solve the stack depth problem, since nested
blocks are processed using recursion anyway. Because of this, we need to keep
our recursion limit fairly small (~50 levels deep).

We address these problems by using coroutines (aka async functions) and having
a partially software stack. Thus we can write the code in a recursive fashion
without worrying about the Python's recursion limit.

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

Author: gbonik

src/cuda/tile/_coroutine_util.py

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+from contextlib import ExitStack
+from dataclasses import dataclass
+from typing import Awaitable
+
+
+# Run a coroutine using a software stack to bypass the Python's recursion limit.
+# Use resume_after() to break the call chain and push a new frame to the software stack.
+def run_coroutine(awaitable: Awaitable):
+    ret = None
+    exc_info = None
+    stack = []
+    with ExitStack() as es:
+        try:
+            stack.append(awaitable.__await__())
+            while stack:
+                top = stack[-1]
+                try:
+                    continuation = top.send(ret) if exc_info is None else top.throw(*exc_info)
+                except StopIteration as s:
+                    ret = s.value
+                    exc_info = None
+                    stack.pop()
+                except Exception:
+                    ret = None
+                    exc_info = sys.exc_info()
+                    stack.pop()
+                else:
+                    ret = exc_info = None
+                    stack.append(continuation.__await__())
+            if exc_info is None:
+                return ret
+            else:
+                raise exc_info[1]
+        finally:
+            for c in stack:
+                es.callback(c.close)
+
+
+# Replace `await foo()` with `await resume_after(foo())` to bypass the recursion limit.
+@dataclass
+class resume_after:
+    awaitable: Awaitable
+
+    def __await__(self):
+        return (yield self.awaitable)

src/cuda/tile/_ir/ir.py

@@ -596,6 +596,15 @@ def change_loop_info(self, new_info: ControlFlowInfo):
         finally:
             self.loop_info = old
 
+    @contextmanager
+    def change_scope(self, new_scope: Scope):
+        old = self.scope
+        self.scope = new_scope
+        try:
+            yield
+        finally:
+            self.scope = old
+
     def __enter__(self):
         assert not self._entered
         self._prev_builder = _current_builder.builder

src/cuda/tile/_ir/op_impl.py

@@ -57,7 +57,7 @@ def wrapper(*args, **kwargs):
                 return func(*args, **kwargs)
             finally:
                 _current_stub.stub_and_args = old
-
+        wrapper._is_coroutine = inspect.iscoroutinefunction(func)
         op_implementations[stub] = wrapper
         return orig_func
 

src/cuda/tile/_ir/ops.py

@@ -161,7 +161,7 @@ def format_var(var):
 
 
 @impl(hir.loop)
-def loop_impl(body: hir.Block, iterable: Var):
+async def loop_impl(body: hir.Block, iterable: Var):
     from .._passes.hir2ir import dispatch_hir_block
 
     range_ty = require_optional_range_type(iterable)
@@ -171,7 +171,7 @@ def loop_impl(body: hir.Block, iterable: Var):
         # have a "break" at the end of the loop body, and no other break/continue statements.
         info = ControlFlowInfo((), flatten=True)
         with Builder.get_current().change_loop_info(info):
-            dispatch_hir_block(body)
+            await dispatch_hir_block(body)
         return ()
 
     builder = Builder.get_current()
@@ -207,7 +207,7 @@ def loop_impl(body: hir.Block, iterable: Var):
         flat_body_vars = flatten_block_parameters(body_vars)
 
         # Dispatch the body (hir.Block) to populate the new_body (ir.Block) with Operations
-        dispatch_hir_block(body)
+        await dispatch_hir_block(body)
 
     # Propagate type information from Continue/Break to body/result phis
     for jump_info in loop_info.jumps:
@@ -359,11 +359,11 @@ def _to_string_rhs(self) -> str:
         return f"if(cond={self.cond})"
 
 
-def _flatten_branch(branch: hir.Block) -> tuple[Var, ...]:
+async def _flatten_branch(branch: hir.Block) -> tuple[Var, ...]:
     from .._passes.hir2ir import dispatch_hir_block
     info = ControlFlowInfo((), flatten=True)
     with Builder.get_current().change_if_else_info(info):
-        dispatch_hir_block(branch)
+        await dispatch_hir_block(branch)
     if len(info.jumps) == 0:
         return ()
     else:
@@ -372,13 +372,13 @@ def _flatten_branch(branch: hir.Block) -> tuple[Var, ...]:
 
 
 @impl(hir.if_else)
-def if_else_impl(cond: Var, then_block: hir.Block, else_block: hir.Block) -> tuple[Var, ...]:
+async def if_else_impl(cond: Var, then_block: hir.Block, else_block: hir.Block) -> tuple[Var, ...]:
     from .._passes.hir2ir import dispatch_hir_block
 
     require_bool(cond)
     if cond.is_constant():
         branch_taken = then_block if cond.get_constant() else else_block
-        return _flatten_branch(branch_taken)
+        return await _flatten_branch(branch_taken)
 
     # Get the total number of results by adding the number of stored variables.
     # Note: we sort the stored variable names to make the order deterministic.
@@ -387,7 +387,7 @@ def if_else_impl(cond: Var, then_block: hir.Block, else_block: hir.Block) -> tup
     # Convert the "then" branch from HIR to IR
     info = ControlFlowInfo(stored_names)
     with nested_block(then_block.name, then_block.loc, if_else_info=info) as new_then_block:
-        dispatch_hir_block(then_block)
+        await dispatch_hir_block(then_block)
 
     # If "then" branch doesn't yield, transform our if-else into the following:
     #    if cond:
@@ -402,11 +402,11 @@ def if_else_impl(cond: Var, then_block: hir.Block, else_block: hir.Block) -> tup
             end_branch(())
         add_operation(IfElse, (),
                       cond=cond, then_block=new_then_block, else_block=new_else_block)
-        return _flatten_branch(else_block)
+        return await _flatten_branch(else_block)
 
     # Convert the "else" branch from HIR to IR
     with nested_block(else_block.name, else_block.loc, if_else_info=info) as new_else_block:
-        dispatch_hir_block(else_block)
+        await dispatch_hir_block(else_block)
 
     # Do type/constant propagation
     num_results = len(info.jumps[0].outputs)

src/cuda/tile/_passes/hir2ir.py

@@ -1,46 +1,52 @@
 # SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
-
 import inspect
 import sys
 from contextlib import contextmanager
-from dataclasses import dataclass
-from typing import Any, Sequence
+from typing import Any
 
 from .ast2hir import get_function_hir
 from .. import TileTypeError
+from .._coroutine_util import resume_after, run_coroutine
 from .._exception import Loc, TileSyntaxError, TileInternalError, TileError, TileRecursionError
 from .._ir import hir, ir
 from .._ir.ir import Var, IRContext, Argument, Scope, LocalScope, BoundMethodValue
-from .._ir.op_impl import op_implementations, impl
+from .._ir.op_impl import op_implementations
 from .._ir.ops import loosely_typed_const, assign, end_branch, return_, continue_, \
-    break_, flatten_block_parameters
+    break_, flatten_block_parameters, store_var
 from .._ir.type import FunctionTy, BoundMethodTy, DTypeConstructor
 from .._ir.typing_support import get_signature
 
 
-MAX_RECURSION_DEPTH = 50
+MAX_RECURSION_DEPTH = 1000
 
 
 def hir2ir(func_hir: hir.Function,
            args: tuple[Argument, ...],
            ir_ctx: IRContext) -> ir.Block:
+    # Run as a coroutine using a software stack, so that we don't exceed Python's recursion limit.
+    return run_coroutine(_hir2ir_coroutine(func_hir, args, ir_ctx))
+
+
+async def _hir2ir_coroutine(func_hir: hir.Function, args: tuple[Argument, ...], ir_ctx: IRContext):
     scope = _create_scope(func_hir, ir_ctx, call_site=None)
     aggregate_params = [
         scope.local.redefine(param_name, param_loc)
         for param_name, param_loc in zip(func_hir.param_names, func_hir.param_locs, strict=True)
     ]
-    preamble = []
-    for param_name, var, arg in zip(func_hir.param_names, aggregate_params, args, strict=True):
-        var.set_type(arg.type)
-        if arg.is_const:
-            preamble.append(hir.Call((), hir.store_var, (param_name, arg.const_value), (), var.loc))
 
     with ir.Builder(ir_ctx, func_hir.body.loc, scope) as ir_builder:
-        flat_params = flatten_block_parameters(aggregate_params)
         try:
-            _dispatch_hir_block_inner(preamble, func_hir.body, ir_builder)
+            for param_name, var, arg in zip(func_hir.param_names, aggregate_params, args,
+                                            strict=True):
+                var.set_type(arg.type)
+                if arg.is_const:
+                    var = loosely_typed_const(arg.const_value)
+                    store_var(param_name, var, var.loc)
+            flat_params = flatten_block_parameters(aggregate_params)
+
+            await _dispatch_hir_block_inner(func_hir.body, ir_builder)
         except Exception as e:
             if 'CUTILEIR' in ir_ctx.tile_ctx.config.log_keys:
                 highlight_loc = e.loc if hasattr(e, 'loc') else None
@@ -60,54 +66,41 @@ def _create_scope(func_hir: hir.Function, ir_ctx: IRContext, call_site: Loc | No
     return Scope(local_scope, func_hir.frozen_globals, call_site)
 
 
-def dispatch_hir_block(block: hir.Block):
-    _dispatch_hir_block_inner((), block, ir.Builder.get_current())
-
-
-@dataclass
-class _State:
-    done: list[hir.Call]
-    current: hir.Call | None
-    todo_stack: list[hir.Call]
+async def dispatch_hir_block(block: hir.Block, cur_builder: ir.Builder | None = None):
+    if cur_builder is None:
+        cur_builder = ir.Builder.get_current()
+    await _dispatch_hir_block_inner(block, cur_builder)
 
-    @contextmanager
-    def next_call(self):
-        call = self.current = self.todo_stack.pop()
-        yield call
-        # Intentionally not in a "finally" block because we want to preserve the state
-        # for the debug printout in case of an exception.
-        self.current = None
-        self.done.append(call)
 
-
-def _dispatch_hir_block_inner(preamble: Sequence[hir.Call],
-                              block: hir.Block,
-                              builder: ir.Builder):
-    state = _State([], None, list(reversed(block.calls)) + list(reversed(preamble)))
+async def _dispatch_hir_block_inner(block: hir.Block, builder: ir.Builder):
+    cursor = 0  # Pre-initialize to guarantee it's defined in the `except` block
     try:
-        if not _dispatch_hir_calls(state, builder):
-            return ()
+        for cursor, call in enumerate(block.calls):
+            loc = _add_call_site(call.loc, builder)
+            with _wrap_exceptions(loc), builder.change_loc(loc):
+                await _dispatch_call(call, builder)
+            if builder.is_terminated:
+                # The current block has been terminated, e.g. by flattening an if-else
+                # with a constant condition (`if True: break`).
+                return
+        cursor = len(block.calls)
+
         result_vars = tuple(_resolve_operand(x) for x in block.results)
         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:
             hir_params = ", ".join(p.name for p in block.params)
-            hir_lines = [str(c) for c in state.done]
-            cur_idx = len(hir_lines)
-            if state.current is not None:
-                hir_lines.append(str(state.current))
-            hir_lines.extend(str(c) for c in reversed(state.todo_stack))
+            hir_lines = [str(c) for c in block.calls]
             hir_lines.append(block.jump_str())
-            hir_str = "\n".join("{}{}".format("--> " if i == cur_idx else "    ", c)
+            hir_str = "\n".join("{}{}".format("--> " if i == cursor else "    ", c)
                                 for i, c in enumerate(hir_lines))
             print(f"==== HIR for ^{block.name}({hir_params}) ====\n{hir_str}\n", file=sys.stderr)
         raise
 
 
-def _dispatch_hir_jump(jump: hir.Jump,
-                       block_results: tuple[Var, ...]):
+def _dispatch_hir_jump(jump: hir.Jump | None, block_results: tuple[Var, ...]):
     match jump:
         case hir.Jump.END_BRANCH:
             end_branch(block_results)
@@ -120,23 +113,10 @@ def _dispatch_hir_jump(jump: hir.Jump,
         case hir.Jump.RETURN:
             assert len(block_results) == 1
             return_(block_results[0])
+        case None: pass
         case _: assert False
 
 
-def _dispatch_hir_calls(state: _State, cur_builder: ir.Builder) -> bool:
-    while len(state.todo_stack) > 0:
-        with state.next_call() as call:
-            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
-                # with a constant condition (`if True: break`). By returning False,
-                # we signal that the original jump and block results should be ignored.
-                return False
-    return True
-
-
 def _add_call_site(loc: Loc, builder: ir.Builder) -> Loc:
     return loc.with_call_site(builder.scope.call_site)
 
@@ -152,7 +132,7 @@ def _wrap_exceptions(loc: Loc):
             raise TileInternalError(str(e)) from e
 
 
-def _dispatch_call(call: hir.Call, builder: ir.Builder, todo_stack: list[hir.Call]):
+async def _dispatch_call(call: hir.Call, builder: ir.Builder):
     first_idx = len(builder.ops)
     callee_var = _resolve_operand(call.callee)
     callee, self_arg = _get_callee_and_self(callee_var)
@@ -161,7 +141,11 @@ def _dispatch_call(call: hir.Call, builder: ir.Builder, todo_stack: list[hir.Cal
     arg_list = _bind_args(callee, args, kwargs)
 
     if callee in op_implementations:
-        result = op_implementations[callee](*arg_list)
+        impl = op_implementations[callee]
+        result = impl(*arg_list)
+        if impl._is_coroutine:
+            result = await result
+
         if builder.is_terminated:
             # The current block has been terminated, e.g. by flattening an if-else
             # with a constant condition (`if True: break`). Ignore the `result` in this case.
@@ -203,25 +187,20 @@ def _dispatch_call(call: hir.Call, builder: ir.Builder, todo_stack: list[hir.Cal
                                       " functions are not supported")
         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.
-        for callee_retval, caller_res in zip(callee_hir.body.results, call.results):
-            todo_stack.append(hir.Call((caller_res,), hir.identity, (callee_retval,), (), call.loc))
+        # Activate a fresh Scope.
+        new_scope = _create_scope(callee_hir, builder.ir_ctx, call_site=builder.loc)
+        with builder.change_scope(new_scope):
+            # Call store_var() to bind arguments to parameters.
+            for arg, param_name, param_loc in zip(arg_list, callee_hir.param_names,
+                                                  callee_hir.param_locs, strict=True):
+                store_var(param_name, arg, param_loc)
 
-        # Now we create a fresh Scope for the new function and install it on the builder.
-        # We need to reset the builder back to the old scope when we return.
-        # 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, call_site=builder.loc)
+            # Dispatch the function body. Use resume_after() to break the call stack
+            # and make sure we stay within the Python's recursion limit.
+            await resume_after(dispatch_hir_block(callee_hir.body, builder))
 
-        # Now push the function body.
-        todo_stack.extend(reversed(callee_hir.body.calls))
-
-        # Finally, call store_var() to bind arguments to parameters.
-        for arg, param_name, param_loc in zip(arg_list, callee_hir.param_names,
-                                              callee_hir.param_locs, strict=True):
-            todo_stack.append(hir.Call((), hir.store_var, (param_name, arg), (), param_loc))
+        for callee_retval, caller_res in zip(callee_hir.body.results, call.results):
+            assign(callee_retval, caller_res)
 
 
 def _is_freshly_defined(var: Var, builder: ir.Builder, first_idx: int):
@@ -277,13 +256,3 @@ def _bind_args(sig_func, args, kwargs) -> list[Var]:
             assert param.default is not param.empty
             ret.append(loosely_typed_const(param.default))
     return ret
-
-
-def _set_scope(scope): ...
-
-
-@impl(_set_scope)
-def _set_scope_impl(scope):
-    assert isinstance(scope, Scope)
-    builder = ir.Builder.get_current()
-    builder.scope = scope