Add ct.static_iter() & tuple concatenation

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

Author: gbonik

changelog.d/static-iter.md

@@ -0,0 +1,5 @@
+<!--- SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!--- SPDX-License-Identifier: Apache-2.0 -->
+
+- Added `ct.static_iter` keyword that enables compile-time `for` loops.
+- Operator `+` can now be used to concatenate tuples.

docs/source/operations.rst

@@ -199,3 +199,4 @@ Metaprogramming Support
 
    static_assert
    static_eval
+   static_iter

src/cuda/tile/__init__.py

@@ -136,6 +136,7 @@
     sqrt,
     static_assert,
     static_eval,
+    static_iter,
     store,
     sub,
     sum,
@@ -276,6 +277,7 @@
     "sqrt",
     "static_assert",
     "static_eval",
+    "static_iter",
     "store",
     "sub",
     "sum",

src/cuda/tile/_dispatch_mode.py

@@ -37,11 +37,14 @@ def __init__(self, kind: StaticEvalKind):
         self._kind = kind
 
     def call_tile_function_from_host(self, func, args, kwargs):
-        from cuda.tile import static_eval, static_assert
-        if func in (static_eval, static_assert):
+        from cuda.tile import static_eval, static_assert, static_iter
+        if func in (static_eval, static_assert, static_iter):
             what = f"{func.__name__}() cannot be used"
         else:
-            what = "Tile functions cannot be called"
+            func_name = getattr(func, "__name__", "")
+            if len(func_name) > 0:
+                func_name = func_name + ": "
+            what = f"{func_name}Tile functions cannot be called"
 
         where = self._kind._value_
         raise TileStaticEvalError(f"{what} inside {where}.")

src/cuda/tile/_ir/hir.py

@@ -228,6 +228,7 @@ class StaticEvalKind(enum.Enum):
     STATIC_EVAL = "static_eval()"
     STATIC_ASSERT_CONDITION = "static_assert() condition"
     STATIC_ASSERT_MESSAGE = "static_assert() message"
+    STATIC_ITER_ITERABLE = "static_iter() iterable"
 
 
 @dataclass
@@ -238,6 +239,7 @@ class StaticEvalExpression:
 
 def if_else(cond, then_block, else_block, /): ...
 def loop(body, iterable, /): ...  # infinite if `iterable` is None
+def static_foreach(body, items, /): ...
 def build_tuple(*items): ...  # Makes a tuple (i.e. returns `items`)
 def unpack(iterable, expected_len, /): ...
 def identity(x): ...   # Identity function (i.e. returns `x`)

src/cuda/tile/_ir/ops.py

@@ -13,10 +13,11 @@
 from typing_extensions import override
 
 import cuda.tile._stub as ct
-from cuda.tile import _datatype as datatype, TileValueError, TileStaticEvalError
+from cuda.tile import _datatype as datatype
 from cuda.tile import RoundingMode, MemoryOrder, MemoryScope
 from cuda.tile._mutex import tile_mutex
