(cco) impl sdma transport in cco primitives

.github/workflows/ci_cco.yml

@@ -90,6 +90,13 @@ jobs:
             # mpirun --allow-run-as-root -np 2 python 03_flydsl_put/main.py
             # mpirun --allow-run-as-root -np 2 python 06_flydsl_gda_modes/main.py
 
+
+      - name: Run CCO SDMA Python API tests (spawn, up to 8 GPUs)
+        run: |
+          $CT exec $CONTAINER bash -c "\
+            cd $GITHUB_WORKSPACE && \
+            timeout 300 pytest tests/python/cco/test_sdma_api.py -v"
+
       - name: Run ops-v2 dispatch/combine test (EP8, 8 ranks)
         run: |
           $CT exec $CONTAINER bash -c "\

.gitignore

@@ -17,8 +17,10 @@ python/mori/cco/*.so
 python/mori/cco/cco.cpp
 python/mori/umbp_master
 python/mori/spdk_proxy
-python/mori/examples/
-python/mori/benchmarks/
+python/mori/examples/*
+python/mori/benchmarks/*
+python/mori/cco/cco.cpp
+
 
 # local cco experiments (not for commit)
 examples/cco/**/_dump_*.py

benchmark/cco/p2p_get_bw.cpp

@@ -98,6 +98,49 @@ __global__ void ibgda_get_bw(ccoWindowDevice* sendWin, ccoWindowDevice* recvWin,
   gda.flush(ccoCoopBlock{});
 }
 
