Add register_tensor/deregister_tensor to public TorchComm API (#2019)

Summary:
Pull Request resolved: #2019

Add `register_address(void*, size_t)` and `deregister_address(void*)` as virtual methods on `TorchCommBackend` with no-op defaults. Add tensor-level
`register_tensor`/`deregister_tensor` on `TorchComm` that delegates to the virtual methods. Bind as `register_tensor`/`deregister_tensor` in Python.

ABI version bumped from 1.0 to 1.1 (new virtual methods on base class).

Reviewed By: d4l3k

Differential Revision: D100188720

fbshipit-source-id: 0c8ac5a2ae8b5137e8e36d86935a5083ddeead27

Author: arttianezhu

comms/torchcomms/TorchComm.cpp

@@ -512,6 +512,14 @@ int64_t TorchComm::get_device_transport() {
   return impl_->get_device_transport();
 }
 
+void TorchComm::tensor_register(const at::Tensor& tensor) {
+  impl_->tensor_register(tensor);
+}
+
+void TorchComm::tensor_deregister(const at::Tensor& tensor) {
+  impl_->tensor_deregister(tensor);
+}
+
 // Communicator Management
 std::shared_ptr<TorchComm> TorchComm::split(
     const std::vector<int>& ranks,

comms/torchcomms/TorchComm.hpp

@@ -170,6 +170,10 @@ class TorchComm : public std::enable_shared_from_this<TorchComm> {
   // Throws if not supported by the backend.
   int64_t get_device_transport();
 
+  // Memory Registration API
+  void tensor_register(const at::Tensor& tensor);
+  void tensor_deregister(const at::Tensor& tensor);
+
   std::shared_ptr<TorchCommBackend> getBackendImpl() const {
     return impl_;
   }

comms/torchcomms/TorchCommBackend.hpp

@@ -361,6 +361,38 @@ class TorchCommBackend {
         std::string(getCommName()));
   }
 
+  /**
+   * Register a tensor's memory with the backend for optimized data transfer.
+   *
+   * Pre-registers the memory region for zero-copy RDMA or similar transport
+   * optimizations. Backends that support registration must override this.
+   *
+   * The caller is responsible for calling tensor_deregister() before the
+   * tensor is freed. Failing to deregister leaks the backend registration
+   * handle but does not cause a crash — the transport layer will clean up
+   * on communicator finalization.
+   *
+   * @param tensor The tensor whose memory to register.
+   */
+  virtual void tensor_register(const at::Tensor& /*tensor*/) {
+    throw std::runtime_error(
+        "[TorchCommBackend]: tensor_register not implemented for "
+        "communicator:" +
+        std::string(getCommName()));
+  }
+
+  /**
+   * Deregister a tensor's previously registered memory.
+   *
+   * @param tensor The tensor whose memory to deregister.
+   */
+  virtual void tensor_deregister(const at::Tensor& /*tensor*/) {
+    throw std::runtime_error(
+        "[TorchCommBackend]: tensor_deregister not implemented for "
+        "communicator:" +
+        std::string(getCommName()));
+  }
+
  protected:
   void runAbortHooks() {
     for (const auto& [_, hook] : abortHooks_) {

comms/torchcomms/TorchCommPy.cpp

@@ -1446,6 +1446,37 @@ is valid until the communicator is destroyed.
     RuntimeError: If the backend does not support device transport.
 )",
           py::call_guard<py::gil_scoped_release>())
+      .def(
+          "tensor_register",
+          [](TorchComm& self, const at::Tensor& tensor) {
+            self.tensor_register(tensor);
+          },
+          R"(
+Register a tensor's memory with the communication backend.
+
+Pre-registers the memory region for optimized data transfer (e.g.,
+RDMA zero-copy). The caller must call tensor_deregister() before
+freeing the tensor. Omitting deregistration leaks the backend handle
+but does not crash; cleanup occurs on communicator finalization.
+
+Args:
+    tensor: The tensor whose memory to register.
+          )",
+          py::arg("tensor"),
+          py::call_guard<py::gil_scoped_release>())
+      .def(
+          "tensor_deregister",
+          [](TorchComm& self, const at::Tensor& tensor) {
+            self.tensor_deregister(tensor);
+          },
+          R"(
+Deregister a tensor's previously registered memory.
+
+Args:
+    tensor: The tensor whose memory to deregister.
+          )",
+          py::arg("tensor"),
+          py::call_guard<py::gil_scoped_release>())
 
       // Point-to-Point Operations
       .def(

comms/torchcomms/_comms.pyi

@@ -707,6 +707,8 @@ class TorchComm:
         self, callback: Callable[[int, PostHookArgs], None]
     ) -> RemovableHandle: ...
     def register_abort_hook(self, callback: Callable[[], None]) -> RemovableHandle: ...
