add rdma tensor write (#36)

Summary:

- Support rdma tensor transfer operations from python
I took some liberties with the python facing interface, open to feedback if they should be modified

Reviewed By: d4l3k

Differential Revision: D85887151

Author: amirafzali

comms/torchcomms/ncclx/TorchCommNCCLXPy.cpp

@@ -29,6 +29,9 @@ PYBIND11_MODULE(_comms_ncclx, m) {
   py::class_<TorchCommNCCLX, std::shared_ptr<TorchCommNCCLX>>(
       m, "TorchCommNCCLX");
 
+  py::class_<RdmaRemoteBuffer, std::shared_ptr<RdmaRemoteBuffer>>(
+      m, "RdmaRemoteBuffer");
+
   py::class_<RdmaTransport, std::shared_ptr<RdmaTransport>>(m, "RdmaTransport")
       // initialize a new RDMATransport using a custom init fn
       .def(py::init([](at::Device device) {
@@ -45,5 +48,37 @@ PYBIND11_MODULE(_comms_ncclx, m) {
             std::string peerUrlStr = peerUrl.cast<std::string>();
             return static_cast<int>(self.connect(peerUrlStr));
           })
-      .def("connected", &RdmaTransport::connected);
+      .def("connected", &RdmaTransport::connected)
+      .def(
+          "write",
+          [](RdmaTransport& self,
+             const RdmaMemory::View& localBuffer,
+             const RdmaRemoteBuffer& remoteBuffer) {
+            return static_cast<int>(
+                self.write(localBuffer, remoteBuffer, false).get());
+          });
+
+  py::class_<RdmaMemory::View, std::shared_ptr<RdmaMemory::View>>(
+      m, "RdmaMemoryView")
+      .def("size", &RdmaMemory::View::size);
+
+  py::class_<RdmaMemory, std::shared_ptr<RdmaMemory>>(m, "RdmaMemory")
+      .def(py::init([](const at::Tensor& tensor) {
+        TORCH_CHECK(
+            tensor.is_contiguous(),
+            "RdmaMemory currently requires a contiguous tensor");
+        // If CPU memory is passed, use device 0 for NIC discovery
+        const auto device = tensor.get_device() < 0 ? 0 : tensor.get_device();
+        return std::make_shared<RdmaMemory>(
+            tensor.data_ptr(), tensor.nbytes(), device);
+      }))
+      .def(
+          "to_view",
+          [](RdmaMemory& self) {
+            return self.createView(size_t(0), self.length());
+          })
+      .def("to_remote_buffer", [](RdmaMemory& self) {
+        return RdmaRemoteBuffer{
+            const_cast<void*>(self.data()), self.remoteKey()};
+      });
 }

comms/torchcomms/ncclx/_comms_ncclx.pyi

@@ -5,10 +5,23 @@ import torch
 
 class TorchCommNCCLX: ...
 
+class RdmaMemoryView:
+    def size(self) -> int: ...
+
+class RdmaRemoteBuffer: ...
+
+class RdmaMemory:
+    def __init__(self, tensor: torch.Tensor) -> None: ...  # pyre-ignore[11]
+    def to_view(self) -> RdmaMemoryView: ...
+    def to_remote_buffer(self) -> RdmaRemoteBuffer: ...
+
 class RdmaTransport:
     def __init__(self, device: torch.device) -> None: ...  # pyre-ignore[11]
     @staticmethod
     def supported() -> bool: ...
     def bind(self) -> bytes: ...
     def connect(self, peer_url: bytes) -> int: ...
     def connected(self) -> bool: ...
+    def write(
+        self, local_buffer: RdmaMemoryView, remote_buffer: RdmaRemoteBuffer
+    ) -> int: ...

comms/torchcomms/tests/integration/py/TransportTest.py

@@ -6,7 +6,7 @@
 import unittest
 
 import torch
-from torchcomms._comms_ncclx import RdmaTransport
+from torchcomms._comms_ncclx import RdmaMemory, RdmaTransport
 
 
 class TransportTest(unittest.TestCase):
@@ -17,28 +17,30 @@ def setUp(self):
     def test_construct(self) -> None:
         _ = RdmaTransport(torch.device("cuda:0"))
 
-    def test_bind_and_connect(self) -> None:
-        if torch.cuda.device_count() < 2:
-            self.skipTest(
-                f"Test requires at least 2 CUDA devices, found {torch.cuda.device_count()}"
-            )
+    def test_rdma_memory_from_tensor(self) -> None:
+        tensor = torch.arange(1024, dtype=torch.uint8, device="cuda:0")
+        compare_tensor = torch.zeros_like(tensor, device="cuda:1")
 
-        server_device = torch.device("cuda:0")
-        client_device = torch.device("cuda:1")
+        tensor_mem = RdmaMemory(tensor)
+        compare_mem = RdmaMemory(compare_tensor)
 
-        server_transport = RdmaTransport(server_device)
-        client_transport = RdmaTransport(client_device)
+        tensor_view = tensor_mem.to_view()
+        compare_view = compare_mem.to_view()
 
-        server_url = server_transport.bind()
-        client_url = client_transport.bind()
+        self.assertEqual(tensor_view.size(), tensor.nbytes)
+        self.assertAlmostEqual(tensor_view.size(), compare_view.size())
+
+    def bind_and_connect(self, server: RdmaTransport, client: RdmaTransport) -> None:
+        server_url = server.bind()
+        client_url = client.bind()
 
         self.assertIsNotNone(server_url)
         self.assertIsNotNone(client_url)
         self.assertNotEqual(server_url, "")
         self.assertNotEqual(client_url, "")
 
-        server_result = server_transport.connect(client_url)
-        client_result = client_transport.connect(server_url)
+        server_result = server.connect(client_url)
+        client_result = client.connect(server_url)
 
         self.assertEqual(
             server_result, 0, "Server connect should return commSuccess (0)"
@@ -47,8 +49,86 @@ def test_bind_and_connect(self) -> None:
             client_result, 0, "Client connect should return commSuccess (0)"
         )
 
-        self.assertTrue(server_transport.connected())
-        self.assertTrue(client_transport.connected())
+        self.assertTrue(server.connected())
+        self.assertTrue(client.connected())
+
+    def test_bind_and_connect(self) -> None:
+        if torch.cuda.device_count() < 2:
+            self.skipTest(
+                f"Test requires at least 2 CUDA devices, found {torch.cuda.device_count()}"
+            )
+
+        server_device = torch.device("cuda:0")
+        client_device = torch.device("cuda:1")
+
+        server_transport = RdmaTransport(server_device)
+        client_transport = RdmaTransport(client_device)
+
+        self.bind_and_connect(server_transport, client_transport)
+
+    def run_send_recv(
+        self,
+        device1: str,
+        device2: str,
+    ) -> None:
+        transport_device_1 = "cuda:0" if device1 == "cpu" else device1
+        transport_device_2 = "cuda:0" if device2 == "cpu" else device2
+        transport1 = RdmaTransport(torch.device(transport_device_1))
+        transport2 = RdmaTransport(torch.device(transport_device_2))
+
+        self.bind_and_connect(transport1, transport2)
+
+        tensor1 = torch.arange(1024, dtype=torch.uint8, device=device1)
+        tensor2 = torch.zeros_like(tensor1, device=device2)
+
+        self.assertEqual(tensor1.nbytes, tensor2.nbytes)
+
+        tensor1_mem = RdmaMemory(tensor1)
+        tensor2_mem = RdmaMemory(tensor2)
+
+        res = transport1.write(tensor1_mem.to_view(), tensor2_mem.to_remote_buffer())
+
+        self.assertEqual(res, 0)
+        self.assertTrue(torch.allclose(tensor1.cpu(), tensor2.cpu()))
+
+        del transport1
+        del transport2
+        del tensor1_mem
+        del tensor2_mem
+
+    def check_multi_gpu(self) -> None:
+        if torch.cuda.device_count() < 2:
+            self.skipTest(
+                f"Test requires at least 2 CUDA devices, found {torch.cuda.device_count()}"
+            )
+
+    def test_write_gpu_to_gpu(self) -> None:
+        self.check_multi_gpu()
+        self.run_send_recv("cuda:0", "cuda:1")
+
+    def test_write_gpu_to_gpu_2(self) -> None:
+        self.check_multi_gpu()
+        self.run_send_recv("cuda:0", "cuda:0")
+
+    def test_write_cpu_to_gpu(self) -> None:
+        self.check_multi_gpu()
+        self.run_send_recv("cpu", "cuda:1")
+
+    def test_write_cpu_to_gpu_2(self) -> None:
+        self.check_multi_gpu()
+        self.run_send_recv("cpu", "cuda:0")
+
+    def test_write_gpu_to_cpu(self) -> None:
+        self.check_multi_gpu()
+        self.run_send_recv("cuda:1", "cpu")
+
+    def test_write_gpu_to_cpu_2(self) -> None:
+        self.check_multi_gpu()
+        self.run_send_recv("cuda:0", "cpu")
+
+    def test_write_cpu_to_cpu(self) -> None:
+        self.check_multi_gpu()
+        self.run_send_recv("cpu", "cpu")
 
 
 if __name__ == "__main__" and os.environ["TEST_BACKEND"] == "ncclx":

comms/torchcomms/transport/RdmaTransport.h

@@ -98,6 +98,14 @@ class RdmaMemory : folly::MoveOnly {
     return cudaDev_;
   }
 
+  size_t length() const {
+    return len_;
+  }
+
+  const void* data() const {
+    return buf_;
+  }
+
   /*
    * Check if the given buffer and length are contained within this memory
    * region.