Add symmetric memory support on XPU device

src/xccl/CMakeLists.txt

@@ -11,9 +11,11 @@
 file(GLOB xccl_h "*.hpp")
 file(GLOB xccl_cpp "*.cpp")
 list(REMOVE_ITEM xccl_cpp "${CMAKE_CURRENT_SOURCE_DIR}/NanCheck_XPU.cpp")
+list(REMOVE_ITEM xccl_cpp "${CMAKE_CURRENT_SOURCE_DIR}/Signal.cpp")
 
 list(APPEND ATen_XPU_XCCL_SRCS ${xccl_cpp})
 list(APPEND ATen_XPU_SYCL_SRCS "${CMAKE_CURRENT_SOURCE_DIR}/NanCheck_XPU.cpp")
+list(APPEND ATen_XPU_SYCL_SRCS "${CMAKE_CURRENT_SOURCE_DIR}/Signal.cpp")
 
 set(ATen_XPU_XCCL_SRCS ${ATen_XPU_XCCL_SRCS} PARENT_SCOPE)
 set(ATen_XPU_SYCL_SRCS ${ATen_XPU_SYCL_SRCS} PARENT_SCOPE)

src/xccl/Signal.cpp

@@ -0,0 +1,201 @@
+#include <ATen/xpu/XPUContext.h>
+#include <comm/SYCLContext.h>
+#include <xccl/Signal.hpp>
+#include <chrono>
+
+namespace c10d::symmetric_memory {
+
+struct barrierKernel {
+  void operator()(sycl::nd_item<1> item) const {
+    auto thread_id = item.get_local_id(0);
+
+    if (thread_id < world_size) {
+      auto target_rank = thread_id;
+      if (target_rank == rank) {
+        return;
+      }
+      auto put_success = try_put_signal_device<std::memory_order_release>(
+          signal_pads[target_rank] + world_size * channel + rank, timeout_ms);
+      if (!put_success) {
+        SYCL_KERNEL_ASSERT(false);
+      }
+
+      auto wait_success = try_wait_signal_device<std::memory_order_acquire>(
+          signal_pads[rank] + world_size * channel + target_rank, timeout_ms);
+      if (!wait_success) {
+        SYCL_KERNEL_ASSERT(false);
+      }
+    }
+  }
+
+  barrierKernel(
+      uint32_t** signal_pads,
+      int channel,
+      int rank,
+      int world_size,
+      size_t timeout_ms)
+      : signal_pads(signal_pads),
+        channel(channel),
+        rank(rank),
+        world_size(world_size),
+        timeout_ms(timeout_ms) {}
+
+ private:
+  uint32_t** signal_pads;
+  int channel;
+  int rank;
+  int world_size;
+  size_t timeout_ms;
+};
+
+void barrier_impl_xpu(
+    uint32_t** signal_pads,
+    int channel,
+    int rank,
+    int world_size,
+    size_t timeout_ms,
+    at::xpu::XPUStream& stream) {
+  int64_t maxNumThreadsPerBlock = syclMaxWorkGroupSize<barrierKernel>();
+  const size_t numThreadsPerBlock =
+      std::min<size_t>(maxNumThreadsPerBlock, std::max(32, world_size));
+
+  if (!(numThreadsPerBlock > 0)) {
+    return;
+  }
+  int64_t numBlocks = 1;
+  auto global_range = numBlocks * numThreadsPerBlock;
+  auto local_range = numThreadsPerBlock;
+
+  using Kernel = barrierKernel;
+  auto kfn = Kernel(signal_pads, channel, rank, world_size, timeout_ms);
+
+  sycl_kernel_submit(global_range, local_range, stream.queue(), kfn);
+}
+
+struct putSignalKernel {
+  void operator()(sycl::nd_item<1> item) const {
+    auto thread_id = item.get_local_id(0);
+
+    if (thread_id == 0) {
+      auto put_success = try_put_signal_device<std::memory_order_release>(
+          signal_pads[dst_rank] + world_size * channel + rank, 10000000);
+      if (!put_success) {
+        SYCL_KERNEL_ASSERT(false);
+      }
+    }
+  }
+
+  putSignalKernel(
+      uint32_t** signal_pads,
+      int dst_rank,
+      int channel,
+      int rank,
+      int world_size,
+      size_t timeout_ms)
+      : signal_pads(signal_pads),
+        dst_rank(dst_rank),
+        channel(channel),
+        rank(rank),
+        world_size(world_size),
+        timeout_ms(timeout_ms) {}
+
+ private:
+  uint32_t** signal_pads;
+  int dst_rank;
+  int channel;
+  int rank;
+  int world_size;
+  size_t timeout_ms;
+};
+
+void put_signal_impl_xpu(
+    uint32_t** signal_pads,
+    int dst_rank,
+    int channel,
+    int rank,
+    int world_size,
+    size_t timeout_ms,
+    at::xpu::XPUStream& stream) {
+  int64_t maxNumThreadsPerBlock = syclMaxWorkGroupSize<putSignalKernel>();
+  const size_t numThreadsPerBlock = std::min<size_t>(maxNumThreadsPerBlock, 32);
+
+  if (!(numThreadsPerBlock > 0)) {
+    return;
+  }
+
+  int64_t numBlocks = 1;
+  auto global_range = numBlocks * numThreadsPerBlock;
+  auto local_range = numThreadsPerBlock;
+
+  using Kernel = putSignalKernel;
+  auto kfn =
+      Kernel(signal_pads, dst_rank, channel, rank, world_size, timeout_ms);
+
+  sycl_kernel_submit(global_range, local_range, stream.queue(), kfn);
+}
+
+struct waitSignalKernel {
+  void operator()(sycl::nd_item<1> item) const {
+    auto thread_id = item.get_local_id(0);
+
+    if (thread_id == 0) {
+      auto wait_success = try_wait_signal_device<std::memory_order_acquire>(
+          signal_pads[rank] + world_size * channel + src_rank, 10000000);
+      if (!wait_success) {
+        SYCL_KERNEL_ASSERT(false);
+      }
+
+      sycl::atomic_fence(sycl::memory_order_seq_cst, sycl::memory_scope_system);
+    }
+  }
+
+  waitSignalKernel(
+      uint32_t** signal_pads,
+      int src_rank,
+      int channel,
+      int rank,
+      int world_size,
+      size_t timeout_ms)
+      : signal_pads(signal_pads),
+        src_rank(src_rank),
+        channel(channel),
+        rank(rank),
+        world_size(world_size),
+        timeout_ms(timeout_ms) {}
+
+ private:
+  uint32_t** signal_pads;
+  int src_rank;
+  int channel;
+  int rank;
+  int world_size;
+  size_t timeout_ms;
+};
+
+void wait_signal_impl_xpu(
+    uint32_t** signal_pads,
+    int src_rank,
+    int channel,
+    int rank,
+    int world_size,
+    size_t timeout_ms,
+    at::xpu::XPUStream& stream) {
+  int64_t maxNumThreadsPerBlock = syclMaxWorkGroupSize<waitSignalKernel>();
+  const size_t numThreadsPerBlock = std::min<size_t>(maxNumThreadsPerBlock, 32);
+
+  if (!(numThreadsPerBlock > 0)) {
+    return;
+  }
+
+  int64_t numBlocks = 1;
+  auto global_range = numBlocks * numThreadsPerBlock;
+  auto local_range = numThreadsPerBlock;
+
+  using Kernel = waitSignalKernel;
+  auto kfn =
+      Kernel(signal_pads, src_rank, channel, rank, world_size, timeout_ms);
+
+  sycl_kernel_submit(global_range, local_range, stream.queue(), kfn);
+}
+
+} // namespace c10d::symmetric_memory

