support symm memory on XPU devices

Author: zhangxiaoli73

src/xccl/IpcExchange.hpp

@@ -0,0 +1,130 @@
+#pragma once
+
+#include <ATen/ATen.h>
+#include <sycl/sycl.hpp>
+#include <torch/csrc/distributed/c10d/Store.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
+#include <xccl/XPUSymmetricMemoryTypes.hpp>
+
+namespace c10d::symmetric_memory {
+
+// Resource wrapper that owns a (vaddr, allocation handle) pair. Upon
+// destruction, it unmaps the vaddr and releases the allocation handle.
+struct AllocationRef : public c10::intrusive_ptr_target {
+  void* ptr;
+  HandleType handle;
+  size_t block_size;
+  int device_idx;
+  bool local_allocation;
+
+  AllocationRef(
+      void* ptr,
+      HandleType handle,
+      size_t block_size,
+      int device_idx,
+      bool local_allocation);
+
+  ~AllocationRef();
+};
+
+class XPUSymmetricMemory : public SymmetricMemory {
+ public:
+  XPUSymmetricMemory(
+      std::vector<c10::intrusive_ptr<AllocationRef>> alloc_refs,
+      std::vector<void*> buffers,
+      std::vector<void*> signal_pads,
+      HandleType mc_handle,
+      void* mc_addr,
+      size_t buffer_size,
+      int local_device_idx,
+      int rank,
+      int world_size);
+
+  ~XPUSymmetricMemory() override{};
+
+  std::vector<void*> get_buffer_ptrs() override;
+  std::vector<void*> get_signal_pad_ptrs() override;
+  void** get_buffer_ptrs_dev() override;
+  void** get_signal_pad_ptrs_dev() override;
+  size_t get_buffer_size() override;
+  size_t get_signal_pad_size() override;
+
+  bool has_multicast_support() override;
+  void* get_multicast_ptr() override;
+
+  at::Tensor get_buffer(
+      int rank,
+      c10::IntArrayRef sizes,
+      c10::ScalarType dtype,
+      int64_t storage_offset);
+
+  void barrier(int channel, size_t timeout_ms) override;
+  void put_signal(int dst_rank, int channel, size_t timeout_ms) override;
+  void wait_signal(int src_rank, int channel, size_t timeout_ms) override;
+
+  int get_rank() override;
+  int get_world_size() override;
+  c10::Device get_device() override;
+
+  void set_group_name(const std::string& group_name) {
+    group_name_ = group_name;
+  }
+
+ private:
+  std::vector<c10::intrusive_ptr<AllocationRef>> alloc_refs_;
+  std::vector<void*> buffers_;
+  std::vector<void*> signal_pads_;
+  HandleType mc_handle_;
+  void* mc_addr_;
+  size_t buffer_size_;
+  int local_device_idx_;
+  int rank_;
+  int world_size_;
+  void** buffers_dev_;
+  void** signal_pads_dev_;
+  std::string group_name_;
+};
+
+struct Block : public c10::intrusive_ptr_target {
+  c10::intrusive_ptr<AllocationRef> alloc_ref;
+  int device_idx;
+  size_t block_size;
+  size_t buffer_size;
+  size_t signal_pad_offset;
+  std::optional<std::string> default_group_name;
+  std::map<std::string, c10::intrusive_ptr<XPUSymmetricMemory>> symm_mems;
+
+  Block(
+      c10::intrusive_ptr<AllocationRef> alloc_ref,
+      int device_idx,
+      size_t block_size,
+      size_t buffer_size,
+      size_t signal_pad_offset,
+      const std::optional<std::string>& group_name);
+};
+
+class XPUSymmetricMemoryAllocator : public SymmetricMemoryAllocator {
+ public:
+  void* alloc(
+      size_t size,
+      int device_idx,
+      const std::optional<std::string>& group_name) override;
+
+  void free(void* ptr) override;
+  size_t get_alloc_size(void* ptr) override;
+  c10::intrusive_ptr<SymmetricMemory> rendezvous(
+      void* ptr,
+      const std::optional<std::string>& group_name) override;
+  bool has_multicast_support(int device_idx) override;
+  //  void exchange_peer_ipc_mem(sycl::queue& queue, void* ptr);
+  c10::DeviceType supported_device_type() override;
+  std::string name() override;
+
+ private:
+  c10::intrusive_ptr<Block> find_block(void* ptr);
+
+  std::shared_mutex mutex_;
+  std::unordered_map<void*, c10::intrusive_ptr<Block>> ptr_to_block_;
+};
+
+} // namespace c10d::symmetric_memory

src/xccl/XPUSymmetricMemory.cpp

@@ -0,0 +1,8 @@
+#pragma once
+
+namespace c10d::symmetric_memory {
+
+constexpr size_t signal_pad_size = 2048;
+using HandleType = void*;
+
+} // namespace c10d::symmetric_memory

src/xccl/XPUSymmetricMemory.hpp

