Add support for **kwargs

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

Author: gbonik

changelog.d/dict.md

@@ -0,0 +1,2 @@
+- Added limited support for dictionaries, variadic keyword parameters in user-defined functions
+(e.g, `def foo(**kwargs)`), and dictionary unpacking (e.g., `foo(x, **y)`).

src/cuda/tile/_ir/core_ops.py

@@ -5,7 +5,7 @@
 import functools
 import operator
 from dataclasses import dataclass
-from types import MethodType, FunctionType, BuiltinFunctionType
+from types import MethodType, FunctionType, BuiltinFunctionType, MappingProxyType
 from typing import Any, Optional
 
 from typing_extensions import override
@@ -27,7 +27,7 @@
     DataclassInfo, DataclassTy, DataclassValue, BoundMethodValue, BoundMethodTy, InvalidType, \
     ContextManagerTy, ContextManagerLifecycle, LiveCapturedScope, ClosureTy, ClosureValue, \
     RangeIterType, RangeValue, TypeTy, ModuleTy, NONE, SliceType, StringTy, FormattedStringTy, \
-    StringFormat, FormattedStringValue, FormattedPiece
+    StringFormat, FormattedStringValue, FormattedPiece, DictTy, DictValue
 from cuda.tile._ir.typing_support import type_of_constant_python_value, \
     loose_type_of_constant_python_value, get_dataclass_info, as_third_party_dtype_spec
 from cuda.tile._ir2bytecode import BytecodeContext
@@ -65,6 +65,7 @@ def core_impl_registry() -> ImplRegistry:
 @overload_dispatcher(operator.lshift, fixed_args=["<<"])
 @overload_dispatcher(operator.rshift, fixed_args=[">>"])
 @overload_dispatcher(operator.matmul, fixed_args=["@"])
+@overload_dispatcher(hir_stubs.is_contained_in, fixed_args=["'in'"])
 @overload_dispatcher(min, fixed_args=["min"])
 @overload_dispatcher(max, fixed_args=["max"])
 def binop_overload_dispatcher(name: str, x: Var, y: Var):
@@ -76,6 +77,13 @@ def binop_overload_dispatcher(name: str, x: Var, y: Var):
         raise TileTypeError(f"Unsupported operand types for {name}: {x_ty} and {y_ty}")
 
 
+@impl(hir_stubs.is_not_contained_in)
+async def is_not_contained_in_impl(x: Var, y: Var):
+    from .._passes.hir2ir import call_function
+    contained = await call_function(hir_stubs.is_contained_in, x, y)
+    return await call_function(operator.not_, contained)
+
+
 def comparison_operator_impl(registry: ImplRegistry, lhs_ty: type[Type], rhs_ty: type[Type]):
     def decorate(func):
         for name in ("eq", "ne", "lt", "le", "gt", "ge"):
@@ -346,6 +354,71 @@ def len_tuple_impl(x: Var[TupleTy]) -> Var:
     return loosely_typed_const(len(x.get_type()))
 
 
