Add support for custom scan via ct.scan

Signed-off-by: Ziheng Deng <zihengd@nvidia.com>

Author: gunnersdeng

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

docs/source/operations.rst

@@ -78,6 +78,7 @@ Scan
 
    cumsum
    cumprod
+   scan
 
 
 Matmul

src/cuda/tile/__init__.py

@@ -90,6 +90,7 @@
     cosh,
     cumprod,
     cumsum,
+    scan,
     equal,
     exp,
     exp2,
@@ -224,6 +225,7 @@
     "cosh",
     "cumprod",
     "cumsum",
+    "scan",
     "equal",
     "exp",
     "exp2",

src/cuda/tile/_ir/ir.py

@@ -414,7 +414,8 @@ def finalize_loopvar_type(self, body_var: Var):
 
 
 class Builder:
-    def __init__(self, ctx: IRContext, loc: Loc, reduction_body: bool = False):
+    def __init__(self, ctx: IRContext, loc: Loc, reduction_body: bool = False,
+                 scan_body: bool = False):
         self.ir_ctx = ctx
         self.is_terminated = False
         self._loc = loc
@@ -423,14 +424,18 @@ def __init__(self, ctx: IRContext, loc: Loc, reduction_body: bool = False):
         self._prev_builder = None
         self._var_map: Dict[str, Var] = dict()
         self.reduction_body = reduction_body
+        self.scan_body = scan_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")
+        if (self.reduction_body or self.scan_body) and op_class.memory_effect != MemoryEffect.NONE:
+            if self.reduction_body:
+                msg = "Operations with memory effects are not supported inside reduction body"
+            else:
+                msg = "Operations with memory effects are not supported inside scan body"
+            raise TileSyntaxError(msg)
 
         assert not self.is_terminated
         force_type = False
@@ -524,11 +529,12 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
 
 @contextmanager
-def nested_block(loc: Loc, reduction_body: bool = False):
+def nested_block(loc: Loc, reduction_body: bool = False, scan_body: bool = False):
     prev_builder = Builder.get_current()
     block = Block(prev_builder.ir_ctx, loc=loc)
     with Builder(prev_builder.ir_ctx, loc,
-                 reduction_body=reduction_body or prev_builder.reduction_body) as builder:
+                 reduction_body=reduction_body or prev_builder.reduction_body,
+                 scan_body=scan_body or prev_builder.scan_body) as builder:
         yield block
     block.extend(builder.ops)