[mlir] Make the inliner arg/res handlers more concise.

[gysit]

Summary

The revision ensures the newly introduced argument
and result handlers cannot be used for type conversion.
Instead use the existing materializeCallConversion hook to
perform type conversions.

mlir/include/mlir/Transforms/InliningUtils.h

public:
   /// 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
+  /// callee arguments. Returns a value of the same type or the `argument`
+  /// itself if nothing changed. 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.
+  /// provided `builder` are inserted right before the inlined callee region. An
+  /// example use case is the insertion of copies for by value arguments.
   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
+  /// results. Returns a value of the same type or the `result` itself if
+  /// nothing changed. 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.
-  /// NOTE: This hook is invoked after inlining the `callable` region.
+  /// region. An example use case is the insertion of copies for by value
+  /// results. NOTE: This hook is invoked after inlining the `callable` region.
   virtual Value handleResult(OpBuilder &builder, Operation *call,
-                             Operation *callable, Value result, Type targetType,
+                             Operation *callable, Value result,
                              DictionaryAttr resultAttrs) const {
     return result;
   }
 
   /// Process a set of blocks that have been inlined for a call. This callback
   /// is invoked before inlined terminator operations have been processed.
   virtual void processInlinedCallBlocks(
       Operation *call, iterator_range<Region::iterator> inlinedBlocks) const {}
public:
   //===--------------------------------------------------------------------===//
 
   virtual void handleTerminator(Operation *op, Block *newDest) const;
   virtual void handleTerminator(Operation *op,
                                 ArrayRef<Value> valuesToRepl) const;
 
   virtual Value handleArgument(OpBuilder &builder, Operation *call,
                                Operation *callable, Value argument,
-                               Type targetType,
                                DictionaryAttr argumentAttrs) const;
   virtual Value handleResult(OpBuilder &builder, Operation *call,
-                             Operation *callable, Value result, Type targetType,
+                             Operation *callable, Value result,
                              DictionaryAttr resultAttrs) const;
 
   virtual void processInlinedCallBlocks(
       Operation *call, iterator_range<Region::iterator> inlinedBlocks) const;
 };
 
 //===----------------------------------------------------------------------===//
 // Inline Methods.

mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp

void handleTerminator(Operation *op,
     // Replace the values directly with the return operands.
     assert(returnOp.getNumOperands() == valuesToRepl.size());
     for (const auto &[dst, src] :
          llvm::zip(valuesToRepl, returnOp.getOperands()))
       dst.replaceAllUsesWith(src);
   }
 
   Value handleArgument(OpBuilder &builder, Operation *call, Operation *callable,
-                       Value argument, Type targetType,
+                       Value argument,
                        DictionaryAttr argumentAttrs) const final {
     if (std::optional<NamedAttribute> attr =
             argumentAttrs.getNamed(LLVM::LLVMDialect::getByValAttrName())) {
       Type elementType = cast<TypeAttr>(attr->getValue()).getValue();
       unsigned requestedAlignment = 1;
       if (std::optional<NamedAttribute> alignAttr =
               argumentAttrs.getNamed(LLVM::LLVMDialect::getAlignAttrName())) {
         requestedAlignment = cast<IntegerAttr>(alignAttr->getValue())

mlir/lib/Transforms/Utils/InliningUtils.cpp

void InlinerInterface::handleTerminator(Operation *op,
                                         ArrayRef<Value> valuesToRepl) const {
   auto *handler = getInterfaceFor(op);
   assert(handler && "expected valid dialect handler");
   handler->handleTerminator(op, valuesToRepl);
 }
 
 Value InlinerInterface::handleArgument(OpBuilder &builder, Operation *call,
                                        Operation *callable, Value argument,
-                                       Type targetType,
                                        DictionaryAttr argumentAttrs) const {
   auto *handler = getInterfaceFor(callable);
   assert(handler && "expected valid dialect handler");
-  return handler->handleArgument(builder, call, callable, argument, targetType,
+  return handler->handleArgument(builder, call, callable, argument,
                                  argumentAttrs);
 }
 
 Value InlinerInterface::handleResult(OpBuilder &builder, Operation *call,
                                      Operation *callable, Value result,
-                                     Type targetType,
                                      DictionaryAttr resultAttrs) const {
   auto *handler = getInterfaceFor(callable);
   assert(handler && "expected valid dialect handler");
-  return handler->handleResult(builder, call, callable, result, targetType,
-                               resultAttrs);
+  return handler->handleResult(builder, call, callable, result, resultAttrs);
 }
 
 void InlinerInterface::processInlinedCallBlocks(
     Operation *call, iterator_range<Region::iterator> inlinedBlocks) const {
   auto *handler = getInterfaceFor(call);
   assert(handler && "expected valid dialect handler");
   handler->processInlinedCallBlocks(call, inlinedBlocks);
 }
if (ArrayAttr arrayAttr = callable.getCallableArgAttrs()) {
     assert(arrayAttr.size() == argAttrs.size());
     for (auto [idx, attr] : llvm::enumerate(arrayAttr))
       argAttrs[idx] = cast<DictionaryAttr>(attr);
   }
 
   // Run the argument attribute handler for the given argument and attribute.
   for (auto [blockArg, argAttr] :
        llvm::zip(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");
+    Value newArgument = interface.handleArgument(
+        builder, call, callable, mapper.lookup(blockArg), argAttr);
+    assert(newArgument.getType() == mapper.lookup(blockArg).getType() &&
+           "expected the argument type to not change");
 
     // 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,
static void handleResultImpl(InlinerInterface &interface, OpBuilder &builder,
   // Run the result attribute handler for the given result and attribute.
   SmallVector<DictionaryAttr> resultAttributes;
   for (auto [result, resAttr] : llvm::zip(results, resAttrs)) {
     // 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 &&
-           "expected the handled result type to match the target type");
+    Value newResult =
+        interface.handleResult(builder, call, callable, result, resAttr);
+    assert(newResult.getType() == result.getType() &&
+           "expected the result type to not change");
 
     // Replace the result uses except for the ones introduce by the handler.
     result.replaceUsesWithIf(newResult, [&](OpOperand &operand) {
       return resultUsers.count(operand.getOwner());
     });
   }
 }
 

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

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,
+                       Value argument,
                        DictionaryAttr argumentAttrs) const final {
     if (!argumentAttrs.contains("test.handle_argument"))
       return argument;
-    return builder.create<TestTypeChangerOp>(call->getLoc(), targetType,
+    return builder.create<TestTypeChangerOp>(call->getLoc(), argument.getType(),
                                              argument);
   }
 
   Value handleResult(OpBuilder &builder, Operation *call, Operation *callable,
-                     Value result, Type targetType,
-                     DictionaryAttr resultAttrs) const final {
+                     Value result, DictionaryAttr resultAttrs) const final {
     if (!resultAttrs.contains("test.handle_result"))
       return result;
-    return builder.create<TestTypeChangerOp>(call->getLoc(), targetType,
+    return builder.create<TestTypeChangerOp>(call->getLoc(), result.getType(),
                                              result);
   }
 
   void processInlinedCallBlocks(
       Operation *call,
       iterator_range<Region::iterator> inlinedBlocks) const final {
     if (!isa<ConversionCallOp>(call))
       return;