+# ===========================================================================================
+# Dictionary
+# ===========================================================================================
+
+def build_dict(keys: tuple[str, ...], values: tuple[Var, ...]) -> Var:
+    keys = tuple(keys)
+    values = tuple(values)
+    assert len(keys) == len(values)
+
+    ty = DictTy(keys, tuple(x.get_type() for x in values))
+    loose_ty = DictTy(keys, tuple(x.get_loose_type() for x in values))
+    res = make_aggregate(DictValue(values), ty, loose_ty)
+    if all(x.is_constant() for x in values):
+        items = [(k, v.get_constant()) for k, v in zip(keys, values, strict=True)]
+        res.set_constant(MappingProxyType(dict(items)))
+    return res
+
+
+def _find_dict_key_index(key: Var, dict_ty: DictTy) -> int | None:
+    key_ty = key.get_type()
+    if not isinstance(key_ty, StringTy):
+        # Python would happily report that the key is not found when a "wrong" key type is passed,
+        # but we can add a stronger check here.
+        raise TileTypeError(f"Dictionary keys must be strings, not {key_ty}")
+
+    return dict_ty.keys.index(key_ty.value) if key_ty.value in dict_ty.keys else None
+
+
+@impl(hir_stubs.is_contained_in, overload=(WILDCARD, DictTy))
+async def is_contained_in_dict_impl(x: Var, y: Var[DictTy]):
+    return loosely_typed_const(_find_dict_key_index(x, y.get_type()) is not None)
+
+
+@impl(getattr, overload=(DictTy, "get"))
+def getattr_dict_method(object: Var, name: Var):
+    name = require_constant_str(name)
+    unbound_func = getattr(dict, name)
+    return bind_method(object, unbound_func)
+
+
+@impl(operator.getitem, overload=(DictTy, WILDCARD))
+def getitem_dict_impl(object: Var[DictTy], key: Var):
+    idx = _find_dict_key_index(key, object.get_type())
+    if idx is None:
+        raise TileTypeError(f"Key '{key.get_type().value}' not found in dictionary")
+    dict_value = object.get_aggregate()
+    assert isinstance(dict_value, DictValue)
+    return dict_value.values[idx]
+
+
+@impl(dict.get)
+def dict_get_impl(self: Var, key: Var, default: Var):
+    dict_ty = self.get_type()
+    if not isinstance(dict_ty, DictTy):
+        raise TileTypeError(f"dict.get() expects a dictionary as `self`, got {dict_ty}")
+
+    idx = _find_dict_key_index(key, dict_ty)
+    if idx is None:
+        return default
+
+    dict_value = self.get_aggregate()
+    assert isinstance(dict_value, DictValue)
+    return dict_value.values[idx]
+
+
 # ===========================================================================================
 # Dataclass
 # ===========================================================================================

src/cuda/tile/_ir/hir.py

@@ -76,7 +76,7 @@ class Call:
     result: Value | None
     callee: Operand
     args: tuple[Operand | Starred, ...]
-    kwargs: tuple[tuple[str, Operand], ...]
+    kwargs: tuple[tuple[str | None, Operand], ...]  # None means an **unpacked argument
     loc: Loc
 
     def __str__(self):

src/cuda/tile/_ir/hir_stubs.py

@@ -65,3 +65,11 @@ def enter_context(manager, /): ...
 
 @stub
 def pop_context(): ...
+
+
+@stub
+def is_contained_in(x, y, /): ...  # "return x in y"
+
+
+@stub
+def is_not_contained_in(x, y, /): ...  # return "x not in y"

src/cuda/tile/_ir/type.py

@@ -7,7 +7,7 @@
 import os
 from dataclasses import dataclass
 from enum import EnumMeta
-from types import ModuleType, FunctionType, BuiltinFunctionType, MethodType
+from types import ModuleType, FunctionType, BuiltinFunctionType, MethodType, MappingProxyType
 from typing import Any, Callable, Optional, Sequence, Tuple, Iterator, Mapping
 from functools import reduce
 import operator
@@ -269,6 +269,43 @@ def as_tuple(self) -> tuple["Var", ...]:
         return self.items
 
 
+# ============== Dictionary (limited support for **kwargs) ==============
+
+@dataclass(frozen=True)
+class DictTy(Type):
+    keys: tuple[str, ...]
+    value_types: tuple[Type, ...]
+
+    def make_symbol(self, var: "Var"):
+        dict_val = var.get_aggregate()
+        assert isinstance(dict_val, DictValue)
+        items = [(k, var2sym(v)) for k, v in zip(self.keys, dict_val.values, strict=True)]
+        return MappingProxyType(dict(items))
+
+    def is_aggregate(self) -> bool:
+        return True
+
+    def aggregate_item_types(self) -> tuple["Type", ...]:
+        return self.value_types
+
+    def make_aggregate_value(self, values: tuple["Var", ...]) -> "AggregateValue":
+        assert len(values) == len(self.keys)
+        return DictValue(values)
+
+    def __str__(self):
+        return (
+                "dict["
+                + ", ".join(f"{k}: {ty}"
+                            for k, ty in zip(self.keys, self.value_types, strict=True))
+                + "]"
+        )
+
+
+@dataclass
+class DictValue(AggregateValue):
+    values: tuple["Var", ...]
+
+
 # ============== Dataclass ===============
 
 

