Draft for adding reactant-opt and make reactant-clang plugin to be activated

enzyme/BUILD

@@ -65,6 +65,7 @@ cc_library(
         "@llvm-project//llvm:TargetParser",
         "@llvm-project//llvm:TransformUtils",
         "@llvm-project//llvm:config",
+        "@llvm-project//llvm:FrontendOpenMP",
     ],
     alwayslink = 1,
 )
@@ -78,6 +79,19 @@ expand_template(
     template = "@llvm-project//llvm:cmake/modules/llvm-driver-template.cpp.in",
 )
 
+cc_binary(
+    name = "reactant-opt",
+    srcs = ["reactant-opt-driver.cpp"],
+    deps = [
+        ":ReactantStatic",
+        "@llvm-project//llvm:AllTargetsAsmParsers",
+        "@llvm-project//llvm:AllTargetsCodeGens",
+        "@llvm-project//llvm:IRReader",
+        "@llvm-project//llvm:Passes",
+        "@llvm-project//llvm:Support",
+    ],
+)
+
 cc_binary(
     name = "reactant-clang",
     srcs = ["reactant-clang-driver.cpp"],
@@ -90,6 +104,7 @@ cc_binary(
         ":ReactantStatic",
         "@llvm-project//clang:clang-driver",
         "@llvm-project//llvm:Support",
+        "@llvm-project//llvm:FrontendOpenMP",
     ],
 )
 
@@ -124,7 +139,11 @@ genrule(
 
 copy_clang_resource_dir(
     name = "reactant-clang-resource",
-    srcs = ["@llvm-project//clang:builtin_headers_gen"],
+    srcs = [
+        "@llvm-project//clang:builtin_headers_gen",
+        "@llvm_openmp_headers//:omp_headers_staged",
+    ],
 )
 
+
 exports_files(["run_lit.sh"])

enzyme/Enzyme/Clang/EnzymeClang.cpp

@@ -72,7 +72,7 @@ for (auto a : args) llvm::errs() << "+ arg: " << a<<"\n";
   }
 
   PluginASTAction::ActionType getActionType() override {
-    return AddBeforeMainAction;
+    return CmdlineBeforeMainAction;
   }
 };
 
@@ -122,6 +122,11 @@ class EnzymePlugin final : public clang::ASTConsumer {
 
 public:
   EnzymePlugin(clang::CompilerInstance &CI) : CI(CI) {
+    llvm::errs() << "EnzymePlugin start \n";
+    if (getenv("NO_REACTANT_PLUGIN")){
+      llvm::errs() << "No registering of plugin - returning\n";
+      return;
+    }
     //FrontendOptions &Opts = CI.getFrontendOpts();
     CodeGenOptions &CGOpts = CI.getCodeGenOpts();
     auto PluginName = "ClangReactant-" + std::to_string(LLVM_VERSION_MAJOR);
@@ -133,29 +138,29 @@ class EnzymePlugin final : public clang::ASTConsumer {
     std::string inFile;
     for (auto in : CI.getFrontendOpts().Inputs) {
     	if (in.isFile()) {
-	 inFile = in.getFile().str();
-      llvm::errs() << " in: " << in.getFile() << "\n";
-	}
+        inFile = in.getFile().str();
+        llvm::errs() << " in: " << in.getFile() << "\n";
+	    }
     }
     if (CI.getLangOpts().CUDAIsDevice) {
 	    std::string file = CI.getFrontendOpts().OutputFile;
 	    file = inFile;
       CGOpts.PassBuilderCallbacks.push_back([=](llvm::PassBuilder & PB) {
-	registerExporter(PB, file);
+          registerExporter(PB, file);
       });
     } else {
-	std::vector<std::string> gpubins;
-	if (CGOpts.CudaGpuBinaryFileName.size()) {
-	  if (inFile.size())
-	    gpubins.push_back(inFile);
-	  //gpubins.push_back(CGOpts.CudaGpuBinaryFileName);
-	}
-        std::string file = CI.getFrontendOpts().OutputFile;
-        CGOpts.PassBuilderCallbacks.push_back([=](llvm::PassBuilder &PB) {
-          registerReactant(PB, gpubins, file);
-        });
+      std::vector<std::string> gpubins;
+      if (CGOpts.CudaGpuBinaryFileName.size()) {
+        if (inFile.size())
+          gpubins.push_back(inFile);
+        //gpubins.push_back(CGOpts.CudaGpuBinaryFileName);
+      }
+            std::string file = CI.getFrontendOpts().OutputFile;
+            CGOpts.PassBuilderCallbacks.push_back([=](llvm::PassBuilder &PB) {
+              registerReactant(PB, gpubins, file);
+            });
     }
-
+        
     CI.getPreprocessorOpts().Includes.push_back("/enzyme/enzyme/version");
 
     std::string PredefineBuffer;

enzyme/WORKSPACE

@@ -131,6 +131,9 @@ llvm("llvm-raw")
 load("@llvm-raw//utils/bazel:configure.bzl", "llvm_configure")
 llvm_configure(name = "llvm-project", targets = LLVM_TARGETS)
 
+load("//:llvm_openmp_repo.bzl", "llvm_openmp_headers")
+llvm_openmp_headers(name = "llvm_openmp_headers")
+
 
 load("@xla//:workspace1.bzl", "xla_workspace1")
 xla_workspace1()

enzyme/llvm_openmp_repo.bzl

@@ -0,0 +1,52 @@
+# llvm_openmp_repo.bzl
+def _llvm_openmp_headers_impl(rctx):
+    # Navigate to llvm-raw root via CMakeLists.txt (stable anchor in all LLVM versions)
+    llvm_root = rctx.path(Label("@llvm-raw//:CMakeLists.txt")).dirname
+    omp_include = str(llvm_root) + "/openmp/runtime/src/include"
+
+    rctx.execute(["mkdir", "-p", "staging/include"])
+
+    # Generate omp.h from the .var template
+    result = rctx.execute(["sh", "-c", """
+        set -e
+        sed \
+          -e 's/@LIBOMP_VERSION_MAJOR@/5/g' \
+          -e 's/@LIBOMP_VERSION_MINOR@/2/g' \
+          -e 's/@LIBOMP_VERSION_BUILD@/20231101/g' \
+          -e 's/@LIBOMP_LIB_FILE@/libomp.so/g' \
+          -e 's/@LIBOMP_HAVE_QUAD_PRECISION@/0/g' \
+          -e 's/@LIBOMP_QUAD_PRECISION_TYPEOF_SPECIFIER@//g' \
+          -e 's/@LIBOMP_HAVE_HIDDEN_HELPER_TASK@/1/g' \
+          -e 's/@LIBOMP_HAVE_OMPT_SUPPORT@/0/g' \
+          -e 's/@LIBOMP_OMPT_OPTIONAL@/0/g' \
+          -e 's/@LIBOMP_USE_ITT_NOTIFY@/0/g' \
+          '{omp_include}/omp.h.var' > staging/include/omp.h
+    """.format(omp_include = omp_include)])
+
+    if result.return_code != 0:
+        fail("Failed to generate omp.h: " + result.stderr)
+
+    # Copy direct headers — use || true so missing ones don't fail
+    # omp-tools.h: OpenMP Tools Interface (OMPT)
+    # ompx.h: OpenMP extensions (LLVM-specific, may not exist in older versions)
+    rctx.execute(["sh", "-c", """
+        set -e
+        for h in omp-tools.h ompx.h omp_lib.h; do
+            src='{omp_include}/'"$h"
+            test -f "$src" && cp "$src" 'staging/include/'"$h" || true
+        done
+    """.format(omp_include = omp_include)])
+
+    # Emit the BUILD so the staged files are accessible
+    rctx.file("BUILD.bazel", """
+filegroup(
+    name = "omp_headers_staged",
+    srcs = glob(["staging/include/**"]),
+    visibility = ["//visibility:public"],
+)
+""")
+
+llvm_openmp_headers = repository_rule(
+    implementation = _llvm_openmp_headers_impl,
+    attrs = {},
+)

enzyme/reactant-opt-driver.cpp

@@ -0,0 +1,57 @@
+#include "llvm/IRReader/IRReader.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/TargetSelect.h"
+
+using namespace llvm;
+
+static cl::opt<std::string> InputFile(cl::Positional, cl::desc("<input bitcode>"), cl::Required);
+
+extern "C" void registerReactant(PassBuilder &PB, std::vector<std::string> gpubins, std::string outfile);
+
+int main(int argc, char **argv) {
+    InitLLVM X(argc, argv);
+    InitializeAllTargets();
+    InitializeAllTargetMCs();
+    InitializeAllAsmPrinters();
+    InitializeAllAsmParsers();
+
+    cl::ParseCommandLineOptions(argc, argv, "reactant-opt\n");
+
+    LLVMContext Ctx;
+    Ctx.setDiscardValueNames(false);
+    SMDiagnostic Err;
+    auto M = parseIRFile(InputFile, Err, Ctx);
+    if (!M) {
+        Err.print(argv[0], errs());
+        return 1;
+    }
+
+    PassBuilder PB;
+    LoopAnalysisManager LAM;
+    FunctionAnalysisManager FAM;
+    CGSCCAnalysisManager CGAM;
+    ModuleAnalysisManager MAM;
+
+    registerReactant(PB, {}, "");
+
+    PB.registerModuleAnalyses(MAM);
+    PB.registerFunctionAnalyses(FAM);
+    PB.registerLoopAnalyses(LAM);
+    PB.registerCGSCCAnalyses(CGAM);
+    PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
+
+    ModulePassManager MPM;
+    MPM.addPass(createModuleToFunctionPassAdaptor(PassManager<Function>()));
+
+    // parse and run the pipeline
+    if (auto Err = PB.parsePassPipeline(MPM, "reactant")) {
+        errs() << "Error: " << toString(std::move(Err)) << "\n";
+        return 1;
+    }
+
+    MPM.run(*M, MAM);
+    return 0;
+}