api: add replace_by_if helper to control ssa value replacement

tests/test_ir.py

@@ -985,3 +985,17 @@ class MyOperation(Operation):
         name = "dialect.op"
 
     assert MyOperation.dialect_name() == "dialect"
+
+
+def test_replace_by_if():
+    a = TestSSAValue(i32)
+    b = test.TestOp((a,))
+    c = test.TestOp((a,))
+
+    assert set(u.operation for u in a.uses) == {b, c}
+
+    d = TestSSAValue(i32)
+    a.replace_by_if(d, lambda u: u.operation is not c)
+
+    assert set(u.operation for u in a.uses) == {c}
+    assert set(u.operation for u in d.uses) == {b}

xdsl/ir/core.py

@@ -2,7 +2,15 @@
 
 import re
 from abc import ABC, abstractmethod
-from collections.abc import Hashable, Iterable, Iterator, Mapping, Reversible, Sequence
+from collections.abc import (
+    Callable,
+    Hashable,
+    Iterable,
+    Iterator,
+    Mapping,
+    Reversible,
+    Sequence,
+)
 from dataclasses import dataclass, field
 from io import StringIO
 from itertools import chain
@@ -161,7 +169,19 @@ def replace_by(self, value: SSAValue) -> None:
         # carry over name if possible
         if value.name_hint is None:
             value.name_hint = self.name_hint
-        assert len(self.uses) == 0, "unexpected error in xdsl"
+        assert not self.uses, "unexpected error in xdsl"
+
+    def replace_by_if(self, value: SSAValue, test: Callable[[Use], bool]):
+        """
+        Replace the value by another value in all its uses that pass the given test
+        function.
+        """
+        for use in self.uses.copy():
+            if test(use):
+                use.operation.operands[use.index] = value
+        # carry over name if possible
+        if value.name_hint is None:
+            value.name_hint = self.name_hint
 
     def erase(self, safe_erase: bool = True) -> None:
         """

xdsl/transforms/convert_memref_stream_to_snitch_stream.py

@@ -185,8 +185,9 @@ def match_and_rewrite(
                 cast_op := builtin.UnrealizedConversionCastOp.get((arg,), (arg.type,)),
                 InsertPoint.at_start(new_body),
             )
-            arg.replace_by(cast_op.results[0])
-            cast_op.operands = (arg,)
+            arg.replace_by_if(
+                cast_op.results[0], lambda use: use.operation is not cast_op
+            )
             rewriter.modify_value_type(arg, stream_type)
 
 

xdsl/transforms/convert_qref_to_qssa.py

@@ -30,17 +30,11 @@ def match_and_rewrite(self, op: QRefBase, rewriter: PatternRewriter):
         if not new_op.is_gate:
             return
 
-        # For gates we replace any other occurences of the original operands
-        # with the results of the new operation
+        # For gates we replace any other occurences of the original operands with the
+        # results of the new operation, except for when used by the new operation.
 
-        old_operands = tuple(new_op.operands)
-
-        # We do this by replacing all uses of the operand...
         for operand, result in zip(op.operands, new_op.results):
-            operand.replace_by(result)
-
-        # and then resetting the operands of the new operation
-        new_op.operands = old_operands
+            operand.replace_by_if(result, lambda use: use.operation is not new_op)
 
 
 class ConvertQRefToQssa(ModulePass):