Remove fn body pass

llvm_ir_dataset_utils/tools/RemoveFunctionBodyPass/CMakeLists.txt

@@ -0,0 +1,36 @@
+cmake_minimum_required(VERSION 3.20)
+
+project(RemoveFunctionBodyPass)
+
+set(LT_LLVM_INSTALL_DIR "" CACHE PATH "LLVM installation directory")
+list(APPEND CMAKE_PREFIX_PATH "${LT_LLVM_INSTALL_DIR}/lib/cmake/llvm/")
+
+set (CMAKE_CXX_STANDARD 17 CACHE STRING "")
+
+find_package(LLVM REQUIRED CONFIG)
+
+include_directories(SYSTEM ${LLVM_INCLUDE_DIRS})
+
+if(NOT LLVM_ENABLE_RTTI)
+  set(CMKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-rtti")
+endif()
+
+message(STATUS "Found LLVM ${LLVM_PACKAGE_VERSION}")
+message(STATUS "Using LLVMConfig.cmake in: ${LLVM_DIR}")
+
+add_library(RemoveFunctionBody SHARED RemoveFunctionBody.cpp)
+target_link_libraries(RemoveFunctionBody "$<$<PLATFORM_ID:Darwin>:-undefined dynamic_lookup>")
+
+# separate_arguments(LLVM_DEFINITIONS_LIST NATIVE_COMMAND ${LLVM_DEFINITIONS})
+# add_definitions(${LLVM_DEFINITIONS_LIST})
+# include_directories(${LLVM_INCLUDE_DIRS})
+#
+# include_directories(${CMAKE_SOURCE_DIR}/include)
+#
+# # Link with passes library
+# llvm_map_components_to_libnames(llvm_libs passes)
+# add_subdirectory(RemoveFunctionBodyPass)
+#
+# # add_llvm_pass_plugin(RemoveFunctionBody ./src/RemoveFunctionBody.cpp)
+
+

llvm_ir_dataset_utils/tools/RemoveFunctionBodyPass/RemoveFunctionBody.cpp

@@ -0,0 +1,86 @@
+#include <llvm/Config/llvm-config-x86_64.h>
+#include <llvm/IR/PassManager.h>
+#include <llvm/IR/Value.h>
+#include <llvm/Support/Compiler.h>
+
+#include "llvm/IR/LegacyPassManager.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Passes/PassPlugin.h"
+#include "llvm/Support/CommandLine.h"
+using namespace llvm;
+using namespace std;
+
+namespace {
+cl::list<int> Lists("index", cl::desc("Specify function index"), cl::OneOrMore);
+struct RemoveFunctionBodyPass : PassInfoMixin<RemoveFunctionBodyPass> {
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {
+    if (Lists.size() != 1) {
+      errs() << "There should be only 1 argument!\n";
+      return PreservedAnalyses::all();
+    }
+
+    int Index = -2;
+    for (auto arg : Lists) {
+      Index = arg;
+    }
+    int n_funcs = M.getFunctionList().size();
+    if (Index == -1) {
+      errs() << "n=" << n_funcs << '\n';
+      return PreservedAnalyses::all();
+    }
+    if (Index >= n_funcs || Index < -1) {
+      errs() << "Index is out of range! "
+             << "Index=" << Index << " n=" << n_funcs << "\n";
+      return PreservedAnalyses::all();
+    }
+
+    int count = -1;
+    for (Function &F : M) {
+      Function *fp = &F;
+      if (count == Index - 1) {
+        if (fp->isDeclaration()) {
+          return PreservedAnalyses::all();
+        }
+        std::vector<Use *> useList;
+        for (auto &use : fp->uses()) {
+          useList.push_back(&use);
+        }
+        for (auto &use : useList) {
+          if (isa<GlobalAlias>(use->getUser())) {
+            GlobalAlias *ga = cast<GlobalAlias>(use->getUser());
+            ga->replaceAllUsesWith(fp);
+            ga->eraseFromParent();
+          }
+        }
+        fp->deleteBody();
+        fp->setComdat(NULL);
+      }
+      count++;
+    }
+    return PreservedAnalyses::none();
+  }
+};
+} // namespace
+
+/* New PM Registration */
+llvm::PassPluginLibraryInfo getRemoveFunctionBodyPluginInfo() {
+  return {LLVM_PLUGIN_API_VERSION, "RemoveFunctionBody", LLVM_VERSION_STRING,
+          [](PassBuilder &PB) {
+            PB.registerPipelineParsingCallback(
+                [](StringRef Name, llvm::ModulePassManager &MPM,
+                   ArrayRef<llvm::PassBuilder::PipelineElement>) {
+                  if (Name == "remove-fn-body") {
+                    MPM.addPass(RemoveFunctionBodyPass());
+                    return true;
+                  }
+                  return false;
+                });
+          }};
+}
+
+// Core interface for pass plugins. To guarantee 'opt'
+// will be able to recognize the pass
+extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
+llvmGetPassPluginInfo() {
+  return getRemoveFunctionBodyPluginInfo();
+}

llvm_ir_dataset_utils/tools/RemoveFunctionBodyPass/install

@@ -0,0 +1,3 @@
+mkdir build && cd build
+cmake -GNinja ..
+ninja