Skip to content

[OpenReg] Add OSX/Windows Support for OpenReg #159441

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 9 commits into
base: main
Choose a base branch
from
+368 −118
Open

[OpenReg] Add OSX/Windows Support for OpenReg #159441

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
1689f44
[OpenReg] Add OSX/Windows Support for OpenReg
FFFrog Jul 25, 2025
a9194c6
Update Windows compilation for OpenReg
can-gaa-hou Jul 29, 2025
5ee0ad5
Fix some issues
can-gaa-hou Jul 30, 2025
4c41abc
Add torch_python_library
can-gaa-hou Jul 31, 2025
ff890b3
Update openreg testcase for Windows
can-gaa-hou Aug 1, 2025
71dc569
Removing libprotobuf from cmake
can-gaa-hou Aug 1, 2025
95279ed
Fix some issues
can-gaa-hou Aug 1, 2025
7381aa4
Fix lint error
can-gaa-hou Aug 1, 2025
65f2cc0
Fix some issues in comments
can-gaa-hou Aug 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 16 additions & 13 deletions test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,35 +4,38 @@ project(TORCH_OPENREG CXX C)

include(GNUInstallDirs)
include(CheckCXXCompilerFlag)
include(CMakeDependentOption)

set(CMAKE_SKIP_BUILD_RPATH FALSE)
set(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)
set(CMAKE_INSTALL_RPATH_USE_LINK_PATH FALSE)
set(CMAKE_INSTALL_RPATH "$ORIGIN/lib/:$ORIGIN/")

set(LINUX TRUE)
set(CMAKE_INSTALL_MESSAGE NEVER)
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_C_STANDARD 11)
set(CMAKE_CXX_EXTENSIONS OFF)

set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
set(CMAKE_SKIP_BUILD_RPATH FALSE)
set(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)
set(CMAKE_INSTALL_RPATH_USE_LINK_PATH FALSE)
set(CMAKE_CXX_VISIBILITY_PRESET hidden)

if(APPLE)
set(CMAKE_INSTALL_RPATH "@loader_path/lib;@loader_path")
elseif(UNIX)
set(CMAKE_INSTALL_RPATH "$ORIGIN/lib:$ORIGIN")
elseif(WIN32)
set(CMAKE_INSTALL_RPATH "")
endif()
set(CMAKE_INSTALL_LIBDIR lib)

add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=1)
set(CMAKE_INSTALL_MESSAGE NEVER)

set(Torch_DIR ${PYTORCH_INSTALL_DIR}/share/cmake/Torch)
find_package(Torch REQUIRED)
include_directories(${PYTORCH_INSTALL_DIR}/include)

if(DEFINED PYTHON_INCLUDE_DIR)
include_directories(${PYTHON_INCLUDE_DIR})
else()
message(FATAL_ERROR "Cannot find Python directory")
endif()

include(${PROJECT_SOURCE_DIR}/cmake/TorchPythonTargets.cmake)
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey @albanD, sorry to bother you.

PyTorch provides some CMake targets, such as torch_cpu_library and torch_cuda_library, but it doesn't provide a CMake target for torch_python, which is very useful for out-of-tree backends, as they require some Python bindings. We currently wrap a CMake target for torch_python in the openreg project, but it would be even better if Torch itself could provide one.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ho that is interesting.
I would have expected that our torch.utils.cpp_extension uses a cmake target to build these cpp extensions. But it doesn't?

Copy link
Collaborator Author

@FFFrog FFFrog Aug 5, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As far as I know, basic functions in torch.utils.cpp_extension such as CPPExtension use -ltorch_python and specify library_dir directly to build the extension instead of using cmake targets.

pytorch/torch/utils/cpp_extension.py

Lines 1204 to 1210 in 1ca8388

if not kwargs.get('py_limited_api', False):
# torch_python uses more than the python limited api
libraries.append('torch_python')
if IS_WINDOWS:
libraries.append("sleef")
kwargs['libraries'] = libraries

In addition, requirements of torch.utils.cpp_extension and integration of the new accelerator are slightly different; the core functions of the former are basically compiled into the extension library itself (such as _C.so), while the latter extension library itself is just an intermediate layer from Python to C++, and the core functions are all in the accelerator's own torch_binding.so. torch.uitls.cpp_extension is more suitable for PyTorch users to extend some functions of torch (such as Pluggable Allocator), but it is not suitable for the latter. A highly integrated CMake object like torch_python_library is what the latter needs.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok!
Adding a new target sounds good then.
cc @swolchok can you give a quick look at the new cmake target? Is that the right style that we want going forward?


add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/openreg)
add_subdirectory(${PROJECT_SOURCE_DIR}/csrc)
add_subdirectory(${PROJECT_SOURCE_DIR}/torch_openreg/csrc)
20 changes: 20 additions & 0 deletions test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
if(WIN32)
set(TORCH_PYTHON_IMPORTED_LOCATION "${PYTORCH_INSTALL_DIR}/lib/torch_python.lib")
else()
set(TORCH_PYTHON_IMPORTED_LOCATION "${PYTORCH_INSTALL_DIR}/lib/libtorch_python.so")
endif()

add_library(torch_python SHARED IMPORTED)

set_target_properties(torch_python PROPERTIES
INTERFACE_INCLUDE_DIRECTORIES "${PYTORCH_INSTALL_DIR}/include"
INTERFACE_LINK_LIBRARIES "c10;torch_cpu"
IMPORTED_LOCATION "${TORCH_PYTHON_IMPORTED_LOCATION}"
)

