Add ct.reduce()

Signed-off-by: Greg Bonik <[email protected]>

Author: gbonik

changelog.d/custom-reduction.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 custom reduction via `ct.reduce()`.

src/cuda/tile/__init__.py

@@ -120,6 +120,7 @@
     pow,
     printf,
     prod,
+    reduce,
     reshape,
     rsqrt,
     scatter,
@@ -250,6 +251,7 @@
     "pow",
     "printf",
     "prod",
+    "reduce",
     "reshape",
     "rsqrt",
     "scatter",

src/cuda/tile/_ir/ir.py

@@ -21,6 +21,7 @@
 from cuda.tile._exception import (
     TileTypeError, Loc, TileInternalError
 )
+from .. import TileSyntaxError
 from .._cext import TileContext
 from .._context import TileContextConfig
 
@@ -317,9 +318,20 @@ def terminator(cls):
     return cls
 
 
-def has_side_effects(cls):
-    cls._has_side_effects = True
-    return cls
+class MemoryEffect(enum.IntEnum):
+    # Int value assigned here is meaningful.
+    # It implies the relative strength of memory effects.
+    # For example, NONE < LOAD < STORE.
+    NONE = 0
+    LOAD = 1
+    STORE = 2
+
+
+def memory_effect(eff: MemoryEffect):
+    def decorate(cls):
+        cls.memory_effect = eff
+        return cls
+    return decorate
 
 
 def has_multiple_results(cls):
@@ -399,19 +411,24 @@ def finalize_loopvar_type(self, body_var: Var):
 
 
 class Builder:
-    def __init__(self, ctx: IRContext, loc: Loc):
+    def __init__(self, ctx: IRContext, loc: Loc, reduction_body: bool = False):
         self.ir_ctx = ctx
         self.is_terminated = False
         self._loc = loc
         self._ops = []
         self._entered = False
         self._prev_builder = None
         self._var_map: Dict[str, Var] = dict()
+        self.reduction_body = reduction_body
 
     def add_operation(self, op_class,
                       result_ty: Type | None | Tuple[Type | None, ...],
                       attrs_and_operands,
                       result: Var | Sequence[Var] | None = None) -> Var | Tuple[Var, ...]:
+        if self.reduction_body and op_class.memory_effect != MemoryEffect.NONE:
+            raise TileSyntaxError("Operations with memory effects are not supported"
+                                  " inside reduction body")
+
         assert not self.is_terminated
         force_type = False
         if isinstance(result_ty, tuple):
@@ -504,10 +521,11 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
 
 @contextmanager
-def nested_block(loc: Loc):
+def nested_block(loc: Loc, reduction_body: bool = False):
     prev_builder = Builder.get_current()
     block = Block(prev_builder.ir_ctx, loc=loc)
-    with Builder(prev_builder.ir_ctx, loc) as builder:
+    with Builder(prev_builder.ir_ctx, loc,
+                 reduction_body=reduction_body or prev_builder.reduction_body) as builder:
         yield block
     block.extend(builder.ops)
 
@@ -520,7 +538,7 @@ class _CurrentBuilder(threading.local):
 
 
 class Operation:
-    _has_side_effects = False
+    memory_effect = MemoryEffect.NONE
     _multiple_results = False
 
     def __init__(
@@ -588,10 +606,6 @@ def all_inputs(self) -> Iterator[Var]:
     def is_terminator(self) -> bool:
         return self._is_terminator
 
-    @property
-    def has_side_effects(self) -> bool:
-        return self._has_side_effects
-
     def _add_operand(self, name: str, var: Var | Tuple[Var, ...]):
         if isinstance(var, Var) and var.is_aggregate() and self.op != "assign":
             # Don't allow aggregate values as operands, except for arrays and lists.

src/cuda/tile/_ir/op_impl.py

@@ -15,10 +15,11 @@
 from cuda.tile._exception import TileTypeError
 from cuda.tile._ir.ops_utils import get_dtype
 
-from .ir import Var
 from .typing_support import datatype, get_signature
+from .ir import Var, TupleValue
 from .type import TupleTy, TileTy, DTypeSpec, EnumTy, StringTy, ArrayTy, SliceType, \
-    ListTy, PointerTy, LooselyTypedScalar, RangeIterType
+    ListTy, PointerTy, LooselyTypedScalar, RangeIterType, FunctionTy, ClosureTy, BoundMethodTy, \
+    DTypeConstructor
 from .. import _datatype
 
 
@@ -46,18 +47,31 @@ def decorate(func):
 
         func_sig = get_signature(func)
         _verify_params_match(stub_sig, func_sig)
-
-        @functools.wraps(func)
-        def wrapper(*args, **kwargs):
-            # Memorize the stub and the args so that we can automatically
-            # provide context for error messages.
-            old = _current_stub.stub_and_args
-            _current_stub.stub_and_args = (stub, stub_sig, func_sig, args, kwargs)
-            try:
-                return func(*args, **kwargs)
-            finally:
-                _current_stub.stub_and_args = old
-        wrapper._is_coroutine = inspect.iscoroutinefunction(func)
+        is_coroutine = inspect.iscoroutinefunction(func)
+        if is_coroutine:
+            @functools.wraps(func)
+            async def wrapper(*args, **kwargs):
+                # Memorize the stub and the args so that we can automatically
+                # provide context for error messages.
+                old = _current_stub.stub_and_args
+                _current_stub.stub_and_args = (stub, stub_sig, func_sig, args, kwargs)
+                try:
+                    return await func(*args, **kwargs)
+                finally:
+                    _current_stub.stub_and_args = old
+        else:
+            @functools.wraps(func)
+            def wrapper(*args, **kwargs):
+                # Memorize the stub and the args so that we can automatically
+                # provide context for error messages.
+                old = _current_stub.stub_and_args
+                _current_stub.stub_and_args = (stub, stub_sig, func_sig, args, kwargs)
+                try:
+                    return func(*args, **kwargs)
+                finally:
+                    _current_stub.stub_and_args = old
+
+        wrapper._is_coroutine = is_coroutine
         op_implementations[stub] = wrapper
         return orig_func
 
@@ -96,6 +110,36 @@ def require_constant_bool(var: Var) -> bool:
     return var.get_constant()
 
 
+def require_constant_scalar(var: Var) -> bool | int | float:
+    ty = var.get_type()
+    if not isinstance(ty, DType):
+        raise _make_type_error(f"Expected a scalar constant, but given value has type {ty}", var)
+    if not var.is_constant():
+        raise _make_type_error(f"Expected a constant, but given value has non-constant type {ty}",
+                               var)
+    ret = var.get_constant()
+    assert isinstance(ret, bool | int | float)
+    return ret
+
+
+def require_constant_scalar_tuple(var: Var) -> tuple[bool | int | float, ...]:
+    ty = require_tuple_type(var)
+    ret = []
+    tuple_val = var.get_aggregate()
+    assert isinstance(tuple_val, TupleValue)
+    for i, (item_ty, item) in enumerate(zip(ty.value_types, tuple_val.items, strict=True)):
+        if not isinstance(item_ty, DType):
+            raise _make_type_error(f"Expected a tuple of scalar constants,"
+                                   f" but item at position #{i} has type {item_ty}", var)
+        if not item.is_constant():
+            raise _make_type_error(f"Expected a tuple of scalar constants,"
+                                   f" but item at position #{i} has non-constant type {ty}", var)
+        value = item.get_constant()
+        assert isinstance(value, bool | int | float)
+        ret.append(value)
+    return tuple(ret)
+
+
 def require_optional_constant_bool(var: Var) -> Optional[bool]:
     if var.is_constant() and var.get_constant() is None:
         return None
@@ -244,6 +288,16 @@ def require_tile_type(var: Var) -> TileTy:
     return ty
 
 
+def require_tile_or_tile_tuple_type(var: Var) -> TileTy | TupleTy:
+    ty = var.get_type()
+    if isinstance(ty, TileTy):
+        return ty
+    if isinstance(ty, TupleTy) and all(isinstance(x, TileTy) for x in ty.value_types):
+        return ty
+    raise _make_type_error(f"Expected a tile or a tuple of tiles, but given value has type {ty}",
+                           var)
+
+
 def require_tile_or_scalar_type(var: Var) -> TileTy | DType | PointerTy:
     ty = var.get_type()
     if not isinstance(ty, TileTy | DType | PointerTy):
@@ -359,6 +413,13 @@ def require_index_or_index_tuple_type(var: Var,
     return ty
 
 
+def require_callable_type(var: Var) -> FunctionTy | BoundMethodTy | ClosureTy | DTypeConstructor:
+    ty = var.get_type()
+    if not isinstance(ty, FunctionTy | BoundMethodTy | ClosureTy | DTypeConstructor):
+        raise _make_type_error(f"Expected a callable object, but given value has type {ty}", var)
+    return ty
+
+
 class PrintfValidator:
     # c-format string has the following: %[flags][width][.precision][length]specifier
     # we only support a subset which makes sense in the tile context