[cl] Support nested conditionals and break_loop in task-scheduling domain loops

Author: Cong Zhang

experimental/task_scheduling/README.md

@@ -30,7 +30,82 @@ The manager also computes a warp-group-rounded initial register budget and threa
 
 Automatic value routing hides storage positions from callback authors. Each callback receives its routed inputs after the `StageInfo` argument and returns the values declared by its captured work method. The frozen `DeviceStep` performs the corresponding value-stack operations, while `DeviceTask.make_context(tasks_inputs)` creates the immutable execution context. The tuple representation remains an internal compiler detail; no placeholder or named routing slot is exposed to resource authors.
 
-`outputs=N` creates `N` independent lexical `ScheduleValue` instances. A work call returns one value directly or a tuple that can be unpacked normally, and its device callback returns the same number of runtime values. A value produced inside a conditional or domain loop cannot escape that scope accidentally. A domain loop(don't support while senmantic and contional early break) that needs state across iterations uses `domain_loop(start, end, step, body, *initial_values)`, with one callback parameter per initial value. The schedule builder invokes `body(*iter_values)` once while capturing the loop body, so every resource call remains a schedule step. The callback must return one work-call output per input; those outputs become both the next-iteration values and the post-loop results. The device adapter preserves zero-trip pass-through and materializes the backedge without exposing routing positions to callbacks. Work methods read the current device offset from `stage_info.loop_offset`; it is not a callback parameter or a routed value. Loop bodies may use `first_iter()`, `last_iter()`, and `every(period, start=...)` to capture iteration-guarded regions.
+`outputs=N` creates `N` independent lexical `ScheduleValue` instances. A work call returns one value directly or a tuple that can be unpacked normally, and its device callback returns the same number of runtime values. A value produced inside a conditional or domain loop cannot escape that scope accidentally. A domain loop that needs state across iterations uses `domain_loop(start, end, step, body, *initial_values)`, with one callback parameter per initial value. The schedule builder invokes `body(*iter_values)` once while capturing the loop body, so every resource call remains a schedule step. The callback must return one work-call output per input; those outputs become both the next-iteration values and the post-loop results. The device adapter preserves zero-trip pass-through and materializes the backedge without exposing routing positions to callbacks. Work methods read the current device offset from `stage_info.loop_offset`; it is not a callback parameter or a routed value. Loop bodies may use `first_iter()`, `last_iter()`, and `every(period, start=...)` to capture iteration-guarded regions.
+
+`when_true()` and `when_false()` regions may nest, including inside iteration
+guards. An inner region executes only when every enclosing guard is active.
+Predicates and other work outputs produced inside a branch may be used by its
+nested regions, but cannot escape to a parent or sibling scope. Reusing the same
+routed predicate preserves true/false correlation in host analysis.
+
+The runnable `tutorial/01_copy_basics/04_copy_tma_nested_conditional.py` example
+extends the conditional TMA copy with all four two-level true/false combinations
+and a third-level condition. It checks the copied tensor and exact per-row branch
+markers, including the absence of writes from inactive branches:
+
+```bash
+python experimental/task_scheduling/tutorial/01_copy_basics/04_copy_tma_nested_conditional.py --rows-cols 9,512
+```
+
+`break_loop()` exits the enclosing domain loop immediately, including from a
+nested conditional. It skips the rest of that iteration and resumes after the
+loop. The context-manager handle provides the same operation:
+
+```python
+with ts.domain_loop(num_tiles) as loop:
+    done = resource.is_done()
+    with ts.when_true(done):
+        loop.break_loop()
+    resource.process()
+```
+
+Use `ts.break_loop()` in a functional loop body. A bare break skips the normal
+backedge: returned loop-carried results retain their values from the last
+completed iteration (or the initial values if the first iteration breaks).
+Pass explicit results to return values computed during the interrupted iteration:
+
+```python
+initial = resource.init_state()
+
+def body(state):
+    updated = resource.advance_state(state)
+    done = resource.is_done(updated)
+    with ts.when_true(done):
+        ts.break_loop(updated)
+    return updated
+
+result = ts.domain_loop(0, num_tiles, 1, body, initial)
+resource.consume_state(result)
+```
+
+For multiple carried inputs, use `ts.break_loop(next_count, next_total)` in the
+same order as the functional loop's inputs. Supply all carried results or none.
+Each explicit argument must be a visible routed `ScheduleValue` with the same
+pipeline-stage provenance and a compatible CUDA type as its carried input;
+ordinary constants must first be returned by a work method. Values created
+inside a nested branch may be returned directly by a break in that branch.
+The exit values are captured before branch-local routes are discarded. Wrong
+arity, foreign-schedule values, and values escaping a sibling/child scope are
+rejected during capture; CUDA types are checked during device compilation.
+An untaken break uses the normal body return values, and a zero-trip loop still
+returns its initial values.
+
+Side effects and pipeline-state updates performed before the break are retained.
+A break exits only the domain loop, so an enclosing `work_tile_loop()` continues.
+`last_iter()` still means the last iteration of the declared range; it is not an
+exit hook. Breaks outside a domain loop and expired loop handles are rejected.
+Unbounded while loops remain unsupported.
+
+The nested-copy tutorial also accepts `--stop-row 5`. Both producer and consumer
+break before acquiring/waiting for row 5, and the example verifies that the
+uncopied output and trace remain zero. A pipeline schedule must ensure that its
+producer and consumer exit consistently and finish any acquired work before
+exiting; `break_loop()` does not implicitly release or drain an interrupted step.
+
+Host expansion honors breaks for its selected opaque-condition assignment and
+representative loop bounds, including zero-trip loops. Opaque assignments remain
+fixed during each explored execution; these checks do not prove safety for all
+possible iteration-varying runtime predicate sequences.
 
 `StageInfo` contains the current `stage_idx`, `phase`, selected full `barrier`, zero-based iteration count, loop offset/bounds, work label, and owning `ExecutionContext`; loop offset/bounds are `None` for peeled work outside a domain loop. This lets pipeline payload work use `stage_info.barrier` without knowing how barrier arrays are stored by the device manager.
 

experimental/task_scheduling/src/task_scheduling/__init__.py

@@ -22,6 +22,7 @@
 )
 from .exhaustive_checker import check_all_interleavings, expand_task
 from .ir import (
+    BreakLoopIR,
     ConditionalIR,
     DependencyEdgeIR,
     DomainLoopIR,
@@ -58,6 +59,8 @@
     producer_work,
 )
 from .schedule_builder import (
+    BreakLoop,
+    break_loop,
     ConditionalBlock,
     DomainLoop,
     DynamicDomainBound,
@@ -96,6 +99,9 @@
 
 
 __all__ = [
+    "BreakLoop",
+    "BreakLoopIR",
+    "break_loop",
     "BarrierAllocation",
     "BarrierAllocator",
     "check_all_interleavings",

experimental/task_scheduling/src/task_scheduling/exhaustive_checker.py

@@ -18,11 +18,13 @@
     guard_fires,
 )
 from .schedule_builder import (
+    BreakLoop,
     ConditionalBlock,
     DomainLoop,
     Step,
     WorkTileLoop,
     _iter_nodes,
+    validate_break_placement,
 )
 from .resources import WorkQueue
 from .task import Task
@@ -63,11 +65,14 @@ def expand_task(
     dynamic_domain_fallback: int = 1,
     representative_domain: bool = False,
 ) -> list[FlatScheduleOp]:
+    validate_break_placement(task.schedule.body)
     result = []
     domain_entry = itertools.count()
 
     def emit(nodes, iteration=None, count=None, phase="O", skipped=False):
         for node in nodes:
+            if isinstance(node, BreakLoop):
+                return True
             if isinstance(node, Step):
                 # Static work-queue stages only update task-local scheduler
                 # state; unlike a CLC queue, they do not synchronize tasks.
@@ -96,10 +101,11 @@ def emit(nodes, iteration=None, count=None, phase="O", skipped=False):
                 stride = _bound(node.step, task, 1)
                 if stride <= 0:
                     raise ValueError("exhaustive checker requires positive domain step")
-                values = tuple(range(start, end, stride)) or (start,)
+                values = tuple(range(start, end, stride))
                 next(domain_entry)
                 for index, _ in enumerate(values):
-                    emit(node.body, index, len(values), "", skipped)
+                    if emit(node.body, index, len(values), "", skipped):
+                        break
             elif isinstance(node, WorkTileLoop):
                 for _ in range(num_tiles):
                     emit(node.body, iteration, count, phase, skipped_tile)
@@ -126,7 +132,9 @@ def emit(nodes, iteration=None, count=None, phase="O", skipped=False):
                         if guard is LAST_ITER
                         else phase
                     )
