[mlir] Lookup the latest value with a legal type when remapping values.

[pifon2a]

Summary

The current condition implies that the target materialization will be
called even if the type is the new operand type is legal, but slightly
different. For example, if there is a bufferization pattern that changes memref
layout, then target materialization for an illegal type (TensorType) would be
called.

[herhut]

mlir/lib/Transforms/Utils/DialectConversion.cpp

1090–1091

I think this could be if (converter && !converter.isLegal(newOperandType)) and then the computation of desiredtype could be moved inside this block.

1090–1091

This line finds the most recent remapped value that has the desired type. Instead, to be in line with the new below code, it should find the most recent remapped value that has a legal type. This is also the only use of this special search functionality,

[pifon2a]

Address the comments.

[herhut]

mlir/lib/Transforms/Utils/DialectConversion.cpp

108–116

Is this still needed? Could we rip out the desiredType support here?

167

Maybe just call lookupOrDefault here?

173

nit: legalValue?

1054–1084

Can we do an early return here to reduce nesting?

[pifon2a]

Address more comments.

[pifon2a]

Thanks @herhut , now it looks much better.

[rriddle]

mlir/lib/Transforms/Utils/DialectConversion.cpp

1087

Sink this down to where it is used.

1101

Remove the semi-colon here.

[pifon2a]

Address the comments

[herhut]

Thanks!

mlir/lib/Transforms/Utils/DialectConversion.cpp

 // ConversionValueMapping
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class wraps a BlockAndValueMapping to provide recursive lookup
 /// functionality, i.e. we will traverse if the mapped value also has a mapping.
 struct ConversionValueMapping {
   /// Lookup a mapped value within the map. If a mapping for the provided value
-  /// does not exist then return the provided value. If `desiredType` is
-  /// non-null, returns the most recently mapped value with that type. If an
-  /// operand of that type does not exist, defaults to normal behavior.
-  Value lookupOrDefault(Value from, Type desiredType = nullptr) const;
+  /// does not exist then return the provided value.
+  Value lookupOrDefault(Value from) const;
+
+  /// Lookup the latest legal value within the map. If a mapping for the
+  /// provided value does not exist then return the provided value. If
+  /// `converter` is non-null, returns the most recently mapped value with the
+  /// legal type. If an operand of that type does not exist, defaults to normal
+  /// behavior.
+  Value lookupLatestLegal(Value from, TypeConverter *converter) const;
 
   /// Lookup a mapped value within the map, or return null if a mapping does not
   /// exist. If a mapping exists, this follows the same behavior of
   /// `lookupOrDefault`.
   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); }
 
 private:
   /// Current value mappings.
   BlockAndValueMapping mapping;
 };
 } // end anonymous namespace
 
-Value ConversionValueMapping::lookupOrDefault(Value from,
-                                              Type desiredType) const {
-  // If there was no desired type, simply find the leaf value.
-  if (!desiredType) {
-    // If this value had a valid mapping, unmap that value as well in the case
-    // that it was also replaced.
-    while (auto mappedValue = mapping.lookupOrNull(from))
-      from = mappedValue;
-    return from;
-  }
+Value ConversionValueMapping::lookupOrDefault(Value from) const {
 
-  // Otherwise, try to find the deepest value that has the desired type.
-  Value desiredValue;
+Value ConversionValueMapping::lookupLatestLegal(
+    Value from, TypeConverter *converter) const {
+  if (!converter)
+    return lookupOrDefault(from);
+
+  // Otherwise, try to find the deepest value that has the legal type.
+  Value legalValue;
   do {
-    if (from.getType() == desiredType)
-      desiredValue = from;
+    if (converter->isLegal(from.getType()))
+      legalValue = from;
 
     Value mappedValue = mapping.lookupOrNull(from);
     if (!mappedValue)
       break;
     from = mappedValue;
   } while (true);
 
   // If the desired value was found use it, otherwise default to the leaf value.
-  return desiredValue ? desiredValue : from;
+  return legalValue ? legalValue : from;
 }
 
 Value ConversionValueMapping::lookupOrNull(Value from) const {
   Value result = lookupOrDefault(from);
   return result == from ? nullptr : result;
 }
 
 //===----------------------------------------------------------------------===//
 // ArgConverter
 //===----------------------------------------------------------------------===//
 namespace {
 /// This class provides a simple interface for converting the types of block
 /// arguments. This is done by creating a new block that contains the new legal
 /// types and extracting the block that contains the old illegal types to allow
 /// for undoing pending rewrites in the case of failure.
 struct ArgConverter {
   ArgConverter(PatternRewriter &rewriter) : rewriter(rewriter) {}
 
   /// This structure contains the information pertaining to an argument that has
   /// been converted.
LogicalResult ConversionPatternRewriterImpl::remapValues(
     Operation::operand_range operands, SmallVectorImpl<Value> &remapped) {
   remapped.reserve(llvm::size(operands));
 
   SmallVector<Type, 1> legalTypes;
   for (auto it : llvm::enumerate(operands)) {
     Value operand = it.value();
     Type origType = operand.getType();
 
-    // If a converter was provided, get the desired legal types for this
-    // operand.
-    Type desiredType;
+    Value newOperand = mapping.lookupLatestLegal(operand, converter);
+
+    // Handle the case where the conversion was 1->1 and the new operand type
+    // isn't legal.
+    Type newOperandType = newOperand.getType();
     if (converter) {
-      // If there is no legal conversion, fail to match this pattern.
+      if (!converter->isLegal(newOperandType)) {
         legalTypes.clear();
+
+        // If there is no legal conversion, fail to match this pattern.
         if (failed(converter->convertType(origType, legalTypes))) {
           return notifyMatchFailure(loc, [=](Diagnostic &diag) {
             diag << "unable to convert type for operand #" << it.index()
                  << ", type was " << origType;
           });
         }
         // TODO: There currently isn't any mechanism to do 1->N type conversion
-      // via the PatternRewriter replacement API, so for now we just ignore it.
-      if (legalTypes.size() == 1)
-        desiredType = legalTypes.front();
-    } else {
-      // TODO: What we should do here is just set `desiredType` to `origType`
-      // and then handle the necessary type conversions after the conversion
-      // process has finished. Unfortunately a lot of patterns currently rely on
-      // receiving the new operands even if the types change, so we keep the
-      // original behavior here for now until all of the patterns relying on
-      // this get updated.
+        // via the PatternRewriter replacement API, so for now we just ignore
+        // it.
+        if (legalTypes.size() != 1) {
+          remapped.push_back(newOperand);
+          continue;
         }
-    Value newOperand = mapping.lookupOrDefault(operand, desiredType);
-
-    // Handle the case where the conversion was 1->1 and the new operand type
-    // isn't legal.
-    Type newOperandType = newOperand.getType();
-    if (converter && desiredType && newOperandType != desiredType) {
-      // Attempt to materialize a conversion for this new value.
+        Type desiredType = legalTypes.front();
         newOperand = converter->materializeTargetConversion(
             rewriter, loc, desiredType, newOperand);
         if (!newOperand) {
           return notifyMatchFailure(loc, [=](Diagnostic &diag) {
             diag << "unable to materialize a conversion for "
                     "operand #"
                  << it.index() << ", from " << newOperandType << " to "
                  << desiredType;
           });
         }
       }
+    } else {
+      // TODO: What we should do here is just set `desiredType` to `origType`
+      // and then handle the necessary type conversions after the conversion
+      // process has finished. Unfortunately a lot of patterns currently rely on
+      // receiving the new operands even if the types change, so we keep the
+      // original behavior here for now until all of the patterns relying on
+      // this get updated.
+      // Attempt to materialize a conversion for this new value.
+    }
     remapped.push_back(newOperand);
   }
   return success();
 }
 
 bool ConversionPatternRewriterImpl::isOpIgnored(Operation *op) const {
   // Check to see if this operation was replaced or its parent ignored.
   return replacements.count(op) || ignoredOps.count(op->getParentOp());
 }
 
 void ConversionPatternRewriterImpl::markNestedOpsIgnored(Operation *op) {
   // Walk this operation and collect nested operations that define non-empty
   // regions. We mark such operations as 'ignored' so that we know we don't have
   // to convert them, or their nested ops.
   if (op->getNumRegions() == 0)
     return;
   op->walk([&](Operation *op) {