Fix pdl-to-smt

Author: math-fehr

xdsl_smt/passes/pdl_to_smt.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass, field
 from typing import Callable, ClassVar, Sequence
-from xdsl.dialects.builtin import ModuleOp, IntegerType, StringAttr, UnitAttr
+from xdsl.dialects.builtin import ModuleOp, IntegerType, StringAttr, UnitAttr, ArrayAttr
 from xdsl.dialects.pdl import (
     ApplyNativeConstraintOp,
     ApplyNativeRewriteOp,
@@ -290,8 +290,12 @@ def match_and_rewrite(self, op: ReplaceOp, rewriter: PatternRewriter):
             replacing_value = replacing_values[0]
 
         if self.refinement is None:
+            replaced_types = op.attributes["__replaced_types"]
+            replacing_types = op.attributes["__replacing_types"]
+            assert isa(replaced_types, ArrayAttr[Attribute])
+            assert isa(replacing_types, ArrayAttr[Attribute])
             refinement = find_refinement_semantics(
-                replaced_value.type, replacing_value.type
+                replaced_types.data[0], replacing_types.data[0]
             )
         else:
             refinement = self.refinement
@@ -529,6 +533,46 @@ def match_and_rewrite(self, op: Operation, rewriter: PatternRewriter):
             ), "Operations used as computations in PDL should not have effects"
 
 
+def annotate_replace_op(pattern: PatternOp):
+    """
+    Annotate all `pdl.replace` operations in the given pattern with the types of the
+    replaced and replacing values.
+    """
+    values_to_types: dict[SSAValue, tuple[Attribute, ...]] = {}
+    for op in pattern.walk():
+        if isinstance(op, OperandOp):
+            if op.value_type is None:
+                raise Exception("Cannot handle non-typed operands")
+            owner = op.value_type.owner
+            assert isinstance(owner, TypeOp)
+            assert owner.constantType is not None
+            values_to_types[op.value] = (owner.constantType,)
+        if isinstance(op, OperationOp):
+            values: list[Attribute] = []
+            for type_value in op.type_values:
+                owner = type_value.owner
+                assert isinstance(owner, TypeOp)
+                assert owner.constantType is not None
+                values.append(owner.constantType)
+            values_to_types[op.op] = tuple(values)
+        if isinstance(op, ResultOp):
+            values_to_types[op.val] = (
+                values_to_types[op.parent_][op.index.value.data],
+            )
+        if isinstance(op, ReplaceOp):
+            replaced_value_types = values_to_types[op.op_value]
+            if len(op.repl_values) != 0:
+                replacing_value_types: list[Attribute] = []
+                for repl_value in op.repl_values:
+                    replacing_value_types.append(values_to_types[repl_value][0])
+                replacing_types = tuple(replacing_value_types)
+            else:
+                assert op.repl_operation is not None
+                replacing_types = values_to_types[op.repl_operation]
+            op.attributes["__replaced_types"] = ArrayAttr(replaced_value_types)
+            op.attributes["__replacing_types"] = ArrayAttr(replacing_types)
+
+
 @dataclass
 class PDLToSMTLowerer:
     native_rewrites: dict[
@@ -621,6 +665,7 @@ def lower_to_smt(self, module: ModuleOp, ctx: Context) -> None:
         new_state_op = DeclareConstOp(StateType())
         Rewriter.insert_op(new_state_op, insert_point)
         rewrite_context = PDLToSMTRewriteContext(new_state_op.res, new_state_op.res)
+        annotate_replace_op(pattern)
 
         walker = PatternRewriteWalker(
             GreedyRewritePatternApplier(