src/xccl/Signal.hpp

@@ -0,0 +1,88 @@
+#pragma once
+
+#include <atomic>
+
+#include <ATen/native/xpu/sycl/MemoryAccess.h>
+#include <comm/SYCLContext.h>
+
+namespace c10d::symmetric_memory {
+
+using at::native::memory::get_alignment;
+
+// =============================================================================
+// Signal primitives using store/load + atomic_fence
+// (sycl::atomic_ref is not supported, use explicit fence instead)
+// =============================================================================
+
+// Store value with release fence (for put_signal)
+// Order: store first, then release fence to flush the store
+inline void store_release(uint32_t* addr, uint32_t val) {
+  *addr = val;
+  sycl::atomic_fence(sycl::memory_order::release, sycl::memory_scope::system);
+}
+
+// Load value with acquire fence (for get_signal/wait_signal)
+// Order: acquire fence first, then load to see the latest value
+inline uint32_t load_acquire(uint32_t* addr) {
+  sycl::atomic_fence(sycl::memory_order::acquire, sycl::memory_scope::system);
+  uint32_t val = *addr;
+  return val;
+}
+
+// =============================================================================
+// Put signal: wait until addr == 0, then set to 1 (release semantics)
+// =============================================================================
+
+template <std::memory_order Sem>
+bool try_put_signal_device(uint32_t* addr, size_t timeout_ms) {
+  // Wait until the slot is free (value == 0)
+  while (load_acquire(addr) != 0) {
+    // Spin wait (no timeout check as IGC issue)
+    continue;
+  }
+  // Set signal to 1 with release semantics
+  store_release(addr, 1);
+  return true;
+}
+
+// =============================================================================
+// Wait signal: wait until addr == 1, then set to 0 (acquire semantics)
+// =============================================================================
+template <std::memory_order Sem>
+bool try_wait_signal_device(uint32_t* addr, size_t timeout_ms) {
+  // Wait until signal is set (value == 1)
+  while (load_acquire(addr) != 1) {
+    // Spin wait (no timeout check as IGC issue)
+    continue;
+  }
+  // Clear signal to 0 with release semantics
+  store_release(addr, 0);
+  return true;
+}
+
+void barrier_impl_xpu(
+    uint32_t** signal_pads,
+    int channel,
+    int rank,
+    int world_size,
+    size_t timeout_ms,
+    at::xpu::XPUStream& stream);
+
+void put_signal_impl_xpu(
+    uint32_t** signal_pads,
+    int dst_rank,
+    int channel,
+    int rank,
+    int world_size,
+    size_t timeout_ms,
+    at::xpu::XPUStream& stream);
+
+void wait_signal_impl_xpu(
+    uint32_t** signal_pads,
+    int src_rank,
+    int channel,
+    int rank,
+    int world_size,
+    size_t timeout_ms,
+    at::xpu::XPUStream& stream);
+} // namespace c10d::symmetric_memory