-from cuda.tile._exception import TileTypeError, TileSyntaxError, TileError, TileStaticAssertionError
+from cuda.tile._exception import TileTypeError, TileSyntaxError, TileError, \
+    TileStaticAssertionError, TileStaticEvalError, TileValueError
 from cuda.tile._ir.ir import (
     Operation, Var, Loc, Block,
     add_operation, Builder,
@@ -1008,7 +1009,6 @@ def binary_arithmetic(fn: str, x: Var, y: Var, rounding_mode: Optional[RoundingM
 @impl(ct.floordiv, fixed_args=["floordiv"])
 @impl(ct.cdiv, fixed_args=["cdiv"])
 @impl(ct.pow, fixed_args=["pow"])
-@impl(operator.add, fixed_args=["add"])
 @impl(operator.sub, fixed_args=["sub"])
 @impl(operator.mul, fixed_args=["mul"])
 @impl(operator.floordiv, fixed_args=["floordiv"])
@@ -1020,6 +1020,15 @@ def binary_arithmetic_impl(fn: str, x: Var, y: Var) -> Var:
     return binary_arithmetic(fn, x, y)
 
 
+@impl(operator.add)
+def add_impl(x: Var, y: Var) -> Var:
+    if isinstance(x.get_type(), TupleTy) and isinstance(y.get_type(), TupleTy):
+        x_items = x.get_aggregate().items
+        y_items = y.get_aggregate().items
+        return build_tuple(x_items + y_items)
+    return binary_arithmetic("add", x, y)
+
+
 @impl(ct.minimum, fixed_args=["min"])
 @impl(ct.maximum, fixed_args=["max"])
 def binary_arithmetic_impl_with_ftz(fn: str, x: Var, y: Var, flush_to_zero: Var) -> Var:
@@ -3901,6 +3910,12 @@ def static_assert_impl(condition: Var, message: Var):
                           " e.g. cuda.tile.static_assert() or ct.static_assert().")
 
 
+@impl(ct.static_iter)
+def static_iter_impl(iterable: Var):
+    raise TileSyntaxError("static_iter() must be used directly by name,"
+                          " e.g. cuda.tile.static_iter() or ct.static_iter().")
+
+
 @impl(hir.do_static_eval)
 def do_static_eval_impl(expr: hir.StaticEvalExpression,
                         local_var_values: tuple[Var, ...]) -> Var:
@@ -3923,10 +3938,52 @@ def do_static_eval_impl(expr: hir.StaticEvalExpression,
         if result is None:
             result = ""
         return loosely_typed_const(str(result))
+    elif expr.kind == hir.StaticEvalKind.STATIC_ITER_ITERABLE:
+        items = _drain_static_iter_iterable(result)
+        return build_tuple(tuple(items))
     else:
         return sym2var(result)
 
 
+_STATIC_ITER_MAX_ITERATIONS = 1000
+
+
+def _drain_static_iter_iterable(iterable) -> list[Var]:
+    try:
+        it = iter(iterable)
+    except Exception as e:
+        msg = str(e)
+        if len(msg) > 0:
+            msg = ": " + msg
+        raise TileTypeError(f"Invalid static_iter() iterable{msg}")
+
+    items = []
+    for i in range(_STATIC_ITER_MAX_ITERATIONS + 1):
+        try:
+            x = next(it)
+        except StopIteration:
+            break
+        except Exception as e:
+            msg = str(e)
+            if len(msg) > 0:
+                msg = ": " + msg
+            raise TileTypeError(f"Error was raised while obtaining item #{i}"
+                                f" from the static_iter() iterable{msg}")
+
+        try:
+            var = sym2var(x)
+        except TileTypeError as e:
+            raise TileStaticEvalError(
+                f"Invalid item #{i} of static_iter() iterable: {str(e)}")
+
+        items.append(var)
+    else:
+        raise TileStaticEvalError(f"Maximum number of iterations"
+                                  f" ({_STATIC_ITER_MAX_ITERATIONS}) has been reached"
+                                  f" while unpacking the static_iter() iterable")
+    return items
+
+
 @impl(hir.do_static_assert)
 async def do_static_assert_impl(condition: Var, message_block: hir.Block) -> None:
     if not condition.is_constant():
@@ -3951,6 +4008,19 @@ async def do_static_assert_impl(condition: Var, message_block: hir.Block) -> Non
     raise TileStaticAssertionError(message)
 
 
+@impl(hir.static_foreach)
+async def static_foreach_impl(body: hir.Block, items: Var):
+    scope = Scope.get_current()
+
+    tuple_val = items.get_aggregate()
+    assert isinstance(tuple_val, TupleValue)
+
+    for item in tuple_val.items:
+        scope.hir2ir_varmap[body.params[0].id] = item
+        from .._passes.hir2ir import dispatch_hir_block
+        await dispatch_hir_block(body)
+
+
 def var2sym(var: Var) -> Any:
     if var.is_constant():
         return var.get_constant()

src/cuda/tile/_passes/ast2hir.py

@@ -17,7 +17,7 @@
 from cuda.tile._ir import hir
 from cuda.tile._ir.type import ClosureDefaultPlaceholder
 from cuda.tile._passes.ast_util import ast_get_all_local_names
-from cuda.tile._stub import static_eval, static_assert
+from cuda.tile._stub import static_eval, static_assert, static_iter
 
 
 @lru_cache
@@ -145,6 +145,7 @@ def _get_function_hir_inner(func_def: ast.FunctionDef | ast.Lambda, signature: i
 
 class LoopKind(Enum):
     FOR = auto()
+    STATIC_FOR = auto()
     WHILE = auto()
 
 
@@ -247,12 +248,13 @@ def decorate(f):
 _expr_handlers: Dict[Type[ast.AST], Callable] = {}
 
 
+_KEYWORD_LIKE_FUNCS = (static_eval, static_assert, static_iter)
+_KEYWORD_LIKE_FUNC_NAMES = ("static_eval", "static_assert", "static_iter")
+
+
 @_register(_expr_handlers, ast.Call)
 def _call_expr(call: ast.Call, ctx: _Context) -> hir.Value:
-    kwd_func = _parse_keyword_like_func(call.func,
-                                        (static_eval, static_assert),
-                                        ("static_eval", "static_assert"),
-                                        ctx)
+    kwd_func = _parse_keyword_like_func(call.func, ctx)
     if kwd_func is not None:
         if kwd_func == "static_eval":
             if len(call.args) != 1 or len(call.keywords) != 0:
@@ -272,8 +274,11 @@ def _call_expr(call: ast.Call, ctx: _Context) -> hir.Value:
             condition = _call_static_eval(call.args[0],
                                           hir.StaticEvalKind.STATIC_ASSERT_CONDITION, ctx)
             return ctx.call(hir.do_static_assert, (condition, message_block))
+        elif kwd_func == "static_iter":
+            raise TileSyntaxError("static_iter() is only allowed as iterable in a `for` loop,"
+                                  " i.e. `for i in ct.static_iter(...)`")
         else:
-            assert False
+            raise TileSyntaxError(f"{kwd_func} is not expected here")
     else:
         callee = _expr(call.func, ctx)
         args = tuple(_expr(a, ctx) for a in call.args)
@@ -342,17 +347,14 @@ def _eval_ast_expr(expr: ast.expr, ctx: _Context):
     return eval(code, dict(ctx.frozen_globals), {})
 
 
-def _parse_keyword_like_func(expr: ast.expr,
-                             kwd_funcs: tuple[Callable, ...],
-                             kwd_func_names: tuple[str, ...],
-                             ctx: _Context) -> str | None:
+def _parse_keyword_like_func(expr: ast.expr, ctx: _Context) -> str | None:
     if isinstance(expr, ast.Name):
         if (expr.id not in ctx.local_names
-                and ctx.frozen_globals.get(expr.id) in kwd_funcs):
-            idx = kwd_funcs.index(ctx.frozen_globals.get(expr.id))
-            return kwd_func_names[idx]
+                and ctx.frozen_globals.get(expr.id) in _KEYWORD_LIKE_FUNCS):
+            idx = _KEYWORD_LIKE_FUNCS.index(ctx.frozen_globals.get(expr.id))
+            return _KEYWORD_LIKE_FUNC_NAMES[idx]
     elif isinstance(expr, ast.Attribute):
-        if expr.attr in kwd_func_names and _is_cuda_tile_module(expr.value, ctx):
+        if expr.attr in _KEYWORD_LIKE_FUNC_NAMES and _is_cuda_tile_module(expr.value, ctx):
             return expr.attr
     return None
 
@@ -583,21 +585,40 @@ def _for_stmt(stmt: ast.For, ctx: _Context):
     if len(stmt.orelse) > 0:
         raise ctx.syntax_error("'for-else' is not supported", loc=stmt.orelse[0])
 
-    iterable = _expr(stmt.iter, ctx)
-    if not isinstance(stmt.target, ast.Name):
-        raise ctx.unsupported_syntax(stmt.target)
+    static_iter_expr = _get_static_iter_expr(stmt.iter, ctx)
+    if static_iter_expr is None:
+        kind = LoopKind.FOR
+        op = hir.loop
+        iterable = _expr(stmt.iter, ctx)
+    else:
+        kind = LoopKind.STATIC_FOR
+        op = hir.static_foreach
+        with ctx.change_loc(static_iter_expr):
+            iterable = _call_static_eval(static_iter_expr,
+                                         hir.StaticEvalKind.STATIC_ITER_ITERABLE, ctx)
 
-    ctx.parent_loops.append(LoopKind.FOR)
+    ctx.parent_loops.append(kind)
     induction_var = ctx.make_value()
     with ctx.new_block(params=(induction_var,)) as body_block:
-        with ctx.change_loc(stmt.target):
-            ctx.store(stmt.target.id, induction_var)
+        _do_assign(induction_var, stmt.target, ctx)
         _stmt_list(stmt.body, ctx)
-        if body_block.jump is None:
+        if body_block.jump is None and static_iter_expr is None:
             ctx.set_block_jump(hir.Jump.CONTINUE)
     ctx.parent_loops.pop()
 
-    ctx.call_void(hir.loop, (body_block, iterable))
+    ctx.call_void(op, (body_block, iterable))
+
+
+def _get_static_iter_expr(expr: ast.expr, ctx: _Context) -> ast.expr | None:
+    if not isinstance(expr, ast.Call):
+        return None
+    if _parse_keyword_like_func(expr.func, ctx) != "static_iter":
+        return None
+
+    if len(expr.args) != 1 or len(expr.keywords) != 0:
+        raise ctx.syntax_error("static_iter() expects a single expression")
+
+    return expr.args[0]
 
 
 def _bool_expr(expr: ast.AST, ctx: _Context) -> hir.Value:
@@ -724,19 +745,21 @@ def _if_stmt(stmt: ast.If, ctx: _Context) -> None:
 
 @_register(_stmt_handlers, ast.Continue)
 def _continue_stmt(stmt: ast.Continue, ctx: _Context) -> None:
+    if ctx.parent_loops and ctx.parent_loops[-1] is LoopKind.STATIC_FOR:
+        raise ctx.syntax_error("Continue in a for loop with static_iter() is not supported")
     ctx.set_block_jump(hir.Jump.CONTINUE)
 
 
 @_register(_stmt_handlers, ast.Break)
 def _break_stmt(stmt: ast.Break, ctx: _Context) -> None:
-    if ctx.parent_loops and ctx.parent_loops[-1] is LoopKind.FOR:
+    if ctx.parent_loops and ctx.parent_loops[-1] in (LoopKind.FOR, LoopKind.STATIC_FOR):
         raise ctx.syntax_error("Break in a for loop is not supported")
     ctx.set_block_jump(hir.Jump.BREAK)
 
 
 @_register(_stmt_handlers, ast.Return)
 def _return_stmt(stmt: ast.Return, ctx: _Context) -> None:
-    if ctx.parent_loops and ctx.parent_loops[-1] is LoopKind.FOR:
+    if ctx.parent_loops and ctx.parent_loops[-1] in (LoopKind.FOR, LoopKind.STATIC_FOR):
         raise ctx.syntax_error("Returning from a for loop is not supported")
 
     return_val = None if stmt.value is None else _expr(stmt.value, ctx)

src/cuda/tile/_stub.py

@@ -2083,6 +2083,32 @@ def static_assert(condition, message=None, /):
     """
 
 
+@function
+def static_iter(iterable):
+    """Iterates at compile time.
+
+    Can only be used as the iterable of a `for` loop::
+
+        for ... in ct.static_iter(...):
+            ...
+
+    The surrounded expression is evaluated using the same rules as :py:func:`static_eval`:
+    it can reference global and local variables, and use the full Python syntax,
+    but must not perform any run-time operations.
+
+    The expression must return a Python iterable, whose length must not exceed some
+    pre-defined number of iterations (currently, 1000). Before any further processing is done,
+    the contents of the iterable are saved to a temporary list, and each item is checked
+    to be valid, as if it were a result of a :py:func:`static_eval` expression
+    (i.e., it must be a supported compile-time constant value or a proxy object
+    for a dynamic value such as a tile).
+
+    Finally, for each item of the iterable, the loop body is inlined, with the induction variable(s)
+    bound to the item. The `break`, `continue`, and `return` statements are not allowed
+    inside a `static_iter` loop.
+    """
+
+
 # ==== Private stubs ====