fix(ptodsl): normalize integer boolean branch merges (#1126)

ptodsl/docs/user_guide/05-control-flow.md

@@ -213,6 +213,11 @@ conditional closes, `br.val` is the SSA-merged result seen by downstream code.
 This surface avoids explicit result-type declarations and explicit
 `pto.yield_(...)` in user code while still keeping the merge contract explicit.
 
+When one branch yields `pto.i1` and the other yields an integer-like scalar,
+the integer value is normalized to `i1` with nonzero truthiness (`0` is false;
+any nonzero value is true). Other incompatible branch result types still
+require an explicit conversion.
+
 ## 5.4 `pto.const_expr` and tracing
 
 `pto.const_expr` parameters (Section 3.6) are compile-time constants. They are fixed at `.compile()` time and cannot change between launches of the same compiled kernel. Because their values are known during tracing, they interact naturally with Python control flow:
@@ -277,6 +282,27 @@ def ast_rewrite_branch_kernel():
 The assigned value `total` is live after the branch, so PTODSL rewrites the
 branch into a `pto.if_` with automatic merge.
 
+Static list slots can also be live across a rewritten branch. The subscript
+index must be an integer that can be resolved during AST rewriting, including
+compile-time constants and `pto.const_expr` values:
+
+```python
+@pto.jit(target="a5")
+def ast_rewrite_static_slot_kernel(*, SLOT: pto.const_expr = 0):
+    values = [pto.const(0, dtype=pto.i32)]
+
+    if pto.const(1, dtype=pto.i1):
+        values[SLOT] = pto.const(1, dtype=pto.i32)
+
+    if values[SLOT]:
+        pto.pipe_barrier(pto.Pipe.ALL)
+```
+
+The rewritten slot is merged as an `scf.if` result, just like a named scalar
+value. Dynamic indices and container aliases remain unsupported; use an
+explicit `pto.if_`/`pto.for_` state value or a real buffer load/store for those
+cases.
+
 If a live-out value is assigned in only one branch, PTODSL keeps the old value
 on the missing branch:
 

ptodsl/ptodsl/_ast_rewrite.py

@@ -804,6 +804,42 @@ def visit_Subscript(self, node):
         return self.generic_visit(node)
 
 
+class _SlotValueRewriter(ast.NodeTransformer):
+    """Replace selected static list slots with scalar branch state names."""
+
+    def __init__(self, slot_values, static_env, static_iters=None):
+        self._slot_values = dict(slot_values)
+        self._static_env = static_env
+        self._static_iters = dict(static_iters or {})
+
+    def visit_For(self, node):
+        if _is_pto_attr_call(node.iter, "static_range") and isinstance(node.target, ast.Name):
+            values = _try_eval_static_range(node.iter, self._static_env, self._static_iters)
+            old = self._static_iters.get(node.target.id)
+            if values is not None:
+                self._static_iters[node.target.id] = values
+            try:
+                node.body = [self.visit(stmt) for stmt in node.body]
+            finally:
+                if values is not None:
+                    if old is None:
+                        self._static_iters.pop(node.target.id, None)
+                    else:
+                        self._static_iters[node.target.id] = old
+            node.orelse = [self.visit(stmt) for stmt in node.orelse]
+            return node
+        return self.generic_visit(node)
+
+    def visit_Subscript(self, node):
+        slots = _resolve_subscript_slots(node, self._static_env, self._static_iters, require_static=False)
+        if len(slots) == 1:
+            slot = next(iter(slots))
+            value_name = self._slot_values.get(slot)
+            if value_name is not None:
+                return ast.copy_location(_name(value_name, node.ctx), node)
+        return self.generic_visit(node)
+
+
 class _ControlFlowRewriter:
     def __init__(self, static_env=None):
         self._static_env = dict(static_env or {})
@@ -940,16 +976,13 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con
         cond_name = self._fresh("cond")
         then_info = _name_info(stmt.body)
         else_info = _name_info(stmt.orelse)
+        then_slot_info = _slot_info(stmt.body, self._static_env, static_iters)
+        else_slot_info = _slot_info(stmt.orelse, self._static_env, static_iters)
         assigned_slots = (
-            _slot_info(stmt.body, self._static_env, static_iters).stores
-            | _slot_info(stmt.orelse, self._static_env, static_iters).stores
+            then_slot_info.stores
+            | else_slot_info.stores
         )
-        if live_after_slots & assigned_slots:
-            slots = ", ".join(slot.display for slot in sorted(live_after_slots & assigned_slots))
-            raise PTODSLAstRewriteError(
-                "ast_rewrite=True does not support automatic branch merges for static subscript slots yet; "
-                f"rewrite {slots} with explicit scalar temporaries"
-            )
+        merge_slots = tuple(sorted(live_after_slots & assigned_slots))
         assigned_any = then_info.stores | else_info.stores
         merge_names = tuple(sorted(live_after & assigned_any))
         old_value_names = {
@@ -959,17 +992,18 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con
         }
 
         branch_live_after = set(live_after) | set(merge_names)
+        branch_live_after_slots = set(live_after_slots) | set(merge_slots)
         then_body = self.rewrite_block(
             stmt.body,
             live_after=branch_live_after,
-            live_after_slots=live_after_slots,
+            live_after_slots=branch_live_after_slots,
             allow_loop_control=False,
             static_iters=static_iters,
         )
         else_body = self.rewrite_block(
             stmt.orelse,
             live_after=branch_live_after,
-            live_after_slots=live_after_slots,
+            live_after_slots=branch_live_after_slots,
             allow_loop_control=False,
             static_iters=static_iters,
         )
@@ -980,15 +1014,45 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con
         )
         branch_name = self._fresh("br")
 
