[mlir][Inliner] Support recursion in Inliner

[javedabsar]

Summary

This fixes Bug https://github.com/llvm/llvm-project/issues/53492
and uses InlineHistory to track recursive inlining.

[rriddle]

Thanks for looking into this! Have you considered/looked into using the same approach as LLVM? i.e. keeping an inline history to track what's already been inlined?

[javedabsar]

In D127072#3559111, @rriddle wrote:

Thanks for looking into this! Have you considered/looked into using the same approach as LLVM? i.e. keeping an inline history to track what's already been inlined?

Yes that seems a good idea. Will try to rework this using InlineHistoryIDs

[javedabsar]

change based on review comments.
Use InlineHistory.

[javedabsar]

gentle reminder

[mehdi_amini]

River is out this week, please ping again on Monday!

[javedabsar]

In D127072#3601560, @mehdi_amini wrote:

River is out this week, please ping again on Monday!

Ah ok. Thanks @mehdi_amini

[rriddle]

Thanks for updating! (Sorry, coming back from OOO)

mlir/lib/Transforms/Inliner.cpp

395–398

Please move these functions outside of the anonymous namespace.

396–400

This doesn't look scalable, given that it assumes the existence of a "callee" attribute. Can you instead use op.getCallableForCallee()? The result of that is either a SymbolRefAttr (which is easily dumpable), or a Value (which could just fallback to the _unnamed_callee_ you have here).

403–404

Can you reflow these comments? it doesn't look like it wraps at 80.

405

406

486

Can we use Optional<int> instead of int here? It would remove the magic -1 handling.

[javedabsar]

Fix based on review feedback.

[javedabsar]

Thanks for the review

[rriddle]

mlir/lib/Transforms/Inliner.cpp

372–373

Do we need to anchor on FlatSymbolRefAttr? We should be able to print any SymbolRefAttr, e.g. using debugString(in mlir/Support/DebugStringHelper.h).

382

Can you use MutableArrayRef here instead?

394

Unrelated?

492

Please cache the end iterator here.

540

You'll likely need a (void)hisStr after this to avoid unused variable warnings.

Also, can you extend the name of this to be more explicit on what this is doing?

mlir/test/Transforms/bugfix53492_inlining.mlir

3 ↗

(On Diff #440682)

Can you rename the file (or just merge it into inlining.mlir) and the functions in this file to better represent what this is testing, as opposed to referencing the bug? The fix is for inlining recursive functions, which is descriptive enough that I don't think we need/should reference the bug here.

[javedabsar]

Fixes based on review comments.

[javedabsar]

rename test

[javedabsar]

mlir/lib/Transforms/Inliner.cpp

492

Actually the 'calls.size' changes in this loop. Below it calls inlineCall (in Utils/InliningUtils.cpp) which calls interface.processInlinedBlocks which then calls collectCallOps that adds new inlined calls. Since the calls is reference to inliner.calls it changes.

[javedabsar]

fix

[javedabsar]

mlir/lib/Transforms/Inliner.cpp

492

Ignore above. I had misunderstood your comment earlier.

[rriddle]

Thank you for tackling this! Very appreciated.

mlir/test/Transforms/inlining-recursive.mlir

4

Honestly I think we could just drop this issue reference.

8

nit: Can you indent these checks?

mlir/lib/Transforms/Inliner.cpp

 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 #include "mlir/Analysis/CallGraph.h"
 #include "mlir/IR/Threading.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Pass/PassManager.h"
+#include "mlir/Support/DebugStringHelper.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/SCCIterator.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "inlining"
 
 using namespace mlir;
for (Operation &op : *block) {
       }
     }
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Inliner
 //===----------------------------------------------------------------------===//