src/cuda/tile/_ir/typing_support.py

@@ -121,6 +121,7 @@ def __missing__(self, key: TypeKey) -> TypeHandler:
     bool: lambda x=False, /: None,
     print: lambda *args, sep=' ', end='\n': None,
     dataclasses.replace: dataclasses.replace,
+    dict.get: dict.get,
 }
 
 

src/cuda/tile/_passes/ast2hir.py

@@ -423,6 +423,7 @@ def _binop_expr(binop: ast.BinOp, ctx: _Context) -> hir.Value:
 _cmp_map = {
     ast.Eq: operator.eq, ast.NotEq: operator.ne, ast.Lt: operator.lt, ast.LtE: operator.le,
     ast.Gt: operator.gt, ast.GtE: operator.ge, ast.Is: operator.is_, ast.IsNot: operator.is_not,
+    ast.In: hir_stubs.is_contained_in, ast.NotIn: hir_stubs.is_not_contained_in
 }
 
 
@@ -975,11 +976,6 @@ def _get_all_parameters(func_def: ast.FunctionDef | ast.Lambda, ctx: _Context) -
         for a in (func_def.args.vararg, func_def.args.kwarg):
             if a is not None:
                 raise ctx.syntax_error("Variadic kernel parameters are not supported", a)
-    else:
-        if func_def.args.kwarg is not None:
-            raise ctx.syntax_error(
-                "Variadic keyword parameters in user-defined functions are not supported",
-                func_def.args.kwarg)
 
     all_args = []
     for arg in func_def.args.posonlyargs:
@@ -990,6 +986,8 @@ def _get_all_parameters(func_def: ast.FunctionDef | ast.Lambda, ctx: _Context) -
         all_args.append(func_def.args.vararg)
     for arg in func_def.args.kwonlyargs:
         all_args.append(arg)
+    if func_def.args.kwarg is not None:
+        all_args.append(func_def.args.kwarg)
     return all_args
 
 

src/cuda/tile/_passes/hir2ir.py

@@ -11,21 +11,20 @@
 from .. import TileTypeError
 from .._coroutine_util import resume_after, run_coroutine
 from .._datatype import PointerInfo
-from .._exception import Loc, FunctionDesc, TileSyntaxError, TileInternalError, TileError, \
-    TileRecursionError
+from .._exception import Loc, FunctionDesc, TileInternalError, TileError, TileRecursionError
 from .._execution import is_stub
 from .._ir import hir, ir
 from .._ir.ir import Var, IRContext
 from .._ir.op_impl import ImplRegistry
 from .._ir.control_flow_ops import end_branch, return_, continue_, break_
 from .._ir.core_ops import (
-    loosely_typed_const, build_dataclass_instance, build_tuple, sym2var, store_var
+    loosely_typed_const, build_dataclass_instance, build_tuple, sym2var, store_var, build_dict
 )
 from .._ir.arithmetic_ops import dtype_constructor
 from .._ir.scope import Scope, LocalScope, IntMap
 from .._ir.type import FunctionTy, BoundMethodTy, DTypeConstructor, ClosureTy, \
     ClosureDefaultPlaceholder, StringFormat, TypeTy, TupleTy, BoundMethodValue, TupleValue, \
-    ClosureValue
+    ClosureValue, DictTy, DictValue
 from .._ir.typing_support import get_signature, is_supported_builtin_func, \
     get_dataclass_info
 
@@ -167,7 +166,20 @@ async def _dispatch_call(hir_call: hir.Call, scope: Scope):
             args.extend(tup_value.items)
         else:
             args.append(_resolve_operand(x, scope))
-    kwargs = {k: _resolve_operand(v, scope) for k, v in hir_call.kwargs}
+    kwargs = {}
+    for k, v in hir_call.kwargs:
+        resolved_val = _resolve_operand(v, scope)
+        if k is None:
+            assert isinstance(resolved_val, Var)
+            dict_ty = resolved_val.get_type()
+            if not isinstance(dict_ty, DictTy):
+                raise TileTypeError(f"Expected a dictionary after **, got {dict_ty}")
+            dict_value = resolved_val.get_aggregate()
+            assert isinstance(dict_value, DictValue)
+            for item_key, item_value in zip(dict_ty.keys, dict_value.values, strict=True):
+                kwargs[item_key] = item_value
+        else:
+            kwargs[k] = resolved_val
     retval = await call(callee_var, args, kwargs)
     if hir_call.result is not None and retval is not None:
         scope.hir2ir_varmap[hir_call.result.id] = retval
