Keep self-reading complete copies in program order

ignition/preparation/compilation.py

@@ -502,8 +502,10 @@ def complete_write(arg):
         mutable_keys = {key_for(arg) for arg in mutated}
         # Opaque operands retain their prepared alignment and contiguous layout.
         source = node.args[1].meta['val'] if target is torch.ops.aten.copy_.default else None
+        # A source that reads the destination region would be gathered in place by the writeback.
         if (target is torch.ops.aten.copy_.default
                 and source.is_contiguous() and source.storage_offset() == 0
+                and mutable_keys.isdisjoint(accesses[node.args[1]])
                 and not any(last_write.get(key, -1) > positions[node] for key in accesses[node.args[1]])
                 and all(complete_write(arg) for arg in mutated)):
             continue

tests/preparation/test_invocation_views.py

@@ -121,6 +121,19 @@ def test_copy_lowering_preserves_destination_for_offset_or_strided_sources(offse
     assert any(call.target is torch.ops.aten.copy_.default for call in calls) is opaque
 
 
+@pytest.mark.parametrize('overlapping', [False, True])
+def test_copy_lowering_preserves_destination_for_sources_gathered_from_it(overlapping):
+    from ignition.preparation.compilation import ForwardTracer, _lower_storage
+
+    output = torch.zeros((4, 8), dtype=torch.int32)
+    physical = output[:2]
+    destination = output if overlapping else torch.zeros_like(output)
+    indices = torch.zeros(4, dtype=torch.int64)
+    module = ForwardTracer()(lambda: destination.copy_(physical[indices]))
+    calls = _lower_storage(module)
+    assert any(call.target is torch.ops.aten.copy_.default for call in calls) is overlapping
+
+
 class CountedShortConvolution:
     def __init__(self, session):
         from b12x.sequence import short_conv

tests/test_compilation.py

@@ -664,6 +664,32 @@ def forward():
         torch.testing.assert_close(snapshot, torch.full_like(snapshot, value), rtol=0, atol=0)
 
 
+@pytest.mark.skipif(os.environ.get("FLARE_RUN_CUDA_TESTS") != "1", reason="explicit CUDA assignment required")
+@torch.no_grad()
+def test_row_expansion_reads_overlapping_rows_before_writing():
+    # Packed rows expand into sample slots of the same buffer, as vocabulary and hidden heads do.
+    rows, width = 6, 129280
+    output = torch.empty((2 * rows, width), device="cuda")
+    physical = output[:rows]
+    source = torch.zeros_like(physical)
+    indices = torch.zeros(2 * rows, dtype=torch.int64, device="cuda")
+
+    def forward():
+        physical.copy_(source)
+        output.copy_(physical[indices])
+
+    compiled = compile_forward(forward)
+    compiled()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        compiled()
+    for shift in range(1, rows):
+        source.copy_(torch.randn_like(source))
+        indices.copy_((torch.arange(2 * rows, device="cuda") + shift) % rows)
+        graph.replay()
+        torch.testing.assert_close(output, source[indices], rtol=0, atol=0)
+
+
 @pytest.mark.skipif(os.environ.get("FLARE_RUN_CUDA_TESTS") != "1", reason="explicit CUDA assignment required")
 @pytest.mark.parametrize("operation", ("unbind", "split"))
 @torch.no_grad()