[mlir] Apply source materialization in case of transitive conversion

[ftynse]

Summary

In dialect conversion infrastructure, source materialization applies as part of
the finalization procedure to results of the newly produced operations that
replace previously existing values with values having a different type.
However, such operations may be created to replace operations created in other
patterns. At this point, it is possible that the results of the _original_
operation are still in use and have mismatching types, but the results of the
_intermediate_ operation that performed the type change are not in use leading
to the absence of source materialization. For example,

%0 = dialect.produce : !dialect.A
dialect.use %0 : !dialect.A

can be replaced with

%0 = dialect.other : !dialect.A
%1 = dialect.produce : !dialect.A  // replaced, scheduled for removal
dialect.use %1 : !dialect.A

and then with

%0 = dialect.final : !dialect.B
%1 = dialect.other : !dialect.A    // replaced, scheduled for removal
%2 = dialect.produce : !dialect.A  // replaced, scheduled for removal
dialect.use %2 : !dialect.A

in the same rewriting, but only the %1->%0 replacement is currently considered.

Change the logic in dialect conversion to look up all values that were replaced
by the given value and performing source materialization if any of those values
is still in use with mismatching types. This is performed by computing the
inverse value replacement mapping. This arguably expensive manipulation is
performed only if there were some type-changing replacements. An alternative
could be to consider all replaced operations and not only those that resulted
in type changes, but it would harm pattern-level composability: the pattern
that performed the non-type-changing replacement would have to be made aware of
the type converter in order to call the materialization hook.

[ftynse]

Drop unused stuff

[rriddle]

If I understand correctly, this revision ensures that we check intermediate values (that did not have a result type change) that still have live uses?

As for the cost of computing the inverse mapping, have you checked if there is an impact to the compile time of any of our large type changing conversions?

If there is no compile time hit, I'm fine with this as is for now.

mlir/lib/Transforms/Utils/DialectConversion.cpp

2228–2264

Can you use an Optional and construct it inside of the loop if necessary? i.e. something like ^

[ftynse]

Review

[ftynse]

In D95626#2534522, @rriddle wrote:

If I understand correctly, this revision ensures that we check intermediate values (that did not have a result type change) that still have live uses?

Yes.

As for the cost of computing the inverse mapping, have you checked if there is an impact to the compile time of any of our large type changing conversions?

If there is no compile time hit, I'm fine with this as is for now.

We are not hitting this code path in any conversion before the child commit, which triggers segfaults on the error path without this change. So, with the optional computation of the inverse map you suggested, there is no effect. I tried with a synthetic test consisting of 10k repetitions of the success test, and see around 9% total time increase in mlir-opt. However, profiling indicates that we spend more time creating cast operations than we spend computing the inverse map.

mlir/include/mlir/IR/BlockAndValueMapping.h

template <typename T> T lookup(T from) const {
     auto result = lookupOrNull(from);
     assert(result && "expected 'from' to be contained within the map");
     return result;
   }
 
   /// Clears all mappings held by the mapper.
   void clear() { valueMap.clear(); }
 