+        slot_value_names = {
+            # BranchHandle deliberately rejects private attribute names. Keep
+            # the generated branch field public while retaining a unique
+            # compiler-generated local name for the rewritten slot value.
+            slot: (
+                f"pto_ast_slot_{slot.base}_"
+                f"{'neg' if slot.index < 0 else ''}{abs(slot.index)}_{self._counter}"
+            )
+            for slot in merge_slots
+        }
+        self._counter += len(slot_value_names)
+        old_slot_value_names = {
+            slot: self._fresh(
+                f"old_slot_{slot.base}_"
+                f"{'neg' if slot.index < 0 else ''}{abs(slot.index)}"
+            )
+            for slot in merge_slots
+        }
         dynamic_then_body = copy.deepcopy(then_body)
         dynamic_else_body = copy.deepcopy(else_body)
-        if merge_names:
+        if slot_value_names:
+            dynamic_then_body = [
+                _SlotValueRewriter(slot_value_names, self._static_env, static_iters).visit(stmt)
+                for stmt in dynamic_then_body
+            ]
+            dynamic_else_body = [
+                _SlotValueRewriter(slot_value_names, self._static_env, static_iters).visit(stmt)
+                for stmt in dynamic_else_body
+            ]
+        if merge_names or slot_value_names:
             dynamic_then_body.append(
                 self._branch_assign(
                     branch_name,
                     merge_names,
                     old_value_names=old_value_names,
                     assigned_names=then_info.stores,
+                    slot_value_names=slot_value_names,
+                    old_slot_value_names=old_slot_value_names,
+                    assigned_slots=then_slot_info.stores,
                 )
             )
             dynamic_else_body.append(
@@ -997,6 +1061,9 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con
                     merge_names,
                     old_value_names=old_value_names,
                     assigned_names=else_info.stores,
+                    slot_value_names=slot_value_names,
+                    old_slot_value_names=old_slot_value_names,
+                    assigned_slots=else_slot_info.stores,
                 )
             )
 
@@ -1050,6 +1117,17 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con
             )
             for name in merge_names
         )
