================ @@ -0,0 +1,464 @@ +//===- tools/plugins-shlib/pypass.cpp -------------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "llvm/ADT/ScopeExit.h" +#include "llvm/Passes/PassBuilder.h" +#include "llvm/Passes/PassPlugin.h" +#include "llvm/Support/DynamicLibrary.h" +#include "llvm/Support/raw_ostream.h" + +#include <algorithm> +#include <cstdlib> +#include <filesystem> +#include <memory> +#include <optional> +#include <string> + +using namespace llvm; + +static cl::opt<std::string> + DylibPath("pypass-dylib", cl::desc("Path to the Python shared library"), + cl::init("")); + +static cl::opt<std::string> + ScriptPath("pypass-script", cl::desc("Path to the Python script to run"), + cl::init("")); + +static std::string findPython() { + if (!DylibPath.empty()) + return DylibPath; + if (const char *Path = std::getenv("LLVM_PYPASS_DYLIB")) + return std::string(Path); + // TODO: Run Python from PATH and use a script to query the shared lib + return std::string{}; +} + +static std::string findScript() { + if (!ScriptPath.empty()) + return ScriptPath; + if (const char *Path = std::getenv("LLVM_PYPASS_SCRIPT")) + return std::string(Path); + return std::string{}; +} + +struct PythonAPI { + using Py_InitializeEx_t = void(int); + using Py_Initialize_t = void(void); + using Py_FinalizeEx_t = int(void); + using Py_Finalize_t = void(void); + using PyDict_GetItemString_t = void *(void *, const char *); + using PyGILStateEnsure_t = int(); + using PyGILStateRelease_t = void(int); + using PyImport_AddModule_t = void *(const char *); + using PyLong_FromVoidPtr_t = void *(void *); + using PyUnicode_FromString_t = void *(const char *); + using PyModule_GetDict_t = void *(void *); + using PyObject_CallObject_t = void *(void *, void *); + using PyObject_IsTrue_t = int(void *); + using PyRun_SimpleString_t = int(const char *); + using PyTuple_SetItem_t = int(void *, long, void *); + using PyTuple_New_t = void *(long); + using PyTypeObject_t = void *; + + // pylifecycle.h + Py_InitializeEx_t *Py_InitializeEx; + Py_Initialize_t *Py_Initialize; + Py_FinalizeEx_t *Py_FinalizeEx; + Py_Finalize_t *Py_Finalize; + + // pythonrun.h + PyRun_SimpleString_t *PyRun_SimpleString; + + // pystate.h + PyGILStateEnsure_t *PyGILState_Ensure; + PyGILStateRelease_t *PyGILState_Release; + + // import.h + PyImport_AddModule_t *PyImport_AddModule; + + // object.h + PyObject_IsTrue_t *PyObject_IsTrue; + + // moduleobject.h + PyModule_GetDict_t *PyModule_GetDict; + + // dictobject.h + PyDict_GetItemString_t *PyDict_GetItemString; + + // abstract.h + PyObject_CallObject_t *PyObject_CallObject; + + // longobject.h + PyLong_FromVoidPtr_t *PyLong_FromVoidPtr; + + // unicodeobject.h + PyUnicode_FromString_t *PyUnicode_FromString; + + // tupleobject.h + PyTuple_SetItem_t *PyTuple_SetItem; + PyTuple_New_t *PyTuple_New; + + PythonAPI() : Ready(false) { + if (!loadDylib(findPython())) + return; + if (!resolveSymbols()) + return; + if (Py_InitializeEx) { + Py_InitializeEx(0); + } else { + Py_Initialize(); + } + Ready = true; + } + + ~PythonAPI() { + if (std::atomic_exchange(&Ready, false)) { + if (Py_FinalizeEx) { + Py_FinalizeEx(); + } else { + Py_Finalize(); + } + } + } + + bool loadDylib(std::string Path) { + std::string Err; + Dylib = sys::DynamicLibrary::getPermanentLibrary(Path.c_str(), &Err); + if (!Dylib.isValid()) { + errs() << "dlopen for '" << Path << "' failed: " << Err << "\n"; + return false; + } + + return true; + } + + bool resolveSymbols() { + bool Success = true; + Success &= resolve("Py_InitializeEx", &Py_InitializeEx); + Success &= resolve("Py_Initialize", &Py_Initialize); + Success &= resolve("Py_FinalizeEx", &Py_FinalizeEx); + Success &= resolve("Py_Finalize", &Py_Finalize); + Success &= resolve("PyGILState_Ensure", &PyGILState_Ensure); + Success &= resolve("PyGILState_Release", &PyGILState_Release); + Success &= resolve("PyRun_SimpleString", &PyRun_SimpleString); + Success &= resolve("PyImport_AddModule", &PyImport_AddModule); + Success &= resolve("PyModule_GetDict", &PyModule_GetDict); + Success &= resolve("PyDict_GetItemString", &PyDict_GetItemString); + Success &= resolve("PyObject_CallObject", &PyObject_CallObject); + Success &= resolve("PyObject_IsTrue", &PyObject_IsTrue); + Success &= resolve("PyLong_FromVoidPtr", &PyLong_FromVoidPtr); + Success &= resolve("PyUnicode_FromString", &PyUnicode_FromString); + Success &= resolve("PyTuple_SetItem", &PyTuple_SetItem); + Success &= resolve("PyTuple_New", &PyTuple_New); + return Success; + } + + bool isReady() const { return Ready; } + + bool loadScript(const std::string &ScriptPath) const { + std::string LoadCmd; + raw_string_ostream(LoadCmd) + << "import runpy\n" + << "globals().update(runpy.run_path('" << ScriptPath << "'))"; + + if (PyRun_SimpleString(LoadCmd.c_str()) != 0) { + errs() << "Failed to load script: " << ScriptPath << "\n"; + return false; + } + + return true; + } + + bool addImportSearchPath(std::string Path) const { + std::string LoadCmd; + raw_string_ostream(LoadCmd) << "import sys\n" + << "sys.path.append('" << Path << "')"; + // Interpreter is not thread-safe + auto GIL = make_scope_exit( + [this, Lock = PyGILState_Ensure()]() { PyGILState_Release(Lock); }); + if (PyRun_SimpleString(LoadCmd.c_str()) != 0) { + errs() << "Failed to add import search path: " << Path << "\n"; + return false; + } + + return true; + } + + void *addModule(const char *Name) const { + void *Mod = PyImport_AddModule(Name); + return PyModule_GetDict(Mod); + } + + // Very simple interface to execute a Python function + bool invoke(void *Mod, const char *Name, void *Args = nullptr) const { + // If the function doesn't exist, we assume no + void *Fn = PyDict_GetItemString(Mod, Name); + if (!Fn) + return false; + // Interpreter is not thread-safe + auto GIL = make_scope_exit( + [this, Lock = PyGILState_Ensure()]() { PyGILState_Release(Lock); }); + // If we get no result, there was an error in Python + void *Result = PyObject_CallObject(Fn, Args); + if (!Result) { + errs() << "PyPassContext error: " << Name << "() failed\n"; + return false; + } + // If the result is truthy, then it's a yes + return PyObject_IsTrue(Result); + } + +private: + sys::DynamicLibrary Dylib; + std::atomic<bool> Ready; + + template <typename FnTy> bool resolve(const char *Name, FnTy **Var) { + assert(Dylib.isValid() && "dlopen shared library first"); + assert(*Var == nullptr && "Resolve symbols once"); + if (void *FnPtr = Dylib.getAddressOfSymbol(Name)) { + *Var = reinterpret_cast<FnTy *>(FnPtr); + return true; + } + errs() << "Missing required CPython API symbol '" << Name + << "' in: " << DylibPath << "\n"; + return false; + }; +}; + +// Python interface is initialized on first access and it is shared across all +// threads. It can be used like a state-less thread-safe object. +const PythonAPI &getPyAPI() { + static const PythonAPI PyAPI; + return PyAPI; +} + +class PyPassContext { +public: + PyPassContext() : PyAPI(getPyAPI()) {} + + bool loadScript(const std::string &Path) { + if (!PyAPI.isReady()) + return false; + + // Make relative paths resolve naturally in import statements + std::string Dir = std::filesystem::path(Path).parent_path().u8string(); + if (!PyAPI.addImportSearchPath(Dir)) + return false; + + if (!PyAPI.loadScript(Path)) + return false; + + PyGlobals = PyAPI.addModule("__main__"); + PyBuiltins = PyAPI.addModule("builtins"); + if (!PyGlobals || !PyBuiltins) + return false; + + return PyAPI.PyDict_GetItemString(PyGlobals, "run"); + } + + bool registerEP(std::string Name) { + std::string EP = "register" + Name; + return PyAPI.invoke(PyGlobals, EP.c_str()); + } + + bool run(void *Entity, void *Ctx, const char *Stage) { + void *Args = PyAPI.PyTuple_New(3); + if (!Args) + return false; + if (PyAPI.PyTuple_SetItem(Args, 0, PyAPI.PyLong_FromVoidPtr(Entity)) != 0) + return false; + if (PyAPI.PyTuple_SetItem(Args, 1, PyAPI.PyLong_FromVoidPtr(Ctx)) != 0) + return false; + if (PyAPI.PyTuple_SetItem(Args, 2, PyAPI.PyUnicode_FromString(Stage)) != 0) + return false; + + // TODO: Should we expose PyPassContext and/or arguments from specific + // entry-points like OptLevel, ThinOrFullLTOPhase or InnerPipeline? + + return PyAPI.invoke(PyGlobals, "run", Args); + } + +private: + const PythonAPI &PyAPI; + void *PyGlobals; + void *PyBuiltins; +}; + +struct PyPass : PassInfoMixin<PyPass> { + PyPass(std::shared_ptr<PyPassContext> Context, StringRef EP) + : Stage(EP.str()), Context(std::move(Context)) {} + + PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) { + LLVMModuleRef Mod = wrap(&M); + LLVMContextRef Ctx = wrap(&M.getContext()); + bool Changed = Context->run(Mod, Ctx, Stage.c_str()); + return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all(); + } + + PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM) { + LLVMValueRef Fn = wrap(&F); + LLVMContextRef Ctx = wrap(&F.getContext()); + bool Changed = Context->run(Fn, Ctx, Stage.c_str()); + return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all(); + } + + std::string Stage; + std::shared_ptr<PyPassContext> Context; + std::optional<OptimizationLevel> OptLevel; + std::optional<ThinOrFullLTOPhase> LTOPhase; + std::optional<ArrayRef<llvm::PassBuilder::PipelineElement>> InnerPipeline; +}; + +static void registerCallbacks(PassBuilder &PB) { + // Context is shared across all entry-points in this pipeline + auto Context = std::make_shared<PyPassContext>(); + + // All entry-point callbacks run the same script + std::string ScriptPath = findScript(); + if (!Context->loadScript(ScriptPath)) + return; + + // Create one PyPass instance per entry-point + if (Context->registerEP("PipelineStartEPCallback")) + PB.registerPipelineStartEPCallback( + [Context](ModulePassManager &MPM, OptimizationLevel Opt) { + PyPass P(Context, "PipelineStartEPCallback"); + P.OptLevel = Opt; + MPM.addPass(std::move(P)); + return true; + }); + + if (Context->registerEP("PipelineEarlySimplificationEPCallback")) + PB.registerPipelineEarlySimplificationEPCallback( + [Context](ModulePassManager &MPM, OptimizationLevel Opt, + ThinOrFullLTOPhase Phase) { + PyPass P(Context, "PipelineEarlySimplificationEPCallback"); + P.OptLevel = Opt; + P.LTOPhase = Phase; + MPM.addPass(std::move(P)); + return true; + }); + + if (Context->registerEP("OptimizerEarlyEPCallback")) + PB.registerOptimizerEarlyEPCallback([Context](ModulePassManager &MPM, + OptimizationLevel Opt, + ThinOrFullLTOPhase Phase) { + PyPass P(Context, "OptimizerEarlyEPCallback"); + P.OptLevel = Opt; + P.LTOPhase = Phase; + MPM.addPass(std::move(P)); + return true; + }); + + if (Context->registerEP("OptimizerLastEPCallback")) + PB.registerOptimizerLastEPCallback([Context](ModulePassManager &MPM, + OptimizationLevel Opt, + ThinOrFullLTOPhase Phase) { + PyPass P(Context, "OptimizerLastEPCallback"); + P.OptLevel = Opt; + P.LTOPhase = Phase; + MPM.addPass(std::move(P)); + return true; + }); + ---------------- weliveindetail wrote:
We have these repetitions everywhere in the codebase. It's not nice, but it is consistent at least. I think if we keep adding entry-points, we should find a better overall solution at some point. https://github.com/llvm/llvm-project/pull/171111 _______________________________________________ cfe-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
