From d97c305696cf52c55952d3ff6776f16b2a54b593 Mon Sep 17 00:00:00 2001
From: feiwen zhu <mzhu@nvidia.com>
Date: Wed, 2 Sep 2026 08:54:58 +0000
Subject: [PATCH 1/9] test: cover reduce and scan capture isolation

Signed-off-by: feiwen zhu <mzhu@nvidia.com>
---
 test/test_code_motion.py | 32 +++++++++++++++++++++++++++++++-
 test/test_reduction.py   | 20 +++++++++++++++++---
 test/test_scan.py        | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 80 insertions(+), 4 deletions(-)

diff --git a/test/test_code_motion.py b/test/test_code_motion.py
index 6963488c..dfdeb80c 100644
--- a/test/test_code_motion.py
+++ b/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
 from cuda.tile._ir.arithmetic_ops import Unary
 from cuda.tile._compile import compile_tile
 
@@ -249,3 +250,32 @@ 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 test_reduce_body_is_licm_barrier():
+    @ct.kernel
+    def kernel(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)
+
+    x = torch.zeros((16, 16), dtype=torch.int32, device="cuda")
+    y = torch.zeros((16,), dtype=torch.int32, device="cuda")
+    sig = ct.compilation.KernelSignature.from_kernel_args(
+        kernel, (x, y), CallingConvention.cutile_python_v1()
+    )
+    [root_block] = compile_tile(
+        kernel._pyfunc, [sig], return_final_ir=True, return_cubin=False
+    ).final_ir
+
+    [reduce] = [op for op in root_block.traverse() if isinstance(op, TileReduce)]
+    modulo_ops = [
+        op
+        for op in reduce.body.operations
+        if isinstance(op, TypedConst) and op.value == 5
+    ]
+    assert len(modulo_ops) == 1
+    [modulo] = modulo_ops
+
+    assert modulo in reduce.body.operations
+    assert modulo not in root_block.operations
diff --git a/test/test_reduction.py b/test/test_reduction.py
index 75123613..f53f43d8 100644
--- a/test/test_reduction.py
+++ b/test/test_reduction.py
@@ -563,6 +563,21 @@ def kernel(x, y, yi):
     assert_equal(yi, yi_ref)
 
 
+def test_custom_reduction_with_constant_capture():
+    @ct.kernel
+    def kernel(x, y):
+        modulo = 5
+        xt = ct.load(x, (0, 0), (16, 16))
+        yt = ct.reduce(xt, -1, lambda a, b: (a + b) % modulo, 0)
+        ct.store(y, (0,), yt)
+
+    x = torch.arange(256, dtype=torch.int32, device="cuda").reshape(16, 16)
+    ref = torch.sum(x, -1, dtype=torch.int32) % 5
+    y = torch.zeros((16,), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, y))
+    assert_equal(y, ref)
+
+
 def test_custom_reduction_with_capture():
     @ct.kernel
     def kernel(x, p, y):
@@ -573,10 +588,9 @@ def kernel(x, p, y):
 
     x = torch.arange(256, dtype=torch.int32, device="cuda").reshape(16, 16)
     p = torch.tensor(5, dtype=torch.int32, device="cuda")
-    ref = torch.sum(x, -1, dtype=torch.int32) % 5
     y = torch.zeros((16,), dtype=torch.int32, device="cuda")
-    ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, p, y))
-    assert_equal(y, ref)
+    with pytest.raises(TileSyntaxError, match="captures runtime value 'modulo'"):
+        ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, p, y))
 
 
 def test_custom_reduction_welford():
diff --git a/test/test_scan.py b/test/test_scan.py
index e1c1d888..6cdd3340 100644
--- a/test/test_scan.py
+++ b/test/test_scan.py
@@ -288,6 +288,38 @@ def kernel(x, y):
     torch.testing.assert_close(y, ref)
 
 
