feat(shmem): add a CPU host-proxy RDMA transport alongside IBGDA

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; }
   bool IsRailOnly() const { return railOnly; }
 
   // Returns the initial RDMA endpoint set. Empty until BuildInitialEndpoints()
@@ -187,6 +197,7 @@ class Context {
   int numQpPerPe{4};
   bool sdmaEnabled{false};
   bool p2pDisabled{false};
+  bool proxyEnabled{false};
   bool railOnly{false};
   std::string myHostname;
   std::vector<PeerInfo> peerInfos;
@@ -195,6 +206,9 @@ class Context {
 
   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/ibverbs/ibverbs.hpp

@@ -39,11 +39,19 @@ class IBVerbsDeviceContext : public RdmaDeviceContext {
                                uint32_t qpId = 0) override;
   bool DestroyRdmaEndpointNoThrow(const RdmaEndpoint&) noexcept override;
 
+  struct ProxyRecvInfo {
+    void* buf;
+    uint32_t lkey;
+    uint32_t count;
+  };
+  ProxyRecvInfo GetProxyRecvInfo(uint32_t qpn) const;
+
  private:
   mutable std::mutex poolMu;
   std::unordered_map<void*, ibv_cq*> cqPool;
   std::unordered_map<uint32_t, ibv_qp*> qpPool;
   std::vector<ibv_comp_channel*> compChPool;
+  std::unordered_map<uint32_t, ProxyRecvInfo> proxyRecvInfo;
 };
 
 class IBVerbsDevice : 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,212 @@
+// Copyright © Advanced Micro Devices, Inc. All rights reserved.
+//
+// MIT License
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+// 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,
+                                                const void* src, 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].dst_addr = dst_addr;
+  ring->cmds[slot].length = length;
+  ring->cmds[slot].lkey = lkey;
+  ring->cmds[slot].rkey = rkey;
+  ring->cmds[slot].flags = 1;
+
+  if (src != nullptr && length > 0 && length <= PROXY_MAX_INLINE_DATA) {
+    const uint64_t* s64 = reinterpret_cast<const uint64_t*>(src);
+    volatile uint64_t* d64 = reinterpret_cast<volatile uint64_t*>(ring->cmds[slot].inline_data);
+    for (uint32_t i = 0; i < (length + 7) / 8; i++) d64[i] = s64[i];
+    ring->cmds[slot].inline_tag = PROXY_INLINE_SCALAR_WRITE;
+    ring->cmds[slot].inline_len = length;
+  } else {
+    ring->cmds[slot].inline_tag = PROXY_INLINE_NONE;
+    ring->cmds[slot].inline_len = 0;
+  }
+
+  __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 = PROXY_FLAGS_DEFAULT;
+
+  __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 = PROXY_FLAGS_FETCH_REQUIRED;
+  // Store slot index in inline_data so the remote can send it back in the reply
+  *reinterpret_cast<volatile uint64_t*>(&ring->cmds[slot].inline_data[0]) =
+      static_cast<uint64_t>(slot);
+  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,86 @@
+// Copyright © Advanced Micro Devices, Inc. All rights reserved.
+//
+// MIT License
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+// 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[4];
+};
+
+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};
+  bool use_native_atomics{false};
+};
+
+class ProxyThread {
+ public:
+  ProxyThread() = default;
+  ~ProxyThread();
+
+  void Init(ProxyRing* ring, std::vector<ProxyQpHandle> qps, int gpuId = 0, uintptr_t heapBase = 0,
+            uintptr_t heapEnd = 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};
+  uintptr_t heap_base_{0};
+  uintptr_t heap_end_{0};
+};
+
+}  // namespace core
+}  // namespace mori
+
+#endif  // !defined(__HIPCC__) && !defined(__CUDACC__)