+  /// Returns a new mapper containing the inverse mapping.
+  BlockAndValueMapping getInverse() const {
+    BlockAndValueMapping result;
+    for (const auto &pair : valueMap)
+      result.valueMap.try_emplace(pair.second, pair.first);
+    return result;
+  }
+
 private:
   /// Utility lookupOrValue that looks up an existing key or returns the
   /// provided value.
   Block *lookupOrValue(Block *from, Block *value) const {
     auto it = valueMap.find(from);
     return it != valueMap.end() ? reinterpret_cast<Block *>(it->second) : value;
   }
   Value lookupOrValue(Value from, Value value) const {

mlir/lib/Transforms/Utils/DialectConversion.cpp

struct ConversionValueMapping {
   Value lookupOrNull(Value from) const;
 
   /// Map a value to the one provided.
   void map(Value oldVal, Value newVal) { mapping.map(oldVal, newVal); }
 
   /// Drop the last mapping for the given value.
   void erase(Value value) { mapping.erase(value); }
 
+  /// Returns the inverse raw value mapping (without recursive query support).
+  BlockAndValueMapping getInverse() const { return mapping.getInverse(); }
+
 private:
   /// Current value mappings.
   BlockAndValueMapping mapping;
 };
 } // end anonymous namespace
 
 Value ConversionValueMapping::lookupOrDefault(Value from,
                                               Type desiredType) const {
legalizeErasedResult(Operation *op, OpResult result,
                        ConversionPatternRewriterImpl &rewriterImpl);
 
   /// Legalize an operation result that was replaced with a value of a different
   /// type.
   LogicalResult
   legalizeChangedResultType(Operation *op, OpResult result, Value newValue,
                             TypeConverter *replConverter,
                             ConversionPatternRewriter &rewriter,
-                            ConversionPatternRewriterImpl &rewriterImpl);
+                            ConversionPatternRewriterImpl &rewriterImpl,
+                            const BlockAndValueMapping &inverseMapping);
 
   /// The legalizer to use when converting operations.
   OperationLegalizer opLegalizer;
 
   /// The conversion mode to use when legalizing operations.
   OpConversionMode mode;
 
   /// A set of pre-existing operations. When mode == OpConversionMode::Analysis,

 LogicalResult
 OperationConverter::finalize(ConversionPatternRewriter &rewriter) {
   ConversionPatternRewriterImpl &rewriterImpl = rewriter.getImpl();
 
   // Legalize converted block arguments.
   if (failed(legalizeConvertedArgumentTypes(rewriter, rewriterImpl)))
     return failure();
 
+  if (rewriterImpl.operationsWithChangedResults.empty())
+    return success();
+
+  Optional<BlockAndValueMapping> inverseMapping;
+
   // Process requested operation replacements.
   for (unsigned i = 0, e = rewriterImpl.operationsWithChangedResults.size();
        i != e; ++i) {
     unsigned replIdx = rewriterImpl.operationsWithChangedResults[i];
     auto &repl = *(rewriterImpl.replacements.begin() + replIdx);
     for (OpResult result : repl.first->getResults()) {
       Value newValue = rewriterImpl.mapping.lookupOrNull(result);
 
       // If the operation result was replaced with null, all of the uses of this
       // value should be replaced.
       if (!newValue) {
         if (failed(legalizeErasedResult(repl.first, result, rewriterImpl)))
           return failure();
         continue;
       }
 
       // Otherwise, check to see if the type of the result changed.
       if (result.getType() == newValue.getType())
         continue;
 
+      // Compute the inverse mapping only if it is really needed.
+      if (!inverseMapping)
+        inverseMapping = rewriterImpl.mapping.getInverse();
+
       // Legalize this result.
       rewriter.setInsertionPoint(repl.first);
       if (failed(legalizeChangedResultType(repl.first, result, newValue,
                                            repl.second.converter, rewriter,
-                                           rewriterImpl)))
+                                           rewriterImpl, *inverseMapping)))
         return failure();
 
       // Update the end iterator for this loop in the case it was updated
       // when legalizing generated conversion operations.
       e = rewriterImpl.operationsWithChangedResults.size();
     }
   }
   return success();
 }
 
 LogicalResult OperationConverter::legalizeConvertedArgumentTypes(
if (liveUserIt != result.user_end()) {
     diag.attachNote(liveUserIt->getLoc())
         << "found live user of result #" << result.getResultNumber() << ": "
         << *liveUserIt;
     return failure();
   }
   return success();
 }
 
