forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
nvrtc.cpp
44 lines (36 loc) · 891 Bytes
/
nvrtc.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
#include "torch/csrc/python_headers.h"
static PyObject* module;
static PyMethodDef TorchNvrtcMethods[] = {
{nullptr, nullptr, 0, nullptr}
};
#if PY_MAJOR_VERSION != 2
static struct PyModuleDef torchnvrtcmodule = {
PyModuleDef_HEAD_INIT,
"torch._nvrtc",
nullptr,
-1,
TorchNvrtcMethods
};
#endif
#if PY_MAJOR_VERSION == 2
PyMODINIT_FUNC init_nvrtc(void)
#else
PyMODINIT_FUNC PyInit__nvrtc(void)
#endif
{
#if PY_MAJOR_VERSION == 2
#define ASSERT_TRUE(cmd) if (!(cmd)) {PyErr_SetString(PyExc_ImportError, "initialization error in torch._nvrtc"); return;}
#else
#define ASSERT_TRUE(cmd) if (!(cmd)) return nullptr
#endif
#if PY_MAJOR_VERSION == 2
ASSERT_TRUE(module = Py_InitModule("torch._nvrtc", TorchNvrtcMethods));
#else
ASSERT_TRUE(module = PyModule_Create(&torchnvrtcmodule));
#endif
#if PY_MAJOR_VERSION == 2
#else
return module;
#endif
#undef ASSERT_TRUE
}