[MLIR][PDL] Generalize result type verification

[sfuniak]

Summary

Presently the result type verification checks if the type is used by a pdl::OperationOp inside the matcher. This is unnecessarily restrictive; the type could come from a pdl::OperandOp or pdl::OperandsOp` and still be inferrable.

[sfuniak]

Rebased.

mlir/lib/Dialect/PDL/IR/PDL.cpp

for (const auto &it : llvm::enumerate(resultTypes)) {
     Operation *resultTypeOp = it.value().getDefiningOp();
     assert(resultTypeOp && "expected valid result type operation");
 
     // If the op was defined by a `apply_native_rewrite`, it is guaranteed to be
     // usable.
     if (isa<ApplyNativeRewriteOp>(resultTypeOp))
       continue;
 
-    // If the type operation was defined in the matcher and constrains the
-    // result of an input operation, it can be used.
-    auto constrainsInputOp = [rewriterBlock](Operation *user) {
-      return user->getBlock() != rewriterBlock && isa<OperationOp>(user);
+    // If the type operation was defined in the matcher and constrains an
+    // operand or the result of an input operation, it can be used.
+    auto constrainsInput = [rewriterBlock](Operation *user) {
+      return user->getBlock() != rewriterBlock &&
+             isa<OperandOp, OperandsOp, OperationOp>(user);
     };
     if (TypeOp typeOp = dyn_cast<TypeOp>(resultTypeOp)) {
-      if (typeOp.type() || llvm::any_of(typeOp->getUsers(), constrainsInputOp))
+      if (typeOp.type() || llvm::any_of(typeOp->getUsers(), constrainsInput))
         continue;
     } else if (TypesOp typeOp = dyn_cast<TypesOp>(resultTypeOp)) {
-      if (typeOp.types() || llvm::any_of(typeOp->getUsers(), constrainsInputOp))
+      if (typeOp.types() || llvm::any_of(typeOp->getUsers(), constrainsInput))
         continue;
     }
 
     return op
         .emitOpError("must have inferable or constrained result types when "
                      "nested within `pdl.rewrite`")
         .attachNote()
         .append("result type #", it.index(), " was not constrained");

mlir/test/Dialect/PDL/ops.mlir

pdl.rewrite %root {
     %newOp = pdl.operation "foo.op" -> (%type1, %type3 : !pdl.type, !pdl.type)
     pdl.replace %root with %newOp
   }
 }
 
 // -----
 
 // Check that the result type of an operation within a rewrite can be inferred
-// from types used within the match block.
+// from the result types of an operation within the match block.
 pdl.pattern @infer_type_from_type_used_in_match : benefit(1) {
   %type1 = pdl.type : i32
   %type2 = pdl.type
   %root = pdl.operation -> (%type1, %type2 : !pdl.type, !pdl.type)
   pdl.rewrite %root {
     %newOp = pdl.operation "foo.op" -> (%type1, %type2 : !pdl.type, !pdl.type)
   }
 }
 
 // -----
 
 // Check that the result type of an operation within a rewrite can be inferred
-// from types used within the match block.
+// from the result types of an operation within the match block.
 pdl.pattern @infer_type_from_type_used_in_match : benefit(1) {
   %types = pdl.types
   %root = pdl.operation -> (%types : !pdl.range<type>)
   pdl.rewrite %root {
     %otherTypes = pdl.types : [i32, i64]
     %newOp = pdl.operation "foo.op" -> (%types, %otherTypes : !pdl.range<type>, !pdl.range<type>)
   }
 }
 
 // -----
 
+// Check that the result type of an operation within a rewrite can be inferred
+// from the type of an operand within the match block.
+pdl.pattern @infer_type_from_type_used_in_match : benefit(1) {
+  %type1 = pdl.type
+  %type2 = pdl.type
+  %operand1 = pdl.operand : %type1
+  %operand2 = pdl.operand : %type2
+  %root = pdl.operation (%operand1, %operand2 : !pdl.value, !pdl.value)
+  pdl.rewrite %root {
+    %newOp = pdl.operation "foo.op" -> (%type1, %type2 : !pdl.type, !pdl.type)
+  }
+}
+
+// -----
+
+// Check that the result type of an operation within a rewrite can be inferred
+// from the types of operands within the match block.
+pdl.pattern @infer_type_from_type_used_in_match : benefit(1) {
+  %types = pdl.types
+  %operands = pdl.operands : %types
+  %root = pdl.operation (%operands : !pdl.range<value>)
+  pdl.rewrite %root {
+    %newOp = pdl.operation "foo.op" -> (%types : !pdl.range<type>)
+  }
+}
+
+// -----
+
 pdl.pattern @apply_rewrite_with_no_results : benefit(1) {
   %root = pdl.operation
   pdl.rewrite %root {
     pdl.apply_native_rewrite "NativeRewrite"(%root : !pdl.operation)
   }
 }
 
 // -----
 
 pdl.pattern @attribute_with_dict : benefit(1) {
   %root = pdl.operation
   pdl.rewrite %root {
     %attr = pdl.attribute {some_unit_attr} attributes {pdl.special_attribute}
     pdl.apply_native_rewrite "NativeRewrite"(%attr : !pdl.attribute)
   }
 }