[mlir][DialectConversion] Legalize all live argument conversions

[rriddle]

Summary

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.

[ftynse]

Nice, thanks!

mlir/lib/Transforms/Utils/DialectConversion.cpp

LogicalResult ArgConverter::materializeLiveConversions(
     function_ref<Operation *(Value)> findLiveUser) {
   for (auto &info : conversionInfo) {
     Block *newBlock = info.first;
     ConvertedBlockInfo &blockInfo = info.second;
     Block *origBlock = blockInfo.origBlock;
 
     // 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);
       if (mapping.lookupOrNull(origArg, origArg.getType()))
         continue;
       Operation *liveUser = findLiveUser(origArg);
       if (!liveUser)
         continue;

mlir/test/Transforms/test-legalize-type-conversion.mlir

func @test_block_argument_not_converted() {
   "test.unsupported_block_arg_type"() ({
     // NOTE: The test pass does not convert `index` types.
     // CHECK: ^bb0({{.*}}: index):
     ^bb0(%0 : index):
       "test.return"(%0) : (index) -> ()
   }) : () -> ()
   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
+}

mlir/test/lib/Dialect/Test/TestOps.td

def TestRemappedValueRegionOp : TEST_Op<"remapped_value_region",
   let regions = (region SizedRegion<1>:$body);
   let results = (outs Variadic<AnyType>:$result);
 }
 
 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.
 //===----------------------------------------------------------------------===//
 
 def ParseIntegerLiteralOp : TEST_Op<"parse_integer_literal"> {
   let results = (outs Variadic<Index>:$results);
   let parser = [{ return ::parse$cppClass(parser, result); }];
   let printer = [{ return ::print(p, *this); }];

mlir/test/lib/Dialect/Test/TestPatterns.cpp

struct TestSignatureConversionUndo
   LogicalResult
   matchAndRewrite(TestSignatureConversionUndoOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
     (void)rewriter.convertRegionTypes(&op->getRegion(0), *getTypeConverter());
     return failure();
   }
 };
 
+/// 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
     : public OpConversionPattern<TestTypeConsumerOp> {
   using OpConversionPattern<TestTypeConsumerOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(TestTypeConsumerOp op, OpAdaptor adaptor,
void runOnOperation() override {
     target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
       return converter.isSignatureLegal(op.getType()) &&
              converter.isLegal(&op.getBody());
     });
     target.addDynamicallyLegalOp<TestCastOp>([&](TestCastOp op) {
       // 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);
 
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };