From 316984056cfa67e44ee16791ebd4c5f002cac37c Mon Sep 17 00:00:00 2001 From: Jane Xu Date: Thu, 23 Jan 2025 13:46:11 -0800 Subject: [PATCH 1/3] Create dummymodule _C instead of using load_library(blahso) --- extension_cpp/__init__.py | 9 +-------- extension_cpp/csrc/muladd.cpp | 23 ++++++++++++++++++++++- setup.py | 1 + 3 files changed, 24 insertions(+), 9 deletions(-) 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..ade0380 100644 --- a/extension_cpp/csrc/muladd.cpp +++ b/extension_cpp/csrc/muladd.cpp @@ -1,7 +1,28 @@ -#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 associated with this extension + built from this file, so that all the TORCH_LIBRARY calls 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..fd6f7a9 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", ], "nvcc": [ "-O3" if not debug_mode else "-O0", From 8c8adcc28dca49babcf99ffb06aea77c025140aa Mon Sep 17 00:00:00 2001 From: Jane Xu Date: Thu, 23 Jan 2025 13:59:30 -0800 Subject: [PATCH 2/3] rICHARD explained that TORCH_LIBRARY calls are called because theyre static initializers --- extension_cpp/csrc/muladd.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/extension_cpp/csrc/muladd.cpp b/extension_cpp/csrc/muladd.cpp index ade0380..7b3415f 100644 --- a/extension_cpp/csrc/muladd.cpp +++ b/extension_cpp/csrc/muladd.cpp @@ -7,8 +7,9 @@ extern "C" { /* Creates a dummy empty _C module that can be imported from Python. - The import from Python will load the .so associated with this extension - built from this file, so that all the TORCH_LIBRARY calls below are run.*/ + 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 = { From 58ac996f12a7c1438bcda4d856b6541c5b8a2ec4 Mon Sep 17 00:00:00 2001 From: Jane Xu Date: Fri, 24 Jan 2025 08:53:02 -0800 Subject: [PATCH 3/3] Add comment --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index fd6f7a9..0dde1e4 100644 --- a/setup.py +++ b/setup.py @@ -38,7 +38,7 @@ def get_extensions(): "cxx": [ "-O3" if not debug_mode else "-O0", "-fdiagnostics-color=always", - "-DPy_LIMITED_API=0x03090000", + "-DPy_LIMITED_API=0x03090000", # min CPython version 3.9 ], "nvcc": [ "-O3" if not debug_mode else "-O0",