+def test_custom_scan_with_constant_capture():
+    @ct.kernel
+    def kernel(x, y):
+        scale = 2
+        xt = ct.load(x, (0, 0), (16, 16))
+        yt = ct.scan(xt, axis=-1, func=lambda a, b: (a + b) % scale, identity=0)
+        ct.store(y, (0, 0), yt)
+
+    x = torch.arange(256, dtype=torch.int32, device="cuda").reshape(16, 16)
+    ref = torch.cumsum(x, -1, dtype=torch.int32) % 2
+    y = torch.zeros((16, 16), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, y))
+    torch.testing.assert_close(y, ref)
+
+
+def test_custom_scan_with_runtime_capture():
+    @ct.kernel
+    def kernel(x, p, y):
+        modulo = ct.gather(p, ())
+        xt = ct.load(x, (0, 0), (16, 16))
+        yt = ct.scan(xt, axis=-1, func=lambda a, b: (a + b) % modulo, identity=0)
+        ct.store(y, (0, 0), yt)
+
+    x = torch.arange(256, dtype=torch.int32, device="cuda").reshape(16, 16)
+    p = torch.tensor(5, dtype=torch.int32, device="cuda")
+    y = torch.zeros((16, 16), dtype=torch.int32, device="cuda")
+    with pytest.raises(
+        TileSyntaxError, match="scan body captures runtime value 'modulo'"
+    ):
+        ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, p, y))
+
+
 def test_custom_scan_none_axis():
     @ct.kernel
     def kernel(x, y):

From c998f1d2a89c16e96d2caaecf429db078d7ad49d Mon Sep 17 00:00:00 2001
From: feiwen zhu <mzhu@nvidia.com>
Date: Wed, 2 Sep 2026 09:17:02 +0000
Subject: [PATCH 2/9] fix: isolate reduce and scan captures

Signed-off-by: feiwen zhu <mzhu@nvidia.com>
---
 src/cuda/tile/_compile.py                    |   8 ++
 src/cuda/tile/_passes/isolate_reduce_scan.py | 128 +++++++++++++++++++
 2 files changed, 136 insertions(+)
 create mode 100644 src/cuda/tile/_passes/isolate_reduce_scan.py

diff --git a/src/cuda/tile/_compile.py b/src/cuda/tile/_compile.py
index 974c42bb..2930a3d5 100644
--- a/src/cuda/tile/_compile.py
+++ b/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 (
+    _remember_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:
+    _remember_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)
 
     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
 
diff --git a/src/cuda/tile/_passes/isolate_reduce_scan.py b/src/cuda/tile/_passes/isolate_reduce_scan.py
new file mode 100644
index 00000000..ee124354
--- /dev/null
+++ b/src/cuda/tile/_passes/isolate_reduce_scan.py
@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+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]]:
+    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 _remember_reduce_scan_capture_names(root_block: Block) -> None:
+    """Record source-level capture names before Assign operations are removed."""
+    definitions = _definitions(root_block)
+    names = {}
+    for region_op in root_block.traverse():
+        if not isinstance(region_op, TileReduce | TileScan):
+            continue
+        for value, _ in _captures(region_op.body):
+            canonical_value = value
+            defining_op = definitions.get(canonical_value.name)
+            while isinstance(defining_op, Assign):
+                canonical_value = defining_op.value
+                defining_op = definitions.get(canonical_value.name)
+            key = (region_op.op, region_op.loc, canonical_value.name)
+            names[key] = value.get_original_name()
+    root_block.ctx._reduce_scan_capture_names = names
+
+
+def _original_capture_name(region_op: TileReduce | TileScan, value: Var) -> str:
+    names = getattr(value.ctx, "_reduce_scan_capture_names", {})
+    key = (region_op.op, region_op.loc, value.name)
+    return names.get(key, value.get_original_name())
+
+
+def legalize_reduce_scan_captures(root_block: Block) -> None:
+    """Rematerialize constant captures and reject runtime captures."""
+    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:
+                original_name = _original_capture_name(region_op, value)
+                raise TileSyntaxError(
+                    f"{_kind(region_op)} body captures runtime value '{original_name}'. "
+                    "Only function arguments and compile-time constants are supported.",
+                    consuming_loc,
+                )
+
+            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 available region-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)

From 3f530f17b68b510543d7b51d876796e49c682d98 Mon Sep 17 00:00:00 2001
From: feiwen zhu <mzhu@nvidia.com>
Date: Wed, 2 Sep 2026 09:30:55 +0000
Subject: [PATCH 3/9] fix: reject shaped reduction captures

Signed-off-by: feiwen zhu <mzhu@nvidia.com>
---
 src/cuda/tile/_passes/isolate_reduce_scan.py | 12 ++++++++-
 test/test_reduction.py                       | 28 +++++++++++++++++++-
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/src/cuda/tile/_passes/isolate_reduce_scan.py b/src/cuda/tile/_passes/isolate_reduce_scan.py
index ee124354..e9c7e736 100644
--- a/src/cuda/tile/_passes/isolate_reduce_scan.py
+++ b/src/cuda/tile/_passes/isolate_reduce_scan.py
@@ -6,6 +6,7 @@
 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