-                    emit(node.body, iteration, count, guard_phase, skipped)
+                    if emit(node.body, iteration, count, guard_phase, skipped):
+                        return True
+        return False
 
     emit(task.schedule.body)
     return result

experimental/task_scheduling/src/task_scheduling/ir.py

@@ -68,6 +68,13 @@ class StepIR:
     stage_resource_id: int | None = None
 
 
+@dataclass(frozen=True)
+class BreakLoopIR:
+    """Exit the enclosing DomainLoopIR without taking its backedge."""
+
+    exit_values: tuple[int, ...] = ()
+
+
 @dataclass(frozen=True)
 class ConditionalIR:
     body: tuple["NodeIR", ...]
@@ -95,7 +102,7 @@ class WorkTileLoopIR:
     skip_if: object | None = None
 
 
-NodeIR = StepIR | ConditionalIR | DomainLoopIR | WorkTileLoopIR
+NodeIR = StepIR | BreakLoopIR | ConditionalIR | DomainLoopIR | WorkTileLoopIR
 
 
 @dataclass(frozen=True)
@@ -136,7 +143,7 @@ def resource_by_id(self) -> Mapping[int, ResourceIR]:
 def _iter_ir_nodes(nodes: tuple[NodeIR, ...]):
     for node in nodes:
         yield node
