[Openreg][PrivateUse1] Refactor csrc files of Pytorch_openreg (#151004)

FFFrog · pytorchmergebot · commit 2653498ff34d · 2025-04-12T17:22:28.000Z
I want to format and refactor the csrc file of pytorch_openreg. To make the code review clearer and easier to understand, I divide the code refactoring into two parts: - Part 1: Code formatting - Part 2: Code refactoring and optimization (Next PR) Pull Request resolved: #151004 Approved by: https://github.com/albanD ghstack dependencies: #151000
diff --git a/test/cpp_extensions/open_registration_extension/pytorch_openreg/csrc/Module.cpp b/test/cpp_extensions/open_registration_extension/pytorch_openreg/csrc/Module.cpp
@@ -1,24 +1,18 @@
-#include <OpenReg.h>
+#include "OpenReg.h"
 
-// Make this a proper CPython module
 static struct PyModuleDef openreg_C_module = {
     PyModuleDef_HEAD_INIT,
     "pytorch_openreg._C",
     nullptr,
     -1,
-    nullptr,
-    nullptr,
-    nullptr,
-    nullptr,
-    nullptr
 };
 
 PyMODINIT_FUNC PyInit__C(void) {
-    PyObject* mod = PyModule_Create(&openreg_C_module);
+  PyObject* mod = PyModule_Create(&openreg_C_module);
 
-    py::object openreg_mod = py::module_::import("pytorch_openreg");
-    // Only borrowed from the python side!
-    openreg::set_impl_factory(openreg_mod.attr("impl_factory").ptr());
+  py::object openreg_mod = py::module_::import("pytorch_openreg");
+  // Only borrowed from the python side!
+  openreg::set_impl_factory(openreg_mod.attr("impl_factory").ptr());
 
-    return mod;
+  return mod;
 }
diff --git a/test/cpp_extensions/open_registration_extension/pytorch_openreg/csrc/OpenReg.h b/test/cpp_extensions/open_registration_extension/pytorch_openreg/csrc/OpenReg.h
@@ -1,11 +1,10 @@
 #pragma once
-// Shared header for OpenReg module
 
 #include <torch/csrc/utils/pybind.h>
 
 namespace openreg {
 
-    void set_impl_factory(PyObject* factory);
-    py::function get_method(const char* name);
+void set_impl_factory(PyObject* factory);
+py::function get_method(const char* name);
 
-}
+} // namespace openreg
diff --git a/test/cpp_extensions/open_registration_extension/pytorch_openreg/csrc/OpenRegHooks.cpp b/test/cpp_extensions/open_registration_extension/pytorch_openreg/csrc/OpenRegHooks.cpp
@@ -1,10 +1,10 @@
-#include <OpenReg.h>
+#include "OpenReg.h"
 
+#include <ATen/CPUGeneratorImpl.h>
+#include <ATen/core/GeneratorForPrivateuseone.h>
 #include <ATen/detail/PrivateUse1HooksInterface.h>
+#include <c10/core/Device.h>
 #include <c10/core/impl/DeviceGuardImplInterface.h>
-#include "ATen/CPUGeneratorImpl.h"
-#include "ATen/core/GeneratorForPrivateuseone.h"
-#include "c10/core/Device.h"
 #include <c10/util/CallOnce.h>
 
 #include <iostream>
@@ -24,7 +24,8 @@ struct HostAllocator final : at::Allocator {
     py::gil_scoped_acquire acquire;
     void* data = nullptr;
     if (nbytes > 0) {
-      data = reinterpret_cast<void*>(get_method("hostMalloc")(nbytes).cast<host_ptr_t>());
+      data = reinterpret_cast<void*>(
+          get_method("hostMalloc")(nbytes).cast<host_ptr_t>());
       TORCH_CHECK(data, "Failed to allocator ", nbytes, " bytes on host.");
     }
     return {data, data, &ReportAndDelete, at::Device(at::kCPU)};
@@ -37,7 +38,8 @@ struct HostAllocator final : at::Allocator {
     py::gil_scoped_acquire acquire;
     TORCH_CHECK(
         get_method("hostFree")(reinterpret_cast<host_ptr_t>(ptr)).cast<bool>(),
-        "Failed to free memory pointer at ", ptr);
+        "Failed to free memory pointer at ",
+        ptr);
   }
 
   at::DeleterFnPtr raw_deleter() const override {
@@ -46,7 +48,10 @@ struct HostAllocator final : at::Allocator {
 
   void copy_data(void* dest, const void* src, std::size_t count) const final {
     py::gil_scoped_acquire acquire;
-    get_method("hostCopyData")(reinterpret_cast<host_ptr_t>(dest), reinterpret_cast<host_ptr_t>(src), count);
+    get_method("hostCopyData")(
+        reinterpret_cast<host_ptr_t>(dest),
+        reinterpret_cast<host_ptr_t>(src),
+        count);
   }
 };
 static HostAllocator global_host_alloc;
@@ -104,7 +109,8 @@ struct OpenRegHooksInterface : public at::PrivateUse1HooksInterface {
 
   bool isPinnedPtr(const void* data) const override {
     py::gil_scoped_acquire acquire;
-    return get_method("isPinnedPtr")(reinterpret_cast<host_ptr_t>(data)).cast<bool>();
+    return get_method("isPinnedPtr")(reinterpret_cast<host_ptr_t>(data))
+        .cast<bool>();
   }
 
   const at::Generator& getDefaultGenerator(
@@ -126,15 +132,26 @@ struct OpenRegHooksInterface : public at::PrivateUse1HooksInterface {
 };
 
 int register_hook() {
-  at::RegisterPrivateUse1HooksInterface(new OpenRegHooksInterface(OpenRegHooksArgs{}));
+  at::RegisterPrivateUse1HooksInterface(
+      new OpenRegHooksInterface(OpenRegHooksArgs{}));
   return 0;
 }
 int temp_register_hook = register_hook();
 
-TORCH_DECLARE_REGISTRY(PrivateUse1HooksRegistry, OpenRegHooksInterface, OpenRegHooksArgs);
-C10_DEFINE_REGISTRY(PrivateUse1HooksRegistry, OpenRegHooksInterface, OpenRegHooksArgs);
-// Using Create function to get PrivateUse1HooksInterface point from PrivateUse1HooksRegistry class.
-C10_REGISTER_TYPED_CLASS(PrivateUse1HooksRegistry, "OpenRegHooks", OpenRegHooksInterface);
+TORCH_DECLARE_REGISTRY(
+    PrivateUse1HooksRegistry,
+    OpenRegHooksInterface,
+    OpenRegHooksArgs);
+C10_DEFINE_REGISTRY(
+    PrivateUse1HooksRegistry,
+    OpenRegHooksInterface,
+    OpenRegHooksArgs);
+// Using Create function to get PrivateUse1HooksInterface point from
+// PrivateUse1HooksRegistry class.
+C10_REGISTER_TYPED_CLASS(
+    PrivateUse1HooksRegistry,
+    "OpenRegHooks",
+    OpenRegHooksInterface);
 
 // Device guard registration
 struct OpenRegGuardImpl final : public c10::impl::DeviceGuardImplInterface {
@@ -158,7 +175,8 @@ struct OpenRegGuardImpl final : public c10::impl::DeviceGuardImplInterface {
   c10::Device exchangeDevice(c10::Device d) const override {
     TORCH_INTERNAL_ASSERT(d.is_privateuseone());
     py::gil_scoped_acquire acquire;
-    auto old_device_index = get_method("exchangeDevice")(d.index()).cast<c10::DeviceIndex>();
+    auto old_device_index =
+        get_method("exchangeDevice")(d.index()).cast<c10::DeviceIndex>();
     return c10::Device(static_type, old_device_index);
   }
 
@@ -207,9 +225,12 @@ struct OpenRegGuardImpl final : public c10::impl::DeviceGuardImplInterface {
   /**
    * Get a stream from the global pool for a given device.
    */
-  c10::Stream getStreamFromGlobalPool(c10::Device d, bool isHighPriority = false) const override {
+  c10::Stream getStreamFromGlobalPool(
+      c10::Device d,
+      bool isHighPriority = false) const override {
     py::gil_scoped_acquire acquire;
-    return get_method("getStreamFromGlobalPool")(d.index(), isHighPriority).cast<c10::Stream>();
+    return get_method("getStreamFromGlobalPool")(d.index(), isHighPriority)
+        .cast<c10::Stream>();
   }
 
   /**
@@ -219,7 +240,8 @@ struct OpenRegGuardImpl final : public c10::impl::DeviceGuardImplInterface {
    */
   c10::Stream getNewStream(c10::Device d, int priority = 0) const override {
     py::gil_scoped_acquire acquire;
-    auto stream_id = get_method("getNewStream")(d.index(), priority).cast<c10::StreamId>();
+    auto stream_id =
+        get_method("getNewStream")(d.index(), priority).cast<c10::StreamId>();
     return c10::Stream(c10::Stream::UNSAFE, d, stream_id);
   }
 
@@ -322,26 +344,31 @@ struct OpenRegGuardImpl final : public c10::impl::DeviceGuardImplInterface {
    * being used on the given stream, and that it should thus avoid recycling the
    * DataPtr until all work on that stream is done.
    */
-  void recordDataPtrOnStream(const c10::DataPtr& data_ptr, const c10::Stream& stream)
-      const override {
+  void recordDataPtrOnStream(
+      const c10::DataPtr& data_ptr,
+      const c10::Stream& stream) const override {
     py::gil_scoped_acquire acquire;
     get_method("recordDataPtrOnStream")(data_ptr, stream);
   }
 
   /**
    * Fetch the elapsed time between two recorded events.
    */
-  double elapsedTime(void* event1, void* event2, const c10::DeviceIndex device_index)
-      const override {
+  double elapsedTime(
+      void* event1,
+      void* event2,
+      const c10::DeviceIndex device_index) const override {
     py::gil_scoped_acquire acquire;
-    return get_method("elapsedTime")((int64_t)event1, (int64_t)event2, device_index).cast<double>();
+    return get_method("elapsedTime")(
+               (int64_t)event1, (int64_t)event2, device_index)
+        .cast<double>();
   }
 };
 
 // Register our device guard
 C10_REGISTER_GUARD_IMPL(PrivateUse1, OpenRegGuardImpl);
 
-} // anonymous namspaces
+} // namespace
 
 // Setter for the python dictionary with implementations
 void set_impl_factory(PyObject* factory) {
@@ -352,4 +379,4 @@ py::function get_method(const char* name) {
   auto factory = py::cast<py::function>(py_factory);
   return factory(name);
 }
-} // openreg
+} // namespace openreg
diff --git a/test/cpp_extensions/open_registration_extension/pytorch_openreg/csrc/OpenRegMem.cpp b/test/cpp_extensions/open_registration_extension/pytorch_openreg/csrc/OpenRegMem.cpp
@@ -1,13 +1,13 @@
-#include <OpenReg.h>
+#include "OpenReg.h"
 
-#include <torch/library.h>
-#include <c10/core/Allocator.h>
-#include <ATen/detail/PrivateUse1HooksInterface.h>
 #include <ATen/EmptyTensor.h>
-#include <c10/util/ArrayRef.h>
-#include <c10/core/TensorOptions.h>
+#include <ATen/detail/PrivateUse1HooksInterface.h>
 #include <ATen/ops/as_strided_cpu_dispatch.h>
 #include <ATen/ops/set_cpu_dispatch.h>
+#include <c10/core/Allocator.h>
+#include <c10/core/TensorOptions.h>
+#include <c10/util/ArrayRef.h>
+#include <torch/library.h>
 
 namespace openreg {
 
@@ -22,11 +22,14 @@ struct OpenRegAllocator final : at::Allocator {
   at::DataPtr allocate(size_t nbytes) override {
     py::gil_scoped_acquire acquire;
     auto curr_device_idx = get_method("getDevice")().cast<c10::DeviceIndex>();
-    auto curr_device = c10::Device(c10::DeviceType::PrivateUse1, curr_device_idx);
+    auto curr_device =
+        c10::Device(c10::DeviceType::PrivateUse1, curr_device_idx);
     void* data = nullptr;
     if (nbytes > 0) {
-        data = reinterpret_cast<void*>(get_method("malloc")(nbytes).cast<openreg_ptr_t>());
-        TORCH_CHECK(data, "Failed to allocator ", nbytes, " bytes on openreg device.");
+      data = reinterpret_cast<void*>(
+          get_method("malloc")(nbytes).cast<openreg_ptr_t>());
+      TORCH_CHECK(
+          data, "Failed to allocator ", nbytes, " bytes on openreg device.");
     }
     return {data, data, &ReportAndDelete, curr_device};
   }
@@ -62,15 +65,17 @@ struct OpenRegAllocator final : at::Allocator {
 
   void copy_data(void* dest, const void* src, std::size_t count) const final {
     py::gil_scoped_acquire acquire;
-    get_method("copy_data")(reinterpret_cast<openreg_ptr_t>(dest), reinterpret_cast<openreg_ptr_t>(src), count);
+    get_method("copy_data")(
+        reinterpret_cast<openreg_ptr_t>(dest),
+        reinterpret_cast<openreg_ptr_t>(src),
+        count);
   }
 };
 
 // Register our dummy allocator
 static OpenRegAllocator global_openreg_alloc;
 REGISTER_ALLOCATOR(c10::DeviceType::PrivateUse1, &global_openreg_alloc);
 
-
 // Empty op needs C++ code and cannot be handled by python side fallback
 at::Tensor empty_openreg(
     c10::IntArrayRef size,
@@ -82,8 +87,12 @@ at::Tensor empty_openreg(
   const auto device = c10::device_or_default(device_opt);
   const auto dtype = c10::dtype_or_default(dtype_opt);
   TORCH_CHECK(device.is_privateuseone());
-  TORCH_CHECK(c10::layout_or_default(layout_opt) == c10::Layout::Strided, "Non strided layout not supported");
-  TORCH_CHECK(!c10::pinned_memory_or_default(pin_memory_opt), "Pin memory can only be on CPU");
+  TORCH_CHECK(
+      c10::layout_or_default(layout_opt) == c10::Layout::Strided,
+      "Non strided layout not supported");
+  TORCH_CHECK(
+      !c10::pinned_memory_or_default(pin_memory_opt),
+      "Pin memory can only be on CPU");
   const c10::DeviceGuard device_guard(device);
   constexpr c10::DispatchKeySet pu1_dks(c10::DispatchKey::PrivateUse1);
   return at::detail::empty_generic(
@@ -100,8 +109,12 @@ at::Tensor empty_strided_openreg(
   const auto device = c10::device_or_default(device_opt);
   const auto dtype = c10::dtype_or_default(dtype_opt);
   TORCH_CHECK(device.is_privateuseone());
-  TORCH_CHECK(c10::layout_or_default(layout_opt) == c10::Layout::Strided, "Non strided layout not supported");
-  TORCH_CHECK(!c10::pinned_memory_or_default(pin_memory_opt), "Pin memory can only be on CPU");
+  TORCH_CHECK(
+      c10::layout_or_default(layout_opt) == c10::Layout::Strided,
+      "Non strided layout not supported");
+  TORCH_CHECK(
+      !c10::pinned_memory_or_default(pin_memory_opt),
+      "Pin memory can only be on CPU");
   const c10::DeviceGuard device_guard(device);
   constexpr c10::DispatchKeySet pu1_dks(c10::DispatchKey::PrivateUse1);
   return at::detail::empty_strided_generic(
@@ -113,8 +126,8 @@ at::Tensor as_strided_openreg(
     c10::IntArrayRef size,
     c10::IntArrayRef stride,
     std::optional<int64_t> storage_offset_) {
-    // Metadata-only change so we re-use the cpu impl
-    return at::cpu::as_strided(self, size, stride, storage_offset_);
+  // Metadata-only change so we re-use the cpu impl
+  return at::cpu::as_strided(self, size, stride, storage_offset_);
 }
 
 at::Tensor& set_openreg(
@@ -123,17 +136,16 @@ at::Tensor& set_openreg(
     int64_t storage_offset,
     c10::IntArrayRef size,
     c10::IntArrayRef stride) {
-    return at::cpu::set_(result, storage, storage_offset, size, stride);
+  return at::cpu::set_(result, storage, storage_offset, size, stride);
 }
 
-
 TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
-    m.impl("empty.memory_format", empty_openreg);
-    m.impl("empty_strided", empty_strided_openreg);
-    m.impl("as_strided", as_strided_openreg);
-    m.impl("set_.source_Storage_storage_offset", set_openreg);
+  m.impl("empty.memory_format", empty_openreg);
+  m.impl("empty_strided", empty_strided_openreg);
+  m.impl("as_strided", as_strided_openreg);
+  m.impl("set_.source_Storage_storage_offset", set_openreg);
 }
 
-} // anonymous namspaces
+} // namespace
 
-} // openreg
+} // namespace openreg