+from cuda.tile._ir.type import TensorLikeTy
 
 
 def _definitions(root_block: Block) -> dict[str, Operation]:
@@ -52,7 +53,7 @@ def _remember_reduce_scan_capture_names(root_block: Block) -> None:
                 canonical_value = defining_op.value
                 defining_op = definitions.get(canonical_value.name)
             key = (region_op.op, region_op.loc, canonical_value.name)
-            names[key] = value.get_original_name()
+            names.setdefault(key, value.get_original_name())
     root_block.ctx._reduce_scan_capture_names = names
 
 
@@ -86,6 +87,15 @@ def legalize_reduce_scan_captures(root_block: Block) -> None:
                     consuming_loc,
                 )
 
+            value_type = value.get_type()
+            if not isinstance(value_type, TensorLikeTy) or value_type.tensor_shape() != ():
+                original_name = _original_capture_name(region_op, value)
+                raise TileSyntaxError(
+                    f"{_kind(region_op)} body captures shaped compile-time constant "
+                    f"'{original_name}'. Only scalar compile-time constants are supported.",
+                    consuming_loc,
+                )
+
             local_value = mapper.clone_var(value)
             constants.append(TypedConst(
                 value=constant_value,
diff --git a/test/test_reduction.py b/test/test_reduction.py
index f53f43d8..85687346 100644
--- a/test/test_reduction.py
+++ b/test/test_reduction.py
@@ -578,12 +578,38 @@ def kernel(x, y):
     assert_equal(y, ref)
 
 
+def test_custom_reduction_with_shaped_constant_capture():
+    @ct.kernel
+    def kernel(x, y):
+        modulo = ct.full((1,), 5, dtype=ct.int32)
+        xt = ct.load(x, (0, 0), (16, 16))
+
+        def combine(a, b):
+            scalar_modulo = ct.extract(modulo, index=(0,), shape=())
+            return (a + b) % scalar_modulo
+
+        yt = ct.reduce(xt, -1, combine, 0)
+        ct.store(y, (0,), yt)
+
+    x = torch.arange(256, dtype=torch.int32, device="cuda").reshape(16, 16)
+    y = torch.zeros((16,), dtype=torch.int32, device="cuda")
+    with pytest.raises(
+        TileSyntaxError,
+        match="captures shaped compile-time constant 'modulo'.*"
+              "Only scalar compile-time constants are supported",
+    ):
+        ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, y))
+
+
 def test_custom_reduction_with_capture():
     @ct.kernel
     def kernel(x, p, y):
         modulo = ct.gather(p, ())
+        same_modulo = modulo
         xt = ct.load(x, (0, 0), (16, 16))
-        yt = ct.reduce(xt, -1, lambda a, b: (a + b) % modulo, 0)
+        yt = ct.reduce(
+            xt, -1, lambda a, b: ((a + b) % modulo) % same_modulo, 0
+        )
         ct.store(y, (0,), yt)
 
     x = torch.arange(256, dtype=torch.int32, device="cuda").reshape(16, 16)

From 0a77f1817c33d1794059db2d254c80efcfcc04fa Mon Sep 17 00:00:00 2001
From: feiwen zhu <mzhu@nvidia.com>
Date: Wed, 2 Sep 2026 09:38:56 +0000
Subject: [PATCH 4/9] fix: keep aggregate bodies isolated during LICM

Signed-off-by: feiwen zhu <mzhu@nvidia.com>
---
 src/cuda/tile/_passes/code_motion.py |  2 +-
 test/test_code_motion.py             | 40 +++++++++++++++++++++++++++-
 test/test_scan.py                    |  5 +++-
 3 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/src/cuda/tile/_passes/code_motion.py b/src/cuda/tile/_passes/code_motion.py
index 004b0c5e..c8a9b519 100644
--- a/src/cuda/tile/_passes/code_motion.py
+++ b/src/cuda/tile/_passes/code_motion.py
@@ -86,7 +86,7 @@ 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)
+            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
diff --git a/test/test_code_motion.py b/test/test_code_motion.py
index dfdeb80c..373263bd 100644
--- a/test/test_code_motion.py
+++ b/test/test_code_motion.py
@@ -11,7 +11,7 @@
 from cuda.tile.compilation import CallingConvention
 from cuda.tile._ir.ir import Operation
 from cuda.tile._ir.core_ops import TypedConst
-from cuda.tile._ir.ops import Loop, IfElse, TileExtract, TileReduce
+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
 
@@ -197,6 +197,22 @@ 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[0]):
+        yt = ct.reduce(xt, -1, lambda a, b: a + b, 0)
+        ct.store(y, (i, 0), yt)
+
+
+@ct.kernel
+def entire_scan_op_yes(x, y):
+    xt = ct.load(x, (0, 0), (16, 16))
+    for i in range(y.shape[0]):
+        yt = ct.scan(xt, -1, lambda a, b: a + b, 0)
+        ct.store(y, (i, 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]
@@ -279,3 +295,25 @@ def kernel(x, y):
 
     assert modulo in reduce.body.operations
     assert modulo not in root_block.operations
+
+
+@pytest.mark.parametrize(
+    "kernel, op_type, x_shape, y_shape",
+    [
+        (entire_reduce_op_yes, TileReduce, (16, 16), (3, 16)),
+        (entire_scan_op_yes, TileScan, (16, 16), (3, 16, 16)),
+    ],
+)
+def test_entire_aggregate_op_can_be_hoisted(kernel, op_type, x_shape, y_shape):
+    x = torch.zeros(x_shape, dtype=torch.float32, device="cuda")
+    y = torch.zeros(y_shape, dtype=torch.float32, device="cuda")
+    sig = ct.compilation.KernelSignature.from_kernel_args(
+        kernel, (x, y), CallingConvention.cutile_python_v1()
+    )
+    [root_block] = compile_tile(
+        kernel._pyfunc, [sig], return_final_ir=True, return_cubin=False
+    ).final_ir
+
+    [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)
diff --git a/test/test_scan.py b/test/test_scan.py
index 6cdd3340..5c3f3cf8 100644
--- a/test/test_scan.py
+++ b/test/test_scan.py
@@ -297,7 +297,10 @@ def kernel(x, y):
         ct.store(y, (0, 0), yt)
 
     x = torch.arange(256, dtype=torch.int32, device="cuda").reshape(16, 16)
-    ref = torch.cumsum(x, -1, dtype=torch.int32) % 2
+    ref = torch.empty_like(x)
+    ref[:, 0] = x[:, 0]
+    for i in range(1, x.shape[1]):
+        ref[:, i] = (ref[:, i - 1] + x[:, i]) % 2
     y = torch.zeros((16, 16), dtype=torch.int32, device="cuda")
     ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, y))
     torch.testing.assert_close(y, ref)

From aa12bbc93147c986a1e9712b0e6c0ec7afc7818b Mon Sep 17 00:00:00 2001
From: feiwen zhu <mzhu@nvidia.com>
Date: Wed, 2 Sep 2026 09:49:18 +0000
Subject: [PATCH 5/9] docs: clarify reduce and scan capture rules

Signed-off-by: feiwen zhu <mzhu@nvidia.com>
---
 changelog.d/reduce-scan-isolation.md | 3 +++
 src/cuda/tile/_stub.py               | 8 ++++++--
 2 files changed, 9 insertions(+), 2 deletions(-)
 create mode 100644 changelog.d/reduce-scan-isolation.md

diff --git a/changelog.d/reduce-scan-isolation.md b/changelog.d/reduce-scan-isolation.md
new file mode 100644
index 00000000..4ef4371e
--- /dev/null
+++ b/changelog.d/reduce-scan-isolation.md
@@ -0,0 +1,3 @@
+- Custom ``ct.reduce()`` and ``ct.scan()`` callbacks now keep captured scalar compile-time
+  constants inside their isolated bodies. Unsupported runtime and shaped captures are diagnosed
+  during cuTile Python compilation, before TileIR is invoked.
diff --git a/src/cuda/tile/_stub.py b/src/cuda/tile/_stub.py
index 8502d92f..8d615cf6 100644
--- a/src/cuda/tile/_stub.py
+++ b/src/cuda/tile/_stub.py
@@ -3049,7 +3049,9 @@ 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 may capture
+            scalar compile-time constants from its enclosing scope. Capturing runtime values or
+            shaped constants 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 +3162,9 @@ 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 may capture
+            scalar compile-time constants from its enclosing scope. Capturing runtime values or
+            shaped constants 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.

