fix(ptodsl): preserve sequential conditional values

ptodsl/ptodsl/_ast_rewrite.py

@@ -140,11 +140,25 @@ def _visit_section_body(self, stmts):
         old_env = self._env
         old_names = self._local_names
         old_outer_bindings = self._section_outer_bindings
+        entry_binding_count = len(self.section_entry_bindings)
         self._env = {}
         self._local_names = _name_info(stmts).stores
         self._section_outer_bindings = set(self._known_bindings)
         try:
-            return [self.visit(stmt) for stmt in stmts]
+            body = [self.visit(stmt) for stmt in stmts]
+            # Materialize outer values under their section-local aliases before
+            # any runtime control flow. Subsequent branch merges can then read
+            # the alias at the current program point instead of always falling
+            # back to the section entry value.
+            entry_bindings = list(self.section_entry_bindings.items())[entry_binding_count:]
+            initializers = [
+                ast.Assign(
+                    targets=[_name(alias, ast.Store())],
+                    value=_name(outer_name),
+                )
+                for alias, outer_name in entry_bindings
+            ]
+            return initializers + body
         finally:
             self._env = old_env
             self._local_names = old_names
@@ -970,13 +984,13 @@ def _fresh(self, prefix: str) -> str:
         self._counter += 1
         return value
 
-    def _section_entry_value(self, name):
+    def _current_value(self, name):
         if name in self._section_uninitialized_aliases:
             raise PTODSLAstRewriteError(
                 "ast_rewrite=True runtime if reads a section-local value before it is initialized; "
                 f"initialize {name!r} before the conditional"
             )
-        return _name(self._section_entry_bindings.get(name, name))
+        return _name(name)
 
     def rewrite_block(self, stmts, *, live_after, live_after_slots=None, allow_loop_control=False, static_iters=None):
         rewritten_reversed = []
@@ -1232,7 +1246,7 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con
         result.extend(
             ast.Assign(
                 targets=[_name(old_name, ast.Store())],
-                value=self._section_entry_value(name),
+                value=self._current_value(name),
             )
             for name, old_name in old_value_names.items()
         )
@@ -1372,7 +1386,7 @@ def _rewrite_for(self, stmt, *, live_after, live_after_slots=None, allow_loop_co
                     keywords=[
                         ast.keyword(
                             arg=name,
-                            value=_name(self._section_entry_bindings.get(name, name)),
+                            value=self._current_value(name),
                         )
                         for name in loop_carried
                     ] + [

ptodsl/tests/test_section.py

@@ -9,6 +9,8 @@
 
 """Focused tracing coverage for explicit physical section hints."""
 
+import re
+
 from ptodsl import pto
 from ptodsl._ast_rewrite import PTODSLAstRewriteError
 from ptodsl._context import make_context
@@ -199,6 +201,34 @@ def lexical_section_sibling_single_sided_conditional_rebinding_probe():
         pto.wait_flag("MTE2", "S", event_id=n_tile_2)
 
 
+@pto.jit(target="a5", mode="explicit")
+def lexical_section_sequential_single_sided_conditional_probe():
+    m_tile = pto.const(0, dtype=pto.i64)
+    n_tile = pto.const(0, dtype=pto.i64)
+    with pto.section("cube"):
+        if pto.get_block_idx() < 16:
+            m_tile = pto.get_block_idx() & 3
+            n_tile = pto.get_block_idx() // 4
+        if 16 <= pto.get_block_idx():
+            m_tile = (pto.get_block_idx() & 3) + 4
+            n_tile = (pto.get_block_idx() // 4) - 4
+        if 15 < pto.get_block_idx():
+            n_tile = 3 - n_tile
+        pto.wait_flag("S", "MTE2", event_id=m_tile + n_tile)
+
+
+@pto.jit(target="a5", mode="explicit")
+def lexical_section_single_sided_read_before_rebinding_probe():
+    value = pto.const(0, dtype=pto.i64)
+    with pto.section("cube"):
+        if pto.get_block_idx() < 16:
+            previous_value = value
+            value = pto.get_block_idx()
+        else:
+            previous_value = value
+        pto.wait_flag("S", "MTE2", event_id=previous_value + value)
+
+
 @pto.jit(target="a5", mode="explicit")
 def lexical_section_uninitialized_conditional_probe():
     one = pto.const(1, dtype=pto.i32)
@@ -342,6 +372,33 @@ def main() -> None:
         module = Module.parse(sibling_single_sided_text, context)
         module.operation.verify()
 
+    sequential_single_sided_text = lexical_section_sequential_single_sided_conditional_probe.compile().mlir_text()
+    if_results = re.findall(r"^\s*(%\d+)(?::\d+)? = scf\.if", sequential_single_sided_text, re.MULTILINE)
+    assert len(if_results) == 3
+    second_if_text = sequential_single_sided_text.split(f"{if_results[1]}:2 = scf.if", 1)[1]
+    second_if_text = second_if_text.split(f"{if_results[2]} = scf.if", 1)[0]
+    assert re.search(
+        rf"else \{{\s+scf\.yield {re.escape(if_results[0])}#0, {re.escape(if_results[0])}#1 : i64, i64",
+        second_if_text,
+    )
+    third_if_text = sequential_single_sided_text.split(f"{if_results[2]} = scf.if", 1)[1]
+    assert re.search(
+        rf"else \{{\s+scf\.yield {re.escape(if_results[1])}#1 : i64",
+        third_if_text,
+    )
+    with make_context() as context:
+        module = Module.parse(sequential_single_sided_text, context)
+        module.operation.verify()
+
+    read_before_rebinding_text = lexical_section_single_sided_read_before_rebinding_probe.compile().mlir_text()
+    assert re.search(
+        r"scf\.yield %c0_i64, %[\d]+ : i64, i64",
+        read_before_rebinding_text,
+    )
+    with make_context() as context:
+        module = Module.parse(read_before_rebinding_text, context)
+        module.operation.verify()
+
     nested_conditional_text = lexical_section_nested_conditional_rebinding_probe.compile().mlir_text()
     assert nested_conditional_text.count("pto.section.cube {") == 1
     assert nested_conditional_text.count("scf.if") == 2