diff --git a/pyg_lib/csrc/library.cpp b/pyg_lib/csrc/library.cpp new file mode 100644 index 000000000..ecd136071 --- /dev/null +++ b/pyg_lib/csrc/library.cpp @@ -0,0 +1,34 @@ +#include "library.h" + +#ifdef USE_PYTHON +#include +#endif + +#ifdef WITH_CUDA +#include +#endif + +#include + +// If we are in a Windows environment, we need to define +// initialization functions for the _custom_ops extension. +// For PyMODINIT_FUNC to work, we need to include Python.h +#ifdef _WIN32 +#ifdef USE_PYTHON +PyMODINIT_FUNC PyInit__C(void) { return NULL; } +#endif // USE_PYTHON +#endif // _WIN32 + +namespace pyg { + +int64_t cuda_version() { +#ifdef WITH_CUDA + return CUDA_VERSION; +#else + return -1; +#endif +} + +TORCH_LIBRARY_FRAGMENT(pyg, m) { m.def("cuda_version", &cuda_version); } + +} // namespace pyg diff --git a/pyg_lib/csrc/library.h b/pyg_lib/csrc/library.h new file mode 100644 index 000000000..3012184e4 --- /dev/null +++ b/pyg_lib/csrc/library.h @@ -0,0 +1,17 @@ +#pragma once + +#include "macros.h" + +namespace pyg { + +PYG_API int64_t cuda_version(); + +namespace detail { + +extern "C" PYG_INLINE_VARIABLE auto _register_ops = &cuda_version; +#ifdef HINT_MSVC_LINKER_INCLUDE_SYMBOL +#pragma comment(linker, "/include:_register_ops") +#endif + +} // namespace detail +} // namespace pyg diff --git a/pyg_lib/csrc/macros.h b/pyg_lib/csrc/macros.h new file mode 100644 index 000000000..63db86133 --- /dev/null +++ b/pyg_lib/csrc/macros.h @@ -0,0 +1,22 @@ +#pragma once + +#ifdef _WIN32 +#if defined(pyg_EXPORTS) +#define PYG_API __declspec(dllexport) +#else +#define PYG_API __declspec(dllimport) +#endif +#else +#define PYG_API +#endif + +#if (defined __cpp_inline_variables) || __cplusplus >= 201703L +#define PYG_INLINE_VARIABLE inline +#else +#ifdef _MSC_VER +#define PYG_INLINE_VARIABLE __declspec(selectany) +#define HINT_MSVC_LINKER_INCLUDE_SYMBOL +#else +#define PYG_INLINE_VARIABLE __attribute__((weak)) +#endif +#endif