From 1689f44fa02d38abbc074648ccc6d61910e13be9 Mon Sep 17 00:00:00 2001
From: FFFrog <ljw1101.vip@gmail.com>
Date: Fri, 25 Jul 2025 11:01:48 +0800
Subject: [PATCH 1/9] [OpenReg] Add OSX/Windows Support for OpenReg

As the title stated.

- Abstract platform-specific APIs
- Add OSX/Windows support

Co-authored-by: can-gaa-hou <jiahaochen535@gmail.com>
---
 .../torch_openreg/CMakeLists.txt              |  24 ++--
 .../torch_openreg/csrc/CMakeLists.txt         |   2 +-
 .../torch_openreg/setup.py                    |  25 +++-
 .../third_party/openreg/csrc/memory.cpp       | 122 ++++++++----------
 .../third_party/openreg/csrc/memory.h         |  98 ++++++++++++++
 .../torch_openreg/csrc/CMakeLists.txt         |   5 +
 test/run_test.py                              |   5 -
 7 files changed, 194 insertions(+), 87 deletions(-)
 create mode 100644 test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.h

diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
index 73163b8cb1ae..7ea5375a6ed3 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
@@ -4,28 +4,26 @@ 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_INSTALL_LIBDIR lib)
+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)
 
-add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=1)
+if(APPLE)
+  set(CMAKE_INSTALL_RPATH "@loader_path/lib;@loader_path")
+elseif(UNIX)
+  set(CMAKE_INSTALL_RPATH "$ORIGIN/lib:$ORIGIN")
+endif()
+set(CMAKE_INSTALL_LIBDIR lib)
+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})
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
index 077f4cf3b640..2bbd7fb0ba12 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
@@ -6,7 +6,7 @@ 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})
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py b/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
index 07d31e73d76b..35a18fe964e3 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
@@ -1,5 +1,6 @@
 import multiprocessing
 import os
+import platform
 import shutil
 import subprocess
 import sys
@@ -9,10 +10,24 @@
 from setuptools import Extension, find_packages, setup
 
 
+# Env Variables
+IS_LINUX = platform.system() == "Linux"
+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
 
@@ -72,11 +87,17 @@ def main():
             extra_compile_args=["-g", "-Wall", "-Werror"],
             libraries=["torch_bindings"],
             library_dirs=[os.path.join(BASE_DIR, "torch_openreg/lib")],
-            extra_link_args=["-Wl,-rpath,$ORIGIN/lib"],
+            extra_link_args=[*make_relative_rpath_args("lib")],
         )
     ]
 