@@ -0,0 +1,76 @@
+#include <sys/socket.h>
+#include <sys/syscall.h>
+#include <sys/un.h>
+#include <unistd.h>
+
+#include <c10/util/error.h>
+
+#include <c10/xpu/XPUCachingAllocator.h>
+#include <level_zero/ze_api.h>
+#include <sycl/sycl.hpp>
+#include <torch/csrc/distributed/c10d/Store.hpp>
+#include <xccl/XPUSymmetricMemoryUtils.hpp>
+
+namespace c10d::symmetric_memory {
+
+std::string getSymmMemBackendXPU() {
+  static auto val = c10::utils::get_env("TORCH_SYMMMEM");
+  if (val.has_value()) {
+    TORCH_CHECK(
+        val.value() == "XPU",
+        "TORCH_SYMMMEM environment variable must be 'XPU'.");
+    return val.value();
+  }
+  return "XPU";
+}
+
+bool device_has_multicast_support(int device_idx) {
+  return false;
+}
+
+bool allow_overlapping_devices() {
+  return false;
+}
+
+void map_block(
+    void** ptr,
+    ze_physical_mem_handle_t handle,
+    size_t size,
+    int device_idx) {
+  sycl::queue current_queue = at::xpu::getCurrentXPUStream().queue();
+  sycl::context sycl_ctx = current_queue.get_context();
+  ze_context_handle_t ze_context =
+      sycl::get_native<sycl::backend::ext_oneapi_level_zero>(sycl_ctx);
+  // 1. Reserve virtual address space
+  void* virtual_ptr = nullptr;
+  ze_result_t status = zeVirtualMemReserve(
+      ze_context, // context
+      nullptr, // let L0 pick virtual address
+      size, // size
+      &virtual_ptr // out: reserved address
+  );
+  TORCH_CHECK(status == ZE_RESULT_SUCCESS, "zeVirtualMemReserve failed");
+
+  // 2. Map physical memory to virtual address
+  status = zeVirtualMemMap(
+      ze_context,
+      virtual_ptr, // virtual memory to map to
+      size,
+      handle, // physical memory handle
+      0, // flags
+      ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE // ze_memory_access_attribute_t
+  );
+  TORCH_CHECK(status == ZE_RESULT_SUCCESS, "zeVirtualMemMap failed");
+
+  // 3. Set access attributes
+  ze_memory_access_attribute_t access = ZE_MEMORY_ACCESS_ATTRIBUTE_READWRITE;
+  status =
+      zeVirtualMemSetAccessAttribute(ze_context, virtual_ptr, size, access);
+  TORCH_CHECK(
+      status == ZE_RESULT_SUCCESS, "zeVirtualMemSetAccessAttribute failed");
+
+  // 4. Return pointer
+  *ptr = virtual_ptr;
+}
+
+} // namespace c10d::symmetric_memory

src/xccl/XPUSymmetricMemoryTypes.hpp

@@ -0,0 +1,89 @@
+#pragma once
+#include <torch/csrc/distributed/c10d/Store.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
+#include <xccl/XPUSymmetricMemoryTypes.hpp>
+
+namespace c10d {
+namespace symmetric_memory {
+
+std::string getSymmMemBackendXPU();
+
+bool device_has_multicast_support(int device_idx);
+
+bool allow_overlapping_devices();
+
+// A set of store-based exchange methods with a preset prefix typically type of
+// the SymmetricMemory.  Most used as static instances at respective
+// SymmetricMemory implementation files.
+class StoreExchange {
+ public:
+  StoreExchange(const std::string& store_prefix)
+      : store_prefix_(store_prefix) {}
+
+  // Put template function in header file so that compiler can easily access it.
+  template <typename T>
+  std::vector<T> all_gather(
+      const c10::intrusive_ptr<c10d::Store>& store,
+      int rank,
+      int world_size,
+      T val) {
+    static_assert(std::is_trivially_copyable_v<T>);
+
+    std::vector<std::string> peer_keys;
+    peer_keys.reserve(world_size);
+    for (int r = 0; r < world_size; ++r) {
+      std::ostringstream oss;
+      oss << store_prefix_ << "/" << seq_id_ << "/" << r;
+      peer_keys.push_back(oss.str());
+    }
+    ++seq_id_;
+
+    {
+      std::vector<uint8_t> payload(
+          reinterpret_cast<uint8_t*>(&val),
+          reinterpret_cast<uint8_t*>(&val) + sizeof(T));
+      store->set(peer_keys[rank], payload);
+    }
+
+    std::vector<T> peer_vals;
+    peer_vals.reserve(world_size);
+    for (int r = 0; r < world_size; ++r) {
+      if (r == rank) {
+        peer_vals.push_back(val);
+        continue;
+      }
+      store->wait({peer_keys[r]});
+      auto payload = store->get(peer_keys[r]);
+      TORCH_CHECK(payload.size() == sizeof(T));
+      T peer_val{};
+      std::memcpy(&peer_val, payload.data(), sizeof(T));
+      peer_vals.push_back(peer_val);
+    }
+    return peer_vals;
+  }
+
+  void barrier(
+      const c10::intrusive_ptr<c10d::Store>& store,
+      int rank,
+      int world_size) {
+    // TODO: implement an efficient one?
+    all_gather(store, rank, world_size, 0);
+  }
+
+ private:
+  const std::string store_prefix_;
+  size_t seq_id_ = 0;
+};
+
+// Returns a pointer of virtual address that is mapped to the physical memory
+// held by the handle.
+// todo: will follow such physical memory handle map with virtual address,
+// when L0 provides physical handle exchange API and we have multicast support.
+void map_block(
+    void** ptr,
+    ze_physical_mem_handle_t handle,
+    size_t size,
+    int device_idx);
+
+} // namespace symmetric_memory
+} // namespace c10d