[lang] Keep custom reduce and scan callback bodies self-contained

changelog.d/reduce-scan-isolation.md

@@ -0,0 +1,5 @@
+- Custom ``ct.reduce()`` and ``ct.scan()`` callbacks now keep captured scalar compile-time
+  constants inside their own bodies, and every value inside a callback must be a scalar tile.
+  Capturing runtime values or non-scalar constants is rejected with a compile-time error.
+  Note: capturing runtime values in these callbacks previously compiled and is no longer
+  supported.

src/cuda/tile/_compile.py

@@ -65,6 +65,11 @@
 from cuda.tile._passes.check_dtype_support import check_dtype_support
 from cuda.tile._passes.dce import dead_code_elimination_pass
 from cuda.tile._passes.materialize_constants import materialize_constants_pass
+from cuda.tile._passes.isolate_reduce_scan import (
+    collect_reduce_scan_capture_names,
+    legalize_reduce_scan_captures,
+    verify_reduce_scan_isolation,
+)
 from cuda.tile._passes.propagate_divby import add_divby_pass
 from cuda.tile._passes.token_order import token_order_pass
 from cutile_cache._cache import MetadataV1, cache_key, cache_lookup, cache_store, evict_lru
@@ -102,12 +107,14 @@ def _transform_ir(func_body: ir.Block,
                   bytecode_version: bc.BytecodeVersion,
                   param_constraints: Sequence[tuple[tuple[ir.Var, ...], ParameterConstraint]]
                   ) -> DataflowResult:
+    capture_names = collect_reduce_scan_capture_names(func_body)
     eliminate_assign_ops(func_body)
     lower_for_with_break(func_body)
     dead_code_elimination_pass(func_body)
     dataflow_result = dataflow_analysis(func_body, param_constraints)
 
     materialize_constants_pass(func_body, dataflow_result)
+    legalize_reduce_scan_captures(func_body, capture_names)
 
     if not CUDA_TILE_TESTING_DISABLE_DIV:
         add_divby_pass(func_body, dataflow_result)
