Introduce BlockRestriction interface for restricting which operations are allowed inside a block

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

Author: gunnersdeng

src/cuda/tile/_ir/ir.py

@@ -20,7 +20,6 @@
 from cuda.tile._exception import (
     TileTypeError, Loc, TileInternalError
 )
-from .. import TileSyntaxError
 from .._context import TileContextConfig
 from cuda.tile._bytecode.version import BytecodeVersion
 
@@ -324,6 +323,14 @@ class MemoryEffect(enum.IntEnum):
     STORE = 2
 
 
+class BlockRestriction:
+    """Interface for restricting which operations are allowed inside a block."""
+
+    def validate_operation(self, op_class: type) -> None:
+        """Raise if the given operation class is not allowed. No restriction by default."""
+        return
+
+
 class Mapper:
     def __init__(self, ctx: IRContext, preserve_vars: bool = False):
         self._ctx = ctx
@@ -396,28 +403,23 @@ def finalize_loopvar_type(self, body_var: Var):
 
 
 class Builder:
-    def __init__(self, ctx: IRContext, loc: Loc, reduction_body: bool = False,
-                 scan_body: bool = False):
+    def __init__(self, ctx: IRContext, loc: Loc,
+                 block_restriction: Optional[BlockRestriction] = None):
         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
-        self.scan_body = scan_body
+        self.block_restriction = block_restriction
 
     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 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)
+        if self.block_restriction is not None:
+            self.block_restriction.validate_operation(op_class)
 
         assert not self.is_terminated
         force_type = False
@@ -512,12 +514,11 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
 
 @contextmanager
-def enter_nested_block(loc: Loc, reduction_body: bool = False, scan_body: bool = False):
+def enter_nested_block(loc: Loc, block_restriction: Optional[BlockRestriction] = None):
     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,
-                 scan_body=scan_body or prev_builder.scan_body) as builder:
+                 block_restriction=block_restriction or prev_builder.block_restriction) as builder:
         yield block
     block.extend(builder.ops)
 

src/cuda/tile/_ir/ops.py

@@ -23,7 +23,7 @@
     add_operation, Builder,
     enter_nested_block, nested_block, PhiState, LoopVarState,
     TupleValue, make_aggregate, RangeValue, BoundMethodValue, ArrayValue, ConstantState,
-    ListValue, ClosureValue, MemoryEffect, attribute, operand,
+    ListValue, ClosureValue, MemoryEffect, attribute, operand, BlockRestriction
 )
 from cuda.tile._ir.load_store_impl import (
     check_load_store_hints,
@@ -306,10 +306,8 @@ async def loop_impl(body: hir.Block, iterable: Var):
 
     # Do this check at the end because this may be an automatically inserted loop
     # around the helper function's body.
-    if builder.reduction_body:
-        raise TileSyntaxError("Loops inside reduction body are not supported")
-    if builder.scan_body:
-        raise TileSyntaxError("Loops inside scan body are not supported")
+    if builder.block_restriction is not None:
+        builder.block_restriction.validate_operation(Loop)
 
 
 def _have_nested_jump(calls: Sequence[hir.Call]) -> bool:
@@ -349,6 +347,26 @@ def _to_string_rhs(self) -> str:
         return f"if(cond={self.cond})"
 
 
+@dataclass
+class ReduceScanRestriction(BlockRestriction):
+    """Restriction for reduction/scan body blocks: no memory effects, loops, or branching."""
+
+    kind: Literal["reduction", "scan"]
+
+    def validate_operation(self, op_class: type) -> None:
+        if getattr(op_class, "memory_effect", MemoryEffect.NONE) != MemoryEffect.NONE:
+            raise TileSyntaxError(
+                f"Operations with memory effects are not supported inside {self.kind} body"
+            )
+        if op_class is Loop:
+            raise TileSyntaxError(f"Loops inside {self.kind} body are not supported")
+        if op_class is IfElse:
+            raise TileSyntaxError(
+                f"Branching inside {self.kind} body is not supported. "
+                f"Consider ct.where() as a workaround."
+            )
+
+
 async def _flatten_branch(branch: hir.Block) -> Var | None:
     from .._passes.hir2ir import dispatch_hir_block
     info = ControlFlowInfo((), flatten=True)
@@ -373,12 +391,8 @@ async def if_else_impl(cond: Var, then_block: hir.Block, else_block: hir.Block)
         return await _flatten_branch(branch_taken)
 
     builder = Builder.get_current()
-    if builder.reduction_body:
-        raise TileSyntaxError("Branching inside reduction body is not supported."
-                              " Consider ct.where() as a workaround.")
-    if builder.scan_body:
-        raise TileSyntaxError("Branching inside scan body is not supported."
-                              " Consider ct.where() as a workaround.")
+    if builder.block_restriction is not None:
+        builder.block_restriction.validate_operation(IfElse)
 
     # Get the total number of results by adding the number of stored variables.
     # Note: we sort the stored variable indices to make the order deterministic.
@@ -2889,7 +2903,7 @@ async def _get_reduce_scan_body_block(
     """Build body block for reduce/scan. Caller passes result_shape; returns
     (body_block, result_types)."""
     builder = Builder.get_current()
-    if builder.reduction_body or builder.scan_body:
+    if isinstance(builder.block_restriction, ReduceScanRestriction):
         raise TileSyntaxError("Nested scan/reduction is not supported")
 
     block_params = []
@@ -2913,8 +2927,7 @@ async def _get_reduce_scan_body_block(
 
     with enter_nested_block(
             builder.loc,
-            reduction_body=(op_name == "reduction"),
-            scan_body=(op_name == "scan")) as body_block:
+            block_restriction=ReduceScanRestriction(op_name)) as body_block:
         body_block.params = tuple(block_params)
         body_results = await body(tuple(lhs_vars), tuple(rhs_vars))
         for body_res, x in zip(body_results, xs, strict=True):

test/test_scan.py

@@ -297,7 +297,8 @@ def test_custom_scan_ifelse_not_supported():
     def kernel(x, y):
         def f(a, b):
             if ct.bid(0) == 0:
-                return a + b
+                # In case of type mismatch, compiler will complain about nested branching
+                return a + b.type()
             else:
                 return (a + b) % 5