[mlir] Replace the materializeCallConversion inlining hook.

[gysit]

Summary

This revision removes the materializeCallConversion hook and replaces
it by a isTypeConvertible hook that checks if an argument or a result
conversion is possible. The actual conversion is then performed using
the recently introduced handleArgument and handleResult hooks.

The revision enables argument and result conversions that create
multiple operations rather than a single cast operation. Additionally,
the handlers also have access to argument and result attributes that
may provide extra information such as if an integer has to be sign or
zero extended.

Depends on D145582

[Dinistro]

I only have one small comment apart from that this LGTM.

mlir/lib/Transforms/Utils/InliningUtils.cpp

493

NIT: you can directly return the result of this function.

[gysit]

Address review comment.

[gysit]

Rebase and extrapolate review comments.

[gysit]

Rebase and remove and remove flang type conversion logic.

[gysit]

This revision uses the argument and result handlers, introduced in https://reviews.llvm.org/D145582, to implement type conversion. One side-effect of this change is that the dialect interface is selected based on the callable rather than based on the call operation. The reason for this change is that the new handlers access to the argument and result attributes that are stored on the callable and not on the call. Selecting the dialect interface based on the callable makes more sense. This impacts Flang's inliner implementation since the FIR dialect implements its own call operation but does not implement a custom function operation. As a result, it is not really possible to implement FIR-specific type conversions with the new approach (except by implementing an FIR function operation).

@clementval would it be ok to temporarily drop the materializeCallConversion without a replacement (I did not find a test that exercises the type conversion logic)? Once the FIR dialect has a custom function operation type conversions would be possible again.

Alternatively, we may also return to selecting the dialect interface based on the call operation. I am not fully convinced this is a good idea since the argument and result attributes really are a property of the callable/function.

@rriddle do you have a strong preference here?

[gysit]

We decided to abandon this revision in favor of https://reviews.llvm.org/D147605 which is simpler and satisfies our requirements. The tricky part with this revision was that we may not call the type conversion hook based on the callable dialect since existing downstream implementations (and flang) perform the type conversion based on the call dialect.

flang/lib/Optimizer/Dialect/FIRDialect.cpp

struct FIRInlinerInterface : public mlir::DialectInlinerInterface {
 
   /// This hook checks to see if the operation `op` is legal to inline into the
   /// given region `reg`.
   bool isLegalToInline(mlir::Operation *op, mlir::Region *reg,
                        bool wouldBeCloned, mlir::IRMapping &map) const final {
     return fir::canLegallyInline(op, reg, wouldBeCloned, map);
   }
 
+  /// This hook checks if `sourceType` is convertible to `targetType`.
+  bool isTypeConvertible(mlir::Operation *call, mlir::Operation *callable,
+                         mlir::Type sourceType, mlir::Type targetType,
+                         mlir::DictionaryAttr argOrResAttrs,
+                         bool isResult) const final {
+    return true;
+  }
+
   /// This hook is called when a terminator operation has been inlined.
   /// We handle the return (a Fortran FUNCTION) by replacing the values
   /// previously returned by the call operation with the operands of the
   /// return.
   void handleTerminator(mlir::Operation *op,
                         llvm::ArrayRef<mlir::Value> valuesToRepl) const final {
     auto returnOp = llvm::cast<mlir::func::ReturnOp>(op);
     assert(returnOp.getNumOperands() == valuesToRepl.size());
     for (const auto &it : llvm::enumerate(returnOp.getOperands()))
       valuesToRepl[it.index()].replaceAllUsesWith(it.value());
   }
 
-  mlir::Operation *materializeCallConversion(mlir::OpBuilder &builder,
-                                             mlir::Value input,
-                                             mlir::Type resultType,
-                                             mlir::Location loc) const final {
-    return builder.create<fir::ConvertOp>(loc, resultType, input);
+  /// This hook converts the `argument` to `targetType` if needed.
+  mlir::Value handleArgument(mlir::OpBuilder &builder, mlir::Operation *call,
+                             mlir::Operation *callable, mlir::Value argument,
+                             mlir::Type targetType,
+                             mlir::DictionaryAttr argumentAttrs) const final {
+    // The convert operation folds if the argument type matches the target type.
+    return builder.createOrFold<fir::ConvertOp>(call->getLoc(), targetType,
+                                                argument);
+  }
+
+  /// This hook converts the `result` to `targetType` if needed.
+  mlir::Value handleResult(mlir::OpBuilder &builder, mlir::Operation *call,
+                           mlir::Operation *callable, mlir::Value result,
+                           mlir::Type targetType,
+                           mlir::DictionaryAttr resultAttrs) const final {
+    // The convert operation folds if the result type matches the target type.
+    return builder.createOrFold<fir::ConvertOp>(call->getLoc(), targetType,
+                                                result);
   }
 };
 } // namespace
 
 fir::FIROpsDialect::FIROpsDialect(mlir::MLIRContext *ctx)
     : mlir::Dialect("fir", ctx, mlir::TypeID::get<FIROpsDialect>()) {
   registerTypes();
   registerAttributes();

mlir/docs/Tutorials/Toy/Ch-4.md

struct ToyInlinerInterface : public DialectInlinerInterface {
   /// This hook checks to see if the given operation is legal to inline into the
   /// given region. For Toy this hook can simply return true, as all Toy
   /// operations are inlinable.
   bool isLegalToInline(Operation *, Region *, bool,
                        IRMapping &) const final {
     return true;
   }
 
-  /// This hook cheks if the given 'src' region can be inlined into the 'dest'
+  /// This hook checks if the given 'src' region can be inlined into the 'dest'
   /// region. The regions here are the bodies of the callable functions. For
   /// Toy, any function can be inlined, so we simply return true.
   bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
                        IRMapping &valueMapping) const final {
     return true;
   }
 
+  /// This hook checks if 'sourceType' and 'targetType' are convertible or if
+  /// an argument or result type mismatch prevents inlining. For Toy, all types
+  /// are convertible, so we simply return true.
+  bool isTypeConvertible(Operation *call, Operation *callable, Type sourceType,
+                         Type targetType, DictionaryAttr argOrResAttrs,
+                         bool isResult) const final {
+    return true;
+  }
+
   /// This hook is called when a terminator operation has been inlined. The only
   /// terminator that we have in the Toy dialect is the return
   /// operation(toy.return). We handle the return by replacing the values
   /// previously returned by the call operation with the operands of the
   /// return.
   void handleTerminator(Operation *op,
                         ArrayRef<Value> valuesToRepl) const final {
     // Only "toy.return" needs to be handled here.
bool CastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   if (!input || !output || input.getElementType() != output.getElementType())
     return false;
   // The shape is required to match if both types are ranked.
   return !input.hasRank() || !output.hasRank() || input == output;
 }
 
 ```
 
-With a proper cast operation, we can now override the necessary hook on the
+With a proper cast operation, we can now override the necessary hooks on the
 ToyInlinerInterface to insert it for us when necessary:
 
 ```c++
 struct ToyInlinerInterface : public DialectInlinerInterface {
   ...
 
-  /// Attempts to materialize a conversion for a type mismatch between a call
-  /// from this dialect, and a callable region. This method should generate an
-  /// operation that takes 'input' as the only operand, and produces a single
-  /// result of 'resultType'. If a conversion can not be generated, nullptr
-  /// should be returned.
-  Operation *materializeCallConversion(OpBuilder &builder, Value input,
-                                       Type resultType,
-                                       Location conversionLoc) const final {
-    return builder.create<CastOp>(conversionLoc, resultType, input);
+  /// Handle possible type mismatches between the arguments of call and
+  /// callable. This handler introduces a cast operation if the argument type
+  /// differs from the target type.
+  Value handleArgument(OpBuilder &builder, Operation *call, Operation *callable,
+                       Value argument, Type targetType,
+                       DictionaryAttr argumentAttrs) const final {
+    if (argument.getType() == targetType)
+      return argument;
+    return builder.create<CastOp>(call->getLoc(), targetType, argument);
+  }
+
+  /// Handle possible type mismatches between the results of callable and call.
+  /// This handler introduces a cast operation if the result type differs from
+  /// the target type.
+  Value handleResult(OpBuilder &builder, Operation *call, Operation *callable,
+                     Value result, Type targetType,
+                     DictionaryAttr resultAttrs) const final {
+    if (result.getType() == targetType)
+      return result;
+    return builder.create<CastOp>(call->getLoc(), targetType, result);
   }
 };
 ```
 
 If we run the working example through the pipeline again, we get the expected:
 
 ```mlir
 toy.func @main() {

mlir/examples/toy/Ch4/mlir/Dialect.cpp

bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
     return true;
   }
 
   // All functions within toy can be inlined.
   bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final {
     return true;
   }
 
+  /// All argument and result types within toy are convertible.
+  bool isTypeConvertible(Operation *call, Operation *callable, Type sourceType,
+                         Type targetType, DictionaryAttr argOrResAttrs,
+                         bool isResult) const final {
+    return true;
+  }
+
   //===--------------------------------------------------------------------===//
   // Transformation Hooks
   //===--------------------------------------------------------------------===//
 
   /// Handle the given inlined terminator(toy.return) by replacing it with a new
   /// operation as necessary.
   void handleTerminator(Operation *op,
                         ArrayRef<Value> valuesToRepl) const final {
     // Only "toy.return" needs to be handled here.
     auto returnOp = cast<ReturnOp>(op);
 
     // Replace the values directly with the return operands.
     assert(returnOp.getNumOperands() == valuesToRepl.size());
     for (const auto &it : llvm::enumerate(returnOp.getOperands()))
       valuesToRepl[it.index()].replaceAllUsesWith(it.value());
   }
 
-  /// Attempts to materialize a conversion for a type mismatch between a call
-  /// from this dialect, and a callable region. This method should generate an
-  /// operation that takes 'input' as the only operand, and produces a single
-  /// result of 'resultType'. If a conversion can not be generated, nullptr
-  /// should be returned.
-  Operation *materializeCallConversion(OpBuilder &builder, Value input,
-                                       Type resultType,
-                                       Location conversionLoc) const final {
-    return builder.create<CastOp>(conversionLoc, resultType, input);
+  /// Handle possible type mismatches between the arguments of call and
+  /// callable. This handler introduces a cast operation if the argument type
+  /// differs from the target type.
+  Value handleArgument(OpBuilder &builder, Operation *call, Operation *callable,
+                       Value argument, Type targetType,
+                       DictionaryAttr argumentAttrs) const final {
+    if (argument.getType() == targetType)
+      return argument;
+    return builder.create<CastOp>(call->getLoc(), targetType, argument);
+  }
+
+  /// Handle possible type mismatches between the results of callable and call.
+  /// This handler introduces a cast operation if the result type differs from
+  /// the target type.
+  Value handleResult(OpBuilder &builder, Operation *call, Operation *callable,
+                     Value result, Type targetType,
+                     DictionaryAttr resultAttrs) const final {
+    if (result.getType() == targetType)
+      return result;
+    return builder.create<CastOp>(call->getLoc(), targetType, result);
   }
 };
 
 //===----------------------------------------------------------------------===//
 // ToyDialect
 //===----------------------------------------------------------------------===//
 
 /// Dialect initialization, the instance will be owned by the context. This is

mlir/examples/toy/Ch5/mlir/Dialect.cpp

bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
     return true;
   }
 
   // All functions within toy can be inlined.
   bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final {
     return true;
   }
 
+  /// All argument and result types within toy are convertible.
+  bool isTypeConvertible(Operation *call, Operation *callable, Type sourceType,
+                         Type targetType, DictionaryAttr argOrResAttrs,
+                         bool isResult) const final {
+    return true;
+  }
+
   //===--------------------------------------------------------------------===//
   // Transformation Hooks
   //===--------------------------------------------------------------------===//
 
   /// Handle the given inlined terminator(toy.return) by replacing it with a new
   /// operation as necessary.
   void handleTerminator(Operation *op,
                         ArrayRef<Value> valuesToRepl) const final {
     // Only "toy.return" needs to be handled here.
     auto returnOp = cast<ReturnOp>(op);
 
     // Replace the values directly with the return operands.
     assert(returnOp.getNumOperands() == valuesToRepl.size());
     for (const auto &it : llvm::enumerate(returnOp.getOperands()))
       valuesToRepl[it.index()].replaceAllUsesWith(it.value());
   }
 
-  /// Attempts to materialize a conversion for a type mismatch between a call
-  /// from this dialect, and a callable region. This method should generate an
-  /// operation that takes 'input' as the only operand, and produces a single
-  /// result of 'resultType'. If a conversion can not be generated, nullptr
-  /// should be returned.
-  Operation *materializeCallConversion(OpBuilder &builder, Value input,
-                                       Type resultType,
-                                       Location conversionLoc) const final {
-    return builder.create<CastOp>(conversionLoc, resultType, input);
+  /// Handle possible type mismatches between the arguments of call and
+  /// callable. This handler introduces a cast operation if the argument type
+  /// differs from the target type.
+  Value handleArgument(OpBuilder &builder, Operation *call, Operation *callable,
+                       Value argument, Type targetType,
+                       DictionaryAttr argumentAttrs) const final {
+    if (argument.getType() == targetType)
+      return argument;
+    return builder.create<CastOp>(call->getLoc(), targetType, argument);
+  }
+
+  /// Handle possible type mismatches between the results of callable and call.
+  /// This handler introduces a cast operation if the result type differs from
+  /// the target type.
+  Value handleResult(OpBuilder &builder, Operation *call, Operation *callable,
+                     Value result, Type targetType,
+                     DictionaryAttr resultAttrs) const final {
+    if (result.getType() == targetType)
+      return result;
+    return builder.create<CastOp>(call->getLoc(), targetType, result);
   }
 };
 
 //===----------------------------------------------------------------------===//
 // ToyDialect
 //===----------------------------------------------------------------------===//
 
 /// Dialect initialization, the instance will be owned by the context. This is

mlir/examples/toy/Ch6/mlir/Dialect.cpp

bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
     return true;
   }
 
   // All functions within toy can be inlined.
   bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final {
     return true;
   }
 
+  /// All argument and result types within toy are convertible.
+  bool isTypeConvertible(Operation *call, Operation *callable, Type sourceType,
+                         Type targetType, DictionaryAttr argOrResAttrs,
+                         bool isResult) const final {
+    return true;
+  }
+
   //===--------------------------------------------------------------------===//
   // Transformation Hooks
   //===--------------------------------------------------------------------===//
 
   /// Handle the given inlined terminator(toy.return) by replacing it with a new
   /// operation as necessary.
   void handleTerminator(Operation *op,
                         ArrayRef<Value> valuesToRepl) const final {
     // Only "toy.return" needs to be handled here.
     auto returnOp = cast<ReturnOp>(op);
 
     // Replace the values directly with the return operands.
     assert(returnOp.getNumOperands() == valuesToRepl.size());
     for (const auto &it : llvm::enumerate(returnOp.getOperands()))
       valuesToRepl[it.index()].replaceAllUsesWith(it.value());
   }
 
-  /// Attempts to materialize a conversion for a type mismatch between a call
-  /// from this dialect, and a callable region. This method should generate an
-  /// operation that takes 'input' as the only operand, and produces a single
-  /// result of 'resultType'. If a conversion can not be generated, nullptr
-  /// should be returned.
-  Operation *materializeCallConversion(OpBuilder &builder, Value input,
-                                       Type resultType,
-                                       Location conversionLoc) const final {
-    return builder.create<CastOp>(conversionLoc, resultType, input);
+  /// Handle possible type mismatches between the arguments of call and
+  /// callable. This handler introduces a cast operation if the argument type
+  /// differs from the target type.
+  Value handleArgument(OpBuilder &builder, Operation *call, Operation *callable,
+                       Value argument, Type targetType,
+                       DictionaryAttr argumentAttrs) const final {
+    if (argument.getType() == targetType)
+      return argument;
+    return builder.create<CastOp>(call->getLoc(), targetType, argument);
+  }
+
+  /// Handle possible type mismatches between the results of callable and call.
+  /// This handler introduces a cast operation if the result type differs from
+  /// the target type.
+  Value handleResult(OpBuilder &builder, Operation *call, Operation *callable,
+                     Value result, Type targetType,
+                     DictionaryAttr resultAttrs) const final {
+    if (result.getType() == targetType)
+      return result;
+    return builder.create<CastOp>(call->getLoc(), targetType, result);
   }
 };
 
 //===----------------------------------------------------------------------===//
 // ToyDialect
 //===----------------------------------------------------------------------===//
 
 /// Dialect initialization, the instance will be owned by the context. This is

mlir/examples/toy/Ch7/mlir/Dialect.cpp

bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
     return true;
   }
 
   // All functions within toy can be inlined.
   bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final {
     return true;
   }
 
+  /// All argument and result types within toy are convertible.
+  bool isTypeConvertible(Operation *call, Operation *callable, Type sourceType,
+                         Type targetType, DictionaryAttr argOrResAttrs,
+                         bool isResult) const final {
+    return true;
+  }
+
   //===--------------------------------------------------------------------===//
   // Transformation Hooks
   //===--------------------------------------------------------------------===//
 
   /// Handle the given inlined terminator(toy.return) by replacing it with a new
   /// operation as necessary.
   void handleTerminator(Operation *op,
                         ArrayRef<Value> valuesToRepl) const final {
     // Only "toy.return" needs to be handled here.
     auto returnOp = cast<ReturnOp>(op);
 
     // Replace the values directly with the return operands.
     assert(returnOp.getNumOperands() == valuesToRepl.size());
     for (const auto &it : llvm::enumerate(returnOp.getOperands()))
       valuesToRepl[it.index()].replaceAllUsesWith(it.value());
   }
 
-  /// Attempts to materialize a conversion for a type mismatch between a call
-  /// from this dialect, and a callable region. This method should generate an
-  /// operation that takes 'input' as the only operand, and produces a single
-  /// result of 'resultType'. If a conversion can not be generated, nullptr
-  /// should be returned.
-  Operation *materializeCallConversion(OpBuilder &builder, Value input,
-                                       Type resultType,
-                                       Location conversionLoc) const final {
-    return builder.create<CastOp>(conversionLoc, resultType, input);
+  /// Handle possible type mismatches between the arguments of call and
+  /// callable. This handler introduces a cast operation if the argument type
+  /// differs from the target type.
+  Value handleArgument(OpBuilder &builder, Operation *call, Operation *callable,
+                       Value argument, Type targetType,
+                       DictionaryAttr argumentAttrs) const final {
+    if (argument.getType() == targetType)
+      return argument;
+    return builder.create<CastOp>(call->getLoc(), targetType, argument);
+  }
+
+  /// Handle possible type mismatches between the results of callable and call.
+  /// This handler introduces a cast operation if the result type differs from
+  /// the target type.
+  Value handleResult(OpBuilder &builder, Operation *call, Operation *callable,
+                     Value result, Type targetType,
+                     DictionaryAttr resultAttrs) const final {
+    if (result.getType() == targetType)
+      return result;
+    return builder.create<CastOp>(call->getLoc(), targetType, result);
   }
 };
 
 //===----------------------------------------------------------------------===//
 // Toy Operations
 //===----------------------------------------------------------------------===//
 
 /// A generalized parser for binary operations. This parses the different forms

mlir/include/mlir/Transforms/InliningUtils.h

public:
   /// in-place(i.e. no duplicates would be created). 'valueMapping' contains any
   /// remapped values from within the 'src' region. This can be used to examine
   /// what values may potentially replace the operands to 'op'.
   virtual bool isLegalToInline(Operation *op, Region *dest, bool wouldBeCloned,
                                IRMapping &valueMapping) const {
     return false;
   }
 
+  /// Returns true if `sourceType`, either a call argument type or a callee
+  /// result type, can be converted to `targetType`, or false otherwise. This
+  /// hook is only called for arguments and results that require type conversion
+  /// for inlining. It executes right before the `isLegalToInline` method, which
+  /// checks if a callable can be inlined in-place of a call. The `call` is
+  /// registered to the current dialect and implements `CallOpInterface`, while
+  /// the `callable` implements `CallableOpInterface`. The `argOrResAttrs`
+  /// dictionary contains the argument or result attributes and is guaranteed to
+  /// be non-null even if no attributes are present.
+  virtual bool isTypeConvertible(Operation *call, Operation *callable,
+                                 Type sourceType, Type targetType,
+                                 DictionaryAttr argOrResAttrs,
+                                 bool isResult) const {
+    return false;
+  }
+
   /// This hook is invoked on an operation that contains regions. It should
   /// return true if the analyzer should recurse within the regions of this
   /// operation when computing legality and cost, false otherwise. The default
   /// implementation returns true.
   virtual bool shouldAnalyzeRecursively(Operation *op) const { return true; }
 
   //===--------------------------------------------------------------------===//
   // Transformation Hooks
public:
   /// operation). The given 'op' will be removed by the caller, after this
   /// function has been called.
   virtual void handleTerminator(Operation *op,
                                 ArrayRef<Value> valuesToReplace) const {
     llvm_unreachable(
         "must implement handleTerminator in the case of one inlined block");
   }
 
-  /// Attempt to materialize a conversion for a type mismatch between a call
-  /// from this dialect, and a callable region. This method should generate an
-  /// operation that takes 'input' as the only operand, and produces a single
-  /// result of 'resultType'. If a conversion can not be generated, nullptr
-  /// should be returned. For example, this hook may be invoked in the following
-  /// scenarios:
-  ///   func @foo(i32) -> i32 { ... }
-  ///
-  ///   // Mismatched input operand
-  ///   ... = foo.call @foo(%input : i16) -> i32
-  ///
-  ///   // Mismatched result type.
-  ///   ... = foo.call @foo(%input : i32) -> i16
-  ///
-  /// NOTE: This hook may be invoked before the 'isLegal' checks above.
-  virtual Operation *materializeCallConversion(OpBuilder &builder, Value input,
-                                               Type resultType,
-                                               Location conversionLoc) const {
-    return nullptr;
-  }
-
   /// Hook to transform the call arguments before using them to replace the
   /// callee arguments. It returns the transformation result or `argument`
   /// itself if the hook did not change anything. The type of the returned value
   /// has to match `targetType`, and the `argumentAttrs` dictionary is non-null
-  /// even if no attribute is present. The hook is called after converting the
-  /// callsite argument types using the materializeCallConversion callback, and
-  /// right before inlining the callee region. Any operations created using the
-  /// provided `builder` are inserted right before the inlined callee region.
-  /// Example use cases are the insertion of copies for by value arguments, or
-  /// integer conversions that require signedness information.
+  /// even if no attributes are present. The hook is called right before
+  /// inlining the callee region. Any operations created using the provided
+  /// `builder` are inserted right before the inlined callee region. Example use
+  /// cases are the insertion of copies for by value arguments, or integer
+  /// conversions that require signedness information.
   virtual Value handleArgument(OpBuilder &builder, Operation *call,
                                Operation *callable, Value argument,
                                Type targetType,
                                DictionaryAttr argumentAttrs) const {
     return argument;
   }
 
   /// Hook to transform the callee results before using them to replace the call
   /// results. It returns the transformation result or the `result` itself if
   /// the hook did not change anything. The type of the returned values has to
   /// match `targetType`, and the `resultAttrs` dictionary is non-null even if
-  /// no attribute is present. The hook is called right before handling
-  /// terminators, and obtains the callee result before converting its type
-  /// using the `materializeCallConversion` callback. Any operations created
-  /// using the provided `builder` are inserted right after the inlined callee
-  /// region. Example use cases are the insertion of copies for by value results
-  /// or integer conversions that require signedness information.
+  /// no attributes are present. The hook is called right before handling
+  /// terminators. Any operations created using the provided `builder` are
+  /// inserted right after the inlined callee region. Example use cases are the
+  /// insertion of copies for by value results or integer conversions that
+  /// require signedness information.
   /// NOTE: This hook is invoked after inlining the `callable` region.
   virtual Value handleResult(OpBuilder &builder, Operation *call,
                              Operation *callable, Value result, Type targetType,
                              DictionaryAttr resultAttrs) const {
     return result;
   }
 
   /// Process a set of blocks that have been inlined for a call. This callback
public:
   //===--------------------------------------------------------------------===//
 
   virtual bool isLegalToInline(Operation *call, Operation *callable,
                                bool wouldBeCloned) const;
   virtual bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
                                IRMapping &valueMapping) const;
   virtual bool isLegalToInline(Operation *op, Region *dest, bool wouldBeCloned,
                                IRMapping &valueMapping) const;
+  virtual bool isTypeConvertible(Operation *call, Operation *callable,
+                                 Type sourceType, Type targetType,
+                                 DictionaryAttr argOrResAttrs,
+                                 bool isResult) const;
   virtual bool shouldAnalyzeRecursively(Operation *op) const;
 
   //===--------------------------------------------------------------------===//
   // Transformation Hooks
   //===--------------------------------------------------------------------===//
 
   virtual void handleTerminator(Operation *op, Block *newDest) const;
   virtual void handleTerminator(Operation *op,

mlir/lib/Transforms/Utils/InliningUtils.cpp

 bool InlinerInterface::isLegalToInline(Operation *op, Region *dest,
                                        bool wouldBeCloned,
                                        IRMapping &valueMapping) const {
   if (auto *handler = getInterfaceFor(op))
     return handler->isLegalToInline(op, dest, wouldBeCloned, valueMapping);
   return false;
 }
 
+bool InlinerInterface::isTypeConvertible(Operation *call, Operation *callable,
+                                         Type sourceType, Type targetType,
+                                         DictionaryAttr argOrResAttrs,
+                                         bool isResult) const {
+  if (auto *handler = getInterfaceFor(call))
+    return handler->isTypeConvertible(call, callable, sourceType, targetType,
+                                      argOrResAttrs, isResult);
+  return false;
+}
+
 bool InlinerInterface::shouldAnalyzeRecursively(Operation *op) const {
   auto *handler = getInterfaceFor(op);
   return handler ? handler->shouldAnalyzeRecursively(op) : true;
 }
 
 /// Handle the given inlined terminator by replacing it with a new operation
 /// as necessary.
 void InlinerInterface::handleTerminator(Operation *op, Block *newDest) const {
static bool isLegalToInline(InlinerInterface &interface, Region *src,
   }
   return true;
 }
 
 //===----------------------------------------------------------------------===//
 // Inline Methods
 //===----------------------------------------------------------------------===//
 
-static void handleArgumentImpl(InlinerInterface &interface, OpBuilder &builder,
-                               CallOpInterface call,
-                               CallableOpInterface callable,
-                               IRMapping &mapper) {
-  if (!call || !callable)
-    return;
-
-  // Unpack the argument attributes if there are any.
+/// Returns the array of argument attribute dictionaries. The attribute
+/// dictionaries are non-null even if no attributes are present.
+static SmallVector<DictionaryAttr>
+getArgumentAttributes(CallableOpInterface callable) {
   SmallVector<DictionaryAttr> argAttrs(
       callable.getCallableRegion()->getNumArguments(),
-      builder.getDictionaryAttr({}));
+      DictionaryAttr::get(callable->getContext(), {}));
   if (ArrayAttr arrayAttr = callable.getCallableArgAttrs()) {
     assert(arrayAttr.size() == argAttrs.size());
     for (auto [idx, attr] : llvm::enumerate(arrayAttr))
       argAttrs[idx] = cast<DictionaryAttr>(attr);
   }
+  return argAttrs;
+}
+
+/// Returns the array of result attribute dictionaries. The attribute
+/// dictionaries are non-null even if no attributes are present.
+static SmallVector<DictionaryAttr>
+getResultAttributes(CallableOpInterface callable) {
+  SmallVector<DictionaryAttr> resAttrs(
+      callable.getCallableResults().size(),
+      DictionaryAttr::get(callable->getContext(), {}));
+  if (ArrayAttr arrayAttr = callable.getCallableResAttrs()) {
+    assert(arrayAttr.size() == resAttrs.size());
+    for (auto [idx, attr] : llvm::enumerate(arrayAttr))
+      resAttrs[idx] = cast<DictionaryAttr>(attr);
+  }
+  return resAttrs;
+}
+
+static void handleArgumentImpl(InlinerInterface &interface, OpBuilder &builder,
+                               CallOpInterface call,
+                               CallableOpInterface callable,
+                               IRMapping &mapper) {
+  if (!call || !callable)
+    return;
+
+  // Unpack the argument attributes.
+  SmallVector<DictionaryAttr> argAttrs = getArgumentAttributes(callable);
 
   // Run the argument attribute handler for the given argument and attribute.
-  for (auto [blockArg, argAttr] :
-       llvm::zip(callable.getCallableRegion()->getArguments(), argAttrs)) {
+  for (auto [blockArg, argAttr] : llvm::zip_equal(
+           callable.getCallableRegion()->getArguments(), argAttrs)) {
     Value newArgument = interface.handleArgument(builder, call, callable,
                                                  mapper.lookup(blockArg),
                                                  blockArg.getType(), argAttr);
     assert(newArgument.getType() == blockArg.getType() &&
            "expected the handled argument type to match the target type");
 
     // Update the mapping to point the new argument returned by the handler.
     mapper.map(blockArg, newArgument);
   }
 }
 
 static void handleResultImpl(InlinerInterface &interface, OpBuilder &builder,
                              CallOpInterface call, CallableOpInterface callable,
                              ValueRange results) {
   if (!call || !callable)
     return;
 
-  // Unpack the result attributes if there are any.
-  SmallVector<DictionaryAttr> resAttrs(results.size(),
-                                       builder.getDictionaryAttr({}));
-  if (ArrayAttr arrayAttr = callable.getCallableResAttrs()) {
-    assert(arrayAttr.size() == resAttrs.size());
-    for (auto [idx, attr] : llvm::enumerate(arrayAttr))
-      resAttrs[idx] = cast<DictionaryAttr>(attr);
-  }
+  // Unpack the result attributes.
+  SmallVector<DictionaryAttr> resAttrs = getResultAttributes(callable);
 
   // Run the result attribute handler for the given result and attribute.
   SmallVector<DictionaryAttr> resultAttributes;
-  for (auto [result, resAttr] : llvm::zip(results, resAttrs)) {
+  for (auto [result, resAttr, callResult] :
+       llvm::zip_equal(results, resAttrs, call->getResults())) {
     // Store the original result users before running the handler.
     DenseSet<Operation *> resultUsers;
     for (Operation *user : result.getUsers())
       resultUsers.insert(user);
 
-    // TODO: Use the type of the call result to replace once the hook can be
-    // used for type conversions. At the moment, all type conversions have to be
-    // done using materializeCallConversion.
-    Type targetType = result.getType();
-
     Value newResult = interface.handleResult(builder, call, callable, result,
-                                             targetType, resAttr);
-    assert(newResult.getType() == targetType &&
+                                             callResult.getType(), resAttr);
+    assert(newResult.getType() == callResult.getType() &&
            "expected the handled result type to match the target type");
 
     // Replace the result uses except for the ones introduce by the handler.
     result.replaceUsesWithIf(newResult, [&](OpOperand &operand) {
       return resultUsers.count(operand.getOwner());
     });
   }
 }
LogicalResult mlir::inlineRegion(InlinerInterface &interface, Region *src,
                                  ValueRange resultsToReplace,
                                  std::optional<Location> inlineLoc,
                                  bool shouldCloneInlinedRegion) {
   return inlineRegionImpl(interface, src, inlineBlock, inlinePoint,
                           inlinedOperands, resultsToReplace, inlineLoc,
                           shouldCloneInlinedRegion);
 }
 
-/// Utility function used to generate a cast operation from the given interface,
-/// or return nullptr if a cast could not be generated.
-static Value materializeConversion(const DialectInlinerInterface *interface,
-                                   SmallVectorImpl<Operation *> &castOps,
-                                   OpBuilder &castBuilder, Value arg, Type type,
-                                   Location conversionLoc) {
-  if (!interface)
-    return nullptr;
-
-  // Check to see if the interface for the call can materialize a conversion.
-  Operation *castOp = interface->materializeCallConversion(castBuilder, arg,
-                                                           type, conversionLoc);
-  if (!castOp)
-    return nullptr;
-  castOps.push_back(castOp);
-
-  // Ensure that the generated cast is correct.
-  assert(castOp->getNumOperands() == 1 && castOp->getOperand(0) == arg &&
-         castOp->getNumResults() == 1 && *castOp->result_type_begin() == type);
-  return castOp->getResult(0);
-}
-
 /// This function inlines a given region, 'src', of a callable operation,
 /// 'callable', into the location defined by the given call operation. This
 /// function returns failure if inlining is not possible, success otherwise. On
 /// failure, no changes are made to the module. 'shouldCloneInlinedRegion'
 /// corresponds to whether the source region should be cloned into the 'call' or
 /// spliced directly.
 LogicalResult mlir::inlineCall(InlinerInterface &interface,
                                CallOpInterface call,
                                CallableOpInterface callable, Region *src,
                                bool shouldCloneInlinedRegion) {
   // We expect the region to have at least one block.
   if (src->empty())
     return failure();
   auto *entryBlock = &src->front();
   ArrayRef<Type> callableResultTypes = callable.getCallableResults();
 
   // Make sure that the number of arguments and results matchup between the call
   // and the region.
   SmallVector<Value, 8> callOperands(call.getArgOperands());
   SmallVector<Value, 8> callResults(call->getResults());
   if (callOperands.size() != entryBlock->getNumArguments() ||
       callResults.size() != callableResultTypes.size())
     return failure();
 
-  // A set of cast operations generated to matchup the signature of the region
-  // with the signature of the call.
-  SmallVector<Operation *, 4> castOps;
-  castOps.reserve(callOperands.size() + callResults.size());
-
-  // Functor used to cleanup generated state on failure.
-  auto cleanupState = [&] {
-    for (auto *op : castOps) {
-      op->getResult(0).replaceAllUsesWith(op->getOperand(0));
-      op->erase();
-    }
+  // Check that argument types are convertible.
+  SmallVector<DictionaryAttr> argAttrs = getArgumentAttributes(callable);
+  for (auto [argumentType, targetType, argumentAttrs] :
+       llvm::zip_equal(call.getArgOperands().getTypes(),
+                       entryBlock->getArgumentTypes(), argAttrs)) {
+    if (argumentType == targetType)
+      continue;
+    if (!interface.isTypeConvertible(call, callable, argumentType, targetType,
+                                     argumentAttrs, false))
       return failure();
-  };
-
-  // Builder used for any conversion operations that need to be materialized.
-  OpBuilder castBuilder(call);
-  Location castLoc = call.getLoc();
-  const auto *callInterface = interface.getInterfaceFor(call->getDialect());
-
-  // Map the provided call operands to the arguments of the region.
-  IRMapping mapper;
-  for (unsigned i = 0, e = callOperands.size(); i != e; ++i) {
-    BlockArgument regionArg = entryBlock->getArgument(i);
-    Value operand = callOperands[i];
+  }
 
-    // If the call operand doesn't match the expected region argument, try to
-    // generate a cast.
-    Type regionArgType = regionArg.getType();
-    if (operand.getType() != regionArgType) {
-      if (!(operand = materializeConversion(callInterface, castOps, castBuilder,
-                                            operand, regionArgType, castLoc)))
-        return cleanupState();
-    }
-    mapper.map(regionArg, operand);
-  }
-
-  // Ensure that the resultant values of the call match the callable.
-  castBuilder.setInsertionPointAfter(call);
-  for (unsigned i = 0, e = callResults.size(); i != e; ++i) {
-    Value callResult = callResults[i];
-    if (callResult.getType() == callableResultTypes[i])
+  // Check that result types are convertible.
+  SmallVector<DictionaryAttr> resAttrs = getResultAttributes(callable);
+  for (auto [resultType, targetType, resultAttrs] :
+       llvm::zip_equal(callableResultTypes, call->getResultTypes(), resAttrs)) {
+    if (resultType == targetType)
       continue;
-
-    // Generate a conversion that will produce the original type, so that the IR
-    // is still valid after the original call gets replaced.
-    Value castResult =
-        materializeConversion(callInterface, castOps, castBuilder, callResult,
-                              callResult.getType(), castLoc);
-    if (!castResult)
-      return cleanupState();
-    callResult.replaceAllUsesWith(castResult);
-    castResult.getDefiningOp()->replaceUsesOfWith(castResult, callResult);
+    if (!interface.isTypeConvertible(call, callable, resultType, targetType,
+                                     resultAttrs, true))
+      return failure();
   }
 
   // Check that it is legal to inline the callable into the call.
   if (!interface.isLegalToInline(call, callable, shouldCloneInlinedRegion))
-    return cleanupState();
+    return failure();
+
+  IRMapping mapper;
+  for (auto [blockArg, operand] :
+       llvm::zip(entryBlock->getArguments(), callOperands))
+    mapper.map(blockArg, operand);
 
   // Attempt to inline the call.
   if (failed(inlineRegionImpl(interface, src, call->getBlock(),
                               ++call->getIterator(), mapper, callResults,
                               callableResultTypes, call.getLoc(),
                               shouldCloneInlinedRegion, call)))
-    return cleanupState();
+    return failure();
   return success();
 }

mlir/test/Transforms/inlining.mlir

 }
 
 // Check that we can handle argument and result attributes.
 func.func @handle_attr_callee_fn_multi_arg(%arg0 : i16, %arg1 : i16 {"test.handle_argument"}) -> (i16 {"test.handle_result"}, i16) {
   %0 = arith.addi %arg0, %arg1 : i16
   %1 = arith.subi %arg0, %arg1 : i16
   return %0, %1 : i16, i16
 }
-func.func @handle_attr_callee_fn(%arg0 : i32 {"test.handle_argument"}) -> (i32 {"test.handle_result"}) {
-  return %arg0 : i32
-}
 
 // CHECK-LABEL: func @inline_handle_attr_call
 // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
 // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]
 func.func @inline_handle_attr_call(%arg0 : i16, %arg1 : i16) -> (i16, i16) {
 
   // CHECK: %[[CHANGE_INPUT:.*]] = "test.type_changer"(%[[ARG1]]) : (i16) -> i16
   // CHECK: %[[SUM:.*]] = arith.addi %[[ARG0]], %[[CHANGE_INPUT]]
   // CHECK: %[[DIFF:.*]] = arith.subi %[[ARG0]], %[[CHANGE_INPUT]]
   // CHECK: %[[CHANGE_RESULT:.*]] = "test.type_changer"(%[[SUM]]) : (i16) -> i16
   // CHECK-NEXT: return %[[CHANGE_RESULT]], %[[DIFF]]
   %res0, %res1 = "test.conversion_call_op"(%arg0, %arg1) { callee=@handle_attr_callee_fn_multi_arg } : (i16, i16) -> (i16, i16)
   return %res0, %res1 : i16, i16
 }
-
-// CHECK-LABEL: func @inline_convert_and_handle_attr_call
-// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
-func.func @inline_convert_and_handle_attr_call(%arg0 : i16) -> (i16) {
-
-  // CHECK: %[[CAST_INPUT:.*]] = "test.cast"(%[[ARG0]]) : (i16) -> i32
-  // CHECK: %[[CHANGE_INPUT:.*]] = "test.type_changer"(%[[CAST_INPUT]]) : (i32) -> i32
-  // CHECK: %[[CHANGE_RESULT:.*]] = "test.type_changer"(%[[CHANGE_INPUT]]) : (i32) -> i32
-  // CHECK: %[[CAST_RESULT:.*]] = "test.cast"(%[[CHANGE_RESULT]]) : (i32) -> i16
-  // CHECK: return %[[CAST_RESULT]]
-  %res = "test.conversion_call_op"(%arg0) { callee=@handle_attr_callee_fn } : (i16) -> (i16)
-  return %res : i16
-}

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

struct TestInlinerInterface : public DialectInlinerInterface {
   // Analysis Hooks
   //===--------------------------------------------------------------------===//
 
   bool isLegalToInline(Operation *call, Operation *callable,
                        bool wouldBeCloned) const final {
     // Don't allow inlining calls that are marked `noinline`.
     return !call->hasAttr("noinline");
   }
+  bool isTypeConvertible(Operation *call, Operation *callable, Type sourceType,
+                         Type targetType, DictionaryAttr argOrResAttrs,
+                         bool isResult) const final {
+    return (sourceType.isSignlessInteger(16) ||
+            sourceType.isSignlessInteger(32)) &&
+           (targetType.isSignlessInteger(16) ||
+            targetType.isSignlessInteger(32));
+  }
   bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final {
     // Inlining into test dialect regions is legal.
     return true;
   }
   bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
     return true;
   }
 
if (!returnOp)
       return;
 
     // Replace the values directly with the return operands.
     assert(returnOp.getNumOperands() == valuesToRepl.size());
     for (const auto &it : llvm::enumerate(returnOp.getOperands()))
       valuesToRepl[it.index()].replaceAllUsesWith(it.value());
   }
 
-  /// Attempt to materialize a conversion for a type mismatch between a call
-  /// from this dialect, and a callable region. This method should generate an
-  /// operation that takes 'input' as the only operand, and produces a single
-  /// result of 'resultType'. If a conversion can not be generated, nullptr
-  /// should be returned.
-  Operation *materializeCallConversion(OpBuilder &builder, Value input,
-                                       Type resultType,
-                                       Location conversionLoc) const final {
-    // Only allow conversion for i16/i32 types.
-    if (!(resultType.isSignlessInteger(16) ||
-          resultType.isSignlessInteger(32)) ||
-        !(input.getType().isSignlessInteger(16) ||
-          input.getType().isSignlessInteger(32)))
-      return nullptr;
-    return builder.create<TestCastOp>(conversionLoc, resultType, input);
-  }
-
   Value handleArgument(OpBuilder &builder, Operation *call, Operation *callable,
                        Value argument, Type targetType,
                        DictionaryAttr argumentAttrs) const final {
-    if (!argumentAttrs.contains("test.handle_argument"))
-      return argument;
+    // Cast if an attribute is present.
+    if (argumentAttrs.contains("test.handle_argument"))
       return builder.create<TestTypeChangerOp>(call->getLoc(), targetType,
                                                argument);
+
+    // Cast only if the types do not match.
+    if (argument.getType() == targetType)
+      return argument;
+    return builder.create<TestCastOp>(call->getLoc(), targetType, argument);
   }
 
   Value handleResult(OpBuilder &builder, Operation *call, Operation *callable,
                      Value result, Type targetType,
                      DictionaryAttr resultAttrs) const final {
-    if (!resultAttrs.contains("test.handle_result"))
-      return result;
+    // Cast if an attribute is present.
+    if (resultAttrs.contains("test.handle_result"))
       return builder.create<TestTypeChangerOp>(call->getLoc(), targetType,
                                                result);
+    // Cast only if the types do not match.
+    if (result.getType() == targetType)
+      return result;
+    return builder.create<TestCastOp>(call->getLoc(), targetType, result);
   }
 
   void processInlinedCallBlocks(
       Operation *call,
       iterator_range<Region::iterator> inlinedBlocks) const final {
     if (!isa<ConversionCallOp>(call))
       return;