registration.h 875 B

12345678910111213141516171819202122
  1. #pragma once
  2. #include <Python.h>
  3. #define _CONCAT(A, B) A##B
  4. #define CONCAT(A, B) _CONCAT(A, B)
  5. #define _STRINGIFY(A) #A
  6. #define STRINGIFY(A) _STRINGIFY(A)
  7. // A version of the TORCH_LIBRARY macro that expands the NAME, i.e. so NAME
  8. // could be a macro instead of a literal token.
  9. #define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE)
  10. // REGISTER_EXTENSION allows the shared library to be loaded and initialized
  11. // via python's import statement.
  12. #define REGISTER_EXTENSION(NAME) \
  13. PyMODINIT_FUNC CONCAT(PyInit_, NAME)() { \
  14. static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, \
  15. STRINGIFY(NAME), nullptr, 0, nullptr}; \
  16. return PyModule_Create(&module); \
  17. }