[llvm-extract] Add option for recursive extraction

[loladiro]

Summary

Particularly, with --delete, this can be very useful for testing
new optimizations on some hotspots, without having to run it on the whole
application. E.g. as such:

llvm-extract app.bc --recursive --rfunc .*hotspot.* > hotspot.bc
llvm-extract app.bc --recursive --delete --rfunc .*hotspot.* > residual.bc
llc -filetype=obj residual.bc > residual.o
# Play with hotspot.bc here
llc -filetype=obj hotspot.bc > hotspot.o
cc -o app residual.o hotspot.o

[davide]

Please add a basic test and a test with mutually recursive calls. The basic logic seems fine to me, some nits (sorry for being pedantic).

tools/llvm-extract/llvm-extract.cpp

236 ↗

(On Diff #94274)

explicit type when non-obvious.

244 ↗

(On Diff #94274)

Can we have a slightly better diagnostic here? Something like Cannot materialize function F + F->getName().

245 ↗

(On Diff #94274)

I generally prefer early continue to avoid many levels of nesting (there are a few in this case), so I'd rewrite the loop as:

for (...) {
  auto *CI = dyn_cast<>();
  if (!CI)
    continue;
  [...]
}

but I don't feel strong about it, so, up to you.

247 ↗

(On Diff #94274)

and probably auto when the type is obvious`.

[davide]

tools/llvm-extract/llvm-extract.cpp

21–23 ↗

(On Diff #94274)

Also this reordering seems unrelated? I think it's a good cleanup but can be committed separately (without pre-commit code review).

[loladiro]

Address review comments

[loladiro]

tools/llvm-extract/llvm-extract.cpp

21–23 ↗

(On Diff #94274)

I run git clang-format as a commit hook and it reorders the include list. Happy to commit separately, but I though the general policy on that was just to put formatting changes into the commit where you're editing the code.

244 ↗

(On Diff #94274)

IIRC, ExitOnError will add diagnostics if the Error object provides them.

[davide]

LGTM

llvm/trunk/test/tools/llvm-extract/recursive.ll

+; RUN: llvm-extract -func=a --recursive %s -S | FileCheck --check-prefix=CHECK-AB %s
+; RUN: llvm-extract -func=a --recursive --delete %s -S | FileCheck --check-prefix=CHECK-CD %s
+; RUN: llvm-extract -func=d --recursive %s -S | FileCheck --check-prefix=CHECK-CD %s
+
+; CHECK-AB: define void @a
+; CHECK-AB: define void @b
+; CHECK-AB-NOT: define void @c
+; CHECK-AB-NOT: define void @d
+
+; CHECK-CD-NOT: define void @a
+; CHECK-CD-NOT: define void @b
+; CHECK-CD: define void @c
+; CHECK-CD: define void @d
+
+define void @a() {
+  call void @b()
+  ret void
+}
+
+define void @b() {
+  ret void
+}
+
+define void @c() {
+  call void @d()
+  ret void
+}
+
+define void @d() {
+  call void @c()
+  ret void
+}

llvm/trunk/tools/llvm-extract/llvm-extract.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Bitcode/BitcodeWriterPass.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/IRPrintingPasses.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IRReader/IRReader.h"
-#include "llvm/IR/LegacyPassManager.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/PrettyStackTrace.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/SourceMgr.h"
OutputFilename("o", cl::desc("Specify output filename"),
                cl::value_desc("filename"), cl::init("-"));
 
 static cl::opt<bool>
 Force("f", cl::desc("Enable binary output on terminals"));
 
 static cl::opt<bool>
 DeleteFn("delete", cl::desc("Delete specified Globals from Module"));
 
+static cl::opt<bool>
+    Recursive("recursive",
+              cl::desc("Recursively extract all called functions"));
+
 // ExtractFuncs - The functions to extract from the module.
 static cl::list<std::string>
 ExtractFuncs("func", cl::desc("Specify function to extract"),
              cl::ZeroOrMore, cl::value_desc("function"));
 
 // ExtractRegExpFuncs - The functions, matched via regular expression, to
 // extract from the module.
 static cl::list<std::string>
if (!match) {
              << ExtractRegExpFuncs[i] << "'!\n";
       return 1;
     }
   }
 
   // Use *argv instead of argv[0] to work around a wrong GCC warning.
   ExitOnError ExitOnErr(std::string(*argv) + ": error reading input: ");
 
+  if (Recursive) {
+    std::vector<llvm::Function *> Workqueue;
+    for (GlobalValue *GV : GVs) {
+      if (auto *F = dyn_cast<Function>(GV)) {
+        Workqueue.push_back(F);
+      }
+    }
+    while (!Workqueue.empty()) {
+      Function *F = &*Workqueue.back();
+      Workqueue.pop_back();
+      ExitOnErr(F->materialize());
+      for (auto &BB : *F) {
+        for (auto &I : BB) {
+          auto *CI = dyn_cast<CallInst>(&I);
+          if (!CI)
+            continue;
+          Function *CF = CI->getCalledFunction();
+          if (!CF)
+            continue;
+          if (CF->isDeclaration() || GVs.count(CF))
+            continue;
+          GVs.insert(CF);
+          Workqueue.push_back(CF);
+        }
+      }
+    }
+  }
+
   auto Materialize = [&](GlobalValue &GV) { ExitOnErr(GV.materialize()); };
 
   // Materialize requisite global values.
   if (!DeleteFn) {
     for (size_t i = 0, e = GVs.size(); i != e; ++i)
       Materialize(*GVs[i]);
   } else {
     // Deleting. Materialize every GV that's *not* in GVs.