-        if not isinstance(node, StepIR):
+        if not isinstance(node, (StepIR, BreakLoopIR)):
             yield from _iter_ir_nodes(node.body)
 
 
@@ -155,7 +162,7 @@ def freeze_program_ir(
     """Normalize capture objects into an identity-free, frozen program view."""
 
     from .enums import Every, FIRST_ITER, LAST_ITER, OpaqueCondition, SKIPPABLE
-    from .schedule_builder import ConditionalBlock, DomainLoop, Step, WorkTileLoop
+    from .schedule_builder import BreakLoop, ConditionalBlock, DomainLoop, Step, WorkTileLoop
 
     resource_ids = {resource: index for index, resource in enumerate(resources)}
     resource_irs = tuple(
@@ -193,6 +200,8 @@ def freeze_guard(guard) -> GuardIR:
         raise TypeError(f"unsupported captured guard {type(guard).__name__}")
 
     def freeze_node(node) -> NodeIR:
+        if isinstance(node, BreakLoop):
+            return BreakLoopIR(tuple(value.value_id for value in node.exit_values))
         if isinstance(node, Step):
             return StepIR(
                 resource_id=resource_ids[node.memory_resource],
@@ -263,6 +272,8 @@ def freeze_node(node) -> NodeIR:
             if isinstance(node, DomainLoopIR):
                 consumed.update(node.initial_values)
                 consumed.update(node.yield_values)
+            elif isinstance(node, BreakLoopIR):
+                consumed.update(node.exit_values)
         task_irs.append(
             TaskIR(
                 task_id=task_id,

experimental/task_scheduling/src/task_scheduling/schedule_builder.py

@@ -46,6 +46,13 @@ class Step:
     unique_id: int
 
 
+@dataclass(kw_only=True, eq=False, frozen=True)
+class BreakLoop:
+    """Exit the enclosing domain loop, skipping the remaining iteration body."""
+
+    exit_values: tuple[ScheduleValue, ...] = ()
+
+
 @dataclass(kw_only=True, eq=False, frozen=True)
 class ConditionalBlock:
     body: tuple["Node", ...] | list["Node"]
@@ -125,7 +132,7 @@ class WorkTileLoop:
     skip_if: Callable[..., object] | None = None
 
 
-Node = Step | ConditionalBlock | DomainLoop | WorkTileLoop
+Node = Step | BreakLoop | ConditionalBlock | DomainLoop | WorkTileLoop
 
 
 @dataclass(kw_only=True, frozen=True)
@@ -151,7 +158,7 @@ def _visit_node(node: Node, visitor: Callable[[object], None]) -> None:
 
 
 def _child_nodes(node: Node) -> tuple[Node, ...] | list[Node]:
-    if isinstance(node, Step):
+    if isinstance(node, (Step, BreakLoop)):
         return ()
     return node.body
 
@@ -186,7 +193,26 @@ def validate_queue_advance_placement(
         elif isinstance(node, WorkTileLoop):
             validate_queue_advance_placement(node.body, parent_is_work_tile_loop=True)
         else:
-            validate_queue_advance_placement(node.body)
+            validate_queue_advance_placement(_child_nodes(node))
+
+
+def validate_break_placement(nodes, *, in_domain_loop=False) -> None:
+    for node in nodes:
+        if isinstance(node, BreakLoop) and not in_domain_loop:
+            raise ScheduleError("break_loop() must be called inside domain_loop()")
+        validate_break_placement(
+            _child_nodes(node),
+            in_domain_loop=in_domain_loop or isinstance(node, DomainLoop),
+        )
+
+
+def contains_break(nodes) -> bool:
+    """Whether this scope can exit its enclosing domain loop."""
+    return any(
+        isinstance(node, BreakLoop)
+        or (isinstance(node, ConditionalBlock) and contains_break(node.body))
+        for node in nodes
+    )
 
 
 def _format_guard(guard: BlockGuard) -> str:
@@ -212,6 +238,11 @@ def _format_schedule(schedule: Schedule, *, show_routes: bool = True) -> str:
 
     def emit(node: Node, indent: int) -> None:
         pad = " " * indent
+        if isinstance(node, BreakLoop):
+            values = ", ".join(f"%{value.value_id}" for value in node.exit_values)
+            suffix = f"values=[{values}]" if values and show_routes else ""
+            lines.append(f"{pad}BreakLoop({suffix})")
+            return
         if isinstance(node, Step):
             label = f", label={node.label!r}" if node.label else ""
             lines.append(
@@ -354,13 +385,6 @@ def make_value(
     def _inside(self, kind: type) -> bool:
         return any(isinstance(item, kind) for item in self.stack)
 
-    def _in_condition(self) -> bool:
-        return any(
-            isinstance(item, ConditionalBlock)
-            and isinstance(item.condition, (IterationPredicate, OpaqueCondition))
-            for item in self.stack
-        )
-
     def open_scope(self, node: ConditionalBlock | DomainLoop | WorkTileLoop) -> None:
         if self.after_work_tile_loop:
             raise ScheduleError("no blocks may follow work_tile_loop()")
@@ -383,8 +407,6 @@ def open_scope(self, node: ConditionalBlock | DomainLoop | WorkTileLoop) -> None
                 DomainLoop
             ):
                 raise ScheduleError("iteration predicates require domain_loop()")
-            if self._in_condition():
-                raise ScheduleError("conditional scheduling blocks cannot be nested")
         self.stack[-1].body.append(node)
         self.stack.append(node)
 
@@ -417,8 +439,11 @@ def finalize(self) -> Schedule:
         if len(self.stack) != 1:
             raise ScheduleError("schedule ended with an unclosed block")
         validate_queue_advance_placement(self.schedule.body)
+        validate_break_placement(self.schedule.body)
 
         def freeze(node: Node) -> Node:
+            if isinstance(node, BreakLoop):
+                return node
             if isinstance(node, Step):
                 return replace(
                     node,
@@ -589,10 +614,12 @@ def _conditional(cond: object, key: Hashable | None, negated: bool):
 
 
 def when_true(cond: object, *, key: Hashable | None = None):
+    """Capture a guarded region, evaluated only when all enclosing guards fire."""
     return _conditional(cond, key, False)
 
 
 def when_false(cond: object, *, key: Hashable | None = None):
+    """Capture a negated guarded region; nested values retain lexical scope."""
     return _conditional(cond, key, True)
 
 
@@ -608,6 +635,30 @@ def first_iter():
     return when_true(FIRST_ITER)
 
 
+def break_loop(*values: ScheduleValue) -> None:
+    """Exit the enclosing domain loop when execution reaches this operation.
+
+    Remaining work and the functional loop's backedge are skipped. Explicit
+    values become the loop's results, in carried-input order. With no arguments,
+    results retain their values at entry to the interrupted iteration.
+    """
+    _require_active_domain_loop("break_loop")
+    builder = _require_active_builder("break_loop()")
+    loop = next(node for node in reversed(builder.stack) if isinstance(node, DomainLoop))
+    if values and len(values) != len(loop.initial_values):
+        raise ScheduleError(
+            "break_loop() must supply one value for every carried input; "
+            f"expected {len(loop.initial_values)}, got {len(values)}"
+        )
+    for (name, initial), value in zip(loop.initial_values.items(), values):
+        if not isinstance(value, ScheduleValue):
+            raise ScheduleError(f"break value {name!r} must be a routed schedule value")
+        builder.validate_value_use(value, name)
+        if value.stage_resource is not initial.stage_resource:
+            raise ScheduleError(f"break value {name!r} changes pipeline stage provenance")
+    builder.stack[-1].body.append(BreakLoop(exit_values=values))
+
+
 def last_iter():
     """Capture a block that runs only on the last domain-loop iteration."""
     _require_active_domain_loop("last_iter")
@@ -680,6 +731,12 @@ def first_iter(self):
         self._check()
         return first_iter()
 
+    def break_loop(self, *values: ScheduleValue) -> None:
+        self._check()
+        if _active_builder.get() is not self.builder:
+            raise ScheduleError("domain-loop handle belongs to another schedule")
+        break_loop(*values)
+
     def last_iter(self):
         self._check()
         return last_iter()