+        dynamic_body.extend(
+            ast.Assign(
+                targets=[_slot_subscript(slot, ast.Store())],
+                value=ast.Attribute(
+                    value=_name(branch_name),
+                    attr=slot_value_names[slot],
+                    ctx=ast.Load(),
+                ),
+            )
+            for slot in merge_slots
+        )
 
         result = [
             ast.Assign(
@@ -1064,6 +1142,20 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con
             )
             for name, old_name in old_value_names.items()
         )
+        result.extend(
+            ast.Assign(
+                targets=[_name(value_name, ast.Store())],
+                value=_slot_subscript(slot),
+            )
+            for slot, value_name in slot_value_names.items()
+        )
+        result.extend(
+            ast.Assign(
+                targets=[_name(old_value_name, ast.Store())],
+                value=_name(slot_value_names[slot]),
+            )
+            for slot, old_value_name in old_slot_value_names.items()
+        )
         result.append(
             ast.copy_location(
                 ast.If(
@@ -1080,18 +1172,36 @@ def _rewrite_if(self, stmt, *, live_after, live_after_slots=None, allow_loop_con
         )
         return result
 
-    def _branch_assign(self, branch_name, names, *, old_value_names, assigned_names):
+    def _branch_assign(
+        self,
+        branch_name,
+        names,
+        *,
+        old_value_names,
+        assigned_names,
+        slot_value_names=(),
+        old_slot_value_names=(),
+        assigned_slots=(),
+    ):
+        keywords = [
+            ast.keyword(
+                arg=name,
+                value=_name(name if name in assigned_names else old_value_names[name]),
+            )
+            for name in names
+        ]
+        keywords.extend(
+            ast.keyword(
+                arg=value_name,
+                value=_name(value_name if slot in assigned_slots else old_slot_value_names[slot]),
+            )
+            for slot, value_name in slot_value_names.items()
+        )
         return ast.Expr(
             value=ast.Call(
                 func=ast.Attribute(value=_name(branch_name), attr="assign", ctx=ast.Load()),
                 args=[],
-                keywords=[
-                    ast.keyword(
-                        arg=name,
-                        value=_name(name if name in assigned_names else old_value_names[name]),
-                    )
-                    for name in names
-                ],
+                keywords=keywords,
             )
         )
 

ptodsl/ptodsl/_control_flow.py

@@ -24,14 +24,15 @@
 
 from ._diagnostics import explicit_mode_required_with_context_error
 from ._runtime_index_ops import coerce_runtime_index
+from ._scalar_adaptation import coerce_runtime_integer_to_i1
 from ._scalar_coercion import coerce_scalar_to_type
 from ._surface_types import const_expr
 from ._tracing.active import current_session, require_active_session
 from ._surface_values import unwrap_surface_value, wrap_like_surface_value, wrap_surface_value
 from ._types import _StructDescriptor
 
 from ptoas.mlir.dialects import pto as _pto, scf
-from ptoas.mlir.ir import InsertionPoint
+from ptoas.mlir.ir import IndexType, InsertionPoint, IntegerType
 
 
 # ── vecscope ──────────────────────────────────────────────────────────────────
@@ -412,6 +413,11 @@ def __init__(self, cond):
 
     def __enter__(self):
         self._cond_value = unwrap_surface_value(self._cond)
+        if not _is_i1_type(self._cond_value.type) and _is_integer_like_type(self._cond_value.type):
+            self._cond_value = coerce_runtime_integer_to_i1(
+                self._cond_value,
+                context="pto.if_(...) condition",
+            )
         self._tmp_if = scf.IfOp(self._cond_value, hasElse=True)
         self._parent_block = _find_parent_block(self._tmp_if)
         self._handle = BranchHandle(self)
@@ -544,6 +550,8 @@ def _validate_merge_spec(self):
                 name,
                 then_value,
                 else_value,
+                then_block=self._tmp_if.then_block,
+                else_block=self._tmp_if.else_block,
             )
             if then_value.type != else_value.type:
                 raise RuntimeError(
@@ -638,11 +646,39 @@ def _is_branch_assign_literal(value) -> bool:
     return isinstance(value, (int, float)) and not isinstance(value, bool)
 
 
-def _reconcile_branch_assignment_values(name, then_value, else_value):
+def _is_i1_type(type_obj) -> bool:
+    return IntegerType.isinstance(type_obj) and IntegerType(type_obj).width == 1
+
+
+def _is_integer_like_type(type_obj) -> bool:
+    return IndexType.isinstance(type_obj) or IntegerType.isinstance(type_obj)
+
+
+def _coerce_integer_to_i1_at(value, *, block, context):
+    with InsertionPoint(block):
+        return coerce_runtime_integer_to_i1(value, context=context)
+
+
+def _reconcile_branch_assignment_values(name, then_value, else_value, *, then_block=None, else_block=None):
     then_is_typed = hasattr(then_value, "type")
     else_is_typed = hasattr(else_value, "type")
 
     if then_is_typed and else_is_typed:
+        then_is_i1 = _is_i1_type(then_value.type)
+        else_is_i1 = _is_i1_type(else_value.type)
+        if then_is_i1 != else_is_i1:
+            if then_is_i1 and _is_integer_like_type(else_value.type):
+                else_value = _coerce_integer_to_i1_at(
+                    else_value,
+                    block=else_block,
+                    context=f"br.assign(...) else branch value for '{name}'",
+                )
+            elif else_is_i1 and _is_integer_like_type(then_value.type):
+                then_value = _coerce_integer_to_i1_at(
+                    then_value,
+                    block=then_block,
+                    context=f"br.assign(...) then branch value for '{name}'",
+                )
         return then_value, else_value
     if then_is_typed:
         return then_value, coerce_scalar_to_type(

ptodsl/ptodsl/_scalar_adaptation.py

@@ -121,6 +121,26 @@ def coerce_runtime_integer_value(value, target_type, *, context: str):
     return coerce_integer_like(value, target_type)
 
 
+def coerce_runtime_integer_to_i1(value, *, context: str):
+    """Convert one runtime integer-like value to ``i1`` using nonzero truthiness."""
+    if not hasattr(value, "type"):
+        raise TypeError(f"{context} expects an integer-like runtime scalar, got {value!r}")
+
+    if IndexType.isinstance(value.type):
+        zero = arith.ConstantOp(IndexType.get(), 0).result
+        return arith.CmpIOp(arith.CmpIPredicate.ne, value, zero).result
+
+    if not IntegerType.isinstance(value.type):
+        raise TypeError(f"{context} expects an integer-like runtime scalar, got {value.type}")
+
+    signless_type = _signless_integer_type(value.type)
+    signless_value = _strip_integer_signedness(value)
+    if IntegerType(signless_type).width == 1:
+        return signless_value
+    zero = arith.ConstantOp(signless_type, 0).result
+    return arith.CmpIOp(arith.CmpIPredicate.ne, signless_value, zero).result
+
+
 def coerce_runtime_i1_value(value, *, context: str):
     """Normalize one authored bool/integer-like value/literal to signless i1."""
     i1_type = IntegerType.get_signless(1)
@@ -136,7 +156,7 @@ def coerce_runtime_i1_value(value, *, context: str):
     kind = classify_runtime_scalar_type(value.type)
     if kind == "float":
         raise TypeError(f"{context} expects a bool or integer-like scalar, got {value.type}")
-    return coerce_integer_like(value, i1_type)
+    return coerce_runtime_integer_to_i1(value, context=context)
 
 
 def normalize_runtime_binary_operands(lhs, rhs):
@@ -272,6 +292,7 @@ def _float_bytewidth(type_obj):
     "classify_runtime_scalar_type",
     "coerce_integer_like",
     "coerce_runtime_i1_value",
+    "coerce_runtime_integer_to_i1",
     "coerce_runtime_index_value",
     "coerce_runtime_integer_value",
     "coerce_scalar_value_to_type",