EP support for rail isolated topology through Host based RDMA

include/mori/application/application_device_types.hpp

@@ -55,7 +55,7 @@ namespace application {
 /*                                       Transport Types                                          */
 /* ---------------------------------------------------------------------------------------------- */
 
-enum TransportType { RDMA = 0, P2P = 1, SDMA = 2 };
+enum TransportType { RDMA = 0, P2P = 1, SDMA = 2, PROXY = 3 };
 
 // Atomic internal buffer configuration. Defined here (device-safe) rather than in
 // the host transport/rdma/rdma.hpp so device kernels (e.g. shmem_ibgda_kernels) can

include/mori/application/context/context.hpp

@@ -97,7 +97,16 @@ class Context {
 
   RdmaContext* GetRdmaContext() const { return rdmaContext.get(); }
   RdmaDeviceContext* GetRdmaDeviceContext() const { return rdmaDeviceContext.get(); }
+  const std::vector<std::unique_ptr<RdmaDeviceContext>>& GetAllRdmaDeviceContexts() const {
+    return allRdmaDeviceContexts;
+  }
   bool RdmaTransportEnabled() const { return GetRdmaDeviceContext() != nullptr; }
+  RdmaDeviceContext* GetRailContext(int peerRank) const {
+    const int nCtx = static_cast<int>(allRdmaDeviceContexts.size());
+    int peerLocalGpu = peerRank % nCtx;
+    int agreedRail = std::max(LocalRankInNode(), peerLocalGpu) % nCtx;
+    return allRdmaDeviceContexts[agreedRail].get();
+  }
 
   // Check if P2P connection is possible with a peer (same node)
   bool CanUseP2P(int destRank) const;
@@ -112,6 +121,7 @@ class Context {
   // in a test function after the workers had already been spawned.
   bool IsSdmaEnabled() const { return sdmaEnabled; }
   bool IsP2PDisabled() const { return p2pDisabled; }
+  bool IsProxyEnabled() const { return proxyEnabled; }
 
   // Returns the initial RDMA endpoint set. Empty until BuildInitialEndpoints()
   // has been called. SHMEM consumes this set; CCO does not (it creates its own
@@ -188,13 +198,17 @@ class Context {
   // Snapshotted at construction; see IsSdmaEnabled() / IsP2PDisabled() above.
   bool sdmaEnabled{false};
   bool p2pDisabled{false};
+  bool proxyEnabled{false};
   std::string myHostname;
   std::vector<PeerInfo> peerInfos;
   std::vector<PeerCapabilities> peerCaps;     // raw capability discovery
   std::vector<TransportType> transportTypes;  // derived via DefaultPolicyResolve
 
   std::unique_ptr<RdmaContext> rdmaContext{nullptr};
   std::unique_ptr<RdmaDeviceContext> rdmaDeviceContext{nullptr};
+  // One context per available RDMA device/port, indexed by NIC index (0..numNics-1).
+  // QP-to-NIC mapping is via the agreed-rail formula, not a direct index.
+  std::vector<std::unique_ptr<RdmaDeviceContext>> allRdmaDeviceContexts;
 
   std::vector<RdmaEndpoint> rdmaEps;
   bool initialEndpointsBuilt{false};

include/mori/application/memory/symmetric_memory.hpp

@@ -81,6 +81,17 @@ class SymmMemManager {
   SymmMemObjPtr GetVMMHeapObj() const { return vmmHeapObj; }
   size_t GetVMMChunkSize() const { return vmmChunkSize; }
 
+  // Per-NIC keys for send-side routing (proxy mode).
+  // perNicLkeys[nic] = lkey for MY buffer registered on nic's PD.
+  // perNicPeerRkeys[nic][peer] = rkey for peer's buffer registered on nic's PD.
+  std::vector<uint32_t> perNicLkeys;
+  std::vector<std::vector<uint32_t>> perNicPeerRkeys;
+
+  // Cached heap rkeys — sub-allocations reuse these instead of doing
+  // redundant ibv_reg_mr + Allgather for each shmem_malloc.
+  uint32_t heapLkey_{0};
+  std::vector<uint32_t> heapRkeys_;
+
   // Common Utilities
   SymmMemObjPtr Get(void* localPtr) const;
   HeapVAManager* GetHeapVAManager() const { return heapVAManager.get(); }

include/mori/application/transport/rdma/providers/ionic/ionic.hpp

@@ -153,12 +153,21 @@ class IonicDeviceContext : public RdmaDeviceContext {
                                         uint64_t resource_type);
   void create_parent_domain(ibv_context* context, struct ibv_pd* pd_orig);
 
+  struct ProxyRecvInfo { void* buf{nullptr}; uint32_t lkey{0}; uint32_t count{0}; };
+  ProxyRecvInfo GetProxyRecvInfo(uint32_t qpn) {
+    auto it = proxyRecvInfo.find(qpn);
+    return (it != proxyRecvInfo.end()) ? it->second : ProxyRecvInfo{};
+  }
+
  private:
+  bool proxyEnabled{false};
   uint32_t pdn;
   struct ibv_pd* pd_uxdma[2];
 
   std::unordered_map<uint32_t, IonicCqContainer*> cqPool;
   std::unordered_map<uint32_t, IonicQpContainer*> qpPool;
+  std::unordered_map<uint32_t, ibv_qp*> proxyQpPool;
+  std::unordered_map<uint32_t, ProxyRecvInfo> proxyRecvInfo;
 };
 
 class IonicDevice : public RdmaDevice {

include/mori/core/transport/rdma/ibverbs_handle.hpp

@@ -42,6 +42,9 @@ struct IBVerbsHandle {
   ibv_cq* cq{nullptr};
   ibv_srq* srq{nullptr};
   ibv_comp_channel* compCh{nullptr};
+  void* recvBuf{nullptr};
+  uint32_t recvLkey{0};
+  uint32_t recvCount{0};
 };
 
 }  // namespace core

include/mori/core/transport/rdma/proxy/proxy_device_primitives.hpp

@@ -0,0 +1,186 @@
+// Copyright © Advanced Micro Devices, Inc. All rights reserved.
+// MIT License
+#pragma once
+
+#include "mori/core/transport/rdma/proxy/proxy_types.hpp"
+
+#ifdef __HIPCC__
+
+namespace mori {
+namespace core {
+
+// Returns sequence number (monotonically increasing). Mask with PROXY_RING_MASK for slot index.
+inline __device__ uint32_t ProxyReserveSlot(volatile ProxyRing* ring) {
+  return __hip_atomic_fetch_add(
+      (uint32_t*)&ring->gpu_head, 1u, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
+}
+
+inline __device__ void ProxyWaitSlotFree(volatile ProxyRing* ring, uint32_t slot) {
+  int spins = 0;
+  while (true) {
+    uint32_t st = __hip_atomic_load(
+        (uint32_t*)&ring->cmds[slot].status, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_SYSTEM);
+    if (st == PROXY_FREE || st == PROXY_COMPLETED) break;
+    if (++spins % 100000 == 0) __builtin_amdgcn_s_sleep(1);
+  }
+}
+
+inline __device__ void ProxyWaitSlotCompleted(volatile ProxyRing* ring, uint32_t slot) {
+  while (true) {
+    uint32_t st = __hip_atomic_load(
+        (uint32_t*)&ring->cmds[slot].status, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_SYSTEM);
+    if (st == PROXY_COMPLETED || st == PROXY_ERROR) break;
+    __builtin_amdgcn_s_sleep(1);
+  }
+}
+
+inline __device__ uint32_t ProxyPostWrite(
+    volatile ProxyRing* ring, uint32_t qp_idx,
+    uint64_t src_addr, uint32_t lkey,
+    uint64_t dst_addr, uint32_t rkey,
+    uint32_t length) {
+  uint32_t seq = ProxyReserveSlot(ring);
+  uint32_t slot = seq & PROXY_RING_MASK;
+  ProxyWaitSlotFree(ring, slot);
+
+  ring->cmds[slot].op = PROXY_RDMA_WRITE;
+  ring->cmds[slot].qp_idx = qp_idx;
+  ring->cmds[slot].src_addr = src_addr;
+  ring->cmds[slot].dst_addr = dst_addr;
+  ring->cmds[slot].length = length;
+  ring->cmds[slot].lkey = lkey;
+  ring->cmds[slot].rkey = rkey;
+  ring->cmds[slot].flags = 1;
+
+  __threadfence_system();
+  ring->cmds[slot].status = PROXY_PENDING;
+  return seq;
+}
+
+inline __device__ uint32_t ProxyPostWriteInline(
+    volatile ProxyRing* ring, uint32_t qp_idx,
+    uint64_t src_addr, uint32_t lkey,
+    uint64_t dst_addr, uint32_t rkey,
+    uint32_t length) {
+  uint32_t seq = ProxyReserveSlot(ring);
+  uint32_t slot = seq & PROXY_RING_MASK;
+  ProxyWaitSlotFree(ring, slot);
+
+  ring->cmds[slot].op = PROXY_RDMA_WRITE_INLINE;
+  ring->cmds[slot].qp_idx = qp_idx;
+  ring->cmds[slot].src_addr = src_addr;
+  ring->cmds[slot].dst_addr = dst_addr;
+  ring->cmds[slot].length = length;
+  ring->cmds[slot].lkey = lkey;
+  ring->cmds[slot].rkey = rkey;
+  ring->cmds[slot].flags = 1;
+
+  __threadfence_system();
+  ring->cmds[slot].status = PROXY_PENDING;
+  return seq;
+}
+
+inline __device__ uint32_t ProxyPostAtomicNonFetch(
+    volatile ProxyRing* ring, uint32_t qp_idx,
+    uint64_t dst_addr, uint32_t rkey,
+    uint64_t add_value, uint32_t lkey,
+    uint64_t ibuf_addr) {
+  uint32_t seq = ProxyReserveSlot(ring);
+  uint32_t slot = seq & PROXY_RING_MASK;
+  ProxyWaitSlotFree(ring, slot);
+
+  ring->cmds[slot].op = PROXY_ATOMIC_FETCH_ADD;
+  ring->cmds[slot].qp_idx = qp_idx;
+  ring->cmds[slot].src_addr = ibuf_addr;
+  ring->cmds[slot].dst_addr = dst_addr;
+  ring->cmds[slot].length = 8;
+  ring->cmds[slot].lkey = lkey;
+  ring->cmds[slot].rkey = rkey;
+  ring->cmds[slot].atomic_arg = add_value;
+  ring->cmds[slot].flags = 1;
+
+  __threadfence_system();
+  ring->cmds[slot].status = PROXY_PENDING;
+  return seq;
+}
+
+// Signal write: RDMA_WRITE of value to remote addr on the SAME NIC path
+// as the preceding data write. Used for signals paired with data
+// (ShmemPutMemNbiSignalThread) to ensure PCIe write ordering.
+inline __device__ uint32_t ProxyPostSignalWrite(
+    volatile ProxyRing* ring, uint32_t qp_idx,
+    uint64_t dst_addr, uint32_t rkey,
+    uint64_t value, uint32_t lkey,
+    uint64_t ibuf_addr) {
+  uint32_t seq = ProxyReserveSlot(ring);
+  uint32_t slot = seq & PROXY_RING_MASK;
+  ProxyWaitSlotFree(ring, slot);
+
+  ring->cmds[slot].op = PROXY_SIGNAL_WRITE;
+  ring->cmds[slot].qp_idx = qp_idx;
+  ring->cmds[slot].src_addr = ibuf_addr;
+  ring->cmds[slot].dst_addr = dst_addr;
+  ring->cmds[slot].length = 8;
+  ring->cmds[slot].lkey = lkey;
+  ring->cmds[slot].rkey = rkey;
+  ring->cmds[slot].atomic_arg = value;
+  ring->cmds[slot].flags = 1;
+
+  __threadfence_system();
+  ring->cmds[slot].status = PROXY_PENDING;
+  return seq;
+}
+
+inline __device__ uint64_t ProxyPostAtomicFetch(
+    volatile ProxyRing* ring, uint32_t qp_idx,
+    uint64_t dst_addr, uint32_t rkey,
+    uint64_t add_value, uint32_t lkey,
+    uint64_t ibuf_addr) {
+  uint32_t seq = ProxyReserveSlot(ring);
+  uint32_t slot = seq & PROXY_RING_MASK;
+  ProxyWaitSlotFree(ring, slot);
+
+  ring->cmds[slot].op = PROXY_ATOMIC_FETCH_ADD;
+  ring->cmds[slot].qp_idx = qp_idx;
+  ring->cmds[slot].src_addr = ibuf_addr;
+  ring->cmds[slot].dst_addr = dst_addr;
+  ring->cmds[slot].length = 8;
+  ring->cmds[slot].lkey = lkey;
+  ring->cmds[slot].rkey = rkey;
+  ring->cmds[slot].atomic_arg = add_value;
+  ring->cmds[slot].flags = 1;
+  ring->cmds[slot].result = 0;
+
+  __threadfence_system();
+  ring->cmds[slot].status = PROXY_PENDING;
+
+  ProxyWaitSlotCompleted(ring, slot);
+  return ring->cmds[slot].result;
+}
+
+// Wait for all ops from [first_seq, first_seq + count) to complete.
+// When count > PROXY_RING_SIZE, slots were reused during submission.
+// ProxyWaitSlotFree already ensured earlier slots completed before reuse,
+// so we only need to wait for the tail — the last PROXY_RING_SIZE slots.
+inline __device__ void ProxyQuiet(volatile ProxyRing* ring, uint32_t first_seq, uint32_t count) {
+  if (count == 0) return;
+  uint32_t start = first_seq;
+  if (count > PROXY_RING_SIZE) {
+    start = first_seq + count - PROXY_RING_SIZE;
+  }
+  uint32_t end = first_seq + count;
+  for (uint32_t seq = start; seq < end; seq++) {
+    uint32_t slot = seq & PROXY_RING_MASK;
+    ProxyWaitSlotCompleted(ring, slot);
+  }
+}
+
+// Range variant for multi-warp callers that know the exact range.
+inline __device__ void ProxyQuietRange(volatile ProxyRing* ring, uint32_t from, uint32_t to) {
+  ProxyQuiet(ring, from, to - from);
+}
+
+}  // namespace core
+}  // namespace mori
+
+#endif  // __HIPCC__

include/mori/core/transport/rdma/proxy/proxy_thread.hpp

@@ -0,0 +1,62 @@
+// Copyright © Advanced Micro Devices, Inc. All rights reserved.
+// MIT License
+#pragma once
+
+#if !defined(__HIPCC__) && !defined(__CUDACC__)
+
+#include <infiniband/verbs.h>
+#include <pthread.h>
+
+#include <atomic>
+#include <vector>
+
+#include "mori/core/transport/rdma/proxy/proxy_types.hpp"
+
+namespace mori {
+namespace core {
+
+struct InlineBuf {
+  uint64_t data[2];
+};
+
+struct ProxyQpHandle {
+  ibv_qp* qp{nullptr};
+  ibv_cq* cq{nullptr};
+  uint32_t lkey_override{0};
+  uint32_t rkey_override{0};
+  void* recv_buf{nullptr};
+  uint32_t recv_lkey{0};
+  uint32_t recv_count{0};
+};
+
+class ProxyThread {
+ public:
+  ProxyThread() = default;
+  ~ProxyThread();
+
+  void Init(ProxyRing* ring, std::vector<ProxyQpHandle> qps, int gpuId = 0);
+  void Start();
+  void Shutdown();
+
+ private:
+  static void* ThreadFunc(void* arg);
+  void MainLoop();
+  void DrainCq(ProxyQpHandle& qph);
+  bool BuildWr(volatile ProxyCmd* cmd, ProxyQpHandle& qph,
+               ibv_send_wr& wr, ibv_sge& sge, uint32_t slot_id,
+               InlineBuf& ibuf);
+
+  ProxyRing* ring_{nullptr};
+  std::vector<ProxyQpHandle> qps_;
+  pthread_t thread_{};
+  std::atomic<bool> running_{false};
+  uint32_t next_slot_{0};
+  uint64_t ops_posted_{0};
+  uint64_t ops_completed_{0};
+  int gpu_id_{0};
+};
+
+}  // namespace core
+}  // namespace mori
+
+#endif  // !defined(__HIPCC__) && !defined(__CUDACC__)