From 4070f305f9a0c488d7177754d77c0b367e8695bf Mon Sep 17 00:00:00 2001
From: River Riddle <riddleriver@gmail.com>
Date: Fri, 5 Nov 2021 18:43:26 +0000
Subject: [PATCH] [mlir][DialectConversion] Legalize all live argument
 conversions

Previously we didn't materialize conversions for arguments in certain
cases as the implicit type propagation was being heavily relied on
by many patterns. Now that those patterns have been fixed to
properly handle type conversions, we can drop the special behavior.

Differential Revision: https://reviews.llvm.org/D113233
---
 .../Transforms/Utils/DialectConversion.cpp    |  8 -----
 .../test-legalize-type-conversion.mlir        | 14 ++++++++
 mlir/test/lib/Dialect/Test/TestOps.td         |  5 +++
 mlir/test/lib/Dialect/Test/TestPatterns.cpp   | 36 ++++++++++++++++++-
 4 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index d90439b6d12d..cea4b2aaf80b 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -682,14 +682,6 @@ LogicalResult ArgConverter::materializeLiveConversions(
 
     // Process the remapping for each of the original arguments.
     for (unsigned i = 0, e = origBlock->getNumArguments(); i != e; ++i) {
-      // FIXME: We should run the below checks even if a type converter wasn't
-      // provided, but a lot of existing lowering rely on the block argument
-      // being blindly replaced. We should rework argument materialization to be
-      // more robust for temporary source materializations, update existing
-      // patterns, and remove these checks.
-      if (!blockInfo.converter && blockInfo.argInfo[i])
-        continue;
-
       // If the type of this argument changed and the argument is still live, we
       // need to materialize a conversion.
       BlockArgument origArg = origBlock->getArgument(i);
diff --git a/mlir/test/Transforms/test-legalize-type-conversion.mlir b/mlir/test/Transforms/test-legalize-type-conversion.mlir
index 4887a87d0156..cfb09f8f272a 100644
--- a/mlir/test/Transforms/test-legalize-type-conversion.mlir
+++ b/mlir/test/Transforms/test-legalize-type-conversion.mlir
@@ -98,3 +98,17 @@ func @test_block_argument_not_converted() {
   }) : () -> ()
   return
 }
+
+// -----
+
+// Make sure argument type changes aren't implicitly forwarded.
+func @test_signature_conversion_no_converter() {
+  "test.signature_conversion_no_converter"() ({
+  // expected-error@below {{failed to materialize conversion for block argument #0 that remained live after conversion}}
+  ^bb0(%arg0: f32):
+    // expected-note@below {{see existing live user here}}
+    "test.type_consumer"(%arg0) : (f32) -> ()
+    "test.return"(%arg0) : (f32) -> ()
+  }) : () -> ()
+  return
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 656ec7ef86b6..c2b408244d08 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1520,6 +1520,11 @@ def TestSignatureConversionUndoOp : TEST_Op<"signature_conversion_undo"> {
   let regions = (region AnyRegion);
 }
 
+def TestSignatureConversionNoConverterOp
+  : TEST_Op<"signature_conversion_no_converter"> {
+  let regions = (region AnyRegion);
+}
+
 //===----------------------------------------------------------------------===//
 // Test parser.
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index c93d233900a2..4df5730b282f 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -950,6 +950,34 @@ struct TestSignatureConversionUndo
   }
 };
 
+/// Call signature conversion without providing a type converter to handle
+/// materializations.
+struct TestTestSignatureConversionNoConverter
+    : public OpConversionPattern<TestSignatureConversionNoConverterOp> {
+  TestTestSignatureConversionNoConverter(TypeConverter &converter,
+                                         MLIRContext *context)
+      : OpConversionPattern<TestSignatureConversionNoConverterOp>(context),
+        converter(converter) {}
+
+  LogicalResult
+  matchAndRewrite(TestSignatureConversionNoConverterOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const final {
+    Region &region = op->getRegion(0);
+    Block *entry = &region.front();
+
+    // Convert the original entry arguments.
+    TypeConverter::SignatureConversion result(entry->getNumArguments());
+    if (failed(
+            converter.convertSignatureArgs(entry->getArgumentTypes(), result)))
+      return failure();
+    rewriter.updateRootInPlace(
+        op, [&] { rewriter.applySignatureConversion(&region, result); });
+    return success();
+  }
+
+  TypeConverter &converter;
+};
+
 /// Just forward the operands to the root op. This is essentially a no-op
 /// pattern that is used to trigger target materialization.
 struct TestTypeConsumerForward
@@ -1041,11 +1069,17 @@ struct TestTypeConversionDriver
       // Allow casts from F64 to F32.
       return (*op.operand_type_begin()).isF64() && op.getType().isF32();
     });
+    target.addDynamicallyLegalOp<TestSignatureConversionNoConverterOp>(
+        [&](TestSignatureConversionNoConverterOp op) {
+          return converter.isLegal(op.getRegion().front().getArgumentTypes());
+        });
 
     // Initialize the set of rewrite patterns.
     RewritePatternSet patterns(&getContext());
     patterns.add<TestTypeConsumerForward, TestTypeConversionProducer,
-                 TestSignatureConversionUndo>(converter, &getContext());
+                 TestSignatureConversionUndo,
+                 TestTestSignatureConversionNoConverter>(converter,
+                                                         &getContext());
     patterns.add<TestTypeConversionAnotherProducer>(&getContext());
     mlir::populateFuncOpTypeConversionPattern(patterns, converter);