Add CFFunctionAnalysisStorePass (#39)

* make instrumentation flexible

* add store pass to O3

* remove useless loop

* track handled blocks

* track handled blocks2

Author: webmiche

llvm/include/llvm/Analysis/CFFunctionAnalysis.h

@@ -9,6 +9,7 @@
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Passes/PassPlugin.h"
 #include "llvm/Support/Debug.h"
+#include <fstream>
 
 namespace llvm {
 
@@ -34,6 +35,20 @@ struct CFFunctionAnalysisPrinterPass
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 };
 
+struct CFFunctionAnalysisStorePass
+    : PassInfoMixin<CFFunctionAnalysisStorePass> {
+
+  std::string Filename;
+
+public:
+  explicit CFFunctionAnalysisStorePass(std::string Filename)
+      : Filename(std::move(Filename)) {}
+
+  explicit CFFunctionAnalysisStorePass() : Filename("called_functions.txt") {}
+
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+};
+
 } // namespace llvm
 
 #endif // LLVM_ANALYSIS_CF_FUNCTION_ANALYSIS_H

llvm/lib/Analysis/CFFunctionAnalysis.cpp

@@ -7,9 +7,7 @@ using namespace llvm;
 // Provide a definition for the static object used to identify passes.
 AnalysisKey CFFunctionAnalysis::Key;
 
-CFFunctionAnalysisInfo CFFunctionAnalysis::run(Module &M,
-                                               ModuleAnalysisManager &AM) {
-
+CFFunctionAnalysisInfo analyse(Module &M) {
   CFFunctionAnalysisInfo CalledFunctions;
   for (auto &F : M) {
     if (F.isDeclaration()) {
@@ -75,6 +73,27 @@ CFFunctionAnalysisInfo CFFunctionAnalysis::run(Module &M,
   return CalledFunctions;
 }
 
+CFFunctionAnalysisInfo CFFunctionAnalysis::run(Module &M,
+                                               ModuleAnalysisManager &AM) {
+
+  CFFunctionAnalysisInfo CalledFunctions;
+
+  // check if called functions file exists
+  std::ifstream file("called_functions.txt");
+  if (file.good()) {
+    std::string line;
+    while (std::getline(file, line)) {
+      char *cstr = new char[line.length() + 1];
+      strcpy(cstr, line.c_str());
+      CalledFunctions.insert(cstr);
+    }
+  } else {
+    CalledFunctions = analyse(M);
+  }
+  file.close();
+  return CalledFunctions;
+}
+
 PreservedAnalyses
 CFFunctionAnalysisPrinterPass::run(Module &M, ModuleAnalysisManager &AM) {
   OS << "Called functions for " << M.getName() << ":\n";
@@ -84,3 +103,25 @@ CFFunctionAnalysisPrinterPass::run(Module &M, ModuleAnalysisManager &AM) {
   OS << "\n";
   return PreservedAnalyses::all();
 }
+
+PreservedAnalyses CFFunctionAnalysisStorePass::run(Module &M,
+                                                   ModuleAnalysisManager &AM) {
+  CFFunctionAnalysisInfo CalledFunctions = AM.getResult<CFFunctionAnalysis>(M);
+
+  // store called functions to called_functions.txt
+  std::ofstream out;
+  out.open(Filename, std::ios::app);
+
+  if (!out) {
+    errs() << "Error: cannot open file " << Filename << "\n";
+    return PreservedAnalyses::none();
+  }
+
+  for (auto &F : CalledFunctions) {
+    out << F.str() << "\n";
+  }
+
+  out.close();
+
+  return PreservedAnalyses::all();
+}

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -17,6 +17,7 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/BasicAliasAnalysis.h"
+#include "llvm/Analysis/CFFunctionAnalysis.h"
 #include "llvm/Analysis/CGSCCPassManager.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/InlineAdvisor.h"
@@ -123,9 +124,9 @@
 #include "llvm/Transforms/Scalar/WarnMissedTransforms.h"
 #include "llvm/Transforms/Utils/AddDiscriminators.h"
 #include "llvm/Transforms/Utils/AssumeBundleBuilder.h"
+#include "llvm/Transforms/Utils/CFFunctionInstrumentation.h"
 #include "llvm/Transforms/Utils/CanonicalizeAliases.h"
 #include "llvm/Transforms/Utils/CountVisits.h"
-#include "llvm/Transforms/Utils/CFFunctionInstrumentation.h"
 #include "llvm/Transforms/Utils/InjectTLIMappings.h"
 #include "llvm/Transforms/Utils/LibCallsShrinkWrap.h"
 #include "llvm/Transforms/Utils/Mem2Reg.h"
@@ -181,9 +182,9 @@ static cl::opt<bool> EnablePostPGOLoopRotation(
     "enable-post-pgo-loop-rotation", cl::init(true), cl::Hidden,
     cl::desc("Run the loop rotation transformation after PGO instrumentation"));
 
-static cl::opt<bool> EnableGlobalAnalyses(
-    "enable-global-analyses", cl::init(true), cl::Hidden,
-    cl::desc("Enable inter-procedural analyses"));
+static cl::opt<bool>
+    EnableGlobalAnalyses("enable-global-analyses", cl::init(true), cl::Hidden,
+                         cl::desc("Enable inter-procedural analyses"));
 
 static cl::opt<bool>
     RunPartialInlining("enable-partial-inlining", cl::init(false), cl::Hidden,
@@ -1084,11 +1085,10 @@ PassBuilder::buildModuleSimplificationPipeline(OptimizationLevel Level,
   // and prior to optimizing globals.
   // FIXME: This position in the pipeline hasn't been carefully considered in
   // years, it should be re-analyzed.
-  MPM.addPass(IPSCCPPass(
-              IPSCCPOptions(/*AllowFuncSpec=*/
-                            Level != OptimizationLevel::Os &&
-                            Level != OptimizationLevel::Oz &&
-                            !isLTOPreLink(Phase))));
+  MPM.addPass(IPSCCPPass(IPSCCPOptions(/*AllowFuncSpec=*/
+                                       Level != OptimizationLevel::Os &&
+                                       Level != OptimizationLevel::Oz &&
+                                       !isLTOPreLink(Phase))));
 
   // Attach metadata to indirect call sites indicating the set of functions
   // they may target at run-time. This should follow IPSCCP.
@@ -1498,6 +1498,7 @@ PassBuilder::buildPerModuleDefaultPipeline(OptimizationLevel Level,
 
   ModulePassManager MPM;
 
+  MPM.addPass(CFFunctionAnalysisStorePass("cffunction-analysis-store.txt"));
   MPM.addPass(CFFunctionInstrumentationPass());
 
   // Convert @llvm.global.annotations to !annotation metadata.
@@ -1559,7 +1560,7 @@ PassBuilder::buildFatLTODefaultPipeline(OptimizationLevel Level, bool ThinLTO,
 ModulePassManager
 PassBuilder::buildThinLTOPreLinkDefaultPipeline(OptimizationLevel Level) {
   if (Level == OptimizationLevel::O0)
-    return buildO0DefaultPipeline(Level, /*LTOPreLink*/true);
+    return buildO0DefaultPipeline(Level, /*LTOPreLink*/ true);
 
   ModulePassManager MPM;
 

llvm/lib/Passes/PassRegistry.def

@@ -42,6 +42,7 @@ MODULE_ALIAS_ANALYSIS("globals-aa", GlobalsAA())
 #ifndef MODULE_PASS
 #define MODULE_PASS(NAME, CREATE_PASS)
 #endif
+MODULE_PASS("store<cffunction-analysis>", CFFunctionAnalysisStorePass())
 MODULE_PASS("print<cffunction-analysis>", CFFunctionAnalysisPrinterPass(errs()))
 MODULE_PASS("cffunction-instrumentation", CFFunctionInstrumentationPass())
 MODULE_PASS("always-inline", AlwaysInlinerPass())

llvm/lib/Transforms/Utils/CFFunctionInstrumentation.cpp

@@ -10,6 +10,9 @@ PreservedAnalyses
 CFFunctionInstrumentationPass::run(Module &M, ModuleAnalysisManager &AM) {
 
   CFFunctionAnalysisInfo CalledFunctions = AM.getResult<CFFunctionAnalysis>(M);
+  LLVM_DEBUG(
+      for (auto &F
+           : CalledFunctions) { dbgs() << "Called function: " << F << "\n"; });
   int permissions_created = 0;
   Value *WritePermission = nullptr;
   Value *FileName = nullptr;
@@ -33,10 +36,17 @@ CFFunctionInstrumentationPass::run(Module &M, ModuleAnalysisManager &AM) {
     StringRef funcFileName = StringRef(fileName);
     // for all return instructions, print the return value to a file with the
     // name of the function
+
+    // store already handled blocks
+    std::set<BasicBlock *> HandledBlocks;
     for (auto &BB : F) {
+      if (HandledBlocks.count(&BB)) {
+        continue;
+      }
       // Do NOT reinstrument the inserted blocks
       if (BB.getName() == "return" || BB.getName() == "print" ||
           BB.getName() == "open") {
+            HandledBlocks.insert(&BB);
         continue;
       }
       if (auto *RI = dyn_cast<ReturnInst>(BB.getTerminator())) {
@@ -113,8 +123,13 @@ CFFunctionInstrumentationPass::run(Module &M, ModuleAnalysisManager &AM) {
 
           // place new BBs in the correct order
           ReturnBB->moveAfter(PrintBB);
+
+          HandledBlocks.insert(AccessBB);
+          HandledBlocks.insert(PrintBB);
+          HandledBlocks.insert(ReturnBB);
         }
       }
+      HandledBlocks.insert(&BB);
     }
   }