Fix inlining multi-block callees with type conversion.

[silvas]

Summary

The previous code result a mismatch between block argument types and
predecessor successor args when a type conversion was needed in a
multiblock case. It was assuming the replaced result types matched the
region result types.

Also, slighly improve the debug output from the inliner and add a
debug counter.

[rriddle]

mlir/test/Transforms/inlining.mlir

134

Please add check lines to this test.

[silvas]

Actually test stuff in test case.

[rriddle]

Thanks for fixing this!

mlir/lib/Transforms/Inliner.cpp

30

This should be camelCase as it defines a variable.

mlir/lib/Transforms/Utils/InliningUtils.cpp

132

Can you use TypeRange here instead? resultsToReplace should also be updated to ValueRange.

250

If you update the input types to ValueRange/TypeRange, you can just do resultsToReplace.getTypes().

[mehdi_amini]

mlir/lib/Transforms/Inliner.cpp

31

Wouldn't this add a static global initializer?

[silvas]

• Switch to ValueRange/TypeRange
• Remove DEBUG_COUNTER :/

[rriddle]

Can you remove the explicit vector materialization here now?
https://github.com/llvm/llvm-project/blob/31282d399b79526303ae4b5d1f73ace8a40c6587/mlir/test/lib/Transforms/TestInlining.cpp#L48

mlir/include/mlir/Transforms/InliningUtils.h

18

Can you use forward declarations instead?

mlir/lib/Transforms/Inliner.cpp

24

This is now unused?

28

These are now unused?

[silvas]

Address feedback

mlir/include/mlir/Transforms/InliningUtils.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TRANSFORMS_INLINING_UTILS_H
 #define MLIR_TRANSFORMS_INLINING_UTILS_H
 
 #include "mlir/IR/DialectInterface.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/Region.h"
 
 namespace mlir {
 
 class Block;
 class BlockAndValueMapping;
 class CallableOpInterface;
 class CallOpInterface;
 class FuncOp;
 class OpBuilder;
 class Operation;
 class Region;
+class TypeRange;
 class Value;
+class ValueRange;
 
 //===----------------------------------------------------------------------===//
 // InlinerInterface
 //===----------------------------------------------------------------------===//
 
 /// This is the interface that must be implemented by the dialects of operations
 /// to be inlined. This interface should only handle the operations of the
 /// given dialect.

 /// failure if it is not possible to inline this function. If the function
 /// returned failure, then no changes to the module have been made.
 ///
 /// The provided 'inlinePoint' must be within a region, and corresponds to the
 /// location where the 'src' region should be inlined. 'mapping' contains any
 /// remapped operands that are used within the region, and *must* include
 /// remappings for the entry arguments to the region. 'resultsToReplace'
 /// corresponds to any results that should be replaced by terminators within the
-/// inlined region. 'inlineLoc' is an optional Location that, if provided, will
-/// be used to update the inlined operations' location information.
-/// 'shouldCloneInlinedRegion' corresponds to whether the source region should
-/// be cloned into the 'inlinePoint' or spliced directly.
+/// inlined region. 'regionResultTypes' specifies the expected return types of
+/// the terminators in the region. 'inlineLoc' is an optional Location that, if
+/// provided, will be used to update the inlined operations' location
+/// information. 'shouldCloneInlinedRegion' corresponds to whether the source
+/// region should be cloned into the 'inlinePoint' or spliced directly.
 LogicalResult inlineRegion(InlinerInterface &interface, Region *src,
                            Operation *inlinePoint, BlockAndValueMapping &mapper,
-                           ArrayRef<Value> resultsToReplace,
+                           ValueRange resultsToReplace,
+                           TypeRange regionResultTypes,
                            Optional<Location> inlineLoc = llvm::None,
                            bool shouldCloneInlinedRegion = true);
 
 /// This function is an overload of the above 'inlineRegion' that allows for
 /// providing the set of operands ('inlinedOperands') that should be used
 /// in-favor of the region arguments when inlining.
 LogicalResult inlineRegion(InlinerInterface &interface, Region *src,
                            Operation *inlinePoint,
-                           ArrayRef<Value> inlinedOperands,
-                           ArrayRef<Value> resultsToReplace,
+                           ValueRange inlinedOperands,
+                           ValueRange resultsToReplace,
                            Optional<Location> inlineLoc = llvm::None,
                            bool shouldCloneInlinedRegion = true);
 
 /// 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 inlineCall(InlinerInterface &interface, CallOpInterface call,
                          CallableOpInterface callable, Region *src,
                          bool shouldCloneInlinedRegion = true);
 
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_INLINING_UTILS_H

mlir/lib/Transforms/Inliner.cpp

 #include "PassDetail.h"
 #include "mlir/Analysis/CallGraph.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/SideEffects.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/SCCIterator.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Parallel.h"
 
 #define DEBUG_TYPE "inlining"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Symbol Use Tracking
 //===----------------------------------------------------------------------===//
 
 /// Returns true if this operation can be discarded if it is a symbol and has no
 /// uses. 'allUsesVisible' corresponds to if the parent symbol table is hidden
 /// from above.
 static bool canDiscardSymbolOnUseEmpty(Operation *op, bool allUsesVisible) {
   if (!SymbolTable::isSymbol(op))
     return false;
inlineCallsInSCC(Inliner &inliner, CGUseList &useList,
   // A set of dead nodes to remove after inlining.
   SmallVector<CallGraphNode *, 1> deadNodes;
 
   // Try to inline each of the call operations. Don't cache the end iterator
   // here as more calls may be added during inlining.
   bool inlinedAnyCalls = false;
   for (unsigned i = 0; i != calls.size(); ++i) {
     ResolvedCall it = calls[i];
+    bool doInline = shouldInline(it);
     LLVM_DEBUG({
-      llvm::dbgs() << "* Considering inlining call: ";
+      if (doInline)
+        llvm::dbgs() << "* Inlining call: ";
+      else
+        llvm::dbgs() << "* Not inlining call: ";
       it.call.dump();
     });
-    if (!shouldInline(it))
+    if (!doInline)
       continue;
     CallOpInterface call = it.call;
     Region *targetRegion = it.targetNode->getCallableRegion();
 
     // If this is the last call to the target node and the node is discardable,
     // then inline it in-place and delete the node if successful.
     bool inlineInPlace = useList.hasOneUseAndDiscardable(it.targetNode);
 

mlir/lib/Transforms/Utils/InliningUtils.cpp

 
 //===----------------------------------------------------------------------===//
 // Inline Methods
 //===----------------------------------------------------------------------===//
 
 LogicalResult mlir::inlineRegion(InlinerInterface &interface, Region *src,
                                  Operation *inlinePoint,
                                  BlockAndValueMapping &mapper,
-                                 ArrayRef<Value> resultsToReplace,
+                                 ValueRange resultsToReplace,
+                                 TypeRange regionResultTypes,
                                  Optional<Location> inlineLoc,
                                  bool shouldCloneInlinedRegion) {
+  assert(resultsToReplace.size() == regionResultTypes.size());
   // We expect the region to have at least one block.
   if (src->empty())
     return failure();
 
   // Check that all of the region arguments have been mapped.
   auto *srcEntryBlock = &src->front();
   if (llvm::any_of(srcEntryBlock->getArguments(),
                    [&](BlockArgument arg) { return !mapper.contains(arg); }))
LogicalResult mlir::inlineRegion(InlinerInterface &interface, Region *src,
 
   // Process the newly inlined blocks.
   interface.processInlinedBlocks(newBlocks);
 
   // Handle the case where only a single block was inlined.
   if (std::next(newBlocks.begin()) == newBlocks.end()) {
     // Have the interface handle the terminator of this block.
     auto *firstBlockTerminator = firstNewBlock->getTerminator();
-    interface.handleTerminator(firstBlockTerminator, resultsToReplace);
+    interface.handleTerminator(firstBlockTerminator,
+                               llvm::to_vector<6>(resultsToReplace));
     firstBlockTerminator->erase();
 
     // Merge the post insert block into the cloned entry block.
     firstNewBlock->getOperations().splice(firstNewBlock->end(),
                                           postInsertBlock->getOperations());
     postInsertBlock->erase();
   } else {
     // Otherwise, there were multiple blocks inlined. Add arguments to the post
     // insertion block to represent the results to replace.
-    for (Value resultToRepl : resultsToReplace) {
-      resultToRepl.replaceAllUsesWith(
-          postInsertBlock->addArgument(resultToRepl.getType()));
+    for (auto resultToRepl : llvm::enumerate(resultsToReplace)) {
+      resultToRepl.value().replaceAllUsesWith(postInsertBlock->addArgument(
+          regionResultTypes[resultToRepl.index()]));
     }
 
     /// Handle the terminators for each of the new blocks.
     for (auto &newBlock : newBlocks)
       interface.handleTerminator(newBlock.getTerminator(), postInsertBlock);
   }
 
   // Splice the instructions of the inlined entry block into the insert block.
   insertBlock->getOperations().splice(insertBlock->end(),
                                       firstNewBlock->getOperations());
   firstNewBlock->erase();
   return success();
 }
 
 /// This function is an overload of the above 'inlineRegion' that allows for
 /// providing the set of operands ('inlinedOperands') that should be used
 /// in-favor of the region arguments when inlining.
 LogicalResult mlir::inlineRegion(InlinerInterface &interface, Region *src,
                                  Operation *inlinePoint,
-                                 ArrayRef<Value> inlinedOperands,
-                                 ArrayRef<Value> resultsToReplace,
+                                 ValueRange inlinedOperands,
+                                 ValueRange resultsToReplace,
                                  Optional<Location> inlineLoc,
                                  bool shouldCloneInlinedRegion) {
   // We expect the region to have at least one block.
   if (src->empty())
     return failure();
 
   auto *entryBlock = &src->front();
   if (inlinedOperands.size() != entryBlock->getNumArguments())
     return failure();
 
   // Map the provided call operands to the arguments of the region.
   BlockAndValueMapping mapper;
   for (unsigned i = 0, e = inlinedOperands.size(); i != e; ++i) {
     // Verify that the types of the provided values match the function argument
     // types.
     BlockArgument regionArg = entryBlock->getArgument(i);
     if (inlinedOperands[i].getType() != regionArg.getType())
       return failure();
     mapper.map(regionArg, inlinedOperands[i]);
   }
 
   // Call into the main region inliner function.
   return inlineRegion(interface, src, inlinePoint, mapper, resultsToReplace,
-                      inlineLoc, shouldCloneInlinedRegion);
+                      resultsToReplace.getTypes(), 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) {
for (unsigned i = 0, e = callResults.size(); i != e; ++i) {
     if (!castResult)
       return cleanupState();
     callResult.replaceAllUsesWith(castResult);
     castResult.getDefiningOp()->replaceUsesOfWith(castResult, callResult);
   }
 
   // Attempt to inline the call.
   if (failed(inlineRegion(interface, src, call, mapper, callResults,
-                          call.getLoc(), shouldCloneInlinedRegion)))
+                          callableResultTypes, call.getLoc(),
+                          shouldCloneInlinedRegion)))
     return cleanupState();
   return success();
 }

mlir/test/Transforms/inlining.mlir

func @inline_convert_call() -> i16 {
 
   // CHECK: %[[CAST_INPUT:.*]] = "test.cast"(%[[INPUT]]) : (i16) -> i32
   // CHECK: %[[CAST_RESULT:.*]] = "test.cast"(%[[CAST_INPUT]]) : (i32) -> i16
   // CHECK-NEXT: return %[[CAST_RESULT]]
   %res = "test.conversion_call_op"(%test_input) { callee=@convert_callee_fn } : (i16) -> (i16)
   return %res : i16
 }
 
+func @convert_callee_fn_multiblock() -> i32 {
+  br ^bb0
+^bb0:
+  %0 = constant 0 : i32
+  return %0 : i32
+}
+
+// CHECK-LABEL: func @inline_convert_result_multiblock
+func @inline_convert_result_multiblock() -> i16 {
+// CHECK:   br ^bb1
+// CHECK: ^bb1:
+// CHECK:   %[[C:.+]] = constant 0 : i32
+// CHECK:   br ^bb2(%[[C]] : i32)
+// CHECK: ^bb2(%[[BBARG:.+]]: i32):
+// CHECK:   %[[CAST_RESULT:.+]] = "test.cast"(%[[BBARG]]) : (i32) -> i16
+// CHECK:   return %[[CAST_RESULT]] : i16
+
+  %res = "test.conversion_call_op"() { callee=@convert_callee_fn_multiblock } : () -> (i16)
+  return %res : i16
+}
+
 // CHECK-LABEL: func @no_inline_convert_call
 func @no_inline_convert_call() {
   // CHECK: "test.conversion_call_op"
   %test_input_i16 = constant 0 : i16
   %test_input_i64 = constant 0 : i64
   "test.conversion_call_op"(%test_input_i16, %test_input_i64) { callee=@convert_callee_fn_multi_arg } : (i16, i64) -> ()
 
   // CHECK: "test.conversion_call_op"

mlir/test/lib/Transforms/TestInlining.cpp

for (auto caller : callers) {
       auto callee = dyn_cast_or_null<FunctionalRegionOp>(
           caller.getCallee().getDefiningOp());
       if (!callee)
         continue;
 
       // Inline the functional region operation, but only clone the internal
       // region if there is more than one use.
       if (failed(inlineRegion(
-              interface, &callee.body(), caller,
-              llvm::to_vector<8>(caller.getArgOperands()),
+              interface, &callee.body(), caller, caller.getArgOperands(),
               SmallVector<Value, 8>(caller.getResults()), caller.getLoc(),
               /*shouldCloneInlinedRegion=*/!callee.getResult().hasOneUse())))
         continue;
 
       // If the inlining was successful then erase the call and callee if
       // possible.
       caller.erase();
       if (callee.use_empty())