+
+#ifndef NDEBUG
+static std::string getNodeName(CallOpInterface op) {
+  if (auto sym = op.getCallableForCallee().dyn_cast<SymbolRefAttr>())
+    return debugString(op);
+  return "_unnamed_callee_";
+}
+#endif
+
+/// Return true if the specified `inlineHistoryID`  indicates an inline history
+/// that already includes `node`.
+static bool inlineHistoryIncludes(
+    CallGraphNode *node, Optional<size_t> inlineHistoryID,
+    MutableArrayRef<std::pair<CallGraphNode *, Optional<size_t>>>
+        inlineHistory) {
+  while (inlineHistoryID.has_value()) {
+    assert(inlineHistoryID.value() < inlineHistory.size() &&
+           "Invalid inline history ID");
+    if (inlineHistory[inlineHistoryID.value()].first == node)
+      return true;
+    inlineHistoryID = inlineHistory[inlineHistoryID.value()].second;
+  }
+  return false;
+}
+
 namespace {
 /// This class provides a specialization of the main inlining interface.
 struct Inliner : public InlinerInterface {
   Inliner(MLIRContext *context, CallGraph &cg,
           SymbolTableCollection &symbolTable)
       : InlinerInterface(context), cg(cg), symbolTable(symbolTable) {}
 
   /// Process a set of blocks that have been inlined. This callback is invoked
   /// *before* inlined terminator operations have been processed.
   void
   processInlinedBlocks(iterator_range<Region::iterator> inlinedBlocks) final {
     // Find the closest callgraph node from the first block.
     CallGraphNode *node;
     Region *region = inlinedBlocks.begin()->getParent();
     while (!(node = cg.lookupNode(region))) {
       region = region->getParentRegion();
       assert(region && "expected valid parent node");
     }
 
     collectCallOps(inlinedBlocks, node, cg, symbolTable, calls,
                    /*traverseNestedCGNodes=*/true);
   }
for (CallGraphNode *node : currentSCC) {
     if (useList.isDead(node)) {
       deadNodes.insert(node);
     } else {
       collectCallOps(*node->getCallableRegion(), node, cg, inliner.symbolTable,
                      calls, /*traverseNestedCGNodes=*/false);
     }
   }
 
+  // When inlining a callee produces new call sites, we want to keep track of
+  // the fact that they were inlined from the callee. This allows us to avoid
+  // infinite inlining.
+  using InlineHistoryT = Optional<size_t>;
+  SmallVector<std::pair<CallGraphNode *, InlineHistoryT>, 8> inlineHistory;
+  std::vector<InlineHistoryT> callHistory(calls.size(), InlineHistoryT{});
+
+  LLVM_DEBUG({
+    llvm::dbgs() << "* Inliner: Initial calls in SCC are: {\n";
+    for (unsigned i = 0, e = calls.size(); i < e; ++i)
+      llvm::dbgs() << "  " << i << ". " << calls[i].call << ",\n";
+    llvm::dbgs() << "}\n";
+  });
+
   // 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) {
+  for (unsigned i = 0; i < calls.size(); ++i) {
     if (deadNodes.contains(calls[i].sourceNode))
       continue;
     ResolvedCall it = calls[i];
-    bool doInline = shouldInline(it);
+
+    InlineHistoryT inlineHistoryID = callHistory[i];
+    bool inHistory =
+        inlineHistoryIncludes(it.targetNode, inlineHistoryID, inlineHistory);
+    bool doInline = !inHistory && shouldInline(it);
     CallOpInterface call = it.call;
     LLVM_DEBUG({
       if (doInline)
-        llvm::dbgs() << "* Inlining call: " << call << "\n";
+        llvm::dbgs() << "* Inlining call: " << i << ". " << call << "\n";
       else
-        llvm::dbgs() << "* Not inlining call: " << call << "\n";
+        llvm::dbgs() << "* Not inlining call: " << i << ". " << call << "\n";
     });
     if (!doInline)
       continue;
+
+    unsigned prevSize = calls.size();
     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);
 
     LogicalResult inlineResult = inlineCall(
         inliner, call, cast<CallableOpInterface>(targetRegion->getParentOp()),
         targetRegion, /*shouldCloneInlinedRegion=*/!inlineInPlace);
     if (failed(inlineResult)) {
       LLVM_DEBUG(llvm::dbgs() << "** Failed to inline\n");
       continue;
     }
     inlinedAnyCalls = true;
 
+    // Create a inline history entry for this inlined call, so that we remember
+    // that new callsites came about due to inlining Callee.
+    InlineHistoryT newInlineHistoryID{inlineHistory.size()};
+    inlineHistory.push_back(std::make_pair(it.targetNode, inlineHistoryID));
+
+    auto historyToString = [](InlineHistoryT h) {
+      return h.has_value() ? std::to_string(h.value()) : "root";
+    };
+    (void)historyToString;
+    LLVM_DEBUG(llvm::dbgs()
+               << "* new inlineHistory entry: " << newInlineHistoryID << ". ["
+               << getNodeName(call) << ", " << historyToString(inlineHistoryID)
+               << "]\n");
+
+    for (unsigned k = prevSize; k != calls.size(); ++k) {
+      callHistory.push_back(newInlineHistoryID);
+      LLVM_DEBUG(llvm::dbgs() << "* new call " << k << " {" << calls[i].call
+                              << "}\n   with historyID = " << newInlineHistoryID
+                              << ", added due to inlining of\n  call {" << call
+                              << "}\n with historyID = "
+                              << historyToString(inlineHistoryID) << "\n");
+    }
+
     // If the inlining was successful, Merge the new uses into the source node.
     useList.dropCallUses(it.sourceNode, call.getOperation(), cg);
     useList.mergeUsesAfterInlining(it.targetNode, it.sourceNode);
 
     // then erase the call.
     call.erase();
 
     // If we inlined in place, mark the node for deletion.

mlir/test/Transforms/inlining-recursive.mlir

+// RUN: mlir-opt %s -inline='default-pipeline=''' | FileCheck %s
+// RUN: mlir-opt %s --mlir-disable-threading -inline='default-pipeline=''' | FileCheck %s
+
+// CHECK-LABEL: func.func @foo0
+func.func @foo0(%arg0 : i32) -> i32 {
+  // CHECK: call @foo1
+  // CHECK: }
+  %0 = arith.constant 0 : i32
+  %1 = arith.cmpi eq, %arg0, %0 : i32
+  cf.cond_br %1, ^exit, ^tail
+^exit:
+  return %0 : i32
+^tail:
+  %3 = call @foo1(%arg0) : (i32) -> i32
+  return %3 : i32
+}
+
+// CHECK-LABEL: func.func @foo1
+func.func @foo1(%arg0 : i32) -> i32 {
+  // CHECK:    call @foo1
+  %0 = arith.constant 1 : i32
+  %1 = arith.subi %arg0, %0 : i32
+  %2 = call @foo0(%1) : (i32) -> i32
+  return %2 : i32
+}