mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
7ff1f3f3f6
This reverts commit 851e89c8e8.
Differential Revision: [D45034526](https://our.internmc.facebook.com/intern/diff/D45034526)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/99275
Approved by: https://github.com/eellison
32 lines
758 B
C++
32 lines
758 B
C++
#if !defined(USE_ROCM) && defined(PYTORCH_EXPANDABLE_SEGMENTS_SUPPORTED)
|
|
#include <c10/cuda/driver_api.h>
|
|
#include <c10/util/Exception.h>
|
|
#include <dlfcn.h>
|
|
#include <iostream>
|
|
namespace c10 {
|
|
namespace cuda {
|
|
|
|
namespace {
|
|
DriverAPI create_driver_api() {
|
|
void* handle = dlopen("libcuda.so", RTLD_LAZY | RTLD_NOLOAD);
|
|
TORCH_INTERNAL_ASSERT(handle);
|
|
DriverAPI r;
|
|
|
|
#define LOOKUP_ENTRY(name) \
|
|
r.name##_ = ((decltype(&name))dlsym(handle, #name)); \
|
|
TORCH_INTERNAL_ASSERT(r.name##_)
|
|
C10_FORALL_DRIVER_API(LOOKUP_ENTRY)
|
|
#undef LOOKUP_ENTRY
|
|
return r;
|
|
}
|
|
} // namespace
|
|
|
|
DriverAPI* DriverAPI::get() {
|
|
static DriverAPI singleton = create_driver_api();
|
|
return &singleton;
|
|
}
|
|
|
|
} // namespace cuda
|
|
} // namespace c10
|
|
|
|
#endif
|