Build torch._C extension module and version.py via CMake (pytorch#180243)

zklaus · pytorchmergebot · commit d35db535674e · 2026-06-04T19:11:35.000Z
Move the build of the `torch._C` Python extension and the install of `torch/version.py` from `setup.py` into `torch/CMakeLists.txt`, so both the existing setuptools backend and the upcoming scikit-build-core backend (pytorch#180247) go through the same code path. The new `Python_add_library(_C MODULE WITH_SOABI ...)` target builds the same `torch/csrc/stub.c` and installs at the wheel root (mapped to `torch/` by `wheel.install-dir`); `setup.py`'s `Extension("torch._C", ...)` block and its companion `extra_compile_args` / `extra_link_args` / `make_relative_rpath_args` plumbing are removed. `version.py` is gitignored, so it is installed explicitly via `install(FILES ... DESTINATION ".")` rather than relying on Python package collection. Adds a `BinaryDistribution` class in `setup.py` that forces `has_ext_modules=True` so `bdist_wheel` uses the binary layout (package_data + data_files only) instead of the purelib layout, which would otherwise dump CMake's `build/lib/` tree -- including test artifacts -- into the wheel root. `package_data` collects the CMake-built extension via the interpreter's exact module name (`f"_C{sysconfig.get_config_var('EXT_SUFFIX')}"`) rather than a `_C*` glob, so a `build/` dir reused across Python versions cannot leak other ABIs' `_C` modules into the wheel. Windows handling: the manual `setup.py` copy of `_C.lib` from `build_temp` to `build_lib` is replaced by the `_C` target's `ARCHIVE DESTINATION "lib"`. For MSVC, the linker is told where to find `pythonXX.lib` via `/LIBPATH` (Python headers use `#pragma comment(lib, "pythonXX")` which expects a bare-name lookup); for Windows venvs, FindPython's library search path is extended with `sys.base_prefix/libs` since the import library lives in the base install, not the venv. `nnapi_backend` gains `TORCH_PYTHON_LINK_FLAGS` so it picks up the same MSVC linker options as `torch_python`. ## Test plan Verification is via artifact-content comparison against the parent (pytorch#180242). Expectation: the wheel must be content-equivalent modulo the one file the PR rebuilds (`torch/_C.cpython-*.so`), which must remain functionally identical (same exports, same imports, same rpath). - [x] **No regression in artifact sizes** across every manywheel / libtorch artifact (sub-0.001 % to 0.001 % delta, zip-timestamp noise consistent across artifacts this PR can't directly affect): | artifact | parent (pytorch#180242) | this PR (pytorch#180243) | delta | | --- | ---: | ---: | ---: | | manywheel-py3_12-rocm7_2 | 6,153,280,701 | 6,153,275,144 | -5,557 (-0.000%) | | manywheel-py3_12-rocm7_1 | 5,788,503,789 | 5,788,499,957 | -3,832 (-0.000%) | | manywheel-py3_12-cuda12_6 | 841,797,521 | 841,797,813 | +292 (+0.000%) | | manywheel-py3_12-cuda13_2 | 624,213,880 | 624,212,127 | -1,753 (-0.000%) | | manywheel-py3_12-cpu | 189,248,452 | 189,248,117 | -335 (-0.000%) | | manywheel-py3_12-xpu | 778,712,445 | 778,711,285 | -1,160 (-0.000%) | | libtorch-rocm7_2-shared-with-deps-release | 12,117,315,856 | 12,117,314,539 | -1,317 (-0.000%) | | libtorch-cpu-shared-with-deps-release | 257,626,615 | 257,624,612 | -2,003 (-0.001%) | - [x] **Wheel manifest equivalence vs parent** (`manywheel-py3_12-rocm7_2`): both wheels contain **26,779 entries**; filename diff is empty -- no files added or removed. Confirms the `BinaryDistribution` layout fix worked (no `build/lib/` spillage into the wheel root). - [x] **Per-file size diff: exactly one file changed.** ``` torch/_C.cpython-312-x86_64-linux-gnu.so 22,960 -> 15,720 (-7,240) ``` Every other one of 26,778 files is byte-identical in size. That is precisely the file the PR rebuilds via CMake. - [x] **`torch._C` is functionally equivalent** between parent and PR. Extracted both `.so` files and compared: - **Defined symbols**: zero-line diff. - **Undefined symbols** (imports from `libtorch_python`): zero-line diff. - **RPATH**: identical (`$ORIGIN:$ORIGIN/lib`). - **NEEDED**: 180242's `.so` includes a redundant `libpthread.so.0` (merged into `libc.so.6` since glibc 2.34); the CMake build drops it. The size shrinkage matches the cleaner NEEDED + tighter section layout from the CMake toolchain. - [x] **`torch/version.py`** is present in both wheels at 339 bytes, byte-identical except for the expected `git_version` SHA (each commit has its own). `__version__`, `debug`, `cuda`, `hip`, `rocm`, `xpu` fields all match -- the CMake-driven generation produces the same output as the setuptools path. - [x] **Exactly one `_C` ABI per wheel.** On a fresh multi-ABI `manywheel-cpu-build` (build_all.sh builds all CPython ABIs in one checkout), each wheel ships only its own `torch/_C.cpython-XY.so`. An earlier greedy `_C*.so` `package_data` glob leaked stale per-ABI stubs (a cp312 wheel carried cp310/cp311/cp312); pinning to the exact `EXT_SUFFIX` fixes it. ``` unzip -l manywheel-py3_13-cpu.whl | grep 'torch/_C\.cpython' # -> only torch/_C.cpython-313-x86_64-linux-gnu.so ``` ### Coverage notes - **NNAPI MSVC link-flag fix** and **Windows venv `pythonXX.lib` discovery fix** are not exercised by the ROCm-Linux artifacts above. They are exercised by the Windows pull-CI (where NNAPI builds run and venv-based discovery is relevant); the green Windows CI is the evidence for those. - ROCm wheels are Linux-only by build matrix (`.github/scripts/generate_binary_build_matrix.py:130-135`, `supported_platforms=["Linux"]`). Co-Generated with [Claude Code](https://claude.com/claude-code) Pull Request resolved: pytorch#180243 Approved by: https://github.com/malfet
diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake
@@ -814,6 +814,21 @@ include_directories(SYSTEM ${EIGEN3_INCLUDE_DIR})
 
 
 if(BUILD_PYTHON)
+  # On Windows venvs, the Python import library (pythonXX.lib) lives in the
+  # base installation's libs/ directory, not in the venv.  Help FindPython
+  # locate it by adding sys.base_prefix/libs to the library search path.
+  if(WIN32 AND Python_EXECUTABLE)
+    execute_process(
+      COMMAND "${Python_EXECUTABLE}" -c "import sys; print(sys.base_prefix)"
+      OUTPUT_VARIABLE _py_base_prefix
+      OUTPUT_STRIP_TRAILING_WHITESPACE
+      ERROR_QUIET
+    )
+    if(_py_base_prefix AND IS_DIRECTORY "${_py_base_prefix}/libs")
+      list(APPEND CMAKE_LIBRARY_PATH "${_py_base_prefix}/libs")
+    endif()
+  endif()
+
   set(PYTHON_COMPONENTS Development.Module)
   if(USE_NUMPY)
     list(APPEND PYTHON_COMPONENTS NumPy)
diff --git a/setup.py b/setup.py
@@ -313,7 +313,7 @@
 from tools.clean import clean as _clean
 from tools.generate_torch_version import get_torch_version
 from tools.setup_helpers.cmake import CMake, CMakeValue
-from tools.setup_helpers.env import build_type, IS_DARWIN, IS_LINUX, IS_WINDOWS
+from tools.setup_helpers.env import IS_DARWIN, IS_LINUX, IS_WINDOWS
 
 
 def str2bool(value: str | None) -> bool:
@@ -973,23 +973,6 @@ def run(self) -> None:
 
         super().run()
 
-        # Copy the essential export library to compile C++ extensions.
-        if IS_WINDOWS:
-            build_temp = Path(self.build_temp)
-            build_lib = Path(self.build_lib)
-
-            ext_filename = self.get_ext_filename("_C")
-            lib_filename = ".".join(ext_filename.split(".")[:-1]) + ".lib"
-
-            export_lib = build_temp / "torch" / "csrc" / lib_filename
-            target_lib = build_lib / "torch" / "lib" / "_C.lib"
-
-            # Create "torch/lib" directory if not exists.
-            # (It is not created yet in "develop" mode.)
-            target_dir = target_lib.parent
-            target_dir.mkdir(parents=True, exist_ok=True)
-            self.copy_file(export_lib, target_lib)
-
     def get_outputs(self) -> list[str]:
         outputs = super().get_outputs()
         outputs.append(os.path.join(self.build_lib, "caffe2"))
@@ -1031,6 +1014,17 @@ def run(self) -> None:
         _clean()
 
 
+class BinaryDistribution(Distribution):
+    # torch._C is built by CMake (torch/CMakeLists.txt), so setuptools has
+    # no ext_modules. Force the wheel to be tagged as a binary distribution
+    # so bdist_wheel uses the binary layout (only files listed via
+    # package_data / data_files are packaged) rather than the purelib
+    # layout, which would dump the entire build/lib/ tree -- including
+    # CMake's test artifacts -- into the wheel at site-packages root.
+    def has_ext_modules(self) -> bool:
+        return True
+
+
 def get_cmake_cache_vars() -> defaultdict[str, CMakeValue]:
     try:
         return defaultdict(lambda: False, cmake.get_cmake_cache_variables())
@@ -1059,58 +1053,8 @@ def configure_extension_build() -> tuple[
     # Configure compile flags
     ################################################################################
 
-    library_dirs: list[str] = [str(TORCH_LIB_DIR)]
     extra_install_requires: list[str] = []
 
-    if IS_WINDOWS:
-        # /NODEFAULTLIB makes sure we only link to DLL runtime
-        # and matches the flags set for protobuf and ONNX
-        extra_link_args: list[str] = ["/NODEFAULTLIB:LIBCMT.LIB"]
-        # /MD links against DLL runtime
-        # and matches the flags set for protobuf and ONNX
-        # /EHsc is about standard C++ exception handling
-        extra_compile_args: list[str] = ["/MD", "/FS", "/EHsc"]
-    else:
-        extra_link_args = []
-        extra_compile_args = [
-            "-Wall",
-            "-Wextra",
-            "-Wno-strict-overflow",
-            "-Wno-unused-parameter",
-            "-Wno-missing-field-initializers",
-            "-Wno-unknown-pragmas",
-            # Python 2.6 requires -fno-strict-aliasing, see
-            # http://legacy.python.org/dev/peps/pep-3123/
-            # We also depend on it in our code (even Python 3).
-            "-fno-strict-aliasing",
-        ]
-
-    main_compile_args: list[str] = []
-    main_libraries: list[str] = ["torch_python"]
-
-    main_link_args: list[str] = []
-    main_sources: list[str] = ["torch/csrc/stub.c"]
-
-    if BUILD_LIBTORCH_WHL:
-        main_libraries = ["torch"]
-        main_sources = []
-
-    if build_type.is_debug():
-        if IS_WINDOWS:
-            extra_compile_args += ["/Z7"]
-            extra_link_args += ["/DEBUG:FULL"]
-        else:
-            extra_compile_args += ["-O0", "-g"]
-            extra_link_args += ["-O0", "-g"]
-
-    if build_type.is_rel_with_deb_info():
-        if IS_WINDOWS:
-            extra_compile_args += ["/Z7"]
-            extra_link_args += ["/DEBUG:FULL"]
-        else:
-            extra_compile_args += ["-g"]
-            extra_link_args += ["-g"]
-
     # pypi cuda package that requires installation of cuda runtime, cudnn and cublas
     # should be included in all wheels uploaded to pypi
     pytorch_extra_install_requires = os.getenv("PYTORCH_EXTRA_INSTALL_REQUIREMENTS")
@@ -1120,39 +1064,12 @@ def configure_extension_build() -> tuple[
             map(str.strip, pytorch_extra_install_requires.split("|"))
         )
 
-    # Cross-compile for M1
-    if IS_DARWIN:
-        macos_target_arch = os.getenv("CMAKE_OSX_ARCHITECTURES", "")
-        if macos_target_arch in ["arm64", "x86_64"]:
-            macos_sysroot_path = os.getenv("CMAKE_OSX_SYSROOT")
-            if macos_sysroot_path is None:
-                macos_sysroot_path = (
-                    subprocess.check_output(
-                        ["xcrun", "--show-sdk-path", "--sdk", "macosx"]
-                    )
-                    .decode("utf-8")
-                    .strip()
-                )
-            extra_compile_args += [
-                "-arch",
-                macos_target_arch,
-                "-isysroot",
-                macos_sysroot_path,
-            ]
-            extra_link_args += ["-arch", macos_target_arch]
-
-    def make_relative_rpath_args(path: str) -> list[str]:
-        if IS_DARWIN:
-            return ["-Wl,-rpath,@loader_path/" + path]
-        elif IS_WINDOWS:
-            return []
-        else:
-            return ["-Wl,-rpath,$ORIGIN/" + path]
-
     ################################################################################
     # Declare extensions and package
     ################################################################################
 
+    # torch._C is now built by CMake (torch/CMakeLists.txt) and installed
+    # into the cmake install prefix, which maps to torch/ in the source tree.
     ext_modules: list[Extension] = []
     # packages that we want to install into site-packages and include them in wheels
     includes = ["torch", "torch.*", "torchgen", "torchgen.*"]
@@ -1163,24 +1080,6 @@ def make_relative_rpath_args(path: str) -> list[str]:
     else:
         excludes.extend(["functorch", "functorch.*"])
     packages = find_packages(include=includes, exclude=excludes)
-    C = Extension(
-        "torch._C",
-        libraries=main_libraries,
-        sources=main_sources,
-        language="c",
-        extra_compile_args=[
-            *main_compile_args,
-            *extra_compile_args,
-        ],
-        include_dirs=[],
-        library_dirs=library_dirs,
-        extra_link_args=[
-            *extra_link_args,
-            *main_link_args,
-            *make_relative_rpath_args("lib"),
-        ],
-    )
-    ext_modules.append(C)
 
     cmdclass = {
         "bdist_wheel": bdist_wheel,
@@ -1272,6 +1171,11 @@ def main() -> None:
 
     torch_package_data = [
         "py.typed",
+        # torch._C is built by CMake (torch/CMakeLists.txt) and installed into
+        # torch/.  Match this interpreter's exact extension suffix, not a glob:
+        # build_all.sh reuses one build/ dir across Python versions, so a glob
+        # would also pick up other ABIs' _C modules left behind in the tree.
+        f"_C{sysconfig.get_config_var('EXT_SUFFIX')}",
         "bin/*",
         "bin/**/*",
         "test/*",
@@ -1378,6 +1282,7 @@ def main() -> None:
     setup(
         name=TORCH_PACKAGE_NAME,
         version=TORCH_VERSION,
+        distclass=BinaryDistribution,
         ext_modules=ext_modules,
         cmdclass=cmdclass,
         packages=packages,
diff --git a/torch/CMakeLists.txt b/torch/CMakeLists.txt
@@ -107,7 +107,16 @@ set(TORCH_PYTHON_LINK_FLAGS "")
 
 if(MSVC)
     string(APPEND TORCH_PYTHON_LINK_FLAGS " /NODEFAULTLIB:LIBCMT.LIB")
-    list(APPEND TORCH_PYTHON_LINK_LIBRARIES ${PYTHON_LIBRARIES} onnx_library)
+    # Python headers use #pragma comment(lib, "pythonXX") which makes MSVC search
+    # for the library by bare name.  Add the directory so the linker can find it.
+    foreach(_pylib IN LISTS Python_LIBRARIES)
+      if(_pylib MATCHES "python[0-9].*\\.lib$")
+        get_filename_component(_pylib_dir "${_pylib}" DIRECTORY)
+        string(APPEND TORCH_PYTHON_LINK_FLAGS " /LIBPATH:\"${_pylib_dir}\"")
+        break()
+      endif()
+    endforeach()
+    list(APPEND TORCH_PYTHON_LINK_LIBRARIES ${Python_LIBRARIES} onnx_library)
     if(NOT CMAKE_BUILD_TYPE MATCHES "Release")
       string(APPEND TORCH_PYTHON_LINK_FLAGS " /DEBUG:FULL")
     endif()
@@ -424,11 +433,40 @@ if(${CMAKE_BUILD_TYPE} STREQUAL "RelWithAssert")
 endif()
 
 if(NOT TORCH_PYTHON_LINK_FLAGS STREQUAL "")
-    set_target_properties(torch_python PROPERTIES LINK_FLAGS ${TORCH_PYTHON_LINK_FLAGS})
+    set_target_properties(torch_python PROPERTIES LINK_FLAGS "${TORCH_PYTHON_LINK_FLAGS}")
 endif()
 
 install(TARGETS torch_python DESTINATION "${TORCH_INSTALL_LIB_DIR}")
 
+# Build the torch._C Python extension module from the thin stub that
+# forwards to torch_python.  This was historically built by setuptools;
+# building it here lets both setuptools and scikit-build-core use the
+# same code path.
+if(BUILD_PYTHON AND NOT BUILD_LIBTORCH_WHL)
+  Python_add_library(_C MODULE WITH_SOABI "${TORCH_SRC_DIR}/csrc/stub.c")
+  target_link_libraries(_C PRIVATE torch_python)
+  set_target_properties(_C PROPERTIES
+    LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/torch"
+    # libtorch_python.so is installed to lib/ relative to _C, so use
+    # $ORIGIN/lib as the install RPATH (the global default is just $ORIGIN).
+    INSTALL_RPATH "${_rpath_portable_origin}/lib"
+  )
+  if(NOT WIN32)
+    set_target_properties(_C PROPERTIES
+      C_VISIBILITY_PRESET "default"
+    )
+  endif()
+  # Install at the install prefix root, which maps to torch/ in the wheel
+  # via wheel.install-dir = "torch". The import library (Windows _C.lib) goes
+  # alongside the other import libs in torch/lib/, matching the layout that
+  # setup.py used before _C moved to CMake.
+  install(TARGETS _C
+    RUNTIME DESTINATION "."
+    LIBRARY DESTINATION "."
+    ARCHIVE DESTINATION "lib"
+  )
+endif()
+
 # Generate torch/version.py from the appropriate CMake cache variables.
 if(${CMAKE_BUILD_TYPE} STREQUAL "Debug")
   set(TORCH_VERSION_DEBUG 1)
@@ -454,6 +492,9 @@ add_custom_target(
   WORKING_DIRECTORY ${TORCH_ROOT}
 )
 add_dependencies(torch_python gen_torch_version)
+# version.py is gitignored (generated), so scikit-build-core won't include it
+# via Python package collection. Install it explicitly via cmake instead.
+install(FILES "${TORCH_SRC_DIR}/version.py" DESTINATION ".")
 
 # Skip building this library under MacOS, since it is currently failing to build on Mac
 # Github issue #61930
@@ -470,6 +511,9 @@ if(NOT ${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
     target_link_libraries(nnapi_backend PRIVATE torch)
   endif()
   target_link_libraries(nnapi_backend PRIVATE torch_python pybind::pybind11 fmt::fmt-header-only)
+  if(MSVC AND TORCH_PYTHON_LINK_FLAGS)
+    set_target_properties(nnapi_backend PROPERTIES LINK_FLAGS "${TORCH_PYTHON_LINK_FLAGS}")
+  endif()
 endif()
 
 set(TORCH_PYTHON_COMPILE_OPTIONS ${TORCH_PYTHON_COMPILE_OPTIONS} PARENT_SCOPE)
