[LLVM][IPO] Add NoInlineFuncCalledOnce and plumb -no-inline-functions-called-once via PassBuilder

Signed-off-by: Karthikdhondi <[email protected]>

Author: Karthikdhondi

llvm/include/llvm/Transforms/IPO/NoInlineFuncCalledOnce.h

@@ -0,0 +1,18 @@
+#ifndef LLVM_TRANSFORMS_IPO_NOINLINEFUNCCALLEDONCE_H
+#define LLVM_TRANSFORMS_IPO_NOINLINEFUNCCALLEDONCE_H
+
+#include "llvm/IR/PassManager.h"
+#include "llvm/Support/CommandLine.h"
+
+namespace llvm {
+
+struct NoInlineFuncCalledOncePass
+    : public PassInfoMixin<NoInlineFuncCalledOncePass> {
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &);
+};
+
+// single definition of the control flag lives in the .cpp (declared here)
+extern cl::opt<bool> EnableNoInlineFuncCalledOnce;
+
+} // namespace llvm
+#endif

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -66,6 +66,7 @@
 #include "llvm/Transforms/IPO/MemProfContextDisambiguation.h"
 #include "llvm/Transforms/IPO/MergeFunctions.h"
 #include "llvm/Transforms/IPO/ModuleInliner.h"
+#include "llvm/Transforms/IPO/NoInlineFuncCalledOnce.h"
 #include "llvm/Transforms/IPO/OpenMPOpt.h"
 #include "llvm/Transforms/IPO/PartialInlining.h"
 #include "llvm/Transforms/IPO/SCCP.h"
@@ -150,6 +151,12 @@
 
 using namespace llvm;
 
+namespace llvm {
+cl::opt<bool> EnableNoInlineFuncCalledOnce(
+    "no-inline-functions-called-once", cl::init(false), cl::Hidden,
+    cl::desc("Mark TU-local functions called exactly once as noinline"));
+} // namespace llvm
+
 static cl::opt<InliningAdvisorMode> UseInlineAdvisor(
     "enable-ml-inliner", cl::init(InliningAdvisorMode::Default), cl::Hidden,
     cl::desc("Enable ML policy for inliner. Currently trained for -Oz only"),
@@ -1274,6 +1281,9 @@ PassBuilder::buildModuleSimplificationPipeline(OptimizationLevel Level,
                  PGOOpt->Action == PGOOptions::SampleUse))
     MPM.addPass(PGOForceFunctionAttrsPass(PGOOpt->ColdOptType));
 
+  if (EnableNoInlineFuncCalledOnce)
+    MPM.addPass(NoInlineFuncCalledOncePass());
+
   MPM.addPass(AlwaysInlinerPass(/*InsertLifetimeIntrinsics=*/true));
 
   if (EnableModuleInliner)
@@ -1447,6 +1457,9 @@ PassBuilder::buildModuleOptimizationPipeline(OptimizationLevel Level,
   const bool LTOPreLink = isLTOPreLink(LTOPhase);
   ModulePassManager MPM;
 
+  if (EnableNoInlineFuncCalledOnce)
+    MPM.addPass(NoInlineFuncCalledOncePass());
+
   // Run partial inlining pass to partially inline functions that have
   // large bodies.
   if (RunPartialInlining)
@@ -1766,6 +1779,9 @@ PassBuilder::buildThinLTOPreLinkDefaultPipeline(OptimizationLevel Level) {
     return MPM;
   }
 
+  if (EnableNoInlineFuncCalledOnce)
+    MPM.addPass(NoInlineFuncCalledOncePass());
+
   // Run partial inlining pass to partially inline functions that have
   // large bodies.
   // FIXME: It isn't clear whether this is really the right place to run this
@@ -2012,6 +2028,9 @@ PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
   // Lower variadic functions for supported targets prior to inlining.
   MPM.addPass(ExpandVariadicsPass(ExpandVariadicsMode::Optimize));
 
+  if (EnableNoInlineFuncCalledOnce)
+    MPM.addPass(NoInlineFuncCalledOncePass());
+
   // Note: historically, the PruneEH pass was run first to deduce nounwind and
   // generally clean up exception handling overhead. It isn't clear this is
   // valuable as the inliner doesn't currently care whether it is inlining an

llvm/lib/Transforms/IPO/CMakeLists.txt

@@ -33,6 +33,7 @@ add_llvm_component_library(LLVMipo
   MemProfContextDisambiguation.cpp
   MergeFunctions.cpp
   ModuleInliner.cpp
+  NoInlineFuncCalledOnce.cpp
   OpenMPOpt.cpp
   PartialInlining.cpp
   SampleContextTracker.cpp

llvm/lib/Transforms/IPO/NoInlineFuncCalledOnce.cpp

@@ -0,0 +1,62 @@
+#include "llvm/Transforms/IPO/NoInlineFuncCalledOnce.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/IR/Attributes.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/Support/CommandLine.h"
+
+using namespace llvm;
+
+PreservedAnalyses NoInlineFuncCalledOncePass::run(Module &M,
+                                                  ModuleAnalysisManager &) {
+  DenseMap<Function *, unsigned> DirectCalls;
+  DenseSet<Function *> Recursive;
+
+  for (Function &F : M)
+    if (!F.isDeclaration() && (F.hasInternalLinkage() || F.hasPrivateLinkage()))
+      DirectCalls[&F] = 0;
+
+  for (Function &Caller : M) {
+    if (Caller.isDeclaration())
+      continue;
+    for (Instruction &I : instructions(Caller)) {
+      auto *CB = dyn_cast<CallBase>(&I);
+      if (!CB)
+        continue;
+      const Value *Op = CB->getCalledOperand()->stripPointerCasts();
+      if (auto *Callee = const_cast<Function *>(dyn_cast<Function>(Op))) {
+        if (!DirectCalls.count(Callee))
+          continue;
+        DirectCalls[Callee] += 1;
+        if (&Caller == Callee)
+          Recursive.insert(Callee);
+      }
+    }
+  }
+
+  bool Changed = false;
+  for (auto &KV : DirectCalls) {
+    Function *F = KV.first;
+    unsigned N = KV.second;
+
+    if (N != 1)
+      continue; // only called-once
+    if (Recursive.count(F))
+      continue; // skip recursion
+    if (F->hasAddressTaken())
+      continue; // skip address-taken
+    if (F->hasFnAttribute(Attribute::AlwaysInline))
+      continue;
+    if (F->hasFnAttribute(Attribute::NoInline))
+      continue;
+
+    F->addFnAttr(Attribute::NoInline);
+    Changed = true;
+  }
+
+  return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
+}