-LogicalResult OperationConverter::legalizeChangedResultType(
-    Operation *op, OpResult result, Value newValue,
-    TypeConverter *replConverter, ConversionPatternRewriter &rewriter,
-    ConversionPatternRewriterImpl &rewriterImpl) {
+/// Finds a user of the given value, or of any other value that the given value
+/// replaced, that was not replaced in the conversion process.
+static Operation *
+findLiveUserOfReplaced(Value value, ConversionPatternRewriterImpl &rewriterImpl,
+                       const BlockAndValueMapping &inverseMapping) {
+  do {
     // Walk the users of this value to see if there are any live users that
     // weren't replaced during conversion.
-  auto liveUserIt = llvm::find_if_not(result.getUsers(), [&](Operation *user) {
+    auto liveUserIt = llvm::find_if_not(value.getUsers(), [&](Operation *user) {
       return rewriterImpl.isOpIgnored(user);
     });
-  if (liveUserIt == result.user_end())
+    if (liveUserIt != value.user_end())
+      return *liveUserIt;
+    value = inverseMapping.lookupOrNull(value);
+  } while (value != nullptr);
+  return nullptr;
+}
+
+LogicalResult OperationConverter::legalizeChangedResultType(
+    Operation *op, OpResult result, Value newValue,
+    TypeConverter *replConverter, ConversionPatternRewriter &rewriter,
+    ConversionPatternRewriterImpl &rewriterImpl,
+    const BlockAndValueMapping &inverseMapping) {
+  Operation *liveUser =
+      findLiveUserOfReplaced(result, rewriterImpl, inverseMapping);
+  if (!liveUser)
     return success();
 
   // If the replacement has a type converter, attempt to materialize a
   // conversion back to the original type.
   if (!replConverter) {
     // TODO: We should emit an error here, similarly to the case where the
     // result is replaced with null. Unfortunately a lot of existing
     // patterns rely on this behavior, so until those patterns are updated
LogicalResult OperationConverter::legalizeChangedResultType(
   Value convertedValue = replConverter->materializeSourceConversion(
       rewriter, op->getLoc(), resultType, newValue);
   if (!convertedValue) {
     InFlightDiagnostic diag = op->emitError()
                               << "failed to materialize conversion for result #"
                               << result.getResultNumber() << " of operation '"
                               << op->getName()
                               << "' that remained live after conversion";
-    diag.attachNote(liveUserIt->getLoc())
-        << "see existing live user here: " << *liveUserIt;
+    diag.attachNote(liveUser->getLoc())
+        << "see existing live user here: " << *liveUser;
     return failure();
   }
 
   // Legalize all of the newly created conversion operations.
   for (int i : llvm::seq<int>(numCreatedOps, rewriterImpl.createdOps.size())) {
     if (failed(opLegalizer.legalize(rewriterImpl.createdOps[i], rewriter))) {
       return op->emitError("failed to legalize conversion operation generated ")
              << "for result #" << result.getResultNumber() << " of operation '"

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

func @test_invalid_result_materialization() {
   %result = "test.type_producer"() : () -> f16
 
   // expected-note@below {{see existing live user here}}
   "foo.return"(%result) : (f16) -> ()
 }
 
 // -----
 
+// CHECK-LABEL: @test_transitive_use_materialization
+func @test_transitive_use_materialization() {
+  // CHECK: %[[V:.*]] = "test.type_producer"() : () -> f64
+  // CHECK: %[[C:.*]] = "test.cast"(%[[V]]) : (f64) -> f32
+  %result = "test.another_type_producer"() : () -> f32
+  // CHECK: "foo.return"(%[[C]])
+  "foo.return"(%result) : (f32) -> ()
+}
+
+// -----
+
+func @test_transitive_use_invalid_materialization() {
+  // expected-error@below {{failed to materialize conversion for result #0 of operation 'test.type_producer' that remained live after conversion}}
+  %result = "test.another_type_producer"() : () -> f16
+  // expected-note@below {{see existing live user here}}
+  "foo.return"(%result) : (f16) -> ()
+}
+
+// -----
+
 func @test_invalid_result_legalization() {
   // expected-error@below {{failed to legalize conversion operation generated for result #0 of operation 'test.type_producer' that remained live after conversion}}
   %result = "test.type_producer"() : () -> i16
   "foo.return"(%result) : (i16) -> ()
 }
 
 // -----
 

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

def TestReturnOp : TEST_Op<"return", [ReturnLike, Terminator]> {
   ];
 }
 def TestCastOp : TEST_Op<"cast">,
   Arguments<(ins Variadic<AnyType>)>, Results<(outs AnyType)>;
 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 TestValidOp : TEST_Op<"valid", [Terminator]>,
   Arguments<(ins Variadic<AnyType>)>;
 
 def TestMergeBlocksOp : TEST_Op<"merge_blocks"> {
   let summary = "merge_blocks operation";
   let description = [{

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

struct TestTypeConsumerForward
   LogicalResult
   matchAndRewrite(TestTypeConsumerOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     rewriter.updateRootInPlace(op, [&] { op->setOperands(operands); });
     return success();
   }
 };
 
+struct TestTypeConversionAnotherProducer
+    : public OpRewritePattern<TestAnotherTypeProducerOp> {
+  using OpRewritePattern<TestAnotherTypeProducerOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(TestAnotherTypeProducerOp op,
+                                PatternRewriter &rewriter) const final {
+    rewriter.replaceOpWithNewOp<TestTypeProducerOp>(op, op.getType());
+    return success();
+  }
+};
+
 struct TestTypeConversionDriver
     : public PassWrapper<TestTypeConversionDriver, OperationPass<ModuleOp>> {
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<TestDialect>();
   }
 
   void runOnOperation() override {
     // Initialize the type converter.
target.addDynamicallyLegalOp<TestCastOp>([&](TestCastOp op) {
       // Allow casts from F64 to F32.
       return (*op.operand_type_begin()).isF64() && op.getType().isF32();
     });
 
     // Initialize the set of rewrite patterns.
     OwningRewritePatternList patterns;
     patterns.insert<TestTypeConsumerForward, TestTypeConversionProducer,
                     TestSignatureConversionUndo>(converter, &getContext());
+    patterns.insert<TestTypeConversionAnotherProducer>(&getContext());
     mlir::populateFuncOpTypeConversionPattern(patterns, &getContext(),
                                               converter);
 
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };