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

32

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.

249

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

[mehdi_amini]

mlir/lib/Transforms/Inliner.cpp

33

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?

29

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;

 /// 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,
+                           ArrayRef<Type> 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,

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/DebugCounter.h"
 #include "llvm/Support/Parallel.h"
 
 #define DEBUG_TYPE "inlining"
 
+using llvm::DebugCounter;
 using namespace mlir;
 
+DEBUG_COUNTER(InlineFunction, "inliner-inline-function",
+              "Controls which inlinings we do.");
+
 //===----------------------------------------------------------------------===//
 // 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) {
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) && DebugCounter::shouldExecute(InlineFunction);
     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,
+                                 ArrayRef<Type> 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); }))
if (std::next(newBlocks.begin()) == newBlocks.end()) {
 
     // 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.
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]);
   }
 
+  auto regionResultTypes = llvm::to_vector<6>(
+      llvm::map_range(resultsToReplace, [](Value v) { return v.getType(); }));
   // Call into the main region inliner function.
   return inlineRegion(interface, src, inlinePoint, mapper, resultsToReplace,
-                      inlineLoc, shouldCloneInlinedRegion);
+                      regionResultTypes, 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"