[mlir][DialectConversion] Enable deeper integration of type conversions

[rriddle]

Summary

This revision adds support for much deeper type conversion integration into the conversion process, and enables auto-generating cast operations when necessary. Type conversions are now largely automatically managed by the conversion infra when using a ConversionPattern with a provided TypeConverter. This removes the need for patterns to do type cast wrapping themselves and moves the burden to the infra. This makes it much easier to perform partial lowerings when type conversions are involved, as any lingering type conversions will be automatically resolved/legalized by the conversion infra.

To support this new integration, a few changes have been made to the type materialization API on TypeConverter. Materialization has been split into three separate categories:

• Argument Materialization: This type of materialization is used when converting the type of block arguments when calling convertRegionTypes. This is useful for contextually inserting additional conversion operations when converting a block argument type, such as when converting the types of a function signature.
• Source Materialization: This type of materialization is used to convert a legal type of the converter into a non-legal type, generally a source type. This may be called when uses of a non-legal type persist after the conversion process has finished.
• Target Materialization: This type of materialization is used to convert a non-legal, or source, type into a legal, or target, type. This type of materialization is used when applying a pattern on an operation, but the types of the operands have not yet been converted.

Depends On D82830

[rriddle]

@ftynse I still need to add tests and rewrite the documentation, but it should be good to start looking over. Thanks.

[ftynse]

Thanks, River! Looks good in general. Two high-level questions:

1. What is the convention on inserting materializations, i.e. who is in charge of inserting one, at which point in the conversion and using which type converter? For example, legalizeChangedResultType checks with the type converter of the user (which it takes from opToConverter, so it is only available for ops created by the infra), but attempts materialization with replConverter that it accepts as an argument, and which may be different.
2. This looks like it can get into infinite recursion or iteration if materialization produces operations that need more materialization in order to be legalized. Is there something the infra does to guard against it?

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

179

Hmm, can we already call this->convertType() from here? DialectCastOp should be constructible if convertType(resultType) == inputs[0] for source conversions and if convertType(inputs[0]) == resultType for target conversions. I was thinking about how one can check this in the verifier of DialectCastOp, but this may actually be a better place.

mlir/lib/Transforms/DialectConversion.cpp

109

Nit: spurious backslash

399

Nit: might be useful to attach a note pointing to the first live user

580

Initialize to nullptr?

1002

assert(legalTypes.size() <= 1) then?

1008

I'm not sure I read this correctly. A lot (all?) patterns rely on receiving new operands to create new operations, which require Values (sometimes, of some specific type) to be passed in. Is the idea here to fail the conversion if there is a type mismatch, no type converter is provided, and we can't find a mapping to the desired (original) type?

2243

What if the user doesn't have a converter, and the types don't match?

[rriddle]

Resolve comments and add tests

[rriddle]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

179

I don't think we can given that users of the type converter can add other type casts, so we aren't guaranteed to have just standard type conversions. At least, that's if we keep DialectCastOp solely between Standard and LLVM.

mlir/lib/Transforms/DialectConversion.cpp

1002

We don't want to assert, because that would prevent 1->N conversions outright. There isn't anything in the infra to detect/protect against it, so the best we can do is what we have been doing; i.e. you need to manage it yourself.

1008

Essentially yes. In the new ideal state, patterns that rely on type conversions must provide a TypeConverter. If one isn't provided, we would always use the original type and rely on the infra for making sure lingering type materializations are inserted when necessary.

2243

Thanks for the catch, refactored to remove the need for this.

[rriddle]

In D82831#2127585, @ftynse wrote:

Thanks, River! Looks good in general. Two high-level questions:

1. What is the convention on inserting materializations, i.e. who is in charge of inserting one, at which point in the conversion and using which type converter? For example, legalizeChangedResultType checks with the type converter of the user (which it takes from opToConverter, so it is only available for ops created by the infra), but attempts materialization with replConverter that it accepts as an argument, and which may be different.

This has been reworked a bit, but ideally only the pattern that replaced the operation should have its type converter used/checked. We are in a bit of an awkward state where a lot of patterns rely on broken behavior w.r.t. type conversions. I was able to remove the checks for user type converters, and there are a few more TODOs sprinkled about to slowly enable more aggressive checks as user patterns get updated.

2. This looks like it can get into infinite recursion or iteration if materialization produces operations that need more materialization in order to be legalized. Is there something the infra does to guard against it?

There is nothing guarding it in this revision, but there should be in a followup. Realistically a generated conversion should be already legal, or trivially legalizable. I'm still playing around with if we should check if it is legal without trying to legalize, but wanted to update more patterns first to decide if that is the desired behavior.

[rriddle]

Rebase

[ftynse]

Thanks, and apologies for the delay! Phabricator wasn't showing this on the front page for some reason (this happened to me in the past, any ideas about the reason for this weird behavior is most welcome) so I missed it. Don't hesitate to ping me if I'm not responding for more than several work days.

This looks much easier to follow after the refactoring! Two more high-level comments, feel free to leave for follow-ups:

• Would it make sense to sprinkle some debug output that connects with the existing output for pattern matches/failures and op creation in patterns?
• Could we start a publicly visible doc that has the list of current assumptions the conversion infra makes, and whether we want to remove some of them? There are out-of-tree users that may rely on some of the deprecated behavior and it will be easier to point them to that than to a diff.

mlir/lib/IR/Value.cpp

82 ↗

(On Diff #277932)

Or maybe blocks? It has always bothered me that we report something like "block with no terminator" while pointing to something else (parent op or first op in the block)

mlir/lib/Transforms/DialectConversion.cpp

2183

Typo: operations

[rriddle]

Resolve comments

[rriddle]

In D82831#2169052, @ftynse wrote:

Thanks, and apologies for the delay! Phabricator wasn't showing this on the front page for some reason (this happened to me in the past, any ideas about the reason for this weird behavior is most welcome) so I missed it. Don't hesitate to ping me if I'm not responding for more than several work days.

No worries, I had plenty of other things to do.

This looks much easier to follow after the refactoring! Two more high-level comments, feel free to leave for follow-ups:

• Would it make sense to sprinkle some debug output that connects with the existing output for pattern matches/failures and op creation in patterns?

I added in a few, but should probably add more for the things that are currently marked with TODOs.

• Could we start a publicly visible doc that has the list of current assumptions the conversion infra makes, and whether we want to remove some of them? There are out-of-tree users that may rely on some of the deprecated behavior and it will be easier to point them to that than to a diff.

I have an in progress rewrite of the DialectConversion doc that should touch on many of these things. I'm going to try and have that out for review early next week, it'll have a section for the deprecated and soon to be replaced behavior.

Thanks again for the review Alex!

mlir/include/mlir/Transforms/DialectConversion.h

public:
   template <typename FnT,
             typename T = typename llvm::function_traits<FnT>::template arg_t<0>>
   void addConversion(FnT &&callback) {
     registerConversion(wrapCallback<T>(std::forward<FnT>(callback)));
   }
 
   /// Register a materialization function, which must be convertible to the
   /// following form:
-  ///   `Optional<Value>(PatternRewriter &, T, ValueRange, Location)`,
+  ///   `Optional<Value>(OpBuilder &, T, ValueRange, Location)`,
   /// where `T` is any subclass of `Type`. This function is responsible for
-  /// creating an operation, using the PatternRewriter and Location provided,
-  /// that "casts" a range of values into a single value of the given type `T`.
-  /// It must return a Value of the converted type on success, an `llvm::None`
-  /// if it failed but other materialization can be attempted, and `nullptr` on
+  /// creating an operation, using the OpBuilder and Location provided, that
+  /// "casts" a range of values into a single value of the given type `T`. It
+  /// must return a Value of the converted type on success, an `llvm::None` if
+  /// it failed but other materialization can be attempted, and `nullptr` on
   /// unrecoverable failure. It will only be called for (sub)types of `T`.
   /// Materialization functions must be provided when a type conversion
   /// results in more than one type, or if a type conversion may persist after
   /// the conversion has finished.
+  ///
+  /// This method registers a materialization that will be called when
+  /// converting an illegal block argument type, to a legal type.
   template <typename FnT,
             typename T = typename llvm::function_traits<FnT>::template arg_t<1>>
-  void addMaterialization(FnT &&callback) {
-    registerMaterialization(
+  void addArgumentMaterialization(FnT &&callback) {
+    argumentMaterializations.emplace_back(
+        wrapMaterialization<T>(std::forward<FnT>(callback)));
+  }
+  /// This method registers a materialization that will be called when
+  /// converting a legal type to an illegal source type. This is used when
+  /// conversions to an illegal type must persist beyond the main conversion.
+  template <typename FnT,
+            typename T = typename llvm::function_traits<FnT>::template arg_t<1>>
+  void addSourceMaterialization(FnT &&callback) {
+    sourceMaterializations.emplace_back(
+        wrapMaterialization<T>(std::forward<FnT>(callback)));
+  }
+  /// This method registers a materialization that will be called when
+  /// converting type from an illegal, or source, type to a legal type.
+  template <typename FnT,
+            typename T = typename llvm::function_traits<FnT>::template arg_t<1>>
+  void addTargetMaterialization(FnT &&callback) {
+    targetMaterializations.emplace_back(
         wrapMaterialization<T>(std::forward<FnT>(callback)));
   }
 
   /// Convert the given type. This function should return failure if no valid
   /// conversion exists, success otherwise. If the new set of types is empty,
   /// the type is removed and any usages of the existing value are expected to
   /// be removed during conversion.
   LogicalResult convertType(Type t, SmallVectorImpl<Type> &results);
LogicalResult convertSignatureArgs(TypeRange types,
                                      unsigned origInputOffset = 0);
 
   /// This function converts the type signature of the given block, by invoking
   /// 'convertSignatureArg' for each argument. This function should return a
   /// valid conversion for the signature on success, None otherwise.
   Optional<SignatureConversion> convertBlockSignature(Block *block);
 
   /// Materialize a conversion from a set of types into one result type by
-  /// generating a cast operation of some kind.
-  Value materializeConversion(PatternRewriter &rewriter, Location loc,
-                              Type resultType, ValueRange inputs);
+  /// generating a cast sequence of some kind. See the respective
+  /// `add*Materialization` for more information on the context for these
+  /// methods.
+  Value materializeArgumentConversion(OpBuilder &builder, Location loc,
+                                      Type resultType, ValueRange inputs) {
+    return materializeConversion(argumentMaterializations, builder, loc,
+                                 resultType, inputs);
+  }
+  Value materializeSourceConversion(OpBuilder &builder, Location loc,
+                                    Type resultType, ValueRange inputs) {
+    return materializeConversion(sourceMaterializations, builder, loc,
+                                 resultType, inputs);
+  }
+  Value materializeTargetConversion(OpBuilder &builder, Location loc,
+                                    Type resultType, ValueRange inputs) {
+    return materializeConversion(targetMaterializations, builder, loc,
+                                 resultType, inputs);
+  }
 
 private:
   /// The signature of the callback used to convert a type. If the new set of
   /// types is empty, the type is removed and any usages of the existing value
   /// are expected to be removed during conversion.
   using ConversionCallbackFn =
       std::function<Optional<LogicalResult>(Type, SmallVectorImpl<Type> &)>;
 
-  using MaterializationCallbackFn = std::function<Optional<Value>(
-      PatternRewriter &, Type, ValueRange, Location)>;
+  /// The signature of the callback used to materialize a conversion.
+  using MaterializationCallbackFn =
+      std::function<Optional<Value>(OpBuilder &, Type, ValueRange, Location)>;
+
+  /// Attempt to materialize a conversion using one of the provided
+  /// materialization functions.
+  Value materializeConversion(
+      MutableArrayRef<MaterializationCallbackFn> materializations,
+      OpBuilder &builder, Location loc, Type resultType, ValueRange inputs);
 
   /// Generate a wrapper for the given callback. This allows for accepting
   /// different callback forms, that all compose into a single version.
   /// With callback of form: `Optional<Type>(T)`
   template <typename T, typename FnT>
   std::enable_if_t<llvm::is_invocable<FnT, T>::value, ConversionCallbackFn>
   wrapCallback(FnT &&callback) {
     return wrapCallback<T>([callback = std::forward<FnT>(callback)](
private:
   }
 
   /// Generate a wrapper for the given materialization callback. The callback
   /// may take any subclass of `Type` and the wrapper will check for the target
   /// type to be of the expected class before calling the callback.
   template <typename T, typename FnT>
   MaterializationCallbackFn wrapMaterialization(FnT &&callback) {
     return [callback = std::forward<FnT>(callback)](
-               PatternRewriter &rewriter, Type resultType, ValueRange inputs,
+               OpBuilder &builder, Type resultType, ValueRange inputs,
                Location loc) -> Optional<Value> {
       if (T derivedType = resultType.dyn_cast<T>())
-        return callback(rewriter, derivedType, inputs, loc);
+        return callback(builder, derivedType, inputs, loc);
       return llvm::None;
     };
   }
 
-  /// Register a materialization.
-  void registerMaterialization(MaterializationCallbackFn &&callback) {
-    materializations.emplace_back(std::move(callback));
-  }
-
   /// The set of registered conversion functions.
   SmallVector<ConversionCallbackFn, 4> conversions;
 
   /// The list of registered materialization functions.
-  SmallVector<MaterializationCallbackFn, 2> materializations;
+  SmallVector<MaterializationCallbackFn, 2> argumentMaterializations;
+  SmallVector<MaterializationCallbackFn, 2> sourceMaterializations;
+  SmallVector<MaterializationCallbackFn, 2> targetMaterializations;
 
   /// A set of cached conversions to avoid recomputing in the common case.
   /// Direct 1-1 conversions are the most common, so this cache stores the
   /// successful 1-1 conversions as well as all failed conversions.
   DenseMap<Type, Type> cachedDirectConversions;
   /// This cache stores the successful 1->N conversions, where N != 1.
   DenseMap<Type, SmallVector<Type, 2>> cachedMultiConversions;
 };
protected:
   /// behavior is what the user actually desired, `MatchAnyOpTypeTag()` should
   /// always be supplied here.
   ConversionPattern(PatternBenefit benefit, TypeConverter &typeConverter,
                     MatchAnyOpTypeTag tag)
       : RewritePattern(benefit, tag), typeConverter(&typeConverter) {}
 
 protected:
   /// An optional type converter for use by this pattern.
-  TypeConverter *typeConverter;
+  TypeConverter *typeConverter = nullptr;
 
 private:
   using RewritePattern::rewrite;
 };
 
 /// OpConversionPattern is a wrapper around ConversionPattern that allows for
 /// matching and rewriting against an instance of a derived operation class as
 /// opposed to a raw Operation.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

LLVMTypeConverter::LLVMTypeConverter(
   addConversion([&](VectorType type) { return convertVectorType(type); });
 
   // LLVMType is legal, so add a pass-through conversion.
   addConversion([](LLVM::LLVMType type) { return type; });
 
   // Materialization for memrefs creates descriptor structs from individual
   // values constituting them, when descriptors are used, i.e. more than one
   // value represents a memref.
-  addMaterialization([&](PatternRewriter &rewriter,
-                         UnrankedMemRefType resultType, ValueRange inputs,
-                         Location loc) -> Optional<Value> {
+  addArgumentMaterialization(
+      [&](OpBuilder &builder, UnrankedMemRefType resultType, ValueRange inputs,
         if (inputs.size() == 1)
           return llvm::None;
-    return UnrankedMemRefDescriptor::pack(rewriter, loc, *this, resultType,
+        return UnrankedMemRefDescriptor::pack(builder, loc, *this, resultType,
                                               inputs);
       });
-  addMaterialization([&](PatternRewriter &rewriter, MemRefType resultType,
-                         ValueRange inputs, Location loc) -> Optional<Value> {
+  addArgumentMaterialization([&](OpBuilder &builder, MemRefType resultType,
+                                 ValueRange inputs,
+                                 Location loc) -> Optional<Value> {
     if (inputs.size() == 1)
       return llvm::None;
-    return MemRefDescriptor::pack(rewriter, loc, *this, resultType, inputs);
+    return MemRefDescriptor::pack(builder, loc, *this, resultType, inputs);
+  });
+  // Add generic source and target materializations to handle cases where
+  // non-LLVM types persist after an LLVM conversion.
+  addSourceMaterialization([&](OpBuilder &builder, Type resultType,
+                               ValueRange inputs,
+                               Location loc) -> Optional<Value> {
+    if (inputs.size() != 1)
+      return llvm::None;
+    // FIXME: These should check LLVM::DialectCastOp can actually be constructed
+    // from the input and result.
+    return builder.create<LLVM::DialectCastOp>(loc, resultType, inputs[0])
+        .getResult();
+  });
+  addTargetMaterialization([&](OpBuilder &builder, Type resultType,
+                               ValueRange inputs,
+                               Location loc) -> Optional<Value> {
+    if (inputs.size() != 1)
+      return llvm::None;
+    // FIXME: These should check LLVM::DialectCastOp can actually be constructed
+    // from the input and result.
+    return builder.create<LLVM::DialectCastOp>(loc, resultType, inputs[0])
+        .getResult();
   });
 }
 
 /// Returns the MLIR context.
 MLIRContext &LLVMTypeConverter::getContext() {
   return *getDialect()->getContext();
 }
 

mlir/lib/Transforms/DialectConversion.cpp

 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/Transforms/Utils.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/SaveAndRestore.h"
 #include "llvm/Support/ScopedPrinter.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 #define DEBUG_TYPE "dialect-conversion"
 
 /// Recursively collect all of the operations to convert from within 'region'.

 //===----------------------------------------------------------------------===//
 // 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.
-  Value lookupOrDefault(Value from) const;
+  /// 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;
 
   /// 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
 
-/// Lookup a mapped value within the map. If a mapping for the provided value
-/// does not exist then return the provided value.
-Value ConversionValueMapping::lookupOrDefault(Value from) const {
+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;
   }
 
+  // Otherwise, try to find the deepest value that has the desired type.
+  Value desiredValue;
+  do {
+    if (from.getType() == desiredType)
+      desiredValue = 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;
+}
+
 //===----------------------------------------------------------------------===//
 // 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 {
   /// undoing them.
   void notifyOpRemoved(Operation *op);
 
   /// Cleanup and undo any generated conversions for the arguments of block.
   /// This method replaces the new block with the original, reverting the IR to
   /// its original state.
   void discardRewrites(Block *block);
 
-  /// Fully replace uses of the old arguments with the new, materializing cast
-  /// operations as necessary.
+  /// Fully replace uses of the old arguments with the new.
   void applyRewrites(ConversionValueMapping &mapping);
 
+  /// Materialize any necessary conversions for converted arguments that have
+  /// live users, using the provided `isDead` for determining if a value is
+  /// dead.
+  LogicalResult materializeLiveConversions(ConversionValueMapping &mapping,
+                                           OpBuilder &builder,
+                                           function_ref<bool(Value)> isDead);
+
   //===--------------------------------------------------------------------===//
   // Conversion
   //===--------------------------------------------------------------------===//
 
   /// Attempt to convert the signature of the given block, if successful a new
   /// block is returned containing the new arguments. Returns `block` if it did
   /// not require conversion.
   FailureOr<Block *> convertSignature(Block *block, TypeConverter &converter,
void ArgConverter::discardRewrites(Block *block) {
   block->erase();
 
   convertedBlocks.erase(origBlock);
   conversionInfo.erase(it);
 }
 
 void ArgConverter::applyRewrites(ConversionValueMapping &mapping) {
   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) {
       Optional<ConvertedArgInfo> &argInfo = blockInfo.argInfo[i];
       BlockArgument origArg = origBlock->getArgument(i);
 
       // Handle the case of a 1->0 value mapping.
       if (!argInfo) {
-        // If a replacement value was given for this argument, use that to
-        // replace all uses.
-        auto argReplacementValue = mapping.lookupOrDefault(origArg);
-        if (argReplacementValue != origArg) {
-          origArg.replaceAllUsesWith(argReplacementValue);
-          continue;
-        }
-        // If there are any dangling uses then replace the argument with one
-        // generated by the type converter. This is necessary as the cast must
-        // persist in the IR after conversion.
-        if (!origArg.use_empty()) {
-          rewriter.setInsertionPointToStart(newBlock);
-          Value newArg = blockInfo.converter->materializeConversion(
-              rewriter, origArg.getLoc(), origArg.getType(), llvm::None);
-          assert(newArg &&
-                 "Couldn't materialize a block argument after 1->0 conversion");
+        Value newArg = mapping.lookupOrDefault(origArg);
+        if (origArg != newArg)
           origArg.replaceAllUsesWith(newArg);
-        }
         continue;
       }
 
       // Otherwise this is a 1->1+ value mapping.
       Value castValue = argInfo->castValue;
       assert(argInfo->newArgSize >= 1 && castValue && "expected 1->1+ mapping");
 
       // If the argument is still used, replace it with the generated cast.
       if (!origArg.use_empty())
         origArg.replaceAllUsesWith(mapping.lookupOrDefault(castValue));
 
       // If all users of the cast were removed, we can drop it. Otherwise, keep
       // the operation alive and let the user handle any remaining usages.
       if (castValue.use_empty() && castValue.getDefiningOp())
         castValue.getDefiningOp()->erase();
     }
   }
 }
 
+LogicalResult
+ArgConverter::materializeLiveConversions(ConversionValueMapping &mapping,
+                                         OpBuilder &builder,
+                                         function_ref<bool(Value)> isDead) {
+  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) {
+      // We only need to handle the case of a 1->0 value mapping. All other
+      // mappings have already materialized the necessary conversions during the
+      // legalization process.
+      if (blockInfo.argInfo[i])
+        continue;
+
+      // If a replacement value wasn't given for this argument and the argument
+      // is still live, we need to materialize a conversion.
+      BlockArgument origArg = origBlock->getArgument(i);
+      auto argReplacementValue = mapping.lookupOrDefault(origArg);
+      if (argReplacementValue != origArg || isDead(origArg))
+        continue;
+
+      rewriter.setInsertionPointToStart(newBlock);
+      Value newArg = blockInfo.converter->materializeSourceConversion(
+          rewriter, origArg.getLoc(), origArg.getType(), llvm::None);
+      if (!newArg) {
+        return emitError(origArg.getLoc())
+               << "failed to materialize conversion for block argument that "
+                  "remained live after conversion, type was "
+               << origArg.getType();
+      }
+      mapping.map(origArg, newArg);
+    }
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Conversion
 
 FailureOr<Block *>
 ArgConverter::convertSignature(Block *block, TypeConverter &converter,
                                ConversionValueMapping &mapping) {
   // Check if the block was already converted. If the block is detached,
   // conservatively assume it is going to be deleted.
if (inputMap->replacementValue) {
       mapping.map(origArg, inputMap->replacementValue);
       continue;
     }
 
     // Otherwise, this is a 1->1+ mapping. Call into the provided type converter
     // to pack the new values. For 1->1 mappings, if there is no materialization
     // provided, use the argument directly instead.
     auto replArgs = newArgs.slice(inputMap->inputNo, inputMap->size);
-    Value newArg = converter.materializeConversion(rewriter, origArg.getLoc(),
-                                                   origArg.getType(), replArgs);
+    Value newArg = converter.materializeArgumentConversion(
+        rewriter, origArg.getLoc(), origArg.getType(), replArgs);
     if (!newArg) {
       assert(replArgs.size() == 1 &&
              "couldn't materialize the result of 1->N conversion");
       newArg = replArgs.front();
     }
     mapping.map(origArg, newArg);
     info.argInfo[i] =
         ConvertedArgInfo(inputMap->inputNo, inputMap->size, newArg);

 private:
   Operation *op;
   LocationAttr loc;
   MutableDictionaryAttr attrs;
   SmallVector<Value, 8> operands;
   SmallVector<Block *, 2> successors;
 };
 
-/// This class represents one requested operation replacement via 'replaceOp'.
+/// This class represents one requested operation replacement via 'replaceOp' or
+/// 'eraseOp`.
 struct OpReplacement {
   OpReplacement() = default;
-  OpReplacement(ValueRange newValues)
-      : newValues(newValues.begin(), newValues.end()) {}
+  OpReplacement(ValueRange newValues, TypeConverter *converter)
+      : newValues(newValues.begin(), newValues.end()), converter(converter) {}
 
+  /// The values used to replace the results of the operation.
   SmallVector<Value, 2> newValues;
+
+  /// An optional type converter that can be used to materialize conversions
+  /// between the new and old values if necessary.
+  TypeConverter *converter;
 };
 
 /// The kind of the block action performed during the rewrite.  Actions can be
 /// undone if the conversion fails.
 enum class BlockActionKind { Create, Erase, Move, Split, TypeConversion };
 
 /// Original position of the given block in its parent region.  We cannot use
 /// a region iterator because it could have been invalidated by other region
struct ConversionPatternRewriterImpl {
   /// Erase any blocks that were unlinked from their regions and stored in block
   /// actions.
   void eraseDanglingBlocks();
 
   /// Undo the block actions (motions, splits) one by one in reverse order until
   /// "numActionsToKeep" actions remains.
   void undoBlockActions(unsigned numActionsToKeep = 0);
 
-  /// Remap the given operands to those with potentially different types.
-  void remapValues(Operation::operand_range operands,
+  /// Remap the given operands to those with potentially different types. The
+  /// provided type converter is used to ensure that the remapped types are
+  /// legal. Returns success if the operands could be remapped, failure
+  /// otherwise.
+  LogicalResult remapValues(Location loc, PatternRewriter &rewriter,
+                            TypeConverter *converter,
+                            Operation::operand_range operands,
                             SmallVectorImpl<Value> &remapped);
 
   /// Returns true if the given operation is ignored, and does not need to be
   /// converted.
   bool isOpIgnored(Operation *op) const;
 
   /// Recursively marks the nested operations under 'op' as ignored. This
   /// removes them from being considered for legalization.
   void markNestedOpsIgnored(Operation *op);
struct ConversionPatternRewriterImpl {
   /// Notifies that the blocks of a region are about to be moved.
   void notifyRegionIsBeingInlinedBefore(Region &region, Region &parent,
                                         Region::iterator before);
 
   /// Notifies that the blocks of a region were cloned into another.
   void notifyRegionWasClonedBefore(iterator_range<Region::iterator> &blocks,
                                    Location origRegionLoc);
 
+  /// Notifies that a pattern match failed for the given reason.
+  LogicalResult
+  notifyMatchFailure(Location loc,
+                     function_ref<void(Diagnostic &)> reasonCallback);
+
   //===--------------------------------------------------------------------===//
   // State
   //===--------------------------------------------------------------------===//
 
   // Mapping between replaced values that differ in type. This happens when
   // replacing a value with one of a different type.
   ConversionValueMapping mapping;
 
struct ConversionPatternRewriterImpl {
   /// A transaction state for each of operations that were updated in-place.
   SmallVector<OperationTransactionState, 4> rootUpdates;
 
   /// A vector of indices to operations that were replaced with values with
   /// different result types than the original operation, e.g. 1->N conversion
   /// of some kind.
   SmallVector<unsigned, 4> operationsWithChangedResults;
 
+  /// A map of operations that were created during conversion to the type
+  /// converter of the pattern that created them.
+  DenseMap<Operation *, TypeConverter *> opToConverter;
+
   /// A default type converter, used when block conversions do not have one
   /// explicitly provided.
   TypeConverter defaultTypeConverter;
 
+  /// The current conversion pattern that is being rewritten, or nullptr if
+  /// called from outside of a conversion pattern rewrite.
+  const ConversionPattern *currentConversionPattern = nullptr;
+
 #ifndef NDEBUG
   /// A set of operations that have pending updates. This tracking isn't
   /// strictly necessary, and is thus only active during debug builds for extra
   /// verification.
   SmallPtrSet<Operation *, 1> pendingRootUpdates;
 
   /// A logger used to emit diagnostics during the conversion process.
   llvm::ScopedPrinter logger{llvm::dbgs()};
void ConversionPatternRewriterImpl::resetState(RewriterState state) {
   for (auto &repl : llvm::drop_begin(replacements, state.numReplacements))
     for (auto result : repl.first->getResults())
       mapping.erase(result);
   while (replacements.size() != state.numReplacements)
     replacements.pop_back();
 
   // Pop all of the newly created operations.
   while (createdOps.size() != state.numCreatedOps) {
+    opToConverter.erase(createdOps.back());
     detachNestedAndErase(createdOps.back());
     createdOps.pop_back();
   }
 
   // Pop all of the recorded ignored operations that are no longer valid.
   while (ignoredOps.size() != state.numIgnoredOperations)
     ignoredOps.pop_back();
 
case BlockActionKind::TypeConversion: {
       argConverter.discardRewrites(action.block);
       break;
     }
     }
   }
   blockActions.resize(numActionsToKeep);
 }
 
-void ConversionPatternRewriterImpl::remapValues(
+LogicalResult ConversionPatternRewriterImpl::remapValues(
+    Location loc, PatternRewriter &rewriter, TypeConverter *converter,
     Operation::operand_range operands, SmallVectorImpl<Value> &remapped) {
   remapped.reserve(llvm::size(operands));
-  for (Value operand : operands)
-    remapped.push_back(mapping.lookupOrDefault(operand));
+
+  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;
+    if (converter) {
+      // If there is no legal conversion, fail to match this pattern.
+      legalTypes.clear();
+      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.
+    }
+    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.
+      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;
+        });
+      }
+    }
+    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) {
void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op,
 
   // Track if any of the results changed, e.g. erased and replaced with null.
   bool resultChanged = false;
 
   // Create mappings for each of the new result values.
   Value newValue, result;
   for (auto it : llvm::zip(newValues, op->getResults())) {
     std::tie(newValue, result) = it;
-    if (!newValue)
+    if (!newValue) {
       resultChanged = true;
-    else
-      mapping.map(result, newValue);
+      continue;
+    }
+    // Remap, and check for any result type changes.
+    resultChanged |= (newValue.getType() != result.getType());
   }
   if (resultChanged)
     operationsWithChangedResults.push_back(replacements.size());
 
   // Record the requested operation replacement.
-  replacements.insert(std::make_pair(op, OpReplacement(newValues)));
+  TypeConverter *converter = nullptr;
+  if (currentConversionPattern)
+    converter = currentConversionPattern->getTypeConverter();
+  replacements.insert(std::make_pair(op, OpReplacement(newValues, converter)));
 
   // Mark this operation as recursively ignored so that we don't need to
   // convert any nested operations.
   markNestedOpsIgnored(op);
 }
 
 void ConversionPatternRewriterImpl::notifyBlockIsBeingErased(Block *block) {
   Region *region = block->getParent();
void ConversionPatternRewriterImpl::notifyRegionWasClonedBefore(
   auto result = computeConversionSet(blocks, origRegionLoc, createdOps);
 
   // This original region has already had its conversion set computed, so there
   // shouldn't be any new failures.
   (void)result;
   assert(succeeded(result) && "expected region to have no unreachable blocks");
 }
 
+LogicalResult ConversionPatternRewriterImpl::notifyMatchFailure(
+    Location loc, function_ref<void(Diagnostic &)> reasonCallback) {
+  LLVM_DEBUG({
+    Diagnostic diag(loc, DiagnosticSeverity::Remark);
+    reasonCallback(diag);
+    logger.startLine() << "** Failure : " << diag.str() << "\n";
+  });
+  return failure();
+}
+
 //===----------------------------------------------------------------------===//
 // ConversionPatternRewriter
 //===----------------------------------------------------------------------===//
 
 ConversionPatternRewriter::ConversionPatternRewriter(MLIRContext *ctx)
     : PatternRewriter(ctx),
       impl(new detail::ConversionPatternRewriterImpl(*this)) {}
 ConversionPatternRewriter::~ConversionPatternRewriter() {}

 
 /// PatternRewriter hook for creating a new operation.
 void ConversionPatternRewriter::notifyOperationInserted(Operation *op) {
   LLVM_DEBUG({
     impl->logger.startLine()
         << "** Insert  : '" << op->getName() << "'(" << op << ")\n";
   });
   impl->createdOps.push_back(op);
+
+  // If this operation was inserted inside of a conversion, track the type
+  // converter for this operation. This will be used to generated conversions if
+  // necessary.
+  if (impl->currentConversionPattern) {
+    if (auto *converter = impl->currentConversionPattern->getTypeConverter())
+      impl->opToConverter.try_emplace(op, converter);
+  }
 }
 
 /// PatternRewriter hook for updating the root operation in-place.
 void ConversionPatternRewriter::startRootUpdate(Operation *op) {
 #ifndef NDEBUG
   impl->pendingRootUpdates.insert(op);
 #endif
   impl->rootUpdates.emplace_back(op);
#endif
   auto &rootUpdates = impl->rootUpdates;
   auto it = llvm::find_if(llvm::reverse(rootUpdates), stateHasOp);
   rootUpdates.erase(rootUpdates.begin() + (rootUpdates.rend() - it));
 }
 
 /// PatternRewriter hook for notifying match failure reasons.
 LogicalResult ConversionPatternRewriter::notifyMatchFailure(
     Operation *op, function_ref<void(Diagnostic &)> reasonCallback) {
-  LLVM_DEBUG({
-    Diagnostic diag(op->getLoc(), DiagnosticSeverity::Remark);
-    reasonCallback(diag);
-    impl->logger.startLine() << "** Failure : " << diag.str() << "\n";
-  });
-  return failure();
+  return impl->notifyMatchFailure(op->getLoc(), reasonCallback);
 }
 
 /// Return a reference to the internal implementation.
 detail::ConversionPatternRewriterImpl &ConversionPatternRewriter::getImpl() {
   return *impl;
 }
 
 //===----------------------------------------------------------------------===//
 // ConversionPattern
 //===----------------------------------------------------------------------===//
 
 /// Attempt to match and rewrite the IR root at the specified operation.
 LogicalResult
 ConversionPattern::matchAndRewrite(Operation *op,
                                    PatternRewriter &rewriter) const {
-  SmallVector<Value, 4> operands;
   auto &dialectRewriter = static_cast<ConversionPatternRewriter &>(rewriter);
-  dialectRewriter.getImpl().remapValues(op->getOperands(), operands);
+  auto &rewriterImpl = dialectRewriter.getImpl();
+
+  // Track the current conversion pattern in the rewriter.
+  assert(!rewriterImpl.currentConversionPattern &&
+         "already inside of a pattern rewrite");
+  llvm::SaveAndRestore<const ConversionPattern *> currentPatternGuard(
+      rewriterImpl.currentConversionPattern, this);
+
+  // Remap the operands of the operation.
+  SmallVector<Value, 4> operands;
+  if (failed(rewriterImpl.remapValues(op->getLoc(), rewriter,
+                                      getTypeConverter(), op->getOperands(),
+                                      operands))) {
+    return failure();
+  }
   return matchAndRewrite(op, operands, dialectRewriter);
 }
 
 //===----------------------------------------------------------------------===//
 // OperationLegalizer
 //===----------------------------------------------------------------------===//
 
 namespace {

 private:
   /// Converts an operation with the given rewriter.
   LogicalResult convert(ConversionPatternRewriter &rewriter, Operation *op);
 
   /// This method is called after the conversion process to legalize any
   /// remaining legalization artifacts and complete the conversion.
   LogicalResult finalize(ConversionPatternRewriter &rewriter);
 
+  /// Legalize the types of converted block arguments.
+  LogicalResult
+  legalizeConvertedArgumentTypes(ConversionPatternRewriter &rewriter,
+                                 ConversionPatternRewriterImpl &rewriterImpl);
+
+  /// Legalize an operation result that was marked as "erased".
+  LogicalResult
+  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);
+
   /// 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,
   /// this is populated with ops found to be legalizable to the target.
LogicalResult OperationConverter::convertOperations(ArrayRef<Operation *> ops) {
   else
     rewriterImpl.applyRewrites();
   return success();
 }
 
 LogicalResult
 OperationConverter::finalize(ConversionPatternRewriter &rewriter) {
   ConversionPatternRewriterImpl &rewriterImpl = rewriter.getImpl();
-  auto isOpDead = [&](Operation *op) { return rewriterImpl.isOpIgnored(op); };
 
-  // Process the operations with changed results.
-  for (unsigned replIdx : rewriterImpl.operationsWithChangedResults) {
+  // Legalize converted block arguments.
+  if (failed(legalizeConvertedArgumentTypes(rewriter, rewriterImpl)))
+    return failure();
+
+  // 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 (auto it : llvm::zip(repl.first->getResults(), repl.second.newValues)) {
-      Value result = std::get<0>(it), newValue = std::get<1>(it);
+      OpResult result = std::get<0>(it);
+      Value newValue = std::get<1>(it);
 
       // If the operation result was replaced with null, all of the uses of this
       // value should be replaced.
       if (!newValue) {
-        auto liveUserIt = llvm::find_if_not(result.getUsers(), isOpDead);
-        if (liveUserIt != result.user_end()) {
-          InFlightDiagnostic diag = repl.first->emitError()
-                                    << "failed to legalize operation '"
-                                    << repl.first->getName()
-                                    << "' marked as erased";
-          diag.attachNote(liveUserIt->getLoc())
-              << "found live user of result #"
-              << result.cast<OpResult>().getResultNumber() << ": "
-              << *liveUserIt;
-          return failure();
-        }
-      }
+        if (failed(legalizeErasedResult(repl.first, result, rewriterImpl)))
+          return failure();
+        continue;
+      }
+
+      // Otherwise, check to see if the type of the result changed.
+      newValue = rewriterImpl.mapping.lookupOrDefault(newValue);
+      if (result.getType() == newValue.getType())
+        continue;
+
+      // Legalize this result.
+      if (failed(legalizeChangedResultType(repl.first, result, newValue,
+                                           repl.second.converter, rewriter,
+                                           rewriterImpl)))
+        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(
+    ConversionPatternRewriter &rewriter,
+    ConversionPatternRewriterImpl &rewriterImpl) {
+  // Functor used to check if all users of a value will be dead after
+  // conversion.
+  auto isValueDead = [&](Value val) {
+    return llvm::all_of(val.getUsers(), [&](Operation *op) {
+      return rewriterImpl.isOpIgnored(op);
+    });
+  };
+
+  // Materialize any necessary conversions for converted block arguments that
+  // are still live.
+  size_t numCreatedOps = rewriterImpl.createdOps.size();
+  if (failed(rewriterImpl.argConverter.materializeLiveConversions(
+          rewriterImpl.mapping, rewriter, isValueDead)))
+    return failure();
+
+  // Legalize any newly created operatoins during argument materialization.
+  for (int i : llvm::seq<int>(numCreatedOps, rewriterImpl.createdOps.size())) {
+    if (failed(opLegalizer.legalize(rewriterImpl.createdOps[i], rewriter))) {
+      return rewriterImpl.createdOps[i]->emitError()
+             << "failed to legalize conversion operation generated for block "
+                "argument that remained live after conversion";
+    }
+  }
+  return success();
+}
+
+LogicalResult OperationConverter::legalizeErasedResult(
+    Operation *op, OpResult result,
+    ConversionPatternRewriterImpl &rewriterImpl) {
+  // If the operation result was replaced with null, all of the uses of this
+  // value should be replaced.
+  auto liveUserIt = llvm::find_if_not(result.getUsers(), [&](Operation *user) {
+    return rewriterImpl.isOpIgnored(user);
+  });
+    InFlightDiagnostic diag = op->emitError("failed to legalize operation '")
+                              << op->getName() << "' marked as erased";
+        << "found live user of result #" << result.getResultNumber() << ": "
+  return success();
+
+LogicalResult OperationConverter::legalizeChangedResultType(
+    Operation *op, OpResult result, Value newValue,
+    TypeConverter *replConverter, ConversionPatternRewriter &rewriter,
+    ConversionPatternRewriterImpl &rewriterImpl) {
+  Type resultType = result.getType(), newType = newValue.getType();
+
+  // Walk the users of this value to see if there are any live users that
+  // weren't replaced during conversion.
+  Operation *liveUser = nullptr;
+  for (Operation *user : result.getUsers()) {
+    if (!rewriterImpl.isOpIgnored(user))
+      continue;
+
+    // Check to see if this live user has a converter, and that the result
+    // of converting the replacement value is legal for this user. This
+    // situation can happen if an intermediate value is used during
+    // conversion with the expectation that it will be replaced with the
+    // proper value by the end of the conversion process.
+    TypeConverter *converter = rewriterImpl.opToConverter.lookup(user);
+    if (converter && converter->convertType(resultType) != newType) {
+      liveUser = user;
+      break;
     }
   }
+  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
+    // we keep the legacy behavior here of just forwarding the new value.
+    return success();
+  }
+
+  // Track the number of created operations so that new ones can be legalized.
+  size_t numCreatedOps = rewriterImpl.createdOps.size();
+
+  // Materialize a conversion for this live result value.
+  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(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 '"
+             << op->getName() << "' that remained live after conversion";
+    }
+  }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Type Conversion
 //===----------------------------------------------------------------------===//
 
 /// Remap an input of the original signature with a new set of types. The
LogicalResult TypeConverter::convertSignatureArgs(TypeRange types,
                                                   SignatureConversion &result,
                                                   unsigned origInputOffset) {
   for (unsigned i = 0, e = types.size(); i != e; ++i)
     if (failed(convertSignatureArg(origInputOffset + i, types[i], result)))
       return failure();
   return success();
 }
 
-Value TypeConverter::materializeConversion(PatternRewriter &rewriter,
-                                           Location loc, Type resultType,
-                                           ValueRange inputs) {
+Value TypeConverter::materializeConversion(
+    MutableArrayRef<MaterializationCallbackFn> materializations,
+    OpBuilder &builder, Location loc, Type resultType, ValueRange inputs) {
   for (MaterializationCallbackFn &fn : llvm::reverse(materializations))
-    if (Optional<Value> result = fn(rewriter, resultType, inputs, loc))
+    if (Optional<Value> result = fn(builder, resultType, inputs, loc))
       return result.getValue();
   return nullptr;
 }
 
 /// This function converts the type signature of the given block, by invoking
 /// 'convertSignatureArg' for each argument. This function should return a valid
 /// conversion for the signature on success, None otherwise.
 auto TypeConverter::convertBlockSignature(Block *block)

mlir/test/Conversion/StandardToLLVM/standard-to-llvm.mlir

func @rsqrt_multidim_vector(%arg0 : vector<4x3xf32>) {
   // CHECK: %[[EXTRACT:.*]] = llvm.extractvalue %arg0[0] : !llvm<"[4 x <3 x float>]">
   // CHECK: %[[ONE:.*]] = llvm.mlir.constant(dense<1.000000e+00> : vector<3xf32>) : !llvm<"<3 x float>">
   // CHECK: %[[SQRT:.*]] = "llvm.intr.sqrt"(%[[EXTRACT]]) : (!llvm<"<3 x float>">) -> !llvm<"<3 x float>">
   // CHECK: %[[DIV:.*]] = llvm.fdiv %[[ONE]], %[[SQRT]] : !llvm<"<3 x float>">
   // CHECK: %[[INSERT:.*]] = llvm.insertvalue %[[DIV]], %0[0] : !llvm<"[4 x <3 x float>]">
   %0 = rsqrt %arg0 : vector<4x3xf32>
   std.return
 }
-
-// -----
-
-// This should not crash. The first operation cannot be converted, so the
-// second should not match. This attempts to convert `return` to `llvm.return`
-// and complains about non-LLVM types.
-func @unknown_source() -> i32 {
-  %0 = "foo"() : () -> i32
-  %1 = addi %0, %0 : i32
-  // expected-error@+1 {{must be LLVM dialect type}}
-  return %1 : i32
-}

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

 };
 } // namespace
 
 namespace {
 struct TestTypeConverter : public TypeConverter {
   using TypeConverter::TypeConverter;
   TestTypeConverter() {
     addConversion(convertType);
-    addMaterialization(materializeCast);
-    addMaterialization(materializeOneToOneCast);
+    addArgumentMaterialization(materializeCast);
+    addArgumentMaterialization(materializeOneToOneCast);
+    addSourceMaterialization(materializeCast);
   }
 
   static LogicalResult convertType(Type t, SmallVectorImpl<Type> &results) {
     // Drop I16 types.
     if (t.isSignlessInteger(16))
       return success();
 
     // Convert I64 to F64.
static LogicalResult convertType(Type t, SmallVectorImpl<Type> &results) {
 
     // Otherwise, convert the type directly.
     results.push_back(t);
     return success();
   }
 
   /// Hook for materializing a conversion. This is necessary because we generate
   /// 1->N type mappings.
-  static Optional<Value> materializeCast(PatternRewriter &rewriter,
-                                         Type resultType, ValueRange inputs,
-                                         Location loc) {
+  static Optional<Value> materializeCast(OpBuilder &builder, Type resultType,
+                                         ValueRange inputs, Location loc) {
     if (inputs.size() == 1)
       return inputs[0];
-    return rewriter.create<TestCastOp>(loc, resultType, inputs).getResult();
+    return builder.create<TestCastOp>(loc, resultType, inputs).getResult();
   }
 
   /// Materialize the cast for one-to-one conversion from i64 to f64.
-  static Optional<Value> materializeOneToOneCast(PatternRewriter &rewriter,
+  static Optional<Value> materializeOneToOneCast(OpBuilder &builder,
                                                  IntegerType resultType,
                                                  ValueRange inputs,
                                                  Location loc) {
     if (resultType.getWidth() == 42 && inputs.size() == 1)
-      return rewriter.create<TestCastOp>(loc, resultType, inputs).getResult();
+      return builder.create<TestCastOp>(loc, resultType, inputs).getResult();
     return llvm::None;
   }
 };
 
 struct TestLegalizePatternDriver
     : public PassWrapper<TestLegalizePatternDriver, OperationPass<ModuleOp>> {
   /// The mode of conversion to use with the driver.
   enum class ConversionMode { Analysis, Full, Partial };