[compiler] isolate mutable state across dynamic if branches

python/flydsl/compiler/ast_rewriter.py

@@ -738,9 +738,11 @@ def scf_if_dispatch(
 
         def _emit_branch(fn, block, label):
             with ir.InsertionPoint(block):
-                branch_result = ReplaceIfWithDispatch._call_branch(fn, result_names, result_values)
+                branch_inputs = _pack_states(state_raw, exemplar)
+                branch_result = ReplaceIfWithDispatch._call_branch(fn, result_names, branch_inputs)
+                branch_map = dict(zip(result_names, branch_inputs))
                 branch_values = ReplaceIfWithDispatch._normalize_branch_result(
-                    branch_result, result_names, result_map, label
+                    branch_result, result_names, branch_map, label
                 )
                 scf.YieldOp(
                     _unpack_branch_outputs(result_names, branch_values, exemplar, label),

tests/system/test_dynamic_controlflow_list_carry_e2e.py

@@ -49,6 +49,20 @@ def _run_if_list(Out: fx.Tensor, flag: fx.Int32, stream: fx.Stream = fx.Stream(N
     _kernel_if_list(Out, flag).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream.value)
 
 
+@flyc.kernel
+def _kernel_if_inplace_list(Out: fx.Tensor, flag: fx.Int32):
+    lst = [fx.Int32(1), fx.Int32(2)]
+    if flag > fx.Int32(0):
+        lst[0] = fx.Int32(10)
+    Out[0] = lst[0]
+    Out[1] = lst[1]
+
+
+@flyc.jit
+def _run_if_inplace_list(Out: fx.Tensor, flag: fx.Int32, stream: fx.Stream = fx.Stream(None)):
+    _kernel_if_inplace_list(Out, flag).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream.value)
+
+
 # ── dynamic for carrying a list ─────────────────────────────────────────────
 
 
@@ -286,6 +300,18 @@ def test_if_list_not_taken(self):
         torch.cuda.synchronize()
         assert out[0].item() == 1 and out[1].item() == 2, out.tolist()
 
+    def test_if_inplace_list_taken(self):
+        out, t_out = _out(2)
+        _run_if_inplace_list(t_out, fx.Int32(1))
+        torch.cuda.synchronize()
+        assert out.tolist() == [10, 2], out.tolist()
+
+    def test_if_inplace_list_not_taken(self):
+        out, t_out = _out(2)
+        _run_if_inplace_list(t_out, fx.Int32(0))
+        torch.cuda.synchronize()
+        assert out.tolist() == [1, 2], out.tolist()
+
     def test_for_list(self):
         out, t_out = _out(2)
         _run_for_list(t_out, fx.Int32(5))  # +1/-1 x5 -> [5, 95]

tests/unit/test_dynamic_controlflow_list_carry.py

@@ -74,6 +74,46 @@ def test_if_carries_list():
         assert "-> (i32, i32)" in text
 
 
+def test_if_inplace_list_updates_do_not_contaminate_branches():
+    with Context(), Location.unknown():
+        module = Module.create()
+        i1 = IntegerType.get_signless(1)
+        with InsertionPoint(module.body):
+            f = func.FuncOp("if_inplace_list", FunctionType.get([i1], []))
+            entry = f.add_entry_block()
+            with InsertionPoint(entry):
+                cond = entry.arguments[0]
+                original_first = _i32(1)
+                lst = [original_first, _i32(2)]
+                branch_inputs = []
+
+                def then_branch(names, branch_lst):
+                    branch_inputs.append(branch_lst)
+                    branch_lst[0] = _i32(10)
+                    return {"lst": branch_lst}
+
+                def else_branch(names, branch_lst):
+                    branch_inputs.append(branch_lst)
+                    assert branch_lst[0].ir_value().owner == original_first.ir_value().owner
+                    return {"lst": branch_lst}
+
+                out = ReplaceIfWithDispatch.scf_if_dispatch(
+                    cond,
+                    then_branch,
+                    else_branch,
+                    result_names=("lst",),
+                    result_values=(lst,),
+                )
+                assert lst[0] is original_first
+                assert branch_inputs[0] is not lst
+                assert branch_inputs[1] is not lst
+                assert branch_inputs[0] is not branch_inputs[1]
+                assert isinstance(out, list) and len(out) == 2
+                func.ReturnOp([])
+
+        assert module.operation.verify()
+
+
 def test_if_carries_nested_list():
     with Context(), Location.unknown():
         module = Module.create()