add_library(torch_python_library INTERFACE IMPORTED)

set_target_properties(torch_python_library PROPERTIES
INTERFACE_INCLUDE_DIRECTORIES "\$<TARGET_PROPERTY:torch_python,INTERFACE_INCLUDE_DIRECTORIES>"
INTERFACE_LINK_LIBRARIES "\$<TARGET_FILE:torch_python>;\$<TARGET_PROPERTY:torch_python,INTERFACE_LINK_LIBRARIES>"
)
8 changes: 6 additions & 2 deletions test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,11 @@ file(GLOB_RECURSE SOURCE_FILES

add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})

target_link_libraries(${LIBRARY_NAME} PRIVATE openreg torch_cpu)
target_link_libraries(${LIBRARY_NAME} PRIVATE torch_cpu_library openreg)
target_include_directories(${LIBRARY_NAME} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})

install(TARGETS ${LIBRARY_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR})
install(TARGETS ${LIBRARY_NAME}
ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}
)
8 changes: 4 additions & 4 deletions ...pp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ int device_count_impl() {
return count;
}

c10::DeviceIndex device_count() noexcept {
OPENREG_EXPORT c10::DeviceIndex device_count() noexcept {
// initialize number of devices only once
static int count = []() {
try {
Expand All @@ -49,17 +49,17 @@ c10::DeviceIndex device_count() noexcept {
return static_cast<c10::DeviceIndex>(count);
}

c10::DeviceIndex current_device() {
OPENREG_EXPORT c10::DeviceIndex current_device() {
c10::DeviceIndex cur_device = -1;
GetDevice(&cur_device);
return cur_device;
}

void set_device(c10::DeviceIndex device) {
OPENREG_EXPORT void set_device(c10::DeviceIndex device) {
SetDevice(device);
}

DeviceIndex ExchangeDevice(DeviceIndex device) {
OPENREG_EXPORT DeviceIndex ExchangeDevice(DeviceIndex device) {
int current_device = -1;
orGetDevice(&current_device);

Expand Down
14 changes: 10 additions & 4 deletions .../cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,22 @@
#pragma once

#ifdef _WIN32
#define OPENREG_EXPORT __declspec(dllexport)
#else
#define OPENREG_EXPORT __attribute__((visibility("default")))
#endif

#include <c10/core/Device.h>
#include <c10/macros/Macros.h>

#include <limits>

namespace c10::openreg {

c10::DeviceIndex device_count() noexcept;
DeviceIndex current_device();
void set_device(c10::DeviceIndex device);
OPENREG_EXPORT c10::DeviceIndex device_count() noexcept;
OPENREG_EXPORT c10::DeviceIndex current_device();
OPENREG_EXPORT void set_device(c10::DeviceIndex device);

DeviceIndex ExchangeDevice(DeviceIndex device);
OPENREG_EXPORT DeviceIndex ExchangeDevice(DeviceIndex device);

} // namespace c10::openreg
60 changes: 56 additions & 4 deletions test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import multiprocessing
import os
import platform
import shutil
import subprocess
import sys
Expand All @@ -9,10 +10,23 @@
from setuptools import Extension, find_packages, setup


# Env Variables
IS_DARWIN = platform.system() == "Darwin"
IS_WINDOWS = platform.system() == "Windows"

BASE_DIR = os.path.dirname(os.path.realpath(__file__))
RUN_BUILD_DEPS = any(arg in {"clean", "dist_info"} for arg in sys.argv)


def make_relative_rpath_args(path):
if IS_DARWIN:
return ["-Wl,-rpath,@loader_path/" + path]
elif IS_WINDOWS:
return []
else:
return ["-Wl,-rpath,$ORIGIN/" + path]


def get_pytorch_dir():
import torch

Expand All @@ -39,9 +53,15 @@ def build_deps():
".",
"--target",
"install",
"--config",
"Release",
"--",
]
build_args += ["-j", str(multiprocessing.cpu_count())]

if IS_WINDOWS:
build_args += ["/m:" + str(multiprocessing.cpu_count())]
else:
build_args += ["-j", str(multiprocessing.cpu_count())]

command = ["cmake"] + build_args
subprocess.check_call(command, cwd=build_dir, env=os.environ)
Expand All @@ -64,19 +84,51 @@ def main():
if not RUN_BUILD_DEPS:
build_deps()

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"] + [
*make_relative_rpath_args("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 = [*make_relative_rpath_args("lib")]
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",
]

ext_modules = [
Extension(
name="torch_openreg._C",
sources=["torch_openreg/csrc/stub.c"],
language="c",
extra_compile_args=["-g", "-Wall", "-Werror"],
extra_compile_args=extra_compile_args,
libraries=["torch_bindings"],
library_dirs=[os.path.join(BASE_DIR, "torch_openreg/lib")],
extra_link_args=["-Wl,-rpath,$ORIGIN/lib"],
extra_link_args=extra_link_args,
)
]

package_data = {"torch_openreg": ["lib/*.so*"]}
package_data = {
"torch_openreg": [
"lib/*.so*",
"lib/*.dylib*",
"lib/*.dll",
"lib/*.lib",
]
}

setup(
packages=find_packages(),
Expand Down
6 changes: 5 additions & 1 deletion ...p_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,4 +8,8 @@ add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})

target_include_directories(${LIBRARY_NAME} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})

install(TARGETS ${LIBRARY_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR})
install(TARGETS ${LIBRARY_NAME}
ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}
)
Loading
Loading