diff --git a/extension_cpp/__init__.py b/extension_cpp/__init__.py index a5b254a..e75150b 100644 --- a/extension_cpp/__init__.py +++ b/extension_cpp/__init__.py @@ -1,10 +1,3 @@ import torch from pathlib import Path - -so_files = list(Path(__file__).parent.glob("_C*.so")) -assert ( - len(so_files) == 1 -), f"Expected one _C*.so file, found {len(so_files)}" -torch.ops.load_library(so_files[0]) - -from . import ops +from . import _C, ops diff --git a/extension_cpp/csrc/muladd.cpp b/extension_cpp/csrc/muladd.cpp index 85f9fce..7b3415f 100644 --- a/extension_cpp/csrc/muladd.cpp +++ b/extension_cpp/csrc/muladd.cpp @@ -1,7 +1,29 @@ -#include +#include +#include +#include +#include #include +extern "C" { + /* Creates a dummy empty _C module that can be imported from Python. + The import from Python will load the .so consisting of this file + in this extension, so that the TORCH_LIBRARY static initializers + below are run. */ + PyObject* PyInit__C(void) + { + static struct PyModuleDef module_def = { + PyModuleDef_HEAD_INIT, + "_C", /* name of module */ + NULL, /* module documentation, may be NULL */ + -1, /* size of per-interpreter state of the module, + or -1 if the module keeps state in global variables. */ + NULL, /* methods */ + }; + return PyModule_Create(&module_def); + } +} + namespace extension_cpp { at::Tensor mymuladd_cpu(const at::Tensor& a, const at::Tensor& b, double c) { diff --git a/setup.py b/setup.py index 6f699bc..0dde1e4 100644 --- a/setup.py +++ b/setup.py @@ -38,6 +38,7 @@ def get_extensions(): "cxx": [ "-O3" if not debug_mode else "-O0", "-fdiagnostics-color=always", + "-DPy_LIMITED_API=0x03090000", # min CPython version 3.9 ], "nvcc": [ "-O3" if not debug_mode else "-O0",