[Mlir-commits] [mlir] b69f10f - [mlir][PDL] Infer result types from a `replace` as the last resort

[River Riddle]

Author: River Riddle
Date: 2022-09-08T11:35:31-07:00
New Revision: b69f10f5a118ec6efe5edfae2deb376b53877170

URL: https://github.com/llvm/llvm-project/commit/b69f10f5a118ec6efe5edfae2deb376b53877170
DIFF: https://github.com/llvm/llvm-project/commit/b69f10f5a118ec6efe5edfae2deb376b53877170.diff

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

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.

Differential Revision: https://reviews.llvm.org/D133141

Added: 
    

Modified: 
    mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp
    mlir/test/Conversion/PDLToPDLInterp/pdl-to-pdl-interp-rewriter.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp b/mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp
index 42ea3c9dc395b..e4cce8d21f779 100644
--- a/mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp
+++ b/mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp
@@ -835,28 +835,7 @@ 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
@@ -895,6 +874,31 @@ void PatternLowering::generateOperationResultTypeRewriter(
     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. 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's in the matcher).
+    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.

diff  --git a/mlir/test/Conversion/PDLToPDLInterp/pdl-to-pdl-interp-rewriter.mlir b/mlir/test/Conversion/PDLToPDLInterp/pdl-to-pdl-interp-rewriter.mlir
index 8d1859692a305..d6e8f4ab322e4 100644
--- a/mlir/test/Conversion/PDLToPDLInterp/pdl-to-pdl-interp-rewriter.mlir
+++ b/mlir/test/Conversion/PDLToPDLInterp/pdl-to-pdl-interp-rewriter.mlir
@@ -86,7 +86,7 @@ module @operation_infer_types_from_replaceop {
     %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
     }
   }