-    package_data = {"torch_openreg": ["lib/*.so*"]}
+    package_data = {
+        "torch_openreg": [
+            "lib/*.so*",
+            "lib/*.dylib*",
+            "lib/*.dll",
+        ]
+    }
 
     setup(
         packages=find_packages(),
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.cpp b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.cpp
index 762cd96d23bb..304b951b00cc 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.cpp
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.cpp
@@ -1,38 +1,9 @@
-#include <include/openreg.h>
+#include "memory.h"
 
-#include <sys/mman.h>
-#include <unistd.h>
-#include <cstdlib>
-#include <cstring>
 #include <map>
 #include <mutex>
 
-namespace openreg {
-namespace internal {
-
-class ScopedMemoryProtector {
- public:
-  ScopedMemoryProtector(const orPointerAttributes& info)
-      : m_info(info), m_protected(false) {
-    if (m_info.type == orMemoryType::orMemoryTypeDevice) {
-      if (mprotect(m_info.pointer, m_info.size, PROT_READ | PROT_WRITE) ==
-          0) {
-        m_protected = true;
-      }
-    }
-  }
-  ~ScopedMemoryProtector() {
-    if (m_protected) {
-      mprotect(m_info.pointer, m_info.size, PROT_NONE);
-    }
-  }
-  ScopedMemoryProtector(const ScopedMemoryProtector&) = delete;
-  ScopedMemoryProtector& operator=(const ScopedMemoryProtector&) = delete;
-
- private:
-  orPointerAttributes m_info;
-  bool m_protected;
-};
+namespace {
 
 class MemoryManager {
  public:
@@ -46,7 +17,7 @@ class MemoryManager {
       return orErrorUnknown;
 
     std::lock_guard<std::mutex> lock(m_mutex);
-    long page_size = sysconf(_SC_PAGESIZE);
+    long page_size = openreg::get_pagesize();
     size_t aligned_size = ((size - 1) / page_size + 1) * page_size;
     void* mem = nullptr;
     int current_device = -1;
@@ -54,21 +25,15 @@ class MemoryManager {
     if (type == orMemoryType::orMemoryTypeDevice) {
       orGetDevice(&current_device);
 
-      mem = mmap(
-          nullptr,
-          aligned_size,
-          PROT_READ | PROT_WRITE,
-          MAP_PRIVATE | MAP_ANONYMOUS,
-          -1,
-          0);
+      mem = openreg::mmap(aligned_size);
       if (mem == MAP_FAILED)
         return orErrorUnknown;
-      if (mprotect(mem, aligned_size, PROT_NONE) != 0) {
-        munmap(mem, aligned_size);
+      if (openreg::mprotect(mem, aligned_size, F_PROT_NONE) != 0) {
+        openreg::munmap(mem, aligned_size);
         return orErrorUnknown;
       }
     } else {
-      if (posix_memalign(&mem, page_size, aligned_size) != 0) {
+      if (openreg::alloc(&mem, page_size, aligned_size) != 0) {
         return orErrorUnknown;
       }
     }
@@ -88,10 +53,10 @@ class MemoryManager {
       return orErrorUnknown;
     const auto& info = it->second;
     if (info.type == orMemoryType::orMemoryTypeDevice) {
-      mprotect(info.pointer, info.size, PROT_READ | PROT_WRITE);
-      munmap(info.pointer, info.size);
+      openreg::mprotect(info.pointer, info.size, F_PROT_READ | F_PROT_WRITE);
+      openreg::munmap(info.pointer, info.size);
     } else {
-      ::free(info.pointer);
+      openreg::free(info.pointer);
     }
     m_registry.erase(it);
     return orSuccess;
@@ -167,7 +132,8 @@ class MemoryManager {
     if (info.type != orMemoryType::orMemoryTypeDevice) {
       return orErrorUnknown;
     }
-    if (mprotect(info.pointer, info.size, PROT_READ | PROT_WRITE) != 0) {
+    if (openreg::mprotect(
+            info.pointer, info.size, F_PROT_READ | F_PROT_WRITE) != 0) {
       return orErrorUnknown;
     }
     return orSuccess;
@@ -179,49 +145,75 @@ class MemoryManager {
     if (info.type != orMemoryType::orMemoryTypeDevice) {
       return orErrorUnknown;
     }
-    if (mprotect(info.pointer, info.size, PROT_NONE) != 0) {
+    if (openreg::mprotect(info.pointer, info.size, F_PROT_NONE) != 0) {
       return orErrorUnknown;
     }
     return orSuccess;
   }
 
  private:
+  class ScopedMemoryProtector {
+   public:
+    ScopedMemoryProtector(const orPointerAttributes& info)
+        : m_info(info), m_protected(false) {
+      if (m_info.type == orMemoryType::orMemoryTypeDevice) {
+        if (openreg::mprotect(
+                m_info.pointer, m_info.size, F_PROT_READ | F_PROT_WRITE) == 0) {
+          m_protected = true;
+        }
+      }
+    }
+    ~ScopedMemoryProtector() {
+      if (m_protected) {
+        openreg::mprotect(m_info.pointer, m_info.size, F_PROT_NONE);
+      }
+    }
+    ScopedMemoryProtector(const ScopedMemoryProtector&) = delete;
+    ScopedMemoryProtector& operator=(const ScopedMemoryProtector&) = delete;
+
+   private:
+    orPointerAttributes m_info;
+    bool m_protected;
+  };
+
   MemoryManager() = default;
+
   orPointerAttributes getPointerInfo(const void* ptr) {
     auto it = m_registry.upper_bound(const_cast<void*>(ptr));
-    if (it == m_registry.begin())
-      return {};
-    --it;
-    const char* p_char = static_cast<const char*>(ptr);
-    const char* base_char = static_cast<const char*>(it->first);
-    if (p_char >= base_char && p_char < (base_char + it->second.size)) {
-      return it->second;
+    if (it != m_registry.begin()) {
+      --it;
+      const char* p_char = static_cast<const char*>(ptr);
+      const char* base_char = static_cast<const char*>(it->first);
+      if (p_char >= base_char && p_char < (base_char + it->second.size)) {
+        return it->second;
+      }
     }
+
     return {};
   }
+
   std::map<void*, orPointerAttributes> m_registry;
   std::mutex m_mutex;
 };
 
-} // namespace internal
-} // namespace openreg
+} // namespace
 
 orError_t orMalloc(void** devPtr, size_t size) {
-  return openreg::internal::MemoryManager::getInstance().allocate(
+  return MemoryManager::getInstance().allocate(
       devPtr, size, orMemoryType::orMemoryTypeDevice);
 }
 
 orError_t orFree(void* devPtr) {
-  return openreg::internal::MemoryManager::getInstance().free(devPtr);
+  return MemoryManager::getInstance().free(devPtr);
 }
 
 orError_t orMallocHost(void** hostPtr, size_t size) {
-  return openreg::internal::MemoryManager::getInstance().allocate(
+  return MemoryManager::getInstance().allocate(
       hostPtr, size, orMemoryType::orMemoryTypeHost);
 }
 
 orError_t orFreeHost(void* hostPtr) {
-  return openreg::internal::MemoryManager::getInstance().free(hostPtr);
+  return MemoryManager::getInstance().free(hostPtr);
 }
 
 orError_t orMemcpy(
@@ -229,21 +221,19 @@ orError_t orMemcpy(
     const void* src,
     size_t count,
     orMemcpyKind kind) {
-  return openreg::internal::MemoryManager::getInstance().memcpy(
-      dst, src, count, kind);
+  return MemoryManager::getInstance().memcpy(dst, src, count, kind);
 }
 
 orError_t orPointerGetAttributes(
     orPointerAttributes* attributes,
     const void* ptr) {
-  return openreg::internal::MemoryManager::getInstance().getPointerAttributes(
-      attributes, ptr);
+  return MemoryManager::getInstance().getPointerAttributes(attributes, ptr);
 }
 
 orError_t orMemoryUnprotect(void* devPtr) {
-  return openreg::internal::MemoryManager::getInstance().unprotect(devPtr);
+  return MemoryManager::getInstance().unprotect(devPtr);
 }
 
 orError_t orMemoryProtect(void* devPtr) {
-  return openreg::internal::MemoryManager::getInstance().protect(devPtr);
+  return MemoryManager::getInstance().protect(devPtr);
 }
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.h b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.h
new file mode 100644
index 000000000000..9de13acc2350
--- /dev/null
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.h
@@ -0,0 +1,98 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdlib>
+#include <cstring>
+
+#include <include/openreg.h>
+
+#if defined(_WIN32)
+#include <windows.h>
+#else
+#include <sys/mman.h>
+#include <unistd.h>
+#endif
+
+#define F_PROT_NONE 0x0
+#define F_PROT_READ 0x1
+#define F_PROT_WRITE 0x2
+
+namespace openreg {
+
+void* mmap(size_t size) {
+#if defined(_WIN32)
+  return VirtualAlloc(nullptr, size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+#else
+  void* addr = ::mmap(
+      nullptr,
+      size,
+      PROT_READ | PROT_WRITE,
+      MAP_PRIVATE | MAP_ANONYMOUS,
+      -1,
+      0);
+  return (addr == MAP_FAILED) ? nullptr : addr;
+#endif
+}
+
+void munmap(void* addr, size_t size) {
+#if defined(_WIN32)
+  VirtualFree(addr, 0, MEM_RELEASE);
+#else
+  ::munmap(addr, size);
+#endif
+}
+
+int mprotect(void* addr, size_t size, int prot) {
+#if defined(_WIN32)
+  DWORD win_prot = 0;
+  DWORD old;
+  if (prot == F_PROT_NONE) {
+    win_prot = PAGE_NOACCESS;
+  } else {
+    win_prot = PAGE_READWRITE;
+  }
+
+  return VirtualProtect(addr, size, win_prot, &old) ? 0 : -1;
+#else
+  int native_prot = 0;
+  if (prot == F_PROT_NONE)
+    native_prot = PROT_NONE;
+  else {
+    if (prot & F_PROT_READ)
+      native_prot |= PROT_READ;
+    if (prot & F_PROT_WRITE)
+      native_prot |= PROT_WRITE;
+  }
+
+  return ::mprotect(addr, size, native_prot);
+#endif
+}
+
+int alloc(void** mem, size_t alignment, size_t size) {
+#ifdef _WIN32
+  *mem = _aligned_malloc(size, alignment);
+  return *mem ? 0 : -1;
+#else
+  return posix_memalign(mem, alignment, size);
+#endif
+}
+
+void free(void* mem) {
+#ifdef _WIN32
+  _aligned_free(mem);
+#else
+  ::free(mem);
+#endif
+}
+
+long get_pagesize() {
+#ifdef _WIN32
+  SYSTEM_INFO si;
+  GetSystemInfo(&si);
+  return static_cast<long>(si.dwPageSize);
+#else
+  return sysconf(_SC_PAGESIZE);
+#endif
+}
+
+} // namespace openreg
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
index 574b5b1c748a..1579c6449a9f 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
@@ -8,5 +8,10 @@ add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
 
 target_link_libraries(${LIBRARY_NAME} PRIVATE torch_python torch_openreg)
 target_link_directories(${LIBRARY_NAME} PRIVATE ${PYTORCH_INSTALL_DIR}/lib)
+target_include_directories(${LIBRARY_NAME} PRIVATE ${PYTORCH_INSTALL_DIR}/include)
+
+if(APPLE)
+  set_target_properties(${LIBRARY_NAME} PROPERTIES LINK_FLAGS "-undefined dynamic_lookup")
+endif()
 
 install(TARGETS ${LIBRARY_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR})
diff --git a/test/run_test.py b/test/run_test.py
index e0bde4e6d52d..4983071ab23f 100755
--- a/test/run_test.py
+++ b/test/run_test.py
@@ -28,7 +28,6 @@
 from torch.testing._internal.common_utils import (
     get_report_path,
     IS_CI,
-    IS_LINUX,
     IS_MACOS,
     retry_shell,
     set_cwd,
@@ -911,10 +910,6 @@ def _test_autoload(test_directory, options, enable=True):
 
 
 def run_test_with_openreg(test_module, test_directory, options):
-    # TODO(FFFrog): Will remove this later when windows/macos are supported.
-    if not IS_LINUX:
-        return 0
-
     openreg_dir = os.path.join(
         test_directory, "cpp_extensions", "open_registration_extension", "torch_openreg"
     )

From a9194c6d9355ead49d95b0fdd37c85f7acca8a16 Mon Sep 17 00:00:00 2001
From: jiahaochen666 <jiahaochen535@gmail.com>
Date: Tue, 29 Jul 2025 19:01:30 +0800
Subject: [PATCH 2/9] Update Windows compilation for OpenReg

---
 .../torch_openreg/CMakeLists.txt              |  2 +
 .../torch_openreg/csrc/CMakeLists.txt         | 14 ++-
 .../csrc/runtime/OpenRegFunctions.cpp         |  8 +-
 .../csrc/runtime/OpenRegFunctions.h           | 18 +++-
 .../torch_openreg/setup.py                    | 15 ++-
 .../third_party/openreg/CMakeLists.txt        | 16 ++-
 .../third_party/openreg/csrc/memory.cpp       |  3 +-
 .../torch_openreg/torch_openreg/__init__.py   | 97 +++++++++++++++++++
 .../torch_openreg/csrc/CMakeLists.txt         | 41 +++++++-
 9 files changed, 198 insertions(+), 16 deletions(-)

diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
index 7ea5375a6ed3..4a1a4a312316 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
@@ -18,6 +18,8 @@ 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)
 set(CMAKE_INSTALL_MESSAGE NEVER)
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
index 2bbd7fb0ba12..42f0ff45d788 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
@@ -9,4 +9,16 @@ add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
 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})
+if (WIN32)
+    target_compile_definitions(${LIBRARY_NAME} PRIVATE OPENREG_EXPORTS)
+endif()
+
+if(WIN32)
+    install(TARGETS ${LIBRARY_NAME} 
+        ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+        LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+        RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    )
+else()
+    install(TARGETS ${LIBRARY_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR})
+endif()
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.cpp b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.cpp
index 240c2d8ce1aa..4d144c3b8e18 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.cpp
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.cpp
@@ -30,7 +30,7 @@ int device_count_impl() {
   return count;
 }
 
-c10::DeviceIndex device_count() noexcept {
+OPENREG_API c10::DeviceIndex device_count() noexcept {
   // initialize number of devices only once
   static int count = []() {
     try {
@@ -49,17 +49,17 @@ c10::DeviceIndex device_count() noexcept {
   return static_cast<c10::DeviceIndex>(count);
 }
 
-c10::DeviceIndex current_device() {
+OPENREG_API c10::DeviceIndex current_device() {
   c10::DeviceIndex cur_device = -1;
   GetDevice(&cur_device);
   return cur_device;
 }
 
-void set_device(c10::DeviceIndex device) {
+OPENREG_API void set_device(c10::DeviceIndex device) {
   SetDevice(device);
 }
 
-DeviceIndex ExchangeDevice(DeviceIndex device) {
+OPENREG_API DeviceIndex ExchangeDevice(DeviceIndex device) {
   int current_device = -1;
   orGetDevice(&current_device);
 
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.h b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.h
index b6b991ff6d3a..b5a1517a7a99 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.h
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.h
@@ -1,5 +1,15 @@
 #pragma once
 
+#ifdef _WIN32
+  #ifdef OPENREG_EXPORTS
+    #define OPENREG_API __declspec(dllexport)
+  #else
+    #define OPENREG_API __declspec(dllimport)
+  #endif
+#else
+  #define OPENREG_API
+#endif
+
 #include <c10/core/Device.h>
 #include <c10/macros/Macros.h>
 
@@ -7,10 +17,10 @@
 
 namespace c10::openreg {
 
-c10::DeviceIndex device_count() noexcept;
-DeviceIndex current_device();
-void set_device(c10::DeviceIndex device);
+OPENREG_API c10::DeviceIndex device_count() noexcept;
+OPENREG_API c10::DeviceIndex current_device();
+OPENREG_API void set_device(c10::DeviceIndex device);
 
-DeviceIndex ExchangeDevice(DeviceIndex device);
+OPENREG_API DeviceIndex ExchangeDevice(DeviceIndex device);
 
 } // namespace c10::openreg
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py b/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
index 35a18fe964e3..d87fb960dee1 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
@@ -54,9 +54,14 @@ 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)
@@ -78,13 +83,18 @@ def run(self):
 def main():
     if not RUN_BUILD_DEPS:
         build_deps()
+    
+    if sys.platform == "win32":
+        extra_compile_args = ["/W3"]
+    else:
+        extra_compile_args = ["-g", "-Wall", "-Werror"]
 
     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=[*make_relative_rpath_args("lib")],
@@ -96,6 +106,7 @@ def main():
             "lib/*.so*",
             "lib/*.dylib*",
             "lib/*.dll",
+            "lib/*.lib",
         ]
     }
 
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
index 7fec109eeb1c..2d86b0a5be52 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
@@ -4,8 +4,20 @@ file(GLOB_RECURSE SOURCE_FILES
     "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp"
 )
 
-add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
+if(WIN32)
+    add_library(${LIBRARY_NAME} STATIC ${SOURCE_FILES})
+else()
+    add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
+endif()
 
 target_include_directories(${LIBRARY_NAME} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
 
-install(TARGETS ${LIBRARY_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR})
+if(WIN32)
+    install(TARGETS ${LIBRARY_NAME} 
+        ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+        LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+        RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    )
+else()
+    install(TARGETS ${LIBRARY_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR})
+endif()
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.cpp b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.cpp
index 304b951b00cc..b169ba8125b2 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.cpp
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.cpp
@@ -26,7 +26,7 @@ class MemoryManager {
       orGetDevice(&current_device);
 
       mem = openreg::mmap(aligned_size);
-      if (mem == MAP_FAILED)
+      if (mem == nullptr)
         return orErrorUnknown;
       if (openreg::mprotect(mem, aligned_size, F_PROT_NONE) != 0) {
         openreg::munmap(mem, aligned_size);
@@ -52,6 +52,7 @@ class MemoryManager {
     if (it == m_registry.end())
       return orErrorUnknown;
     const auto& info = it->second;
+    
     if (info.type == orMemoryType::orMemoryTypeDevice) {
       openreg::mprotect(info.pointer, info.size, F_PROT_READ | F_PROT_WRITE);
       openreg::munmap(info.pointer, info.size);
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py
index 3ed73794b06d..792c7c6d54b2 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py
@@ -1,5 +1,102 @@
+import ctypes
+import glob
+import os
+import platform
+import sys
+import textwrap
+import sysconfig
+
 import torch
 
+if sys.platform == "win32":
+
+    def _load_dll_libraries() -> None:
+
+        py_dll_path = os.path.join(sys.exec_prefix, "Library", "bin")
+        th_dll_path = os.path.join(os.path.dirname(__file__), "lib")
+        usebase_path = os.path.join(sysconfig.get_config_var("userbase"), "Library", "bin")
+        py_root_bin_path = os.path.join(sys.exec_prefix, "bin")
+
+        # When users create a virtualenv that inherits the base environment,
+        # we will need to add the corresponding library directory into
+        # DLL search directories. Otherwise, it will rely on `PATH` which
+        # is dependent on user settings.
+        if sys.exec_prefix != sys.base_exec_prefix:
+            base_py_dll_path = os.path.join(sys.base_exec_prefix, "Library", "bin")
+        else:
+            base_py_dll_path = ""
+
+        dll_paths = [
+            p
+            for p in (
+                th_dll_path,
+                py_dll_path,
+                base_py_dll_path,
+                usebase_path,
+                py_root_bin_path,
+            )
+            if os.path.exists(p)
+        ]
+
+
+        kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True)
+        with_load_library_flags = hasattr(kernel32, "AddDllDirectory")
+        prev_error_mode = kernel32.SetErrorMode(0x0001)
+
+        kernel32.LoadLibraryW.restype = ctypes.c_void_p
+        if with_load_library_flags:
+            kernel32.LoadLibraryExW.restype = ctypes.c_void_p
+
+        for dll_path in dll_paths:
+            os.add_dll_directory(dll_path)
+
+        try:
+            ctypes.CDLL("vcruntime140.dll")
+            ctypes.CDLL("msvcp140.dll")
+            if platform.machine() != "ARM64":
+                ctypes.CDLL("vcruntime140_1.dll")
+        except OSError:
+            print(
+                textwrap.dedent(
+                    """
+                    Microsoft Visual C++ Redistributable is not installed, this may lead to the DLL load failure.
+                    It can be downloaded at https://aka.ms/vs/16/release/vc_redist.x64.exe
+                    """
+                ).strip()
+            )
+
+        dlls = glob.glob(os.path.join(th_dll_path, "*.dll"))
+        path_patched = False
+        for dll in dlls:
+            is_loaded = False
+            if with_load_library_flags:
+                res = kernel32.LoadLibraryExW(dll, None, 0x00001100)
+                last_error = ctypes.get_last_error()
+                if res is None and last_error != 126:
+                    err = ctypes.WinError(last_error)
+                    err.strerror += (
+                        f' Error loading "{dll}" or one of its dependencies.'
+                    )
+                    raise err
+                elif res is not None:
+                    is_loaded = True
+            if not is_loaded:
+                if not path_patched:
+                    os.environ["PATH"] = ";".join(dll_paths + [os.environ["PATH"]])
+                    path_patched = True
+                res = kernel32.LoadLibraryW(dll)
+                if res is None:
+                    err = ctypes.WinError(ctypes.get_last_error())
+                    err.strerror += (
+                        f' Error loading "{dll}" or one of its dependencies.'
+                    )
+                    raise err
+
+        kernel32.SetErrorMode(prev_error_mode)
+
+    _load_dll_libraries()
+    del _load_dll_libraries
+
 import torch_openreg._C  # type: ignore[misc]
 import torch_openreg.openreg
 
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
index 1579c6449a9f..4744bbe7207f 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
@@ -6,7 +6,36 @@ file(GLOB_RECURSE SOURCE_FILES
 
 add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
 
-target_link_libraries(${LIBRARY_NAME} PRIVATE torch_python torch_openreg)
+if(WIN32)
+    # For Windows, we need specific PyTorch libraries
+    find_library(TORCH_PYTHON_LIBRARY torch_python PATHS ${PYTORCH_INSTALL_DIR}/lib REQUIRED)
+    find_library(TORCH_CPU_LIBRARY torch_cpu PATHS ${PYTORCH_INSTALL_DIR}/lib REQUIRED)
+    find_library(C10_LIBRARY c10 PATHS ${PYTORCH_INSTALL_DIR}/lib REQUIRED)
+    
+    target_link_libraries(${LIBRARY_NAME} PRIVATE 
+        ${TORCH_PYTHON_LIBRARY}
+        ${TORCH_CPU_LIBRARY}
+        ${C10_LIBRARY}
+        torch_openreg
+    )
+else()
+    target_link_libraries(${LIBRARY_NAME} PRIVATE torch_python torch_openreg)
+endif()
+
+if(WIN32)
+    if(DEFINED PYTHON_LIBRARY_DIR AND DEFINED PYTHON_LIBRARY)
+        target_link_directories(${LIBRARY_NAME} PRIVATE ${PYTHON_LIBRARY_DIR})
+        target_link_libraries(${LIBRARY_NAME} PRIVATE ${PYTHON_LIBRARY})
+    else()
+        find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
+        target_link_libraries(${LIBRARY_NAME} PRIVATE ${Python3_LIBRARIES})
+    endif()
+endif()
+
+if(WIN32)
+    set_target_properties(${LIBRARY_NAME} PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS ON)
+endif()
+
 target_link_directories(${LIBRARY_NAME} PRIVATE ${PYTORCH_INSTALL_DIR}/lib)
 target_include_directories(${LIBRARY_NAME} PRIVATE ${PYTORCH_INSTALL_DIR}/include)
 
@@ -14,4 +43,12 @@ if(APPLE)
   set_target_properties(${LIBRARY_NAME} PROPERTIES LINK_FLAGS "-undefined dynamic_lookup")
 endif()
 
-install(TARGETS ${LIBRARY_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR})
+if(WIN32)
+    install(TARGETS ${LIBRARY_NAME} 
+        ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+        LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+        RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    )
+else()
+    install(TARGETS ${LIBRARY_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR})
+endif()

From 5ee0ad5cbb5fb0cb3b22ade1b36932b1ffb4a757 Mon Sep 17 00:00:00 2001
From: jiahaochen666 <jiahaochen535@gmail.com>
Date: Wed, 30 Jul 2025 17:11:44 +0800
Subject: [PATCH 3/9] Fix some issues

---
 .../torch_openreg/csrc/CMakeLists.txt         | 18 ++--
 .../csrc/runtime/OpenRegFunctions.cpp         |  8 +-
 .../csrc/runtime/OpenRegFunctions.h           | 16 ++--
 .../torch_openreg/setup.py                    | 56 ++++++++++-
 .../third_party/openreg/CMakeLists.txt        | 22 ++---
 .../third_party/openreg/include/openreg.h     | 28 +++---
 .../torch_openreg/torch_openreg/__init__.py   | 95 +------------------
 .../torch_openreg/torch_openreg/_utils.py     | 92 ++++++++++++++++++
 .../torch_openreg/csrc/CMakeLists.txt         | 31 ++----
 .../torch_openreg/csrc/Module.cpp             |  2 +-
 .../torch_openreg/torch_openreg/csrc/stub.c   | 13 ++-
 11 files changed, 205 insertions(+), 176 deletions(-)
 create mode 100644 test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/_utils.py

diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
index 42f0ff45d788..10162061d504 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
@@ -9,16 +9,10 @@ add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
 target_link_libraries(${LIBRARY_NAME} PRIVATE torch_cpu_library openreg)
 target_include_directories(${LIBRARY_NAME} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
 
-if (WIN32)
-    target_compile_definitions(${LIBRARY_NAME} PRIVATE OPENREG_EXPORTS)
-endif()
+set(CMAKE_CXX_VISIBILITY_PRESET hidden)
 
-if(WIN32)
-    install(TARGETS ${LIBRARY_NAME} 
-        ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
-        LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
-        RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}
-    )
-else()
-    install(TARGETS ${LIBRARY_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR})
-endif()
+install(TARGETS ${LIBRARY_NAME} 
+    ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}
+)
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.cpp b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.cpp
index 4d144c3b8e18..6b928f4ad9cc 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.cpp
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.cpp
@@ -30,7 +30,7 @@ int device_count_impl() {
   return count;
 }
 
-OPENREG_API c10::DeviceIndex device_count() noexcept {
+OPENREG_EXPORT c10::DeviceIndex device_count() noexcept {
   // initialize number of devices only once
   static int count = []() {
     try {
@@ -49,17 +49,17 @@ OPENREG_API c10::DeviceIndex device_count() noexcept {
   return static_cast<c10::DeviceIndex>(count);
 }
 
-OPENREG_API c10::DeviceIndex current_device() {
+OPENREG_EXPORT c10::DeviceIndex current_device() {
   c10::DeviceIndex cur_device = -1;
   GetDevice(&cur_device);
   return cur_device;
 }
 
-OPENREG_API void set_device(c10::DeviceIndex device) {
+OPENREG_EXPORT void set_device(c10::DeviceIndex device) {
   SetDevice(device);
 }
 
-OPENREG_API DeviceIndex ExchangeDevice(DeviceIndex device) {
+OPENREG_EXPORT DeviceIndex ExchangeDevice(DeviceIndex device) {
   int current_device = -1;
   orGetDevice(&current_device);
 
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.h b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.h
index b5a1517a7a99..8d8e9cd1e302 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.h
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegFunctions.h
@@ -1,13 +1,9 @@
 #pragma once
 
 #ifdef _WIN32
-  #ifdef OPENREG_EXPORTS
-    #define OPENREG_API __declspec(dllexport)
-  #else
-    #define OPENREG_API __declspec(dllimport)
-  #endif
+  #define OPENREG_EXPORT __declspec(dllexport)
 #else
-  #define OPENREG_API
+  #define OPENREG_EXPORT __attribute__((visibility("default")))
 #endif
 
 #include <c10/core/Device.h>
@@ -17,10 +13,10 @@
 
 namespace c10::openreg {
 
-OPENREG_API c10::DeviceIndex device_count() noexcept;
-OPENREG_API c10::DeviceIndex current_device();
-OPENREG_API 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);
 
-OPENREG_API DeviceIndex ExchangeDevice(DeviceIndex device);
+OPENREG_EXPORT DeviceIndex ExchangeDevice(DeviceIndex device);
 
 } // namespace c10::openreg
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py b/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
index d87fb960dee1..6c8ad8c5ed1c 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
@@ -19,6 +19,19 @@
 RUN_BUILD_DEPS = any(arg in {"clean", "dist_info"} for arg in sys.argv)
 
 
+def check_env_flag(name, default = ""):
+    return os.getenv(name, default).upper() in ["ON", "1", "YES", "TRUE", "Y"]
+
+
+if "CMAKE_BUILD_TYPE" not in os.environ:
+    if check_env_flag("DEBUG"):
+        os.environ["CMAKE_BUILD_TYPE"] = "Debug"
+    elif check_env_flag("REL_WITH_DEB_INFO"):
+        os.environ["CMAKE_BUILD_TYPE"] = "RelWithDebInfo"
+    else:
+        os.environ["CMAKE_BUILD_TYPE"] = "Release"
+
+
 def make_relative_rpath_args(path):
     if IS_DARWIN:
         return ["-Wl,-rpath,@loader_path/" + path]
@@ -54,7 +67,7 @@ def build_deps():
         ".",
         "--target",
         "install",
-        "--config", "Release",
+        "--config", os.environ["CMAKE_BUILD_TYPE"],
         "--",
     ]
 
@@ -83,11 +96,44 @@ def run(self):
 def main():
     if not RUN_BUILD_DEPS:
         build_deps()
-    
-    if sys.platform == "win32":
-        extra_compile_args = ["/W3"]
+
+    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_compile_args = ["-g", "-Wall", "-Werror"]
+        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",
+        ]
+
+    if os.environ["CMAKE_BUILD_TYPE"] == "Debug":
+        if IS_WINDOWS:
+            extra_compile_args += ["/Z7"]
+            extra_link_args += ["/DEBUG:FULL"]
+        else:
+            extra_compile_args += ["-O0", "-g"]
+            extra_link_args += ["-O0", "-g"]
+    elif os.environ["CMAKE_BUILD_TYPE"] == "RelWithDebInfo":
+        if IS_WINDOWS:
+            extra_compile_args += ["/Z7"]
+            extra_link_args += ["/DEBUG:FULL"]
+        else:
+            extra_compile_args += ["-g"]
+            extra_link_args += ["-g"]
 
     ext_modules = [
         Extension(
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
index 2d86b0a5be52..51269071db49 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
@@ -4,20 +4,14 @@ file(GLOB_RECURSE SOURCE_FILES
     "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp"
 )
 
-if(WIN32)
-    add_library(${LIBRARY_NAME} STATIC ${SOURCE_FILES})
-else()
-    add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
-endif()
+add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
 
 target_include_directories(${LIBRARY_NAME} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
 
-if(WIN32)
-    install(TARGETS ${LIBRARY_NAME} 
-        ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
-        LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
-        RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}
-    )
-else()
-    install(TARGETS ${LIBRARY_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR})
-endif()
+set(CMAKE_CXX_VISIBILITY_PRESET hidden)
+
+install(TARGETS ${LIBRARY_NAME} 
+    ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}
+)
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/include/openreg.h b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/include/openreg.h
index b6b0b3da4295..a5e8b77c421c 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/include/openreg.h
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/include/openreg.h
@@ -2,6 +2,12 @@
 
 #include <cstddef>
 
+#ifdef _WIN32
+  #define OPENREG_EXPORT __declspec(dllexport)
+#else
+  #define OPENREG_EXPORT __attribute__((visibility("default")))
+#endif
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -28,19 +34,19 @@ struct orPointerAttributes {
   size_t size;
 };
 
-orError_t orMalloc(void** devPtr, size_t size);
-orError_t orFree(void* devPtr);
-orError_t orMallocHost(void** hostPtr, size_t size);
-orError_t orFreeHost(void* hostPtr);
-orError_t orMemcpy(void* dst, const void* src, size_t count, orMemcpyKind kind);
-orError_t orMemoryUnprotect(void* devPtr);
-orError_t orMemoryProtect(void* devPtr);
+OPENREG_EXPORT orError_t orMalloc(void** devPtr, size_t size);
+OPENREG_EXPORT orError_t orFree(void* devPtr);
+OPENREG_EXPORT orError_t orMallocHost(void** hostPtr, size_t size);
+OPENREG_EXPORT orError_t orFreeHost(void* hostPtr);
+OPENREG_EXPORT orError_t orMemcpy(void* dst, const void* src, size_t count, orMemcpyKind kind);
+OPENREG_EXPORT orError_t orMemoryUnprotect(void* devPtr);
+OPENREG_EXPORT orError_t orMemoryProtect(void* devPtr);
 
-orError_t orGetDeviceCount(int* count);
-orError_t orSetDevice(int device);
-orError_t orGetDevice(int* device);
+OPENREG_EXPORT orError_t orGetDeviceCount(int* count);
+OPENREG_EXPORT orError_t orSetDevice(int device);
+OPENREG_EXPORT orError_t orGetDevice(int* device);
 
-orError_t orPointerGetAttributes(
+OPENREG_EXPORT orError_t orPointerGetAttributes(
     orPointerAttributes* attributes,
     const void* ptr);
 
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py
index 792c7c6d54b2..444a6cbfbca1 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py
@@ -1,102 +1,13 @@
-import ctypes
-import glob
-import os
-import platform
 import sys
-import textwrap
-import sysconfig
-
 import torch
 
-if sys.platform == "win32":
-
-    def _load_dll_libraries() -> None:
-
-        py_dll_path = os.path.join(sys.exec_prefix, "Library", "bin")
-        th_dll_path = os.path.join(os.path.dirname(__file__), "lib")
-        usebase_path = os.path.join(sysconfig.get_config_var("userbase"), "Library", "bin")
-        py_root_bin_path = os.path.join(sys.exec_prefix, "bin")
-
-        # When users create a virtualenv that inherits the base environment,
-        # we will need to add the corresponding library directory into
-        # DLL search directories. Otherwise, it will rely on `PATH` which
-        # is dependent on user settings.
-        if sys.exec_prefix != sys.base_exec_prefix:
-            base_py_dll_path = os.path.join(sys.base_exec_prefix, "Library", "bin")
-        else:
-            base_py_dll_path = ""
-
-        dll_paths = [
-            p
-            for p in (
-                th_dll_path,
-                py_dll_path,
-                base_py_dll_path,
-                usebase_path,
-                py_root_bin_path,
-            )
-            if os.path.exists(p)
-        ]
-
-
-        kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True)
-        with_load_library_flags = hasattr(kernel32, "AddDllDirectory")
-        prev_error_mode = kernel32.SetErrorMode(0x0001)
-
-        kernel32.LoadLibraryW.restype = ctypes.c_void_p
-        if with_load_library_flags:
-            kernel32.LoadLibraryExW.restype = ctypes.c_void_p
-
-        for dll_path in dll_paths:
-            os.add_dll_directory(dll_path)
-
-        try:
-            ctypes.CDLL("vcruntime140.dll")
-            ctypes.CDLL("msvcp140.dll")
-            if platform.machine() != "ARM64":
-                ctypes.CDLL("vcruntime140_1.dll")
-        except OSError:
-            print(
-                textwrap.dedent(
-                    """
-                    Microsoft Visual C++ Redistributable is not installed, this may lead to the DLL load failure.
-                    It can be downloaded at https://aka.ms/vs/16/release/vc_redist.x64.exe
-                    """
-                ).strip()
-            )
-
-        dlls = glob.glob(os.path.join(th_dll_path, "*.dll"))
-        path_patched = False
-        for dll in dlls:
-            is_loaded = False
-            if with_load_library_flags:
-                res = kernel32.LoadLibraryExW(dll, None, 0x00001100)
-                last_error = ctypes.get_last_error()
-                if res is None and last_error != 126:
-                    err = ctypes.WinError(last_error)
-                    err.strerror += (
-                        f' Error loading "{dll}" or one of its dependencies.'
-                    )
-                    raise err
-                elif res is not None:
-                    is_loaded = True
-            if not is_loaded:
-                if not path_patched:
-                    os.environ["PATH"] = ";".join(dll_paths + [os.environ["PATH"]])
-                    path_patched = True
-                res = kernel32.LoadLibraryW(dll)
-                if res is None:
-                    err = ctypes.WinError(ctypes.get_last_error())
-                    err.strerror += (
-                        f' Error loading "{dll}" or one of its dependencies.'
-                    )
-                    raise err
-
-        kernel32.SetErrorMode(prev_error_mode)
 
+if sys.platform == "win32":
+    from ._utils import _load_dll_libraries
     _load_dll_libraries()
     del _load_dll_libraries
 
+
 import torch_openreg._C  # type: ignore[misc]
 import torch_openreg.openreg
 
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/_utils.py b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/_utils.py
new file mode 100644
index 000000000000..5210af72df81
--- /dev/null
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/_utils.py
@@ -0,0 +1,92 @@
+import ctypes
+import glob
+import os
+import platform
+import sys
+import textwrap
+import sysconfig
+
+
+def _load_dll_libraries() -> None:
+
+    py_dll_path = os.path.join(sys.exec_prefix, "Library", "bin")
+    th_dll_path = os.path.join(os.path.dirname(__file__), "lib")
+    usebase_path = os.path.join(sysconfig.get_config_var("userbase"), "Library", "bin")
+    py_root_bin_path = os.path.join(sys.exec_prefix, "bin")
+
+    # When users create a virtualenv that inherits the base environment,
+    # we will need to add the corresponding library directory into
+    # DLL search directories. Otherwise, it will rely on `PATH` which
+    # is dependent on user settings.
+    if sys.exec_prefix != sys.base_exec_prefix:
+        base_py_dll_path = os.path.join(sys.base_exec_prefix, "Library", "bin")
+    else:
+        base_py_dll_path = ""
+
+    dll_paths = [
+        p
+        for p in (
+            th_dll_path,
+            py_dll_path,
+            base_py_dll_path,
+            usebase_path,
+            py_root_bin_path,
+        )
+        if os.path.exists(p)
+    ]
+
+
+    kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True)
+    with_load_library_flags = hasattr(kernel32, "AddDllDirectory")
+    prev_error_mode = kernel32.SetErrorMode(0x0001)
+
+    kernel32.LoadLibraryW.restype = ctypes.c_void_p
+    if with_load_library_flags:
+        kernel32.LoadLibraryExW.restype = ctypes.c_void_p
+
+    for dll_path in dll_paths:
+        os.add_dll_directory(dll_path)
+
+    try:
+        ctypes.CDLL("vcruntime140.dll")
+        ctypes.CDLL("msvcp140.dll")
+        if platform.machine() != "ARM64":
+            ctypes.CDLL("vcruntime140_1.dll")
+    except OSError:
+        print(
+            textwrap.dedent(
+                """
+                Microsoft Visual C++ Redistributable is not installed, this may lead to the DLL load failure.
+                It can be downloaded at https://aka.ms/vs/16/release/vc_redist.x64.exe
+                """
+            ).strip()
+        )
+
+    dlls = glob.glob(os.path.join(th_dll_path, "*.dll"))
+    path_patched = False
+    for dll in dlls:
+        is_loaded = False
+        if with_load_library_flags:
+            res = kernel32.LoadLibraryExW(dll, None, 0x00001100)
+            last_error = ctypes.get_last_error()
+            if res is None and last_error != 126:
+                err = ctypes.WinError(last_error)
+                err.strerror += (
+                    f' Error loading "{dll}" or one of its dependencies.'
+                )
+                raise err
+            elif res is not None:
+                is_loaded = True
+        if not is_loaded:
+            if not path_patched:
+                os.environ["PATH"] = ";".join(dll_paths + [os.environ["PATH"]])
+                path_patched = True
+            res = kernel32.LoadLibraryW(dll)
+            if res is None:
+                err = ctypes.WinError(ctypes.get_last_error())
+                err.strerror += (
+                    f' Error loading "{dll}" or one of its dependencies.'
+                )
+                raise err
+
+    kernel32.SetErrorMode(prev_error_mode)
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
index 4744bbe7207f..0fd1e7926e47 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
@@ -8,16 +8,7 @@ add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
 
 if(WIN32)
     # For Windows, we need specific PyTorch libraries
-    find_library(TORCH_PYTHON_LIBRARY torch_python PATHS ${PYTORCH_INSTALL_DIR}/lib REQUIRED)
-    find_library(TORCH_CPU_LIBRARY torch_cpu PATHS ${PYTORCH_INSTALL_DIR}/lib REQUIRED)
-    find_library(C10_LIBRARY c10 PATHS ${PYTORCH_INSTALL_DIR}/lib REQUIRED)
-    
-    target_link_libraries(${LIBRARY_NAME} PRIVATE 
-        ${TORCH_PYTHON_LIBRARY}
-        ${TORCH_CPU_LIBRARY}
-        ${C10_LIBRARY}
-        torch_openreg
-    )
+    target_link_libraries(${LIBRARY_NAME} PRIVATE torch_python torch_cpu_library c10 torch_openreg)
 else()
     target_link_libraries(${LIBRARY_NAME} PRIVATE torch_python torch_openreg)
 endif()
@@ -32,10 +23,6 @@ if(WIN32)
     endif()
 endif()
 
-if(WIN32)
-    set_target_properties(${LIBRARY_NAME} PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS ON)
-endif()
-
 target_link_directories(${LIBRARY_NAME} PRIVATE ${PYTORCH_INSTALL_DIR}/lib)
 target_include_directories(${LIBRARY_NAME} PRIVATE ${PYTORCH_INSTALL_DIR}/include)
 
@@ -43,12 +30,10 @@ if(APPLE)
   set_target_properties(${LIBRARY_NAME} PROPERTIES LINK_FLAGS "-undefined dynamic_lookup")
 endif()
 
-if(WIN32)
-    install(TARGETS ${LIBRARY_NAME} 
-        ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
-        LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
-        RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}
-    )
-else()
-    install(TARGETS ${LIBRARY_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR})
-endif()
+set(CMAKE_CXX_VISIBILITY_PRESET hidden)
+
+install(TARGETS ${LIBRARY_NAME} 
+    ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}
+)
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/Module.cpp b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/Module.cpp
index 4acdbfc8e1dc..38c456339003 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/Module.cpp
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/Module.cpp
@@ -90,7 +90,7 @@ static PyMethodDef methods[] = {
  * Therefore, it cannot be named initModule here, otherwise initModule
  * in torch/csrc/Module.cpp will be called, resulting in failure.
  */
-extern "C" PyObject* initOpenRegModule(void) {
+extern "C" OPENREG_EXPORT PyObject* initOpenRegModule(void) {
   static struct PyModuleDef openreg_C_module = {
       PyModuleDef_HEAD_INIT, "torch_openreg._C", nullptr, -1, methods};
   PyObject* mod = PyModule_Create(&openreg_C_module);
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/stub.c b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/stub.c
index cd3eb4fe1ecc..243a43a37e5e 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/stub.c
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/stub.c
@@ -1,13 +1,18 @@
 #include <Python.h>
 
-extern PyObject* initOpenRegModule(void);
+#ifdef _WIN32
+  #define OPENREG_EXPORT __declspec(dllexport)
+#else
+  #define OPENREG_EXPORT __attribute__((visibility("default")))
+#endif
+
+extern OPENREG_EXPORT PyObject* initOpenRegModule(void);
 
-#ifndef _WIN32
 #ifdef __cplusplus
 extern "C"
 #endif
-__attribute__((visibility("default"))) PyObject* PyInit__C(void);
-#endif
+
+OPENREG_EXPORT PyObject* PyInit__C(void);
 
 PyMODINIT_FUNC PyInit__C(void)
 {

From 4c41abce4c9234cd735f8393c25b382ea57afa3e Mon Sep 17 00:00:00 2001
From: jiahaochen666 <jiahaochen535@gmail.com>
Date: Thu, 31 Jul 2025 20:42:10 +0800
Subject: [PATCH 4/9] Add torch_python_library

---
 .../torch_openreg/CMakeLists.txt              |  2 +
 .../cmake/TorchPythonTargets.cmake            | 37 +++++++++++++++++++
 .../torch_openreg/csrc/CMakeLists.txt         |  4 +-
 3 files changed, 41 insertions(+), 2 deletions(-)
 create mode 100644 test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake

diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
index 4a1a4a312316..bf4ee5b660ca 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
@@ -33,6 +33,8 @@ else()
   message(FATAL_ERROR "Cannot find Python directory")
 endif()
 
+include(${PROJECT_SOURCE_DIR}/cmake/TorchPythonTargets.cmake)
+
 add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/openreg)
 add_subdirectory(${PROJECT_SOURCE_DIR}/csrc)
 add_subdirectory(${PROJECT_SOURCE_DIR}/torch_openreg/csrc)
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake b/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake
new file mode 100644
index 000000000000..11315b664a64
--- /dev/null
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake
@@ -0,0 +1,37 @@
+if (WIN32)
+  set(TORCH_PYTHON_IMPORTED_LOCATION "${PYTORCH_INSTALL_DIR}/lib/torch_python.lib")
+  set(TORCH_PYTHON_LINK_LIBS
+    "${PYTORCH_INSTALL_DIR}/lib/libprotobuf.lib"
+    "${PYTORCH_INSTALL_DIR}/lib/c10.lib"
+    "${PYTORCH_INSTALL_DIR}/lib/torch_cpu.lib"
+  )
+  set(TORCH_PYTHON_LIBRARY_LINK
+    "\$<TARGET_FILE:torch_python>;\$<TARGET_PROPERTY:torch_python,INTERFACE_LINK_LIBRARIES>"
+  )
+else()
+  set(TORCH_PYTHON_IMPORTED_LOCATION "${PYTORCH_INSTALL_DIR}/lib/libtorch_python.so")
+  set(TORCH_PYTHON_LINK_LIBS "protobuf::libprotobuf;c10;torch_cpu")
+  set(TORCH_PYTHON_LIBRARY_LINK
+    "-Wl,--no-as-needed,\"\$<TARGET_FILE:torch_python>\" -Wl,--as-needed;"
+    "\$<TARGET_PROPERTY:torch_python,INTERFACE_LINK_LIBRARIES>"
+  )
+endif()
+
+add_library(torch_python SHARED IMPORTED)
+
+set_target_properties(torch_python PROPERTIES
+  INTERFACE_COMPILE_DEFINITIONS "USE_DISTRIBUTED;USE_C10D_GLOO;USE_C10D_MPI;USE_RPC;USE_TENSORPIPE"
+  INTERFACE_INCLUDE_DIRECTORIES "${PYTORCH_INSTALL_DIR}/include"
+  INTERFACE_LINK_LIBRARIES "${TORCH_PYTHON_LINK_LIBS}"
+  IMPORTED_LOCATION "${TORCH_PYTHON_IMPORTED_LOCATION}"
+)
+
+add_library(torch_python_library INTERFACE IMPORTED)
+
+set_target_properties(torch_python_library PROPERTIES
+  INTERFACE_COMPILE_DEFINITIONS "\$<TARGET_PROPERTY:torch_python,INTERFACE_COMPILE_DEFINITIONS>"
+  INTERFACE_COMPILE_OPTIONS "\$<TARGET_PROPERTY:torch_python,INTERFACE_COMPILE_OPTIONS>"
+  INTERFACE_INCLUDE_DIRECTORIES "\$<TARGET_PROPERTY:torch_python,INTERFACE_INCLUDE_DIRECTORIES>"
+  INTERFACE_LINK_LIBRARIES "${TORCH_PYTHON_LIBRARY_LINK}"
+  INTERFACE_SYSTEM_INCLUDE_DIRECTORIES "\$<TARGET_PROPERTY:torch_python,INTERFACE_SYSTEM_INCLUDE_DIRECTORIES>"
+)
\ No newline at end of file
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
index 0fd1e7926e47..949b903bd308 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
@@ -8,9 +8,9 @@ add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
 
 if(WIN32)
     # For Windows, we need specific PyTorch libraries
-    target_link_libraries(${LIBRARY_NAME} PRIVATE torch_python torch_cpu_library c10 torch_openreg)
+    target_link_libraries(${LIBRARY_NAME} PRIVATE torch_python_library torch_cpu_library c10 torch_openreg)
 else()
-    target_link_libraries(${LIBRARY_NAME} PRIVATE torch_python torch_openreg)
+    target_link_libraries(${LIBRARY_NAME} PRIVATE torch_python_library torch_openreg)
 endif()
 
 if(WIN32)

From ff890b34b75e79f0b70685726fa7f6594657210c Mon Sep 17 00:00:00 2001
From: jiahaochen666 <jiahaochen535@gmail.com>
Date: Fri, 1 Aug 2025 11:16:21 +0800
Subject: [PATCH 5/9] Update openreg testcase for Windows

---
 test/test_openreg.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/test/test_openreg.py b/test/test_openreg.py
index cae20b16f479..a2418fbdef20 100644
--- a/test/test_openreg.py
+++ b/test/test_openreg.py
@@ -17,6 +17,7 @@
 from torch.testing._internal.common_utils import (
     run_tests,
     skipIfTorchDynamo,
+    skipIfWindows,
     skipIfXpu,
     TemporaryFileName,
     TestCase,
@@ -284,6 +285,7 @@ def test_manual_seed(self):
         self.assertEqual(torch.openreg.initial_seed(), 2024)  # type: ignore[misc]
 
     # Autograd
+    @skipIfWindows()
     def test_autograd_init(self):
         # Make sure autograd is initialized
         torch.ones(2, requires_grad=True, device="openreg").sum().backward()

From 71dc569a1cef9dfe7a0e7741a7cb3fc7b23d5f95 Mon Sep 17 00:00:00 2001
From: jiahaochen666 <jiahaochen535@gmail.com>
Date: Fri, 1 Aug 2025 14:16:05 +0800
Subject: [PATCH 6/9] Removing libprotobuf from cmake

---
 .../torch_openreg/cmake/TorchPythonTargets.cmake               | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake b/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake
index 11315b664a64..ff69b487dcc9 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake
@@ -1,7 +1,6 @@
 if (WIN32)
   set(TORCH_PYTHON_IMPORTED_LOCATION "${PYTORCH_INSTALL_DIR}/lib/torch_python.lib")
   set(TORCH_PYTHON_LINK_LIBS
-    "${PYTORCH_INSTALL_DIR}/lib/libprotobuf.lib"
     "${PYTORCH_INSTALL_DIR}/lib/c10.lib"
     "${PYTORCH_INSTALL_DIR}/lib/torch_cpu.lib"
   )
@@ -10,7 +9,7 @@ if (WIN32)
   )
 else()
   set(TORCH_PYTHON_IMPORTED_LOCATION "${PYTORCH_INSTALL_DIR}/lib/libtorch_python.so")
-  set(TORCH_PYTHON_LINK_LIBS "protobuf::libprotobuf;c10;torch_cpu")
+  set(TORCH_PYTHON_LINK_LIBS "c10;torch_cpu")
   set(TORCH_PYTHON_LIBRARY_LINK
     "-Wl,--no-as-needed,\"\$<TARGET_FILE:torch_python>\" -Wl,--as-needed;"
     "\$<TARGET_PROPERTY:torch_python,INTERFACE_LINK_LIBRARIES>"

From 95279ed7b01139819db3c3a232f936d0fead1497 Mon Sep 17 00:00:00 2001
From: jiahaochen666 <jiahaochen535@gmail.com>
Date: Fri, 1 Aug 2025 15:45:59 +0800
Subject: [PATCH 7/9] Fix some issues

---
 .../torch_openreg/CMakeLists.txt                  |  1 +
 .../torch_openreg/csrc/CMakeLists.txt             |  2 --
 .../torch_openreg/setup.py                        |  1 -
 .../third_party/openreg/CMakeLists.txt            |  2 --
 .../torch_openreg/csrc/CMakeLists.txt             | 15 +--------------
 5 files changed, 2 insertions(+), 19 deletions(-)

diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
index bf4ee5b660ca..c1cc0eeeb3b1 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/CMakeLists.txt
@@ -13,6 +13,7 @@ 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")
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
index 10162061d504..12d7500b22a0 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
@@ -9,8 +9,6 @@ add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
 target_link_libraries(${LIBRARY_NAME} PRIVATE torch_cpu_library openreg)
 target_include_directories(${LIBRARY_NAME} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
 
-set(CMAKE_CXX_VISIBILITY_PRESET hidden)
-
 install(TARGETS ${LIBRARY_NAME} 
     ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py b/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
index 6c8ad8c5ed1c..fea08dd19d32 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
@@ -11,7 +11,6 @@
 
 
 # Env Variables
-IS_LINUX = platform.system() == "Linux"
 IS_DARWIN = platform.system() == "Darwin"
 IS_WINDOWS = platform.system() == "Windows"
 
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
index 51269071db49..5d9aeb23e2c5 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
@@ -8,8 +8,6 @@ add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
 
 target_include_directories(${LIBRARY_NAME} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
 
-set(CMAKE_CXX_VISIBILITY_PRESET hidden)
-
 install(TARGETS ${LIBRARY_NAME} 
     ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
index 949b903bd308..36ed78b716e3 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
@@ -6,32 +6,19 @@ file(GLOB_RECURSE SOURCE_FILES
 
 add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
 
-if(WIN32)
-    # For Windows, we need specific PyTorch libraries
-    target_link_libraries(${LIBRARY_NAME} PRIVATE torch_python_library torch_cpu_library c10 torch_openreg)
-else()
-    target_link_libraries(${LIBRARY_NAME} PRIVATE torch_python_library torch_openreg)
-endif()
+target_link_libraries(${LIBRARY_NAME} PRIVATE torch_python_library torch_openreg)
 
 if(WIN32)
-    if(DEFINED PYTHON_LIBRARY_DIR AND DEFINED PYTHON_LIBRARY)
-        target_link_directories(${LIBRARY_NAME} PRIVATE ${PYTHON_LIBRARY_DIR})
-        target_link_libraries(${LIBRARY_NAME} PRIVATE ${PYTHON_LIBRARY})
-    else()
         find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
         target_link_libraries(${LIBRARY_NAME} PRIVATE ${Python3_LIBRARIES})
-    endif()
 endif()
 
 target_link_directories(${LIBRARY_NAME} PRIVATE ${PYTORCH_INSTALL_DIR}/lib)
-target_include_directories(${LIBRARY_NAME} PRIVATE ${PYTORCH_INSTALL_DIR}/include)
 
 if(APPLE)
   set_target_properties(${LIBRARY_NAME} PROPERTIES LINK_FLAGS "-undefined dynamic_lookup")
 endif()
 
-set(CMAKE_CXX_VISIBILITY_PRESET hidden)
-
 install(TARGETS ${LIBRARY_NAME} 
     ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}

From 7381aa4ef6befd6a82e5b3e8ae36b02b30200712 Mon Sep 17 00:00:00 2001
From: jiahaochen666 <jiahaochen535@gmail.com>
Date: Fri, 1 Aug 2025 16:51:41 +0800
Subject: [PATCH 8/9] Fix lint error

---
 .../torch_openreg/cmake/TorchPythonTargets.cmake     |  2 +-
 .../torch_openreg/csrc/CMakeLists.txt                |  2 +-
 .../torch_openreg/setup.py                           |  5 +++--
 .../torch_openreg/third_party/openreg/CMakeLists.txt |  2 +-
 .../third_party/openreg/csrc/memory.cpp              |  2 +-
 .../torch_openreg/torch_openreg/__init__.py          |  2 ++
 .../torch_openreg/torch_openreg/_utils.py            | 12 +++---------
 .../torch_openreg/torch_openreg/csrc/CMakeLists.txt  |  2 +-
 8 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake b/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake
index ff69b487dcc9..92aa53c4733f 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake
@@ -1,4 +1,4 @@
-if (WIN32)
+if(WIN32)
   set(TORCH_PYTHON_IMPORTED_LOCATION "${PYTORCH_INSTALL_DIR}/lib/torch_python.lib")
   set(TORCH_PYTHON_LINK_LIBS
     "${PYTORCH_INSTALL_DIR}/lib/c10.lib"
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
index 12d7500b22a0..e2ae2b3f3667 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/csrc/CMakeLists.txt
@@ -9,7 +9,7 @@ add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
 target_link_libraries(${LIBRARY_NAME} PRIVATE torch_cpu_library openreg)
 target_include_directories(${LIBRARY_NAME} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
 
-install(TARGETS ${LIBRARY_NAME} 
+install(TARGETS ${LIBRARY_NAME}
     ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
     RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py b/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
index fea08dd19d32..b57e15608d72 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
@@ -18,7 +18,7 @@
 RUN_BUILD_DEPS = any(arg in {"clean", "dist_info"} for arg in sys.argv)
 
 
-def check_env_flag(name, default = ""):
+def check_env_flag(name, default=""):
     return os.getenv(name, default).upper() in ["ON", "1", "YES", "TRUE", "Y"]
 
 
@@ -66,7 +66,8 @@ def build_deps():
         ".",
         "--target",
         "install",
-        "--config", os.environ["CMAKE_BUILD_TYPE"],
+        "--config",
+        os.environ["CMAKE_BUILD_TYPE"],
         "--",
     ]
 
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
index 5d9aeb23e2c5..5450b49be164 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/CMakeLists.txt
@@ -8,7 +8,7 @@ add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
 
 target_include_directories(${LIBRARY_NAME} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
 
-install(TARGETS ${LIBRARY_NAME} 
+install(TARGETS ${LIBRARY_NAME}
     ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
     RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.cpp b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.cpp
index b169ba8125b2..942b04b3b50a 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.cpp
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/csrc/memory.cpp
@@ -52,7 +52,7 @@ class MemoryManager {
     if (it == m_registry.end())
       return orErrorUnknown;
     const auto& info = it->second;
-    
+
     if (info.type == orMemoryType::orMemoryTypeDevice) {
       openreg::mprotect(info.pointer, info.size, F_PROT_READ | F_PROT_WRITE);
       openreg::munmap(info.pointer, info.size);
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py
index 444a6cbfbca1..45b2343070fe 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/__init__.py
@@ -1,9 +1,11 @@
 import sys
+
 import torch
 
 
 if sys.platform == "win32":
     from ._utils import _load_dll_libraries
+
     _load_dll_libraries()
     del _load_dll_libraries
 
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/_utils.py b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/_utils.py
index 5210af72df81..286a66afd789 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/_utils.py
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/_utils.py
@@ -3,12 +3,11 @@
 import os
 import platform
 import sys
-import textwrap
 import sysconfig
+import textwrap
 
 
 def _load_dll_libraries() -> None:
-
     py_dll_path = os.path.join(sys.exec_prefix, "Library", "bin")
     th_dll_path = os.path.join(os.path.dirname(__file__), "lib")
     usebase_path = os.path.join(sysconfig.get_config_var("userbase"), "Library", "bin")
@@ -35,7 +34,6 @@ def _load_dll_libraries() -> None:
         if os.path.exists(p)
     ]
 
-
     kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True)
     with_load_library_flags = hasattr(kernel32, "AddDllDirectory")
     prev_error_mode = kernel32.SetErrorMode(0x0001)
@@ -71,9 +69,7 @@ def _load_dll_libraries() -> None:
             last_error = ctypes.get_last_error()
             if res is None and last_error != 126:
                 err = ctypes.WinError(last_error)
-                err.strerror += (
-                    f' Error loading "{dll}" or one of its dependencies.'
-                )
+                err.strerror += f' Error loading "{dll}" or one of its dependencies.'
                 raise err
             elif res is not None:
                 is_loaded = True
@@ -84,9 +80,7 @@ def _load_dll_libraries() -> None:
             res = kernel32.LoadLibraryW(dll)
             if res is None:
                 err = ctypes.WinError(ctypes.get_last_error())
-                err.strerror += (
-                    f' Error loading "{dll}" or one of its dependencies.'
-                )
+                err.strerror += f' Error loading "{dll}" or one of its dependencies.'
                 raise err
 
     kernel32.SetErrorMode(prev_error_mode)
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
index 36ed78b716e3..2864e5bd9cf5 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
@@ -19,7 +19,7 @@ if(APPLE)
   set_target_properties(${LIBRARY_NAME} PROPERTIES LINK_FLAGS "-undefined dynamic_lookup")
 endif()
 
-install(TARGETS ${LIBRARY_NAME} 
+install(TARGETS ${LIBRARY_NAME}
     ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
     RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}

From 65f2cc0fda95f10e80dd120f09794a56f8cbbbfe Mon Sep 17 00:00:00 2001
From: jiahaochen666 <jiahaochen535@gmail.com>
Date: Wed, 6 Aug 2025 17:29:49 +0800
Subject: [PATCH 9/9] Fix some issues in comments

---
 .../cmake/TorchPythonTargets.cmake            | 20 +------
 .../torch_openreg/setup.py                    | 38 +++-----------
 .../torch_openreg/torch_openreg/_utils.py     | 52 ++-----------------
 .../torch_openreg/csrc/CMakeLists.txt         | 10 ++--
 4 files changed, 16 insertions(+), 104 deletions(-)

diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake b/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake
index 92aa53c4733f..181dee20e0bb 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake
@@ -1,36 +1,20 @@
 if(WIN32)
   set(TORCH_PYTHON_IMPORTED_LOCATION "${PYTORCH_INSTALL_DIR}/lib/torch_python.lib")
-  set(TORCH_PYTHON_LINK_LIBS
-    "${PYTORCH_INSTALL_DIR}/lib/c10.lib"
-    "${PYTORCH_INSTALL_DIR}/lib/torch_cpu.lib"
-  )
-  set(TORCH_PYTHON_LIBRARY_LINK
-    "\$<TARGET_FILE:torch_python>;\$<TARGET_PROPERTY:torch_python,INTERFACE_LINK_LIBRARIES>"
-  )
 else()
   set(TORCH_PYTHON_IMPORTED_LOCATION "${PYTORCH_INSTALL_DIR}/lib/libtorch_python.so")
-  set(TORCH_PYTHON_LINK_LIBS "c10;torch_cpu")
-  set(TORCH_PYTHON_LIBRARY_LINK
-    "-Wl,--no-as-needed,\"\$<TARGET_FILE:torch_python>\" -Wl,--as-needed;"
-    "\$<TARGET_PROPERTY:torch_python,INTERFACE_LINK_LIBRARIES>"
-  )
 endif()
 
 add_library(torch_python SHARED IMPORTED)
 
 set_target_properties(torch_python PROPERTIES
-  INTERFACE_COMPILE_DEFINITIONS "USE_DISTRIBUTED;USE_C10D_GLOO;USE_C10D_MPI;USE_RPC;USE_TENSORPIPE"
   INTERFACE_INCLUDE_DIRECTORIES "${PYTORCH_INSTALL_DIR}/include"
-  INTERFACE_LINK_LIBRARIES "${TORCH_PYTHON_LINK_LIBS}"
+  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_COMPILE_DEFINITIONS "\$<TARGET_PROPERTY:torch_python,INTERFACE_COMPILE_DEFINITIONS>"
-  INTERFACE_COMPILE_OPTIONS "\$<TARGET_PROPERTY:torch_python,INTERFACE_COMPILE_OPTIONS>"
   INTERFACE_INCLUDE_DIRECTORIES "\$<TARGET_PROPERTY:torch_python,INTERFACE_INCLUDE_DIRECTORIES>"
-  INTERFACE_LINK_LIBRARIES "${TORCH_PYTHON_LIBRARY_LINK}"
-  INTERFACE_SYSTEM_INCLUDE_DIRECTORIES "\$<TARGET_PROPERTY:torch_python,INTERFACE_SYSTEM_INCLUDE_DIRECTORIES>"
+  INTERFACE_LINK_LIBRARIES "\$<TARGET_FILE:torch_python>;\$<TARGET_PROPERTY:torch_python,INTERFACE_LINK_LIBRARIES>"
 )
\ No newline at end of file
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py b/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
index b57e15608d72..48124c079e48 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
@@ -18,19 +18,6 @@
 RUN_BUILD_DEPS = any(arg in {"clean", "dist_info"} for arg in sys.argv)
 
 
-def check_env_flag(name, default=""):
-    return os.getenv(name, default).upper() in ["ON", "1", "YES", "TRUE", "Y"]
-
-
-if "CMAKE_BUILD_TYPE" not in os.environ:
-    if check_env_flag("DEBUG"):
-        os.environ["CMAKE_BUILD_TYPE"] = "Debug"
-    elif check_env_flag("REL_WITH_DEB_INFO"):
-        os.environ["CMAKE_BUILD_TYPE"] = "RelWithDebInfo"
-    else:
-        os.environ["CMAKE_BUILD_TYPE"] = "Release"
-
-
 def make_relative_rpath_args(path):
     if IS_DARWIN:
         return ["-Wl,-rpath,@loader_path/" + path]
@@ -67,7 +54,7 @@ def build_deps():
         "--target",
         "install",
         "--config",
-        os.environ["CMAKE_BUILD_TYPE"],
+        "Release",
         "--",
     ]
 
@@ -100,13 +87,15 @@ def main():
     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"]
+        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 = []
+        extra_link_args = [*make_relative_rpath_args("lib")]
         extra_compile_args = [
             "-Wall",
             "-Wextra",
@@ -120,21 +109,6 @@ def main():
             "-fno-strict-aliasing",
         ]
 
-    if os.environ["CMAKE_BUILD_TYPE"] == "Debug":
-        if IS_WINDOWS:
-            extra_compile_args += ["/Z7"]
-            extra_link_args += ["/DEBUG:FULL"]
-        else:
-            extra_compile_args += ["-O0", "-g"]
-            extra_link_args += ["-O0", "-g"]
-    elif os.environ["CMAKE_BUILD_TYPE"] == "RelWithDebInfo":
-        if IS_WINDOWS:
-            extra_compile_args += ["/Z7"]
-            extra_link_args += ["/DEBUG:FULL"]
-        else:
-            extra_compile_args += ["-g"]
-            extra_link_args += ["-g"]
-
     ext_modules = [
         Extension(
             name="torch_openreg._C",
@@ -143,7 +117,7 @@ def main():
             extra_compile_args=extra_compile_args,
             libraries=["torch_bindings"],
             library_dirs=[os.path.join(BASE_DIR, "torch_openreg/lib")],
-            extra_link_args=[*make_relative_rpath_args("lib")],
+            extra_link_args=extra_link_args,
         )
     ]
 
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/_utils.py b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/_utils.py
index 286a66afd789..1c26f475ba7a 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/_utils.py
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/_utils.py
@@ -1,38 +1,10 @@
 import ctypes
 import glob
 import os
-import platform
-import sys
-import sysconfig
-import textwrap
 
 
 def _load_dll_libraries() -> None:
-    py_dll_path = os.path.join(sys.exec_prefix, "Library", "bin")
-    th_dll_path = os.path.join(os.path.dirname(__file__), "lib")
-    usebase_path = os.path.join(sysconfig.get_config_var("userbase"), "Library", "bin")
-    py_root_bin_path = os.path.join(sys.exec_prefix, "bin")
-
-    # When users create a virtualenv that inherits the base environment,
-    # we will need to add the corresponding library directory into
-    # DLL search directories. Otherwise, it will rely on `PATH` which
-    # is dependent on user settings.
-    if sys.exec_prefix != sys.base_exec_prefix:
-        base_py_dll_path = os.path.join(sys.base_exec_prefix, "Library", "bin")
-    else:
-        base_py_dll_path = ""
-
-    dll_paths = [
-        p
-        for p in (
-            th_dll_path,
-            py_dll_path,
-            base_py_dll_path,
-            usebase_path,
-            py_root_bin_path,
-        )
-        if os.path.exists(p)
-    ]
+    openreg_dll_path = os.path.join(os.path.dirname(__file__), "lib")
 
     kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True)
     with_load_library_flags = hasattr(kernel32, "AddDllDirectory")
@@ -42,25 +14,9 @@ def _load_dll_libraries() -> None:
     if with_load_library_flags:
         kernel32.LoadLibraryExW.restype = ctypes.c_void_p
 
-    for dll_path in dll_paths:
-        os.add_dll_directory(dll_path)
-
-    try:
-        ctypes.CDLL("vcruntime140.dll")
-        ctypes.CDLL("msvcp140.dll")
-        if platform.machine() != "ARM64":
-            ctypes.CDLL("vcruntime140_1.dll")
-    except OSError:
-        print(
-            textwrap.dedent(
-                """
-                Microsoft Visual C++ Redistributable is not installed, this may lead to the DLL load failure.
-                It can be downloaded at https://aka.ms/vs/16/release/vc_redist.x64.exe
-                """
-            ).strip()
-        )
+    os.add_dll_directory(openreg_dll_path)
 
-    dlls = glob.glob(os.path.join(th_dll_path, "*.dll"))
+    dlls = glob.glob(os.path.join(openreg_dll_path, "*.dll"))
     path_patched = False
     for dll in dlls:
         is_loaded = False
@@ -75,7 +31,7 @@ def _load_dll_libraries() -> None:
                 is_loaded = True
         if not is_loaded:
             if not path_patched:
-                os.environ["PATH"] = ";".join(dll_paths + [os.environ["PATH"]])
+                os.environ["PATH"] = ";".join([openreg_dll_path] + [os.environ["PATH"]])
                 path_patched = True
             res = kernel32.LoadLibraryW(dll)
             if res is None:
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
index 2864e5bd9cf5..4ff321c43f2c 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/torch_openreg/csrc/CMakeLists.txt
@@ -9,16 +9,14 @@ add_library(${LIBRARY_NAME} SHARED ${SOURCE_FILES})
 target_link_libraries(${LIBRARY_NAME} PRIVATE torch_python_library torch_openreg)
 
 if(WIN32)
-        find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
-        target_link_libraries(${LIBRARY_NAME} PRIVATE ${Python3_LIBRARIES})
+    find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
+    target_link_libraries(${LIBRARY_NAME} PRIVATE ${Python3_LIBRARIES})
+elseif(APPLE)
+    set_target_properties(${LIBRARY_NAME} PROPERTIES LINK_FLAGS "-undefined dynamic_lookup")
 endif()
 
 target_link_directories(${LIBRARY_NAME} PRIVATE ${PYTORCH_INSTALL_DIR}/lib)
 
-if(APPLE)
-  set_target_properties(${LIBRARY_NAME} PROPERTIES LINK_FLAGS "-undefined dynamic_lookup")
-endif()
-
 install(TARGETS ${LIBRARY_NAME}
     ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}