[3/N][SymmMem] Expose offset field from handle (pytorch#161532)

As titled, so that kernels relying on direct pointers can use base address and `hdl.offset` to access remote memory.

Pull Request resolved: pytorch#161532
Approved by: https://github.com/ngimel
ghstack dependencies: pytorch#161470, pytorch#161471

Author: kwen2501

test/distributed/test_nvshmem.py

@@ -117,6 +117,29 @@ def test_mempool_compute_ops(self) -> None:
         expected = torch.mm(x0, w)
         self.assertEqual(y, expected)
 
+    @skipIfRocm
+    def test_handle_offset(self) -> None:
+        """
+        Test if handle offset is correctly set.
+        """
+        self._init_device()
+        group_name = dist.group.WORLD.group_name
+        symm_mem.enable_symm_mem_for_group(group_name)
+
+        dtype = torch.float
+        numel = 1024
+        allocator = symm_mem.get_mempool_allocator(self.device)
+        mempool = torch.cuda.MemPool(allocator)
+
+        with torch.cuda.use_mem_pool(mempool):
+            x0 = torch.empty(numel, dtype=dtype, device=self.device)
+            x1 = torch.empty_like(x0)
+
+        hdl0 = symm_mem.rendezvous(x0, group=group_name)
+        hdl1 = symm_mem.rendezvous(x1, group=group_name)
+        self.assertEqual(hdl0.offset, 0)
+        self.assertEqual(hdl1.offset, x0.untyped_storage().nbytes())
+
     @skipIfRocm
     def test_nvshmem_put(self) -> None:
         self._init_device()

torch/csrc/distributed/c10d/init.cpp

@@ -1170,6 +1170,7 @@ This class does not support ``__members__`` property.)");
       .def_property_readonly("buffer_size", &SymmetricMemory::get_buffer_size)
       .def_property_readonly(
           "signal_pad_size", &SymmetricMemory::get_signal_pad_size)
+      .def_property_readonly("offset", &SymmetricMemory::get_offset)
       .def(
           "get_buffer",
           &SymmetricMemory::get_buffer,

torch/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cu

@@ -195,6 +195,10 @@ class NVSHMEMSymmetricMemory : public SymmetricMemory {
     return nullptr;
   }
 
+  size_t get_offset() override {
+    return offset_;
+  }
+
   at::Tensor get_buffer(
       int rank,
       c10::IntArrayRef sizes,

torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp

@@ -50,6 +50,10 @@ class TORCH_API SymmetricMemory : public c10::intrusive_ptr_target {
   virtual size_t get_buffer_size() = 0;
   virtual size_t get_signal_pad_size() = 0;
 
+  virtual size_t get_offset() {
+    TORCH_CHECK(false, "NYI");
+  }
+
   virtual bool has_multicast_support() = 0;
   virtual void* get_multicast_ptr() = 0;