From caa74c4d1856e5fdb1aa9728aff7f5bc1c304e0f Mon Sep 17 00:00:00 2001
From: feiwen zhu <mzhu@nvidia.com>
Date: Wed, 2 Sep 2026 13:53:03 +0000
Subject: [PATCH 6/9] refactor: pass capture names explicitly and clarify
 diagnostics

Return the source-level names of captured values from the pre-pass and
hand them to the legalization pass, instead of stashing them on the
IRContext and importing a private helper from _compile.py. Key them by the
value that survives Assign elimination. Also document why reduce/scan
bodies are LICM barriers and why capture legalization must run after
constant materialization, and name the callback's parameters explicitly in
the runtime-capture error.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: feiwen zhu <mzhu@nvidia.com>
---
 src/cuda/tile/_compile.py                    |  6 +-
 src/cuda/tile/_passes/code_motion.py         |  3 +
 src/cuda/tile/_passes/isolate_reduce_scan.py | 62 ++++++++++++--------
 3 files changed, 44 insertions(+), 27 deletions(-)

diff --git a/src/cuda/tile/_compile.py b/src/cuda/tile/_compile.py
index 2930a3d5..ff3c829b 100644
--- a/src/cuda/tile/_compile.py
+++ b/src/cuda/tile/_compile.py
@@ -66,7 +66,7 @@
 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 (
-    _remember_reduce_scan_capture_names,
+    collect_reduce_scan_capture_names,
     legalize_reduce_scan_captures,
     verify_reduce_scan_isolation,
 )
@@ -107,14 +107,14 @@ def _transform_ir(func_body: ir.Block,
                   bytecode_version: bc.BytecodeVersion,
                   param_constraints: Sequence[tuple[tuple[ir.Var, ...], ParameterConstraint]]
                   ) -> DataflowResult:
-    _remember_reduce_scan_capture_names(func_body)
+    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)
+    legalize_reduce_scan_captures(func_body, capture_names)
 
     if not CUDA_TILE_TESTING_DISABLE_DIV:
         add_divby_pass(func_body, dataflow_result)
diff --git a/src/cuda/tile/_passes/code_motion.py b/src/cuda/tile/_passes/code_motion.py
index c8a9b519..41926ec4 100644
--- a/src/cuda/tile/_passes/code_motion.py
+++ b/src/cuda/tile/_passes/code_motion.py
@@ -86,6 +86,9 @@ 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
 
+            # 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.
diff --git a/src/cuda/tile/_passes/isolate_reduce_scan.py b/src/cuda/tile/_passes/isolate_reduce_scan.py
index e9c7e736..d8da54db 100644
--- a/src/cuda/tile/_passes/isolate_reduce_scan.py
+++ b/src/cuda/tile/_passes/isolate_reduce_scan.py
@@ -2,6 +2,14 @@
 #
 # 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
@@ -18,6 +26,8 @@ def _definitions(root_block: Block) -> dict[str, Operation]:
 
 
 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
@@ -39,32 +49,36 @@ def _kind(op: TileReduce | TileScan) -> str:
     return "reduction" if isinstance(op, TileReduce) else "scan"
 
 
-def _remember_reduce_scan_capture_names(root_block: Block) -> None:
-    """Record source-level capture names before Assign operations are removed."""
+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 = {}
+    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 = value
-            defining_op = definitions.get(canonical_value.name)
+            canonical = value
+            defining_op = definitions.get(canonical.name)
             while isinstance(defining_op, Assign):
-                canonical_value = defining_op.value
-                defining_op = definitions.get(canonical_value.name)
-            key = (region_op.op, region_op.loc, canonical_value.name)
-            names.setdefault(key, value.get_original_name())
-    root_block.ctx._reduce_scan_capture_names = names
-
+                canonical = defining_op.value
+                defining_op = definitions.get(canonical.name)
+            names.setdefault(canonical.name, value.get_original_name())
+    return names
 
-def _original_capture_name(region_op: TileReduce | TileScan, value: Var) -> str:
-    names = getattr(value.ctx, "_reduce_scan_capture_names", {})
-    key = (region_op.op, region_op.loc, value.name)
-    return names.get(key, value.get_original_name())
 
+def legalize_reduce_scan_captures(root_block: Block, capture_names: dict[str, str]) -> None:
+    """Rematerialize constant captures inside each body and reject runtime captures.
 
-def legalize_reduce_scan_captures(root_block: Block) -> None:
-    """Rematerialize constant captures 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():
@@ -80,19 +94,19 @@ def legalize_reduce_scan_captures(root_block: Block) -> None:
             elif isinstance(defining_op, TypedConst):
                 constant_value = defining_op.value
             else:
-                original_name = _original_capture_name(region_op, value)
+                name = capture_names.get(value.name, value.get_original_name())
                 raise TileSyntaxError(
-                    f"{_kind(region_op)} body captures runtime value '{original_name}'. "
-                    "Only function arguments and compile-time constants are supported.",
+                    f"{_kind(region_op)} body captures runtime value '{name}'. Only the "
+                    "callback's own parameters and scalar compile-time constants are supported.",
                     consuming_loc,
                 )
 
             value_type = value.get_type()
             if not isinstance(value_type, TensorLikeTy) or value_type.tensor_shape() != ():
-                original_name = _original_capture_name(region_op, value)
+                name = capture_names.get(value.name, value.get_original_name())
                 raise TileSyntaxError(
-                    f"{_kind(region_op)} body captures shaped compile-time constant "
-                    f"'{original_name}'. Only scalar compile-time constants are supported.",
+                    f"{_kind(region_op)} body captures shaped compile-time constant '{name}'. "
+                    "Only scalar compile-time constants are supported.",
                     consuming_loc,
                 )
 
@@ -110,7 +124,7 @@ def legalize_reduce_scan_captures(root_block: Block) -> None:
 
 
 def verify_reduce_scan_isolation(root_block: Block) -> None:
-    """Verify that reduce and scan bodies only use available region-local values."""
+    """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

From ec6a72c13d9d4eff440e7f2c31d002e836eb70e5 Mon Sep 17 00:00:00 2001
From: feiwen zhu <mzhu@nvidia.com>
Date: Wed, 2 Sep 2026 13:54:05 +0000
Subject: [PATCH 7/9] fix: require scalar tiles inside reduce and scan bodies

Check, right after a callback body is built, that every value inside it
is a 0-d tile. This gives one consistent diagnostic for non-scalar
constants whether they are captured from the enclosing scope or created
inside the callback, and it replaces the capture-only shape check in the
legalization pass. Document the restriction in the ct.reduce()/ct.scan()
docstrings.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: feiwen zhu <mzhu@nvidia.com>
---
 src/cuda/tile/_ir/ops.py                     | 27 +++++++++++++++++++-
 src/cuda/tile/_passes/isolate_reduce_scan.py | 12 ++-------
 src/cuda/tile/_stub.py                       | 14 +++++-----
 test/test_reduction.py                       | 24 +++++++++++++++--
 test/test_scan.py                            | 21 +++++++++++++++
 5 files changed, 79 insertions(+), 19 deletions(-)

diff --git a/src/cuda/tile/_ir/ops.py b/src/cuda/tile/_ir/ops.py
index 8a364c14..a09ff60a 100644
--- a/src/cuda/tile/_ir/ops.py
+++ b/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
 
 
diff --git a/src/cuda/tile/_passes/isolate_reduce_scan.py b/src/cuda/tile/_passes/isolate_reduce_scan.py
index d8da54db..ca132c94 100644
--- a/src/cuda/tile/_passes/isolate_reduce_scan.py
+++ b/src/cuda/tile/_passes/isolate_reduce_scan.py
@@ -14,7 +14,6 @@
 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
-from cuda.tile._ir.type import TensorLikeTy
 
 
 def _definitions(root_block: Block) -> dict[str, Operation]:
@@ -101,15 +100,8 @@ def legalize_reduce_scan_captures(root_block: Block, capture_names: dict[str, st
                     consuming_loc,
                 )
 
-            value_type = value.get_type()
-            if not isinstance(value_type, TensorLikeTy) or value_type.tensor_shape() != ():
-                name = capture_names.get(value.name, value.get_original_name())
-                raise TileSyntaxError(
-                    f"{_kind(region_op)} body captures shaped compile-time constant '{name}'. "
-                    "Only 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,
diff --git a/src/cuda/tile/_stub.py b/src/cuda/tile/_stub.py
index 8d615cf6..f052e4c1 100644
--- a/src/cuda/tile/_stub.py
+++ b/src/cuda/tile/_stub.py
@@ -3049,9 +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. The function may capture
-            scalar compile-time constants from its enclosing scope. Capturing runtime values or
-            shaped constants is unsupported and rejected during compilation.
+            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.
@@ -3162,9 +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. The function may capture
-            scalar compile-time constants from its enclosing scope. Capturing runtime values or
-            shaped constants is unsupported and rejected during compilation.
+            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.
diff --git a/test/test_reduction.py b/test/test_reduction.py
index 85687346..4d70e895 100644
--- a/test/test_reduction.py
+++ b/test/test_reduction.py
@@ -595,8 +595,28 @@ def combine(a, b):
     y = torch.zeros((16,), dtype=torch.int32, device="cuda")
     with pytest.raises(
         TileSyntaxError,
-        match="captures shaped compile-time constant 'modulo'.*"
-              "Only scalar compile-time constants are supported",
+        match=r"reduction body must only operate on scalar tiles, but 'modulo' has shape \(1,\)",
+    ):
+        ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, y))
+
+
+def test_custom_reduction_with_shaped_constant_in_body():
+    @ct.kernel
+    def kernel(x, y):
+        xt = ct.load(x, (0, 0), (16, 16))
+
+        def combine(a, b):
+            modulo = ct.full((1,), 5, dtype=ct.int32)
+            return (a + b) % ct.extract(modulo, index=(0,), shape=())
+
+        yt = ct.reduce(xt, -1, combine, 0)
+        ct.store(y, (0,), yt)
+
+    x = torch.arange(256, dtype=torch.int32, device="cuda").reshape(16, 16)
+    y = torch.zeros((16,), dtype=torch.int32, device="cuda")
+    with pytest.raises(
+        TileSyntaxError,
+        match=r"reduction body must only operate on scalar tiles, but 'modulo' has shape \(1,\)",
     ):
         ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, y))
 
diff --git a/test/test_scan.py b/test/test_scan.py
index 5c3f3cf8..5e0cb05d 100644
--- a/test/test_scan.py
+++ b/test/test_scan.py
@@ -306,6 +306,27 @@ def kernel(x, y):
     torch.testing.assert_close(y, ref)
 
 
+def test_custom_scan_with_shaped_constant_capture():
+    @ct.kernel
+    def kernel(x, y):
+        scale = ct.full((1,), 2, dtype=ct.int32)
+        xt = ct.load(x, (0, 0), (16, 16))
+
+        def combine(a, b):
+            return (a + b) % ct.extract(scale, index=(0,), shape=())
+
+        yt = ct.scan(xt, axis=-1, func=combine, identity=0)
+        ct.store(y, (0, 0), yt)
+
+    x = torch.arange(256, dtype=torch.int32, device="cuda").reshape(16, 16)
+    y = torch.zeros((16, 16), dtype=torch.int32, device="cuda")
+    with pytest.raises(
+        TileSyntaxError,
+        match=r"scan body must only operate on scalar tiles, but 'scale' has shape \(1,\)",
+    ):
+        ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, y))
+
+
 def test_custom_scan_with_runtime_capture():
     @ct.kernel
     def kernel(x, p, y):

From b64d242e51042375e3569be717da8dcf07e5035a Mon Sep 17 00:00:00 2001
From: feiwen zhu <mzhu@nvidia.com>
Date: Wed, 2 Sep 2026 13:54:52 +0000
Subject: [PATCH 8/9] test: launch hoisted aggregate kernels and use a true
 scan identity

The kernels checking that a whole reduce/scan can be hoisted stored tiles
whose rank did not match the array; give them matching ranks and launch
them so the hoisted result is verified. Cover scan in the LICM barrier
test as well. The scan constant-capture test used identity 0 together with
'(a + b) % 2', which is not an identity for arbitrary inputs; restrict the
inputs to {0, 1} so the reference does not depend on where the identity is
combined.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: feiwen zhu <mzhu@nvidia.com>
---
 test/test_code_motion.py | 90 ++++++++++++++++++++++++----------------
 test/test_scan.py        |  8 ++--
 2 files changed, 57 insertions(+), 41 deletions(-)

diff --git a/test/test_code_motion.py b/test/test_code_motion.py
index 373263bd..45cf7890 100644
--- a/test/test_code_motion.py
+++ b/test/test_code_motion.py
@@ -200,17 +200,31 @@ def carried_from_nested_loop_no(x, a, t):
 @ct.kernel
 def entire_reduce_op_yes(x, y):
     xt = ct.load(x, (0, 0), (16, 16))
-    for i in range(y.shape[0]):
-        yt = ct.reduce(xt, -1, lambda a, b: a + b, 0)
-        ct.store(y, (i, 0), yt)
+    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[0]):
+    for i in range(y.shape[1] // 16):
         yt = ct.scan(xt, -1, lambda a, b: a + b, 0)
-        ct.store(y, (i, 0, 0), yt)
+        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):
@@ -268,52 +282,56 @@ def test_hoisting(kernel, op_finder, expected_x):
     assert_close(x, ref)
 
 
-def test_reduce_body_is_licm_barrier():
-    @ct.kernel
-    def kernel(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)
-
-    x = torch.zeros((16, 16), dtype=torch.int32, device="cuda")
-    y = torch.zeros((16,), dtype=torch.int32, device="cuda")
+def _final_ir(kernel, args):
     sig = ct.compilation.KernelSignature.from_kernel_args(
-        kernel, (x, y), CallingConvention.cutile_python_v1()
+        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
 
-    [reduce] = [op for op in root_block.traverse() if isinstance(op, TileReduce)]
-    modulo_ops = [
-        op
-        for op in reduce.body.operations
-        if isinstance(op, TypedConst) and op.value == 5
-    ]
-    assert len(modulo_ops) == 1
-    [modulo] = modulo_ops
 
-    assert modulo in reduce.body.operations
-    assert modulo not in root_block.operations
+@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, x_shape, y_shape",
+    "kernel, op_type, y_shape, reference",
     [
-        (entire_reduce_op_yes, TileReduce, (16, 16), (3, 16)),
-        (entire_scan_op_yes, TileScan, (16, 16), (3, 16, 16)),
+        (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, x_shape, y_shape):
-    x = torch.zeros(x_shape, dtype=torch.float32, device="cuda")
+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")
-    sig = ct.compilation.KernelSignature.from_kernel_args(
-        kernel, (x, y), CallingConvention.cutile_python_v1()
-    )
-    [root_block] = compile_tile(
-        kernel._pyfunc, [sig], return_final_ir=True, return_cubin=False
-    ).final_ir
+    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))
diff --git a/test/test_scan.py b/test/test_scan.py
index 5e0cb05d..8483feaf 100644
--- a/test/test_scan.py
+++ b/test/test_scan.py
@@ -296,11 +296,9 @@ def kernel(x, y):
         yt = ct.scan(xt, axis=-1, func=lambda a, b: (a + b) % scale, identity=0)
         ct.store(y, (0, 0), yt)
 
-    x = torch.arange(256, dtype=torch.int32, device="cuda").reshape(16, 16)
-    ref = torch.empty_like(x)
-    ref[:, 0] = x[:, 0]
-    for i in range(1, x.shape[1]):
-        ref[:, i] = (ref[:, i - 1] + x[:, i]) % 2
+    # With inputs in {0, 1}, `(a + b) % 2` is XOR and 0 is a true identity for it.
+    x = (torch.arange(256, dtype=torch.int32, device="cuda") % 2).reshape(16, 16)
+    ref = torch.cumsum(x, -1, dtype=torch.int32) % 2
     y = torch.zeros((16, 16), dtype=torch.int32, device="cuda")
     ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, y))
     torch.testing.assert_close(y, ref)

From 5fbd116a5d640b91b58d48489ab79bc700c9c449 Mon Sep 17 00:00:00 2001
From: feiwen zhu <mzhu@nvidia.com>
Date: Wed, 2 Sep 2026 13:54:52 +0000
Subject: [PATCH 9/9] docs: state the reduce/scan callback behavior change in
 the changelog

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: feiwen zhu <mzhu@nvidia.com>
---
 changelog.d/reduce-scan-isolation.md | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/changelog.d/reduce-scan-isolation.md b/changelog.d/reduce-scan-isolation.md
index 4ef4371e..27945d3b 100644
--- a/changelog.d/reduce-scan-isolation.md
+++ b/changelog.d/reduce-scan-isolation.md
@@ -1,3 +1,5 @@
 - Custom ``ct.reduce()`` and ``ct.scan()`` callbacks now keep captured scalar compile-time
-  constants inside their isolated bodies. Unsupported runtime and shaped captures are diagnosed
-  during cuTile Python compilation, before TileIR is invoked.
+  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.