+// SDMA thread scope: one thread per SDMA queue pulls its slice each iteration;
+// quiet drains completion at the end.
+__global__ void sdma_get_bw(ccoWindowDevice* sendWin, ccoWindowDevice* recvWin, size_t len_doubles,
+                            ccoDevComm devComm, int peerLsa, int iter) {
+  ccoSdma sdma{devComm};
+  const int nq = devComm.sdma.sdmaNumQueue;
+  const int q = threadIdx.x;
+  if (q >= nq) return;
+  const size_t total = len_doubles * sizeof(double);
+  const size_t per = total / static_cast<size_t>(nq);
+  const size_t off = static_cast<size_t>(q) * per;
+  const size_t bytes = (q == nq - 1) ? (total - off) : per;
+  for (int i = 0; i < iter; i++)
+    sdma.get(peerLsa, reinterpret_cast<ccoWindow_t>(recvWin), off,
+             reinterpret_cast<ccoWindow_t>(sendWin), off, bytes, q);
+  sdma.quietQueue(peerLsa, q);
+}
+
+// SDMA warp scope: one warp drives all SDMA queues via a single warp-scope get
+// that splits the transfer across the queue set internally (one lane per queue).
+__global__ void sdma_get_bw_warp(ccoWindowDevice* sendWin, ccoWindowDevice* recvWin,
+                                 size_t len_doubles, ccoDevComm devComm, int peerLsa, int iter) {
+  ccoSdma sdma{devComm};
+  const size_t total = len_doubles * sizeof(double);
+  for (int i = 0; i < iter; i++)
+    sdma.get<ccoCoopWarp>(peerLsa, reinterpret_cast<ccoWindow_t>(recvWin), 0,
+                          reinterpret_cast<ccoWindow_t>(sendWin), 0, total);
+  sdma.quiet<ccoCoopWarp>(peerLsa);
+}
+
+static void launch_sdma(PutScope scope, ccoWindow_t sendWin, ccoWindow_t recvWin,
+                        size_t len_doubles, ccoDevComm devComm, int peerLsa, int count,
+                        int warp_size) {
+  const int nq = devComm.sdma.sdmaNumQueue;
+  if (scope == PutScope::kWarp) {
+    hipLaunchKernelGGL(sdma_get_bw_warp, dim3(1), dim3(warp_size), 0, 0, sendWin, recvWin,
+                       len_doubles, devComm, peerLsa, count);
+  } else {
+    hipLaunchKernelGGL(sdma_get_bw, dim3(1), dim3(nq), 0, 0, sendWin, recvWin, len_doubles, devComm,
+                       peerLsa, count);
+  }
+}
+
 static void launch_lsa(PutScope scope, dim3 grid, dim3 block, ccoWindow_t sendWin,
                        ccoWindow_t recvWin, unsigned int* counter_d, size_t len_doubles,
                        int peerLsa, int count, int warp_size) {
@@ -165,7 +208,9 @@ int main(int argc, char** argv) {
     if (size_bytes % sizeof(double) != 0) continue;
     const size_t len_doubles = size_bytes / sizeof(double);
 
-    if (!size_ok(args.put_scope, size_bytes, args.nblocks, args.threads_per_block,
+    // SDMA splits by queue count and handles uneven tails in-kernel (see put_bw).
+    if (args.transport != Transport::kSdma &&
+        !size_ok(args.put_scope, size_bytes, args.nblocks, args.threads_per_block,
                  ctx.device_warp_size)) {
       if (my_pe == 0) table.push_back(PerfTableRow{size_bytes, true, 0.0});
       ccoBarrierAll(ctx.comm);
@@ -174,7 +219,10 @@ int main(int argc, char** argv) {
 
     if (run_kernels) {
       const float ms = RunWarmupAndTimed(res, args.warmup, args.iters, [&](int count) {
-        if (args.transport == Transport::kLsa) {
+        if (args.transport == Transport::kSdma) {
+          launch_sdma(args.put_scope, ctx.send_win, ctx.recv_win, len_doubles, ctx.devComm,
+                      ctx.peer_lsa_rank, count, ctx.device_warp_size);
+        } else if (args.transport == Transport::kLsa) {
           launch_lsa(args.put_scope, grid, block, ctx.send_win, ctx.recv_win, res.counter_d,
                      len_doubles, ctx.peer_lsa_rank, count, ctx.device_warp_size);
         } else {
@@ -194,10 +242,25 @@ int main(int argc, char** argv) {
 
   ccoBarrierAll(ctx.comm);
   if (my_pe == 0) {
-    PrintPerfTable("p2p_get_bw unidirection", TransportToChar(args.transport),
-                   ScopeToChar(args.put_scope), args.nblocks, args.threads_per_block,
-                   ctx.device_warp_size, args.iters, args.warmup, PerfTableMetric::kBandwidthGbps,
-                   table);
+    // SDMA parallelism is the queue count; see put_bw. thread scope → 1×nq
+    // threads (one message per queue); warp scope → a single warp fanning the
+    // whole transfer across the queues.
+    int print_grid = args.nblocks;
+    int print_block = args.threads_per_block;
+    const char* print_scope = ScopeToChar(args.put_scope);
+    if (args.transport == Transport::kSdma) {
+      print_grid = 1;
+      if (args.put_scope == PutScope::kWarp) {
+        print_block = ctx.device_warp_size;
+        print_scope = "warp";
+      } else {
+        print_block = static_cast<int>(ctx.devComm.sdma.sdmaNumQueue);
+        print_scope = "thread";
+      }
+    }
+    PrintPerfTable("p2p_get_bw unidirection", TransportToChar(args.transport), print_scope,
+                   print_grid, print_block, ctx.device_warp_size, args.iters, args.warmup,
+                   PerfTableMetric::kBandwidthGbps, table);
   }
 
   if (run_kernels) {

benchmark/cco/p2p_get_latency.cpp

@@ -70,6 +70,20 @@ __global__ void ibgda_get_lat(ccoWindowDevice* sendWin, ccoWindowDevice* recvWin
   }
 }
 
+// SDMA: one thread issues a single whole-buffer get on queue 0 + quiet per
+// iteration, so the per-op time tracks the SDMA dispatch + completion round trip.
+__global__ void sdma_get_lat(ccoWindowDevice* sendWin, ccoWindowDevice* recvWin, size_t len_doubles,
+                             ccoDevComm devComm, int peerLsa, int iter) {
+  if (blockIdx.x != 0 || threadIdx.x != 0) return;
+  ccoSdma sdma{devComm};
+  const size_t bytes = len_doubles * sizeof(double);
+  for (int i = 0; i < iter; i++) {
+    sdma.get(peerLsa, reinterpret_cast<ccoWindow_t>(recvWin), 0,
+             reinterpret_cast<ccoWindow_t>(sendWin), 0, bytes, 0);
+    sdma.quietQueue(peerLsa, 0);
+  }
+}
+
 }  // namespace mori::cco::benchmark
 
 int main(int argc, char** argv) {
@@ -114,7 +128,10 @@ int main(int argc, char** argv) {
 
     if (run_kernels) {
       const float ms = RunWarmupAndTimed(res, args.warmup, args.iters, [&](int count) {
-        if (args.transport == Transport::kLsa) {
+        if (args.transport == Transport::kSdma) {
+          hipLaunchKernelGGL(sdma_get_lat, dim3(1), dim3(1), 0, 0, ctx.send_win, ctx.recv_win,
+                             len_doubles, ctx.devComm, ctx.peer_lsa_rank, count);
+        } else if (args.transport == Transport::kLsa) {
           hipLaunchKernelGGL(lsa_get_lat, grid, block, 0, 0, ctx.send_win, ctx.recv_win,
                              len_doubles, ctx.peer_lsa_rank, count);
         } else {
@@ -134,9 +151,16 @@ int main(int argc, char** argv) {
 
   ccoBarrierAll(ctx.comm);
   if (my_pe == 0) {
-    PrintPerfTable("p2p_get_latency unidirection", TransportToChar(args.transport),
-                   ScopeToChar(args.put_scope), 1, block_threads, ctx.device_warp_size, args.iters,
-                   args.warmup, PerfTableMetric::kLatencyUs, table);
+    // SDMA latency uses a single queue / single thread.
+    int print_block = block_threads;
+    const char* print_scope = ScopeToChar(args.put_scope);
+    if (args.transport == Transport::kSdma) {
+      print_block = 1;
+      print_scope = "thread";
+    }
+    PrintPerfTable("p2p_get_latency unidirection", TransportToChar(args.transport), print_scope, 1,
+                   print_block, ctx.device_warp_size, args.iters, args.warmup,
+                   PerfTableMetric::kLatencyUs, table);
   }
 
   if (run_kernels) {

benchmark/cco/p2p_put_bw.cpp

@@ -104,6 +104,49 @@ __global__ void ibgda_put_bw(ccoWindowDevice* sendWin, ccoWindowDevice* recvWin,
   gda.flush(ccoCoopBlock{});
 }
 
+// SDMA thread scope: one thread per SDMA queue copies its slice each iteration;
+// quiet drains completion at the end. Parallelism is the queue count.
+__global__ void sdma_put_bw(ccoWindowDevice* sendWin, ccoWindowDevice* recvWin, size_t len_doubles,
+                            ccoDevComm devComm, int peerLsa, int iter) {
+  ccoSdma sdma{devComm};
+  const int nq = devComm.sdma.sdmaNumQueue;
+  const int q = threadIdx.x;
+  if (q >= nq) return;
+  const size_t total = len_doubles * sizeof(double);
+  const size_t per = total / static_cast<size_t>(nq);
+  const size_t off = static_cast<size_t>(q) * per;
+  const size_t bytes = (q == nq - 1) ? (total - off) : per;
+  for (int i = 0; i < iter; i++)
+    sdma.put(peerLsa, reinterpret_cast<ccoWindow_t>(recvWin), off,
+             reinterpret_cast<ccoWindow_t>(sendWin), off, bytes, q);
+  sdma.quietQueue(peerLsa, q);
+}
+
+// SDMA warp scope: one warp drives all SDMA queues via a single warp-scope put
+// that splits the transfer across the queue set internally (one lane per queue).
+__global__ void sdma_put_bw_warp(ccoWindowDevice* sendWin, ccoWindowDevice* recvWin,
+                                 size_t len_doubles, ccoDevComm devComm, int peerLsa, int iter) {
+  ccoSdma sdma{devComm};
+  const size_t total = len_doubles * sizeof(double);
+  for (int i = 0; i < iter; i++)
+    sdma.put<ccoCoopWarp>(peerLsa, reinterpret_cast<ccoWindow_t>(recvWin), 0,
+                          reinterpret_cast<ccoWindow_t>(sendWin), 0, total);
+  sdma.quiet<ccoCoopWarp>(peerLsa);
+}
+
+static void launch_sdma(PutScope scope, ccoWindow_t sendWin, ccoWindow_t recvWin,
+                        size_t len_doubles, ccoDevComm devComm, int peerLsa, int count,
+                        int warp_size) {
+  const int nq = devComm.sdma.sdmaNumQueue;
+  if (scope == PutScope::kWarp) {
+    hipLaunchKernelGGL(sdma_put_bw_warp, dim3(1), dim3(warp_size), 0, 0, sendWin, recvWin,
+                       len_doubles, devComm, peerLsa, count);
+  } else {
+    hipLaunchKernelGGL(sdma_put_bw, dim3(1), dim3(nq), 0, 0, sendWin, recvWin, len_doubles, devComm,
+                       peerLsa, count);
+  }
+}
+
 static void launch_lsa(PutScope scope, dim3 grid, dim3 block, ccoWindow_t sendWin,
                        ccoWindow_t recvWin, unsigned int* counter_d, size_t len_doubles,
                        int peerLsa, int count, int warp_size) {
@@ -173,7 +216,10 @@ int main(int argc, char** argv) {
     if (size_bytes % sizeof(double) != 0) continue;
     const size_t len_doubles = size_bytes / sizeof(double);
 
-    if (!size_ok(args.put_scope, size_bytes, args.nblocks, args.threads_per_block,
+    // SDMA splits by queue count and handles uneven tails in-kernel, so the
+    // scope/nblocks divisibility rules don't apply — only require 8B alignment.
+    if (args.transport != Transport::kSdma &&
+        !size_ok(args.put_scope, size_bytes, args.nblocks, args.threads_per_block,
                  ctx.device_warp_size)) {
       if (my_pe == 0) table.push_back(PerfTableRow{size_bytes, true, 0.0});
       ccoBarrierAll(ctx.comm);
@@ -182,7 +228,10 @@ int main(int argc, char** argv) {
 
     if (run_kernels) {
       const float ms = RunWarmupAndTimed(res, args.warmup, args.iters, [&](int count) {
-        if (args.transport == Transport::kLsa) {
+        if (args.transport == Transport::kSdma) {
+          launch_sdma(args.put_scope, ctx.send_win, ctx.recv_win, len_doubles, ctx.devComm,
+                      ctx.peer_lsa_rank, count, ctx.device_warp_size);
+        } else if (args.transport == Transport::kLsa) {
           launch_lsa(args.put_scope, grid, block, ctx.send_win, ctx.recv_win, res.counter_d,
                      len_doubles, ctx.peer_lsa_rank, count, ctx.device_warp_size);
         } else {
@@ -202,10 +251,25 @@ int main(int argc, char** argv) {
 
   ccoBarrierAll(ctx.comm);
   if (my_pe == 0) {
-    PrintPerfTable("p2p_put_bw unidirection", TransportToChar(args.transport),
-                   ScopeToChar(args.put_scope), args.nblocks, args.threads_per_block,
-                   ctx.device_warp_size, args.iters, args.warmup, PerfTableMetric::kBandwidthGbps,
-                   table);
+    // SDMA parallelism is the queue count. thread scope launches 1 block × nq
+    // threads (one message per queue → units=nq); warp scope launches a single
+    // warp that fans the whole transfer across the queues (one logical op).
+    int print_grid = args.nblocks;
+    int print_block = args.threads_per_block;
+    const char* print_scope = ScopeToChar(args.put_scope);
+    if (args.transport == Transport::kSdma) {
+      print_grid = 1;
+      if (args.put_scope == PutScope::kWarp) {
+        print_block = ctx.device_warp_size;
+        print_scope = "warp";
+      } else {
+        print_block = static_cast<int>(ctx.devComm.sdma.sdmaNumQueue);
+        print_scope = "thread";
+      }
+    }
+    PrintPerfTable("p2p_put_bw unidirection", TransportToChar(args.transport), print_scope,
+                   print_grid, print_block, ctx.device_warp_size, args.iters, args.warmup,
+                   PerfTableMetric::kBandwidthGbps, table);
   }
 
   if (run_kernels) {

benchmark/cco/p2p_put_latency.cpp

@@ -73,6 +73,20 @@ __global__ void ibgda_put_lat(ccoWindowDevice* sendWin, ccoWindowDevice* recvWin
   }
 }
 
+// SDMA: one thread issues a single whole-buffer put on queue 0 + quiet per
+// iteration, so the per-op time tracks the SDMA dispatch + completion round trip.
+__global__ void sdma_put_lat(ccoWindowDevice* sendWin, ccoWindowDevice* recvWin, size_t len_doubles,
+                             ccoDevComm devComm, int peerLsa, int iter) {
+  if (blockIdx.x != 0 || threadIdx.x != 0) return;
+  ccoSdma sdma{devComm};
+  const size_t bytes = len_doubles * sizeof(double);
+  for (int i = 0; i < iter; i++) {
+    sdma.put(peerLsa, reinterpret_cast<ccoWindow_t>(recvWin), 0,
+             reinterpret_cast<ccoWindow_t>(sendWin), 0, bytes, 0);
+    sdma.quietQueue(peerLsa, 0);
+  }
+}
+
 }  // namespace mori::cco::benchmark
 
 int main(int argc, char** argv) {
@@ -117,7 +131,10 @@ int main(int argc, char** argv) {
 
     if (run_kernels) {
       const float ms = RunWarmupAndTimed(res, args.warmup, args.iters, [&](int count) {
-        if (args.transport == Transport::kLsa) {
+        if (args.transport == Transport::kSdma) {
+          hipLaunchKernelGGL(sdma_put_lat, dim3(1), dim3(1), 0, 0, ctx.send_win, ctx.recv_win,
+                             len_doubles, ctx.devComm, ctx.peer_lsa_rank, count);
+        } else if (args.transport == Transport::kLsa) {
           hipLaunchKernelGGL(lsa_put_lat, grid, block, 0, 0, ctx.send_win, ctx.recv_win,
                              len_doubles, ctx.peer_lsa_rank, count);
         } else {
@@ -137,9 +154,16 @@ int main(int argc, char** argv) {
 
   ccoBarrierAll(ctx.comm);
   if (my_pe == 0) {
-    PrintPerfTable("p2p_put_latency unidirection", TransportToChar(args.transport),
-                   ScopeToChar(args.put_scope), 1, block_threads, ctx.device_warp_size, args.iters,
-                   args.warmup, PerfTableMetric::kLatencyUs, table);
+    // SDMA latency uses a single queue / single thread.
+    int print_block = block_threads;
+    const char* print_scope = ScopeToChar(args.put_scope);
+    if (args.transport == Transport::kSdma) {
+      print_block = 1;
+      print_scope = "thread";
+    }
+    PrintPerfTable("p2p_put_latency unidirection", TransportToChar(args.transport), print_scope, 1,
+                   print_block, ctx.device_warp_size, args.iters, args.warmup,
+                   PerfTableMetric::kLatencyUs, table);
   }
 
   if (run_kernels) {