+    def tensor_register(self, tensor: Any) -> None: ...
+    def tensor_deregister(self, tensor: Any) -> None: ...
 
 def new_comm(
     backend: str,

comms/torchcomms/fake/TorchCommFake.hpp

@@ -4,6 +4,7 @@
 
 #include <comms/torchcomms/TorchCommBackend.hpp>
 #include <comms/torchcomms/TorchWork.hpp>
+#include <unordered_set>
 #include <vector>
 
 namespace torch::comms {
@@ -193,6 +194,19 @@ class TorchCommFake : public TorchCommBackend {
     return abortEnabled_ && aborted_;
   }
 
+  // Memory registration — tracks registered addresses for test verification
+  void tensor_register(const at::Tensor& tensor) override {
+    registered_addrs_.insert(tensor.data_ptr());
+  }
+
+  void tensor_deregister(const at::Tensor& tensor) override {
+    registered_addrs_.erase(tensor.data_ptr());
+  }
+
+  bool is_tensor_registered(const at::Tensor& tensor) const {
+    return registered_addrs_.count(tensor.data_ptr()) > 0;
+  }
+
  private:
   bool initialized_;
   at::Device device_;
@@ -203,6 +217,7 @@ class TorchCommFake : public TorchCommBackend {
   bool abortEnabled_{false};
   bool aborted_{false};
   bool shouldFailReconfigure_{false};
+  std::unordered_set<void*> registered_addrs_;
 };
 
 } // namespace torch::comms

comms/torchcomms/tests/integration/py/RegisterTensorTest.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+# pyre-unsafe
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+import itertools
+import unittest
+
+import torch
+from torchcomms.tests.integration.py.TorchCommTestHelpers import (
+    skip_backend,
+    TorchCommTestWrapper,
+)
+
+
+class RegisterTensorTest(unittest.TestCase):
+    """Test tensor_register/tensor_deregister via CPU broadcast.
+
+    Verifies that pre-registering a CPU tensor with the communicator
+    enables zero-copy RDMA for CPU collectives. Backends that override
+    tensor_register should inherit this test.
+    """
+
+    counts = [4, 1024, 1024 * 1024]
+    dtypes = [torch.float, torch.int, torch.int8]
+
+    def setUp(self):
+        self.wrapper = TorchCommTestWrapper()
+        self.torchcomm = self.wrapper.get_torchcomm()
+        self.rank = self.torchcomm.get_rank()
+        self.num_ranks = self.torchcomm.get_size()
+
+    def tearDown(self):
+        self.torchcomm = None
+        self.wrapper = None
+
+    def _verify_broadcast_results(self, tensor, expected_value, count, dtype):
+        expected = torch.ones(count, dtype=dtype, device="cpu") * expected_value
+        if dtype == torch.float:
+            self.assertTrue(
+                torch.allclose(tensor, expected),
+                f"CPU broadcast tensors not close enough for count={count}",
+            )
+        else:
+            self.assertTrue(
+                torch.equal(tensor, expected),
+                f"CPU broadcast tensors not equal for count={count}",
+            )
+
+    def _cpu_broadcast_with_tensor_register(self, count, dtype):
+        """Test CPU broadcast using TorchComm.tensor_register() public API."""
+        root_rank = 0
+        root_value = 99
+
+        if self.rank == root_rank:
+            tensor = torch.ones(count, dtype=dtype, device="cpu") * root_value
+        else:
+            tensor = torch.zeros(count, dtype=dtype, device="cpu")
+
+        self.torchcomm.tensor_register(tensor)
+
+        try:
+            work = self.torchcomm.broadcast(tensor, root_rank, False)
+            self.assertTrue(work.is_completed())
+            self._verify_broadcast_results(tensor, root_value, count, dtype)
+        finally:
+            self.torchcomm.tensor_deregister(tensor)
+
+    @skip_backend("nccl", msg="tensor_register not implemented for backend: ")
+    @skip_backend("ncclx", msg="tensor_register not implemented for backend: ")
+    @skip_backend("gloo", msg="tensor_register not implemented for backend: ")
+    def test_cpu_broadcast_with_tensor_register(self):
+        """Test CPU broadcast with TorchComm.tensor_register() public API."""
+        for count, dtype in itertools.product(self.counts, self.dtypes):
+            with self.subTest(count=count, dtype=dtype):
+                self._cpu_broadcast_with_tensor_register(count, dtype)
+
+
+if __name__ == "__main__":
+    unittest.main()

comms/torchcomms/tests/unit/cpp/TorchCommRegisterTensorTest.cpp

@@ -0,0 +1,79 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+
+#include <gtest/gtest.h>
+
+#include <ATen/ATen.h>
+#include <comms/torchcomms/TorchComm.hpp>
+#include <comms/torchcomms/TorchCommFactory.hpp>
+#include <comms/torchcomms/fake/TorchCommFake.hpp>
+#include <cstdlib>
+
+namespace torch::comms {
+
+namespace {
+constexpr const char* kBackendName = "fake_test";
+constexpr const char* kBackendEnvKey = "TORCHCOMMS_BACKEND_LIB_PATH_FAKE_TEST";
+} // namespace
+
+class TorchCommRegisterTensorTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    const char* lib_path = std::getenv("FAKE_TEST_BACKEND_LIB_PATH");
+    ASSERT_NE(lib_path, nullptr) << "FAKE_TEST_BACKEND_LIB_PATH not set";
+    setenv(kBackendEnvKey, lib_path, 1);
+
+    comm_ = new_comm(kBackendName, at::Device(at::kCPU), "register_test");
+    ASSERT_NE(comm_, nullptr);
+  }
+
+  void TearDown() override {
+    comm_.reset();
+    unsetenv(kBackendEnvKey);
+  }
+
+  TorchCommFake* getFakeBackend() {
+    return dynamic_cast<TorchCommFake*>(comm_->getBackendImpl().get());
+  }
+
+  std::shared_ptr<TorchComm> comm_;
+};
+
+TEST_F(TorchCommRegisterTensorTest, RegisterDelegatesToBackend) {
+  auto tensor = at::ones({1024}, at::kFloat);
+  auto* backend = getFakeBackend();
+  ASSERT_NE(backend, nullptr);
+
+  EXPECT_FALSE(backend->is_tensor_registered(tensor));
+  comm_->tensor_register(tensor);
+  EXPECT_TRUE(backend->is_tensor_registered(tensor));
+}
+
+TEST_F(TorchCommRegisterTensorTest, DeregisterRemovesRegistration) {
+  auto tensor = at::ones({1024}, at::kFloat);
+  auto* backend = getFakeBackend();
+  ASSERT_NE(backend, nullptr);
+
+  comm_->tensor_register(tensor);
+  EXPECT_TRUE(backend->is_tensor_registered(tensor));
+
+  comm_->tensor_deregister(tensor);
+  EXPECT_FALSE(backend->is_tensor_registered(tensor));
+}
+
+TEST_F(TorchCommRegisterTensorTest, RegisterMultipleTensors) {
+  auto t1 = at::ones({1024}, at::kFloat);
+  auto t2 = at::zeros({2048}, at::kByte);
+  auto* backend = getFakeBackend();
+  ASSERT_NE(backend, nullptr);
+
+  comm_->tensor_register(t1);
+  comm_->tensor_register(t2);
+  EXPECT_TRUE(backend->is_tensor_registered(t1));
+  EXPECT_TRUE(backend->is_tensor_registered(t2));
+
+  comm_->tensor_deregister(t1);
+  EXPECT_FALSE(backend->is_tensor_registered(t1));
+  EXPECT_TRUE(backend->is_tensor_registered(t2));
+}
+
+} // namespace torch::comms