transformations: (pdl) Different type matching between polymorphic an…

…d fixed case (#3405)

Co-authored-by: Sasha Lopoukhine <superlopuh@gmail.com>

tests/filecheck/interpreters/test_pdl_match_type.mlir

@@ -0,0 +1,26 @@
+// RUN: xdsl-opt %s -p apply-pdl | filecheck %s
+
+%x = arith.constant 42: i32
+%y = arith.constant 84: i64
+
+pdl.pattern : benefit(1) {
+    %in_type = pdl.type: i32
+    %value = pdl.attribute: %in_type
+    %constant_op = pdl.operation "arith.constant" {"value" = %value} -> (%in_type: !pdl.type)
+    pdl.rewrite %constant_op {
+      pdl.erase %constant_op
+    }
+}
+
+// CHECK:       builtin.module {
+// CHECK-NEXT:    %y = arith.constant 84 : i64
+// CHECK-NEXT:    pdl.pattern : benefit(1) {
+// CHECK-NEXT:      %in_type = pdl.type : i32
+// CHECK-NEXT:      %value = pdl.attribute : %in_type
+// CHECK-NEXT:      %constant_op = pdl.operation "arith.constant" {"value" = %value} -> (%in_type : !pdl.type)
+// CHECK-NEXT:      pdl.rewrite %constant_op {
+// CHECK-NEXT:        pdl.erase %constant_op
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// CHECK-NEXT:  

tests/filecheck/interpreters/test_pdl_polymorphic.mlir

@@ -0,0 +1,25 @@
+// RUN: xdsl-opt %s -p apply-pdl | filecheck %s
+
+%x = arith.constant 42: i32
+%y = arith.constant 84: i64
+
+pdl.pattern : benefit(1) {
+    %in_type = pdl.type
+    %value = pdl.attribute
+    %constant_op = pdl.operation "arith.constant" {"value" = %value} -> (%in_type: !pdl.type)
+    pdl.rewrite %constant_op {
+      pdl.erase %constant_op
+    }
+}
+
+// CHECK:       builtin.module {
+// CHECK-NEXT:    pdl.pattern : benefit(1) {
+// CHECK-NEXT:      %in_type = pdl.type
+// CHECK-NEXT:      %value = pdl.attribute
+// CHECK-NEXT:      %constant_op = pdl.operation "arith.constant" {"value" = %value} -> (%in_type : !pdl.type)
+// CHECK-NEXT:      pdl.rewrite %constant_op {
+// CHECK-NEXT:        pdl.erase %constant_op
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// CHECK-NEXT:  

tests/interpreters/test_pdl_interpreter.py

@@ -482,6 +482,48 @@ def even_length_string(attr: Attribute) -> bool:
     assert new_input_module_true.is_structurally_equivalent(input_module_true)
 
 
+def test_match_type():
+    matcher = PDLMatcher()
+
+    pdl_op = pdl.TypeOp()
+    ssa_value = pdl_op.result
+    xdsl_value = StringAttr("a")
+
+    # New value
+    assert matcher.match_type(ssa_value, pdl_op, xdsl_value)
+    assert matcher.matching_context == {ssa_value: xdsl_value}
+
+    # Same value
+    assert matcher.match_type(ssa_value, pdl_op, xdsl_value)
+    assert matcher.matching_context == {ssa_value: xdsl_value}
+
+    # Other value
+    assert not matcher.match_type(ssa_value, pdl_op, StringAttr("b"))
+    assert matcher.matching_context == {ssa_value: xdsl_value}
+
+
+def test_match_fixed_type():
+    matcher = PDLMatcher()
+
+    pdl_op = pdl.TypeOp(IntegerType(32))
+    xdsl_value = IntegerType(32)
+    ssa_value = pdl_op.result
+
+    assert matcher.match_type(ssa_value, pdl_op, xdsl_value)
+    assert matcher.matching_context == {ssa_value: xdsl_value}
+
+
+def test_not_match_fixed_type():
+    matcher = PDLMatcher()
+
+    pdl_op = pdl.TypeOp(IntegerType(64))
+    xdsl_value = IntegerType(32)
+    ssa_value = pdl_op.result
+
+    assert not matcher.match_type(ssa_value, pdl_op, xdsl_value)
+    assert matcher.matching_context == {}
+
+
 def test_native_constraint_constant_parameter():
     """
     Check that `pdl.apply_native_constraint` can take constant attribute parameters

xdsl/interpreters/pdl.py

@@ -109,9 +109,11 @@ def match_type(self, ssa_val: SSAValue, pdl_op: pdl.TypeOp, xdsl_attr: Attribute
         if ssa_val in self.matching_context:
             return self.matching_context[ssa_val] == xdsl_attr
 
-        self.matching_context[ssa_val] = xdsl_attr
-
-        return True
+        if pdl_op.constantType is None or pdl_op.constantType == xdsl_attr:
+            self.matching_context[ssa_val] = xdsl_attr
+            return True
+        else:
+            return False
 
     def match_attribute(
         self,