Add support for closures (nested functions and lambdas)

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

Author: gbonik

changelog.d/closures.md

@@ -0,0 +1,4 @@
+<!--- SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!--- SPDX-License-Identifier: Apache-2.0 -->
+
+- Added support for nested functions and lambdas.

src/cuda/tile/_compile.py

@@ -77,7 +77,8 @@ def _get_final_ir(pyfunc, args, tile_context) -> ir.Function:
     ir_ctx = ir.IRContext(tile_context)
     func_hir: hir.Function = get_function_hir(pyfunc, entry_point=True)
 
-    ir_args = _bind_kernel_arguments(func_hir.param_names, args, get_constant_annotations(pyfunc))
+    ir_args = _bind_kernel_arguments(tuple(func_hir.signature.parameters),
+                                     args, get_constant_annotations(pyfunc))
     func_body = hir2ir(func_hir, ir_args, ir_ctx)
     eliminate_assign_ops(func_body)
     dead_code_elimination_pass(func_body)
@@ -162,7 +163,7 @@ def _compiler_crash_dump(func_ir: ir.Function,
 
     artifacts = {
         f"{func_ir.name}.bytecode": bytes(bytecode_buf),
-        f"{func_ir.name}.cutileir": f"{func_ir.to_string(include_loc=False)}\n",
+        f"{func_ir.name}.cutileir": f"{func_ir.body.to_string(include_loc=False)}\n",
         "debug_info.txt": debug_info,
     }
 
@@ -184,7 +185,7 @@ def compile_tile(pyfunc,
 
     if 'CUTILEIR' in context.config.log_keys:
         code = (f"==== CuTile IR for {func_ir.name}==== \n\n"
-                f"{func_ir.to_string(include_loc=False)}\n\n")
+                f"{func_ir.body.to_string(include_loc=False)}\n\n")
         print(f'\n{code}', file=sys.stderr)
 
     sm_arch = get_sm_arch()

src/cuda/tile/_exception.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import dataclasses
 import linecache
+import os.path
 import re
 from dataclasses import dataclass
 from typing import Optional
@@ -11,10 +12,20 @@
 
 @dataclass(eq=False, frozen=True)
 class FunctionDesc:
-    name: str
+    name: str | None
     filename: str
     line: int
 
+    def __str__(self):
+        return f"'{self.name}' @{self.filename}:{self.line}"
+
+    def short_str(self):
+        if self.name is None:
+            base_name = os.path.basename(self.filename)
+            return f"<lambda at {base_name}:{self.line}>"
+        else:
+            return f"<function {self.name}>"
+
 
 @dataclass(slots=True, frozen=True)
 class Loc:

src/cuda/tile/_ir/hir.py

@@ -15,10 +15,11 @@
 
 
 import enum
+import inspect
 import threading
 from dataclasses import dataclass
 from textwrap import indent
-from typing import Any, Set, Mapping
+from typing import Any, Set, OrderedDict
 
 from cuda.tile._exception import Loc, FunctionDesc
 
@@ -119,15 +120,81 @@ def jump_str(self):
         return f"{self.jump._value_}{results_str}  # Line {self.jump_loc.line}"
 
 
-@dataclass
+# Resolved position of a named variable.
+#
+# If `depth` is -1, then this is a global variable, and the `index`
+# points into `hir.Function.frozen_global_names` / `hir.Function.frozen_global_values`.
+# If both `depth` and `index` are -1, then this is an indication that no variable
+# with the given name has been found.
+#
+# If `depth` is non-negative but is less than the depth of the current function,
+# then this is a captured local variable of an enclosing function.
+# In this case, the local variable name can be found in
+# `hir.Function.enclosing_funcs[depth].local_names[index].
+#
+# If `depth` is equal to the depth of the current function, then this is a local variable.
+# In this case, the local variable name can be found in `hir.Function.local_names[index]`.
+@dataclass(frozen=True)
+class ResolvedName:
+    depth: int
+    index: int
+
+
+UNKNOWN_NAME = ResolvedName(-1, -1)
+
+
+@dataclass(eq=False, repr=False)
 class Function:
     desc: FunctionDesc
     body: Block
-    param_names: tuple[str, ...]
+
+    # For nested functions, this signature is synthesized from AST. In this case,
+    # default values of parameters are represented with ClosureDefaultPlaceholder objects.
+    signature: inspect.Signature
+
+    # Names of all local variables defined in this function.
+    local_names: tuple[str, ...]
+
+    # For each function parameter, an index into `local_names`
+    param_local_indices: tuple[int, ...]
+
+    # For each function parameter, its location (same length as `param_local_indices`)
     param_locs: tuple[Loc, ...]
-    frozen_globals: Mapping[str, Any]
+
+    # Names of all global-like variables accessible by this function.
+    frozen_global_names: tuple[str, ...]
+
+    # Values of all global-like variables access by this function.
+    # We assume that they don't change after compilation, hence the name "frozen".
+    # Same length as `frozen_global_names`.
+    frozen_global_values: tuple[Any, ...]
+
+    # Strict upper bound on hir.Value.id inside this function definition,
+    # excluding any nested functions. Can be used to pre-allocate lists etc.
     value_id_upper_bound: int
 
+    # List of all function definitions nested directly inside this function.
+    nested_functions: "tuple[Function, ...]"
+
+    # Names of all variables that are ever loaded by this function
+    loaded_names: tuple[str, ...]
+
+    # For each variable name used in this function (i.e. loaded from or stored to),
+    # its resolved position in the scope.
+    used_names: OrderedDict[str, ResolvedName]
+
+    # Pre-computed view of `used_names`: for each non-negative `depth` strictly less than this
+    # function's depth, the list of all local variable indices captured by this function.
+    # Empty when this is a top-level function.
+    captures_by_depth: tuple[tuple[int, ...], ...]
+
+    # Sequence of enclosing function definitions, outermost first.
+    # Empty when this is a top-level function.
+    enclosing_funcs: "tuple[Function, ...]"
+
+    def __repr__(self):
+        return f"<HIR for function {self.desc}>"
+
 
 @dataclass
 class _OperandFormatter:
@@ -157,3 +224,4 @@ def build_tuple(*items): ...  # Makes a tuple (i.e. returns `items`)
 def identity(x): ...   # Identity function (i.e. returns `x`)
 def store_var(name, value, /): ...  # Store into a named variable
 def load_var(name, /): ...  # Load from a named variable
+def make_closure(func_hir: Function, /, *default_values): ...

src/cuda/tile/_ir/ir.py

@@ -286,6 +286,31 @@ def as_tuple(self) -> tuple[Var, ...]:
         return self.base_ptr, self.length
 
 
+@dataclass
+class ClosureValue(AggregateValue):
+    # Default values of parameters. These need to be carried by the closure's value
+    # because default expressions are evaluated at definition time, not when the closure is called.
+    # Should have the same length as the corresponding `ClosureTy.default_value_types`.
+    default_values: tuple[Var, ...]
+
+    # Tuple of the same length as `ty.func_hir.enclosing_functions`
+    # and `ty.frozen_capture_types_by_depth`, where `ty` is the `ClosureTy` of this closure.
+    #
+    # For each depth `i`, `frozen_captures_by_depth[i]` is either:
+    #   - None: means the enclosing function's LocalScope is still live;
+    #   - tuple[Var, ...]: means the enclosing function's LocalScope is no longer live.
+    #       The tuple contains the final values of the variables captured from the enclosing
+    #       function's scope. Its length should be the same as `ty.func_hir.captures_by_depth`.
+    frozen_captures_by_depth: tuple[tuple[Var, ...] | None, ...]
+
+    def as_tuple(self) -> tuple["Var", ...]:
+        return (
+            *self.default_values,
+            *(v for values in self.frozen_captures_by_depth
+              if values is not None for v in values)
+        )
+
+
 def terminator(cls):
     cls._is_terminator = True
     return cls