[mlir][PDL] Infer result types from a `replace` as the last resort

[rriddle]

Summary

This prevents situations where explicit results types were provided,
which have different types than the operation being replaced. This
is useful for supporting dialect conversion, which will have proper
support added in a followup.

[jpienaar]

Would an unrealized cast be needed where the replaced op has different type? (I forget if that is automatically handled).

mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp

888

This is optimistic check isn't it? E.g., if in a block dominated by this one it would allow it, so only skipping on cheap to check case.

[rriddle]

In D133141#3765387, @jpienaar wrote:

Would an unrealized cast be needed where the replaced op has different type? (I forget if that is automatically handled).

Cast insertion is handled automatically by the rewriter (see the follow up commit for a test that has that).

mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp

888

PDL rewrites only have single block regions, so if the op isn't in the rewriter block (i.e. the current block of the operation) we already know it dominates (i.e. it would be in the matcher).

[jpienaar]

mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp

888

Ah, could you expand comment in 885 to capture this?

mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp

rewriteValues[typeOp] = builder.create<pdl_interp::CreateTypesOp>(
         typeOp.getLoc(), typeOp.getType(), typeAttr);
   }
 }
 
 void PatternLowering::generateOperationResultTypeRewriter(
     pdl::OperationOp op, function_ref<Value(Value)> mapRewriteValue,
     SmallVectorImpl<Value> &types, DenseMap<Value, Value> &rewriteValues,
     bool &hasInferredResultTypes) {
-  // Look for an operation that was replaced by `op`. The result types will be
-  // inferred from the results that were replaced.
   Block *rewriterBlock = op->getBlock();
-  for (OpOperand &use : op.op().getUses()) {
-    // Check that the use corresponds to a ReplaceOp and that it is the
-    // replacement value, not the operation being replaced.
-    pdl::ReplaceOp replOpUser = dyn_cast<pdl::ReplaceOp>(use.getOwner());
-    if (!replOpUser || use.getOperandNumber() == 0)
-      continue;
-    // Make sure the replaced operation was defined before this one.
-    Value replOpVal = replOpUser.operation();
-    Operation *replacedOp = replOpVal.getDefiningOp();
-    if (replacedOp->getBlock() == rewriterBlock &&
-        !replacedOp->isBeforeInBlock(op))
-      continue;
-
-    Value replacedOpResults = builder.create<pdl_interp::GetResultsOp>(
-        replacedOp->getLoc(), mapRewriteValue(replOpVal));
-    types.push_back(builder.create<pdl_interp::GetValueTypeOp>(
-        replacedOp->getLoc(), replacedOpResults));
-    return;
-  }
 
   // Try to handle resolution for each of the result types individually. This is
   // preferred over type inferrence because it will allow for us to use existing
   // types directly, as opposed to trying to rebuild the type list.
   OperandRange resultTypeValues = op.types();
   auto tryResolveResultTypes = [&] {
     types.reserve(resultTypeValues.size());
     for (const auto &it : llvm::enumerate(resultTypeValues)) {
if (!resultTypeValues.empty() && succeeded(tryResolveResultTypes()))
     return;
 
   // Otherwise, check if the operation has type inference support itself.
   if (op.hasTypeInference()) {
     hasInferredResultTypes = true;
     return;
   }
 
+  // Look for an operation that was replaced by `op`. The result types will be
+  // inferred from the results that were replaced.
+  for (OpOperand &use : op.op().getUses()) {
+    // Check that the use corresponds to a ReplaceOp and that it is the
+    // replacement value, not the operation being replaced.
+    pdl::ReplaceOp replOpUser = dyn_cast<pdl::ReplaceOp>(use.getOwner());
+    if (!replOpUser || use.getOperandNumber() == 0)
+      continue;
+    // Make sure the replaced operation was defined before this one.
+    Value replOpVal = replOpUser.operation();
+    Operation *replacedOp = replOpVal.getDefiningOp();
+    if (replacedOp->getBlock() == rewriterBlock &&
+        !replacedOp->isBeforeInBlock(op))
+      continue;
+
+    Value replacedOpResults = builder.create<pdl_interp::GetResultsOp>(
+        replacedOp->getLoc(), mapRewriteValue(replOpVal));
+    types.push_back(builder.create<pdl_interp::GetValueTypeOp>(
+        replacedOp->getLoc(), replacedOpResults));
+    return;
+  }
+
   // If the types could not be inferred from any context and there weren't any
   // explicit result types, assume the user actually meant for the operation to
   // have no results.
   if (resultTypeValues.empty())
     return;
 
   // The verifier asserts that the result types of each pdl.operation can be
   // inferred. If we reach here, there is a bug either in the logic above or

mlir/test/Conversion/PDLToPDLInterp/pdl-to-pdl-interp-rewriter.mlir

module @operation_infer_types_from_replaceop {
   // CHECK:     %[[RESULT_TYPES:.*]] = pdl_interp.get_value_type of %[[RESULTS]]
   // CHECK:     pdl_interp.create_operation "foo.op" -> (%[[RESULT_TYPES]] : !pdl.range<type>)
   pdl.pattern : benefit(1) {
     %rootType = type
     %rootType1 = type
     %root = operation "foo.op" -> (%rootType, %rootType1 : !pdl.type, !pdl.type)
     rewrite %root {
       %newType1 = type
-      %newOp = operation "foo.op" -> (%rootType, %newType1 : !pdl.type, !pdl.type)
+      %newOp = operation "foo.op"
       replace %root with %newOp
     }
   }
 }
 
 // -----
 
 // CHECK-LABEL: module @operation_infer_types_from_otherop_individual_results