[mlir] call target materialization more in dialect conversion

[ftynse]

Summary

During dialect conversion, target materialization is triggered to create
cast-like operations when a type mismatch occurs between the value that
replaces a rewritten operation and the type that another operations expects as
operands processed by the type conversion. First, a dummy cast is inserted to
make sure the pattern application can proceed. The decision to trigger the
user-provided materialization hook is taken later based on the result of the
dummy cast having uses. However, it only has uses if other patterns constructed
new operations using the casted value as operand. If existing (legal)
operations use the replaced value, they may have not been updated to use the
casted value yet. The conversion infra would then delete the dummy cast first,
and then would replace the uses with now-invalid (null in the bast case) value.
When deciding whether to trigger cast materialization, check for liveness the
uses not only of the casted value, but also of all the values that it replaces.

This was discovered in the finalizing bufferize pass that cleans up
mutually-cancelling casts without touching other operations. It is not
impossible that there are other scenarios where the dialect converison infra
could produce invalid operand uses because of dummy casts erased too eagerly.

[ftynse]

Forgot FileCheck.

[springerm]

Thanks for fixing this!

mlir/lib/Transforms/Utils/DialectConversion.cpp

static void computeNecessaryMaterializations(
     SetVector<UnresolvedMaterialization *> &necessaryMaterializations) {
   auto isLive = [&](Value value) {
     auto findFn = [&](Operation *user) {
       auto matIt = materializationOps.find(user);
       if (matIt != materializationOps.end())
         return !necessaryMaterializations.count(matIt->second);
       return rewriterImpl.isOpIgnored(user);
     };
+    // This value may be replacing another value that has a live user.
+    for (Value inv : inverseMapping.lookup(value))
+      if (llvm::find_if_not(inv.getUsers(), findFn) != inv.user_end())
+        return true;
+    // Or have live users itself.
     return llvm::find_if_not(value.getUsers(), findFn) != value.user_end();
   };
 
   llvm::unique_function<Value(Value, Value, Type)> lookupRemappedValue =
       [&](Value invalidRoot, Value value, Type type) {
         // Check to see if the input operation was remapped to a variant of the
         // output.
         Value remappedValue = rewriterImpl.mapping.lookupOrDefault(value, type);

mlir/test/Transforms/test-legalize-target-materialization-no-uses.mlir

+// RUN: mlir-opt -test-target-materialization-with-no-uses %s | FileCheck %s
+
+// The conversion is set up as follows:
+// - type_changer ops are illegal;
+// - type_changer ops are replaced with their operands;
+// - i16 types are converted to i64 by the type conversion;
+// - the rest of the types are legal.
+// The first type_changer is replaced with its operand. For the pattern to
+// apply to the second type_changer, the conversion infra creates a dummy
+// cast operation to cast from the i32 to i64 because the original op takes an
+// (illegal) i16 that became i64. This dummy operation should be replaced by
+// the one produced by the target materialization hook. At the moment when the
+// materialization decision is taken, the i64 replacement of the first type
+// change (the result of the dummy cast) has no uses, but the value it replaces
+// does, so the infra must call the materialization rather than assume the
+// dummy cast to be dead.
+
+// CHECK-LABEL: @foo
+func @foo() {
+  %0 = "test.type_producer"() : () -> i32
+  // CHECK: test.cast
+  // CHECK-NOT: test.type_changer
+  %1 = "test.type_changer"(%0) : (i32) -> i16
+  %2 = "test.type_changer"(%1) : (i16) -> i64
+  "test.type_consumer"(%2) : (i64) -> ()
+  return
+}

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

 def TestInvalidOp : TEST_Op<"invalid", [Terminator]>,
   Arguments<(ins Variadic<AnyType>)>;
 def TestTypeProducerOp : TEST_Op<"type_producer">,
   Results<(outs AnyType)>;
 def TestAnotherTypeProducerOp : TEST_Op<"another_type_producer">,
   Results<(outs AnyType)>;
 def TestTypeConsumerOp : TEST_Op<"type_consumer">,
   Arguments<(ins AnyType)>;
+def TestTypeChangerOp : TEST_Op<"type_changer">,
+  Arguments<(ins AnyType)>, Results<(outs AnyType)>;
 def TestValidOp : TEST_Op<"valid", [Terminator]>,
   Arguments<(ins Variadic<AnyType>)>;
 
 def TestMergeBlocksOp : TEST_Op<"merge_blocks"> {
   let summary = "merge_blocks operation";
   let description = [{
     Test op with multiple blocks that are merged with Dialect Conversion
   }];

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

void runOnOperation() override {
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
+// Test Target Materialization With No Uses
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct ForwardOperandPattern : public OpConversionPattern<TestTypeChangerOp> {
+  using OpConversionPattern<TestTypeChangerOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(TestTypeChangerOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const final {
+    rewriter.replaceOp(op, adaptor.getOperands());
+    return success();
+  }
+};
+
+struct TestTargetMaterializationWithNoUses
+    : public PassWrapper<TestTargetMaterializationWithNoUses,
+                         OperationPass<ModuleOp>> {
+  StringRef getArgument() const final {
+    return "test-target-materialization-with-no-uses";
+  }
+  StringRef getDescription() const final {
+    return "Test a special case of target materialization in DialectConversion";
+  }
+
+  void runOnOperation() override {
+    TypeConverter converter;
+    converter.addConversion([](Type t) { return t; });
+    converter.addConversion([](IntegerType intTy) -> Type {
+      if (intTy.getWidth() == 16)
+        return IntegerType::get(intTy.getContext(), 64);
+      return intTy;
+    });
+    converter.addTargetMaterialization(
+        [](OpBuilder &builder, Type type, ValueRange inputs, Location loc) {
+          return builder.create<TestCastOp>(loc, type, inputs).getResult();
+        });
+
+    ConversionTarget target(getContext());
+    target.addIllegalOp<TestTypeChangerOp>();
+
+    RewritePatternSet patterns(&getContext());
+    patterns.add<ForwardOperandPattern>(converter, &getContext());
+
+    if (failed(applyPartialConversion(getOperation(), target,
+                                      std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
 // Test Block Merging
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// A rewriter pattern that tests that blocks can be merged.
 struct TestMergeBlock : public OpConversionPattern<TestMergeBlocksOp> {
   using OpConversionPattern<TestMergeBlocksOp>::OpConversionPattern;
 
PassRegistration<TestLegalizePatternDriver>([] {
     return std::make_unique<TestLegalizePatternDriver>(legalizerConversionMode);
   });
 
   PassRegistration<TestRemappedValue>();
 
   PassRegistration<TestUnknownRootOpDriver>();
 
   PassRegistration<TestTypeConversionDriver>();
+  PassRegistration<TestTargetMaterializationWithNoUses>();
 
   PassRegistration<TestMergeBlocksPatternDriver>();
   PassRegistration<TestSelectiveReplacementPatternDriver>();
 }
 } // namespace test
 } // namespace mlir