llvm-py-pass/plugin.cpp at main · weliveindetail/llvm-py-pass · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
#include <llvm/IR/Function.h>
#include <llvm/Pass.h>
#include <llvm/Passes/PassBuilder.h>
#include <llvm/Passes/PassPlugin.h>
#include <llvm/Support/FileSystem.h>
#include <llvm/Support/Path.h>
#include <llvm/Transforms/IPO/PassManagerBuilder.h>

#include <llvm-c/Core.h>
#include <llvm-c/Types.h>

#include <cstdio>
#include <cstdlib>
#include <fstream>
#include <sstream>
#include <string>
#include <unistd.h>

using namespace llvm;

extern bool runPyPassOn(const std::string &, LLVMModuleRef, LLVMContextRef);

namespace {

template <typename PredFn> static void visitScriptPaths(PredFn &&Pred) {
  SmallString<256> dir;
  sys::fs::current_path(dir);
  std::vector<std::string> paths;
  while (true) {
    if (Pred(dir.str()))
      return;
    if (!sys::path::has_parent_path(dir))
      break;
    dir = sys::path::parent_path(dir);
  }
}

static std::string Script;
static std::string ScriptPath;

struct PyPass : PassInfoMixin<PyPass> {
  PreservedAnalyses run(Module &M, ModuleAnalysisManager &) {
    LLVMModuleRef Mod = wrap(&M);
    LLVMContextRef Ctx = wrap(&M.getContext());
    fprintf(stderr, "[llvm-py-pass] Running pass on %s\n",
            M.getModuleIdentifier().c_str());
    SmallString<256> CWD;
    sys::fs::current_path(CWD);
    sys::fs::set_current_path(ScriptPath);
    bool Changed = runPyPassOn(Script, Mod, Ctx);
    sys::fs::set_current_path(CWD);
    return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
  }
};

} // namespace

extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
llvmGetPassPluginInfo() {
  if (Script.empty()) {
    fprintf(stderr, "[llvm-py-pass] Obtaining pass plugin info\n");
    visitScriptPaths([&](Twine path) {
      auto file = path + "/run.py";
      if (sys::fs::exists(path)) {
        std::ifstream stream(file.str());
        std::stringstream buffer;
        buffer << stream.rdbuf();
        Script = buffer.str();
        ScriptPath = path.str();
        fprintf(stderr, "[llvm-py-pass] Found %s/run.py\n", ScriptPath.c_str());
        return true; // found, done
      }
      return false; // check next location
    });
  }

  if (Script.empty()) {
    fprintf(stderr, "[llvm-py-pass] Skipped: Cannot find script 'run.py' in "
                    "any of the following locations\n");
    visitScriptPaths([&](Twine path) {
      fprintf(stderr, "  %s\n", path.str().c_str());
      return false;
    });
    return {LLVM_PLUGIN_API_VERSION, "llvm-py-pass", LLVM_VERSION_STRING,
            [](PassBuilder &PB) {}};
  }

  return {LLVM_PLUGIN_API_VERSION, "llvm-py-pass", LLVM_VERSION_STRING,
          [](PassBuilder &PB) {
            PB.registerPipelineEarlySimplificationEPCallback(
                [](ModulePassManager &MPM, OptimizationLevel opt) {
                  fprintf(stderr,
                          "[llvm-py-pass] Adding pass to optimizer pipeline\n");
                  MPM.addPass(PyPass());
                  return true;
                });
            PB.registerPipelineParsingCallback(
                [](StringRef Name, ModulePassManager &MPM,
                   ArrayRef<PassBuilder::PipelineElement> InnerPipeline) {
                  if (Name == "llvm-py-pass") {
                    fprintf(stderr,
                            "[llvm-py-pass] Adding pass llvm-py-pass\n");
                    MPM.addPass(PyPass());
                    return true;
                  }
                  return false;
                });
          }};
}