@@ -128,6 +135,7 @@ def _transform_ir(func_body: ir.Block,
 
     split_loops(func_body)
     dead_code_elimination_pass(func_body)
+    verify_reduce_scan_isolation(func_body)
 
     return dataflow_result
 

src/cuda/tile/_ir/ops.py

@@ -31,6 +31,7 @@
     UNARY_STRICT_FLOAT, UNARY_FLOAT, divmod_tensorlike
 from .cast_ops import implicit_cast
 from .control_flow_ops import Loop, IfElse, control_flow_impl_registry, EndBranch
+from .core_ops import Assign
 from .core_ops import loosely_typed_const, strictly_typed_const, build_tuple, bind_method, \
     sym2var, core_impl_registry, print_impl, TilePrintf, tuple_item
 from .static_eval_ops import static_eval_impl_registry
@@ -61,7 +62,7 @@
 from .type import (
     PartitionViewTy, StridedViewTy, GatherScatterViewTy, TupleTy, TileTy, NoneType, ArrayTy,
     ListTy, Type, LooselyTypedScalar, TokenTy, TiledViewTy,
-    RawArrayMemoryTy, IndexSliceTy,
+    RawArrayMemoryTy, IndexSliceTy, TensorLikeTy,
 )
 from cuda.tile._datatype import (
     DType, is_integral, is_float, is_signed, is_boolean, PointerInfo,
@@ -2212,6 +2213,29 @@ def generate_bytecode(self, ctx: BytecodeContext) -> tuple[bc.Value, ...]:
         return nested_builder.done()
 
 
+def _require_scalar_body(body_block: Block, op_name: Literal["reduction", "scan"]) -> None:
+    """Reject non-scalar tiles inside a reduce/scan body.
+
+    The body combines 0-d elements, and keeping every value inside it 0-d keeps the emitted
+    body self-contained. Non-scalar constants must be reduced to a scalar outside the callback.
+    """
+    # Assign ops are still present at this point; use them to name temporaries after the
+    # variable they were assigned to.
+    aliases = {op.value.name: op.result_var.get_original_name()
+               for op in body_block.operations if isinstance(op, Assign)}
+    for op in body_block.operations:
+        for var in (*op.all_inputs(), *op.result_vars):
+            ty = var.get_type_allow_invalid()
+            if isinstance(ty, TensorLikeTy) and ty.tensor_shape() != ():
+                name = var.get_original_name()
+                if name.startswith("$"):
+                    name = aliases.get(var.name, name)
+                what = "a value" if name.startswith("$") else f"'{name}'"
+                raise TileSyntaxError(
+                    f"{op_name} body must only operate on scalar tiles, but {what} has shape "
+                    f"{ty.tensor_shape()}", op.loc)
+
+
 async def _get_reduce_scan_body_block(
     xs: tuple[Var, ...],
     body: Callable,
@@ -2255,6 +2279,7 @@ async def _get_reduce_scan_body_block(
 
         add_operation_variadic(EndBranch, (), outputs=body_results)
 
+    _require_scalar_body(body_block, op_name)
     return body_block
 
 

src/cuda/tile/_passes/code_motion.py

@@ -86,7 +86,10 @@ def _hoist(block: Block, stack: list[_StackItem], def_depth: dict[str, int], is_
             for var in op.body.params:
                 def_depth[var.name] = depth + 1
 
-            body_res = _hoist(op.body, stack, def_depth, True)
+            # Only loop bodies are hoisting sources. A reduce/scan body is a self-contained
+            # combine function: nothing may move out of it, although the operation as a whole
+            # can still be hoisted together with its body.
+            body_res = _hoist(op.body, stack, def_depth, isinstance(op, Loop))
             if body_res.mobility == _BlockMobility.IMMOVABLE:
                 # Propagate IMMOVABLE to all ancestors.
                 ret.mobility = _BlockMobility.IMMOVABLE

src/cuda/tile/_passes/isolate_reduce_scan.py

@@ -0,0 +1,144 @@
+# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Keep custom reduce and scan callback bodies self-contained.
+
+A reduce or scan callback is a pure combine function of its block parameters, so the body we
+emit for it should not depend on anything from the enclosing scope. Compile-time constants the
+callback uses are materialized inside the body, while captures of runtime values are rejected
+with a source-level diagnostic.
+"""
+
+from cuda.tile._exception import Loc, TileInternalError, TileSyntaxError
+from cuda.tile._ir.core_ops import Assign, TypedConst
+from cuda.tile._ir.ir import Block, Mapper, Operation, Var
+from cuda.tile._ir.ops import TileReduce, TileScan
+
+
+def _definitions(root_block: Block) -> dict[str, Operation]:
+    return {
+        result.name: op
+        for op in root_block.traverse()
+        for result in op.result_vars
+    }
+
+
+def _captures(body: Block) -> list[tuple[Var, Loc]]:
+    """Return the values used in `body` but defined outside of it, with the location of their
+    first use. Reduce and scan bodies contain no nested blocks, so one level suffices."""
+    local_names = {var.name for var in body.params}
+    local_names.update(
+        result.name
+        for op in body.operations
+        for result in op.result_vars
+    )
+
+    captures = []
+    captured_names = set()
+    for op in body.operations:
+        for operand in op.all_inputs():
+            if operand.name not in local_names and operand.name not in captured_names:
+                captures.append((operand, op.loc))
+                captured_names.add(operand.name)
+    return captures
+
+
+def _kind(op: TileReduce | TileScan) -> str:
+    return "reduction" if isinstance(op, TileReduce) else "scan"
+
+
+def collect_reduce_scan_capture_names(root_block: Block) -> dict[str, str]:
+    """Map the values captured by reduce/scan bodies to their source-level names.
+
+    This must run before `eliminate_assign_ops`: a variable such as `m = ct.gather(...)` is an
+    Assign of a temporary, and once the Assign is gone only the temporary's name is left for
+    diagnostics. The keys are the names of the values that remain after Assign elimination.
+    """
+    definitions = _definitions(root_block)
+    names: dict[str, str] = {}
+    for region_op in root_block.traverse():
+        if not isinstance(region_op, TileReduce | TileScan):
+            continue
+        for value, _ in _captures(region_op.body):
+            canonical = value
+            defining_op = definitions.get(canonical.name)
+            while isinstance(defining_op, Assign):
+                canonical = defining_op.value
+                defining_op = definitions.get(canonical.name)
+            names.setdefault(canonical.name, value.get_original_name())
+    return names
+
+
+def legalize_reduce_scan_captures(root_block: Block, capture_names: dict[str, str]) -> None:
+    """Rematerialize constant captures inside each body and reject runtime captures.
+
+    This must run after `materialize_constants_pass`: that pass emits every dataflow-proven
+    constant at the start of the root block, which turns uses inside a callback body into
+    captures. Cloning such constants back into the body keeps the body self-contained.
+    `capture_names` comes from `collect_reduce_scan_capture_names`.
+    """
+    definitions = _definitions(root_block)
+
+    for region_op in root_block.traverse():
+        if not isinstance(region_op, TileReduce | TileScan):
+            continue
+
+        mapper = Mapper(root_block.ctx)
+        constants = []
+        for value, consuming_loc in _captures(region_op.body):
+            defining_op = definitions.get(value.name)
+            if value.is_constant():
+                constant_value = value.get_constant()
+            elif isinstance(defining_op, TypedConst):
+                constant_value = defining_op.value
+            else:
+                name = capture_names.get(value.name, value.get_original_name())
+                raise TileSyntaxError(
+                    f"{_kind(region_op)} body captures runtime value '{name}'. Only the "
+                    "callback's own parameters and scalar compile-time constants are supported.",
+                    consuming_loc,
+                )
+
+            # Shapes were validated when the body was built (see `_require_scalar_body` in
+            # ops.py), so the capture is a scalar and can simply be cloned into the body.
+            local_value = mapper.clone_var(value)
+            constants.append(TypedConst(
+                value=constant_value,
+                result_vars=(local_value,),
+                loc=value.loc,
+            ))
+
+        if constants:
+            for op in region_op.body.operations:
+                op.remap_operands(mapper)
+            region_op.body[:0] = constants
+
+
+def verify_reduce_scan_isolation(root_block: Block) -> None:
+    """Verify that reduce and scan bodies only use their parameters and body-local values."""
+    for region_op in root_block.traverse():
+        if not isinstance(region_op, TileReduce | TileScan):
+            continue
+
+        local_names = {var.name for var in region_op.body.params}
+        local_names.update(
+            result.name
+            for op in region_op.body.operations
+            for result in op.result_vars
+        )
+        available = {var.name for var in region_op.body.params}
+        for op in region_op.body.operations:
+            for operand in op.all_inputs():
+                if operand.name not in available:
+                    original_name = operand.get_original_name()
+                    problem = (
+                        "is used before its definition"
+                        if operand.name in local_names
+                        else "is defined outside the region"
+                    )
+                    raise TileInternalError(
+                        f"{_kind(region_op)} body value '{original_name}' {problem}",
+                        op.loc,
+                    )
+            available.update(result.name for result in op.result_vars)

src/cuda/tile/_stub.py

@@ -3049,7 +3049,10 @@ def reduce(x, /, axis, func, identity, *, keepdims=False):
             `lambda a, b: a + b` or `operator.add` can be used to implement the sum reduction.
             If `x` is a tuple of N tiles, then the function takes 2N tiles and returns a tuple
             of N combined tiles. The first N arguments correspond to one of the groups of values
-            being combined, while the rest correspond to the other.
+            being combined, while the rest correspond to the other. The function must only
+            operate on scalar tiles. It may capture scalar compile-time constants from its
+            enclosing scope; capturing runtime values is unsupported and rejected during
+            compilation.
         identity: a constant scalar or a tuple of constant scalars that specifies the identity
             element of the `func`.
         keepdims (bool): True to keep the axis of size 1, False to remove the reduced axis.
@@ -3160,7 +3163,10 @@ def scan(x, /, axis, func, identity, *, reverse=False):
             `lambda a, b: a + b` or `operator.add` can be used to implement cumsum.
             If `x` is a tuple of N tiles, then the function takes 2N tiles and returns a tuple
             of N combined tiles. The first N arguments correspond to one of the groups of values
-            being combined, while the rest correspond to the other.
+            being combined, while the rest correspond to the other. The function must only
+            operate on scalar tiles. It may capture scalar compile-time constants from its
+            enclosing scope; capturing runtime values is unsupported and rejected during
+            compilation.
         identity: a constant scalar or a tuple of constant scalars that specifies the identity
             element of the `func`.
         reverse (bool): if True, the scan is performed in the reverse direction along the axis.

test/test_code_motion.py

@@ -10,7 +10,8 @@
 import cuda.tile as ct
 from cuda.tile.compilation import CallingConvention
 from cuda.tile._ir.ir import Operation
-from cuda.tile._ir.ops import Loop, IfElse, TileExtract
+from cuda.tile._ir.core_ops import TypedConst
+from cuda.tile._ir.ops import Loop, IfElse, TileExtract, TileReduce, TileScan
 from cuda.tile._ir.arithmetic_ops import Unary
 from cuda.tile._compile import compile_tile
 
@@ -196,6 +197,36 @@ def carried_from_nested_loop_no(x, a, t):
         ct.store(x, i, val)
 
 
+@ct.kernel
+def entire_reduce_op_yes(x, y):
+    xt = ct.load(x, (0, 0), (16, 16))
+    for i in range(y.shape[1]):
+        yt = ct.reduce(xt, -1, lambda a, b: a + b, 0, keepdims=True)
+        ct.store(y, (0, i), yt)
+
+
+@ct.kernel
+def entire_scan_op_yes(x, y):
+    xt = ct.load(x, (0, 0), (16, 16))
+    for i in range(y.shape[1] // 16):
+        yt = ct.scan(xt, -1, lambda a, b: a + b, 0)
+        ct.store(y, (0, i), yt)
+
+
+@ct.kernel
+def reduce_body_modulo(x, y):
+    xt = ct.load(x, (0, 0), (16, 16))
+    yt = ct.reduce(xt, -1, lambda a, b: (a + b) % 5, 0)
+    ct.store(y, (0,), yt)
+
+
+@ct.kernel
+def scan_body_modulo(x, y):
+    xt = ct.load(x, (0, 0), (16, 16))
+    yt = ct.scan(xt, -1, lambda a, b: (a + b) % 5, 0)
+    ct.store(y, (0, 0), yt)
+
+
 def make_cases(tuples):
     return [pytest.param(kernel, op_finder, expected_x, id=kernel._pyfunc.__name__)
             for kernel, op_finder, expected_x in tuples]
@@ -249,3 +280,58 @@ def test_hoisting(kernel, op_finder, expected_x):
     ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, a, 4.0))
     ref = torch.tensor(expected_x, dtype=torch.float32, device="cuda")
     assert_close(x, ref)
+
+
+def _final_ir(kernel, args):
+    sig = ct.compilation.KernelSignature.from_kernel_args(
+        kernel, args, CallingConvention.cutile_python_v1()
+    )
+    [root_block] = compile_tile(
+        kernel._pyfunc, [sig], return_final_ir=True, return_cubin=False
+    ).final_ir
+    return root_block
+
+
+@pytest.mark.parametrize(
+    "kernel, op_type, y_shape",
+    [
+        (reduce_body_modulo, TileReduce, (16,)),
+        (scan_body_modulo, TileScan, (16, 16)),
+    ],
+    ids=["reduce", "scan"],
+)
+def test_aggregate_body_is_licm_barrier(kernel, op_type, y_shape):
+    x = torch.zeros((16, 16), dtype=torch.int32, device="cuda")
+    y = torch.zeros(y_shape, dtype=torch.int32, device="cuda")
+    root_block = _final_ir(kernel, (x, y))
+
+    [aggregate] = [op for op in root_block.traverse() if isinstance(op, op_type)]
+
+    def is_modulo_const(op):
+        return isinstance(op, TypedConst) and op.value == 5
+
+    assert sum(map(is_modulo_const, aggregate.body.operations)) == 1
+    assert not any(map(is_modulo_const, root_block.operations))
+
+
+@pytest.mark.parametrize(
+    "kernel, op_type, y_shape, reference",
+    [
+        (entire_reduce_op_yes, TileReduce, (16, 3),
+         lambda x: x.sum(-1, keepdim=True).expand(16, 3)),
+        (entire_scan_op_yes, TileScan, (16, 48),
+         lambda x: torch.cumsum(x, -1).repeat(1, 3)),
+    ],
+    ids=["reduce", "scan"],
+)
+def test_entire_aggregate_op_can_be_hoisted(kernel, op_type, y_shape, reference):
+    x = torch.arange(256, dtype=torch.float32, device="cuda").reshape(16, 16)
+    y = torch.zeros(y_shape, dtype=torch.float32, device="cuda")
+    root_block = _final_ir(kernel, (x, y))
+
+    [aggregate] = [op for op in root_block.traverse() if isinstance(op, op_type)]
+    [loop] = [op for op in root_block.traverse() if isinstance(op, Loop)]
+    assert not _is_inside_loop(aggregate, loop)
+
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, y))
+    assert_close(y, reference(x))