forked from scikit-hep/awkward
-
Notifications
You must be signed in to change notification settings - Fork 0
/
startup.cpp
45 lines (36 loc) · 1.41 KB
/
startup.cpp
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
// BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
#include <iostream>
#include <string>
#include "awkward/python/startup.h"
namespace ak = awkward;
std::string StartupLibraryPathCallback::library_path() {
// Fetches Path Eagerly
std::string eager_path;
try {
py::object awkward_cuda_kernels = py::module::import(
"awkward_cuda_kernels");
if (py::hasattr(awkward_cuda_kernels, "shared_library_path")) {
py::object library_path_pyobj = py::getattr(awkward_cuda_kernels,
"shared_library_path");
eager_path = library_path_pyobj.cast<std::string>();
}
}
catch (...) {
// If the Python Extension fails to import the module, fall back to the
// previous fetched path
eager_path = library_path_;
}
// Keep the Library Patch as a cache, this will fetch paths even when the shared
// library path is changed but in case of a failure it will revert to the last
// successful library fetch. In case the shared Library ceases to exist, it will
// be handled by the dlopen and dlsym in the C++ layer.
library_path_ = eager_path;
return library_path_;
}
void
make_startup(py::module& m, const std::string& name) {
m.def(name.c_str(), []() -> void {
ak::kernel::lib_callback->add_library_path_callback(
ak::kernel::lib::cuda, std::make_shared<StartupLibraryPathCallback>());
});
}