@@ -178,10 +190,6 @@ async def _call_user_defined(callee_hir: hir.Function,
                              builder: ir.Builder,
                              parent_scopes: tuple[LocalScope, ...] = ()):
     _check_recursive_call(builder.loc)
-    for param_name, param in callee_hir.signature.parameters.items():
-        if param.kind == inspect.Parameter.VAR_KEYWORD:
-            raise TileSyntaxError("Variadic keyword parameters in user-defined"
-                                  " functions are not supported")
 
     # Activate a fresh Scope. Each inlining gets its own concretized
     # FunctionDesc so that DI never merges two specializations whose generated
@@ -197,6 +205,8 @@ async def _call_user_defined(callee_hir: hir.Function,
             if isinstance(arg, tuple):
                 # Handle the *vararg parameter
                 arg = build_tuple(arg)
+            elif isinstance(arg, dict):
+                arg = build_dict(tuple(arg.keys()), tuple(arg.values()))
             store_var(local_idx, arg, param_loc)
 
         # Dispatch the function body. Use resume_after() to break the call stack
@@ -411,6 +421,8 @@ def _bind_args(sig: inspect.Signature, func_name: str, args, kwargs,
             ret.append(bound_args.arguments[name])
         elif param.kind == param.VAR_POSITIONAL:
             ret.append(())
+        elif param.kind == param.VAR_KEYWORD:
+            ret.append({})
         else:
             assert param.default is not param.empty
             if isinstance(param.default, ClosureDefaultPlaceholder):

test/test_dict.py

@@ -0,0 +1,87 @@
+# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+import pytest
+import torch.cuda
+
+import cuda.tile as ct
+from cuda.tile import TileTypeError
+
+
+def test_variadic_kwargs_in_helper_function():
+    def helper(**kwargs):
+        ct.static_assert(kwargs == {"foo": 123, "bar": 456})
+        return 789
+
+    @ct.kernel
+    def kernel():
+        res = helper(foo=123, bar=456)
+        ct.static_assert(res == 789)
+
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, ())
+
+
+def test_forward_variadic_kwargs():
+    def leaf(x, foo, bar):
+        return x * 100 + foo * 10 + bar
+
+    def forward(f, **kwargs):
+        return f(3, **kwargs)
+
+    @ct.kernel
+    def kernel():
+        res = forward(leaf, foo=4, bar=5)
+        ct.static_assert(res == 345)
+
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, ())
+
+
+def test_dict_access():
+    def helper(**kwargs):
+        foo1 = kwargs["foo"]
+        ct.static_assert(foo1 == 123)
+
+        foo2 = kwargs.get("foo")
+        ct.static_assert(foo2 == 123)
+
+        bar1 = kwargs["bar"]
+        ct.static_assert(bar1 == 456)
+
+        bar2 = kwargs.get("bar")
+        ct.static_assert(bar2 == 456)
+
+        qux1 = kwargs.get("qux")
+        ct.static_assert(qux1 is None)
+
+        res1 = "foo" in kwargs
+        ct.static_assert(res1)
+
+        res2 = "foo" not in kwargs
+        ct.static_assert(not res2)
+
+        res3 = "qux" in kwargs
+        ct.static_assert(not res3)
+
+        res4 = "qux" not in kwargs
+        ct.static_assert(res4)
+
+        return 789
+
+    @ct.kernel
+    def kernel():
+        res = helper(foo=123, bar=456)
+        ct.static_assert(res == 789)
+
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, ())
+
+
+def test_dict_getitem_miss():
+    def helper(**kwargs):
+        return kwargs["qux"]
+
+    @ct.kernel
+    def kernel():
+        helper(foo=123, bar=456)
+
+    with pytest.raises(TileTypeError, match="Key 'qux' not found in dictionary"):
+        ct.launch(torch.cuda.current_stream(), (1,), kernel, ())