From 1a134bc460dab381692a1f68e729b06d163b0c1f Mon Sep 17 00:00:00 2001 From: qizzhang Date: Sun, 5 Jul 2026 21:45:23 -0500 Subject: [PATCH 01/10] add initial sdma impl --- benchmark/cco/p2p_get_bw.cpp | 74 ++++++- benchmark/cco/p2p_get_latency.cpp | 32 ++- benchmark/cco/p2p_put_bw.cpp | 75 ++++++- benchmark/cco/p2p_put_latency.cpp | 32 ++- benchmark/cco/util.cpp | 107 ++++++++-- benchmark/cco/util.hpp | 12 +- include/mori/cco/cco.hpp | 84 ++++++++ .../core/transport/sdma/device_primitives.hpp | 20 +- tests/cpp/cco/test_sdma_get.cpp | 186 ++++++++++++++++++ tests/cpp/cco/test_sdma_put.cpp | 184 +++++++++++++++++ 10 files changed, 752 insertions(+), 54 deletions(-) create mode 100644 tests/cpp/cco/test_sdma_get.cpp create mode 100644 tests/cpp/cco/test_sdma_put.cpp diff --git a/benchmark/cco/p2p_get_bw.cpp b/benchmark/cco/p2p_get_bw.cpp index 22342bfb1..533ae162d 100644 --- a/benchmark/cco/p2p_get_bw.cpp +++ b/benchmark/cco/p2p_get_bw.cpp @@ -98,6 +98,48 @@ __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(nq); + const size_t off = static_cast(q) * per; + const size_t bytes = (q == nq - 1) ? (total - off) : per; + for (int i = 0; i < iter; i++) + sdma.get(peerLsa, reinterpret_cast(recvWin), off, + reinterpret_cast(sendWin), off, bytes, q); + sdma.quiet(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(peerLsa, reinterpret_cast(recvWin), 0, + reinterpret_cast(sendWin), 0, total); + sdma.quiet(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 +207,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 +218,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 +241,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(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) { diff --git a/benchmark/cco/p2p_get_latency.cpp b/benchmark/cco/p2p_get_latency.cpp index ce1544120..6076cd9c0 100644 --- a/benchmark/cco/p2p_get_latency.cpp +++ b/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(recvWin), 0, + reinterpret_cast(sendWin), 0, bytes, 0); + sdma.quiet(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) { diff --git a/benchmark/cco/p2p_put_bw.cpp b/benchmark/cco/p2p_put_bw.cpp index 09e185c06..1d1ceaffd 100644 --- a/benchmark/cco/p2p_put_bw.cpp +++ b/benchmark/cco/p2p_put_bw.cpp @@ -104,6 +104,48 @@ __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(nq); + const size_t off = static_cast(q) * per; + const size_t bytes = (q == nq - 1) ? (total - off) : per; + for (int i = 0; i < iter; i++) + sdma.put(peerLsa, reinterpret_cast(recvWin), off, + reinterpret_cast(sendWin), off, bytes, q); + sdma.quiet(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(peerLsa, reinterpret_cast(recvWin), 0, + reinterpret_cast(sendWin), 0, total); + sdma.quiet(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 +215,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 +227,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 +250,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(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) { diff --git a/benchmark/cco/p2p_put_latency.cpp b/benchmark/cco/p2p_put_latency.cpp index a0f4427ce..fd8cce162 100644 --- a/benchmark/cco/p2p_put_latency.cpp +++ b/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(recvWin), 0, + reinterpret_cast(sendWin), 0, bytes, 0); + sdma.quiet(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) { diff --git a/benchmark/cco/util.cpp b/benchmark/cco/util.cpp index cf9c32e37..8d3af5ae8 100644 --- a/benchmark/cco/util.cpp +++ b/benchmark/cco/util.cpp @@ -36,16 +36,22 @@ namespace mori::cco::benchmark { void PrintUsage(const char* program) { std::fprintf(stderr, "Usage: %s [options]\n" - " transport is selected by env MORI_DISABLE_P2P:\n" - " unset / on / 1 -> ibgda (RDMA) [default]\n" - " off / 0 / false -> lsa (intra-node P2P)\n" + " -t transport lsa | sdma | ibgda (default ibgda)\n" + " sets the required env internally:\n" + " sdma -> MORI_ENABLE_SDMA=1\n" + " ibgda -> MORI_DISABLE_P2P=1\n" + " lsa -> MORI_DISABLE_P2P=0\n" + " when -t is omitted the transport falls back to env:\n" + " MORI_ENABLE_SDMA=1 -> sdma [highest priority]\n" + " MORI_DISABLE_P2P unset / on / 1 -> ibgda [default]\n" + " MORI_DISABLE_P2P off / 0 / false -> lsa\n" " -b min_bytes minimum message size\n" " -e max_bytes maximum message size\n" " -f step multiply size by this factor each step\n" " -n iters timed iterations\n" " -w warmup warmup iterations\n" " -c grid_x HIP grid x (blocks)\n" - " -t threads threads per block\n" + " -T threads threads per block\n" " -s scope thread | warp | block | thread_agg (default block)\n" " thread_agg = thread scope + ThreadAggregate (bw only)\n" " -h this help\n", @@ -79,7 +85,7 @@ int ParseArgs(int argc, char** argv, PerfArgs* out_args) { }; int opt = 0; - while ((opt = getopt(argc, argv, "hb:e:f:n:w:c:t:s:")) != -1) { + while ((opt = getopt(argc, argv, "hb:e:f:n:w:c:T:t:s:")) != -1) { switch (opt) { case 'h': return 2; @@ -101,9 +107,21 @@ int ParseArgs(int argc, char** argv, PerfArgs* out_args) { case 'c': out_args->nblocks = std::atoi(optarg); break; - case 't': + case 'T': out_args->threads_per_block = std::atoi(optarg); break; + case 't': + if (std::strcmp(optarg, "lsa") == 0) { + out_args->transport = Transport::kLsa; + } else if (std::strcmp(optarg, "sdma") == 0) { + out_args->transport = Transport::kSdma; + } else if (std::strcmp(optarg, "ibgda") == 0) { + out_args->transport = Transport::kIbgda; + } else { + return 1; + } + out_args->transport_explicit = true; + break; case 's': if (std::strcmp(optarg, "thread") == 0) { out_args->put_scope = PutScope::kThread; @@ -242,16 +260,40 @@ int PerfInit(int argc, char** argv, PerfContext* ctx) { args.min_size = (args.min_size + sizeof(double) - 1) / sizeof(double) * sizeof(double); } - // Transport from MORI_DISABLE_P2P. Default (unset) is IBGDA (P2P disabled); - // an explicit false value ("0"/"false"/"off"/"no") selects LSA. This mirrors - // mori::env::IsEnvVarEnabled's parsing but defaults to enabled when unset. - { - const char* v = std::getenv("MORI_DISABLE_P2P"); - bool p2p_disabled = true; // default: IBGDA - if (v != nullptr && v[0] != '\0') { - p2p_disabled = env::detail::ParseBool(v).value_or(true); + // Transport selection. When -t is given it wins and we export the env the CCO + // comm keys off (MORI_ENABLE_SDMA gates SDMA-queue build in ccoCommCreate; + // MORI_DISABLE_P2P chooses IBGDA vs LSA) before ccoCommCreate runs below. + // Without -t we fall back to those same env vars: MORI_ENABLE_SDMA takes + // precedence, else MORI_DISABLE_P2P picks IBGDA (default) vs LSA. + if (args.transport_explicit) { + switch (args.transport) { + case Transport::kSdma: + setenv("MORI_ENABLE_SDMA", "1", 1); + break; + case Transport::kIbgda: + setenv("MORI_ENABLE_SDMA", "0", 1); + setenv("MORI_DISABLE_P2P", "1", 1); + break; + case Transport::kLsa: + setenv("MORI_ENABLE_SDMA", "0", 1); + setenv("MORI_DISABLE_P2P", "0", 1); + break; + } + } else { + bool sdma_enabled = false; + if (const char* sv = std::getenv("MORI_ENABLE_SDMA")) { + if (sv[0] != '\0') sdma_enabled = env::detail::ParseBool(sv).value_or(false); + } + if (sdma_enabled) { + args.transport = Transport::kSdma; + } else { + const char* v = std::getenv("MORI_DISABLE_P2P"); + bool p2p_disabled = true; // default: IBGDA + if (v != nullptr && v[0] != '\0') { + p2p_disabled = env::detail::ParseBool(v).value_or(true); + } + args.transport = p2p_disabled ? Transport::kIbgda : Transport::kLsa; } - args.transport = p2p_disabled ? Transport::kIbgda : Transport::kLsa; } // Local communicator → local rank → device binding. CCO pins the device at @@ -313,7 +355,10 @@ int PerfInit(int argc, char** argv, PerfContext* ctx) { // DevComm tuned to the chosen transport. ccoDevCommRequirements reqs = CCO_DEV_COMM_REQUIREMENTS_INITIALIZER; - if (args.transport == Transport::kLsa) { + if (args.transport == Transport::kLsa || args.transport == Transport::kSdma) { + // Both intra-node backends need no GDA connectivity. The SDMA signal pool is + // materialized by ccoDevCommCreate whenever the comm has SDMA queues (built + // in ccoCommCreate when MORI_ENABLE_SDMA is on and peers are canSDMA). reqs.gdaConnectionType = CCO_GDA_CONNECTION_NONE; reqs.gdaSignalCount = 0; reqs.gdaCounterCount = 0; @@ -338,13 +383,21 @@ int PerfInit(int argc, char** argv, PerfContext* ctx) { } // Transport feasibility checks. - if (args.transport == Transport::kLsa) { + if (args.transport == Transport::kLsa || args.transport == Transport::kSdma) { if (ctx->devComm.lsaSize < 2) { if (ctx->my_pe == 0) { std::fprintf(stderr, - "LSA transport requires both PEs on the same node (lsaSize=%d). " - "Run -T ibgda for cross-node.\n", - ctx->devComm.lsaSize); + "%s transport requires both PEs on the same node (lsaSize=%d).\n", + TransportToChar(args.transport), ctx->devComm.lsaSize); + } + PerfFinalize(ctx); + return 1; + } + if (args.transport == Transport::kSdma && ctx->devComm.sdma.sdmaNumQueue == 0) { + if (ctx->my_pe == 0) { + std::fprintf(stderr, + "SDMA transport has no queues — set MORI_ENABLE_SDMA=1 and ensure peers are " + "SDMA-capable.\n"); } PerfFinalize(ctx); return 1; @@ -366,7 +419,12 @@ int PerfInit(int argc, char** argv, PerfContext* ctx) { // Announce the resolved transport up front (PE 0 only) so the run is // self-identifying without parsing the table header. if (ctx->my_pe == 0) { - if (args.transport == Transport::kLsa) { + if (args.transport == Transport::kSdma) { + std::printf( + "[cco-bench] transport = SDMA (intra-node copy engine via ccoSdma; lsaSize=%d, " + "numQueue=%u)\n", + ctx->devComm.lsaSize, ctx->devComm.sdma.sdmaNumQueue); + } else if (args.transport == Transport::kLsa) { std::printf("[cco-bench] transport = LSA (intra-node P2P, flat-VA load/store; lsaSize=%d)\n", ctx->devComm.lsaSize); } else { @@ -378,6 +436,13 @@ int PerfInit(int argc, char** argv, PerfContext* ctx) { std::fflush(stdout); } + // CCO setup (window P2P import / DevComm creation) may call + // hipDeviceEnablePeerAccess on links already enabled, returning the benign + // hipErrorPeerAccessAlreadyEnabled and leaving it as the sticky last error. + // Clear it so the first HIP_RUNTIME_CHECK(hipGetLastError()) after a kernel + // launch doesn't misfire on a leftover setup error. + (void)hipGetLastError(); + return 0; } diff --git a/benchmark/cco/util.hpp b/benchmark/cco/util.hpp index 06b96f2b0..2aa2e27b5 100644 --- a/benchmark/cco/util.hpp +++ b/benchmark/cco/util.hpp @@ -27,6 +27,7 @@ #include #include +#include #include #include #include @@ -50,7 +51,8 @@ enum class PutScope { kThread, kWarp, kBlock, kThreadAgg }; // disabled, use RDMA); explicitly disabled (0/false/off/no) → kLsa. // kLsa — intra-node flat-VA load/store (no NIC), requires same-node peer. // kIbgda — cross-node one-sided RDMA via ccoGda. -enum class Transport { kLsa, kIbgda }; +// kSdma — intra-node SDMA copy-engine via ccoSdma; selected by MORI_ENABLE_SDMA. +enum class Transport { kLsa, kIbgda, kSdma }; inline constexpr std::size_t kDefaultMinSize = 8; inline constexpr std::size_t kDefaultMaxSize = 64ULL * 1024ULL * 1024ULL; @@ -78,8 +80,12 @@ struct PerfArgs { int nblocks = kDefaultNumBlocks; int threads_per_block = kDefaultThreadsPerBlock; PutScope put_scope = PutScope::kBlock; - // Default IBGDA; overridden in PerfInit from MORI_DISABLE_P2P. + // Default IBGDA; overridden by the -t flag or, when -t is absent, in PerfInit + // from MORI_ENABLE_SDMA / MORI_DISABLE_P2P. Transport transport = Transport::kIbgda; + // Set when -t explicitly picks a transport: PerfInit then exports the matching + // env (MORI_ENABLE_SDMA / MORI_DISABLE_P2P) instead of auto-detecting. + bool transport_explicit = false; }; // Holds MPI + CCO state for the lifetime of a benchmark run. PerfInit registers @@ -163,6 +169,8 @@ inline const char* TransportToChar(Transport t) { return "lsa"; case Transport::kIbgda: return "ibgda"; + case Transport::kSdma: + return "sdma"; } return "none"; } diff --git a/include/mori/cco/cco.hpp b/include/mori/cco/cco.hpp index 777116e67..4c2b1a7ba 100644 --- a/include/mori/cco/cco.hpp +++ b/include/mori/cco/cco.hpp @@ -877,3 +877,87 @@ int ccoBarrierAll(ccoComm* comm); } // namespace cco } // namespace mori + +/* ════════════════════════════════════════════════════════════════════════════ + * SDMA (intra-node copy-engine) device session. + * + * Moves data via the SDMA queues on ccoDevComm::sdma; peers are addressed by + * LSA rank. put/get are non-blocking (single thread per (peer, queueId)); + * completion is recorded in the local signal pool and awaited with quiet(). + * Device-only — pulls in the SDMA core primitives, so it lives behind the same + * __HIPCC__ guard and is skipped entirely by host translation units. + * ════════════════════════════════════════════════════════════════════════════ */ +#if defined(__HIPCC__) || defined(__CUDACC__) +#include "mori/core/transport/sdma/device_primitives.hpp" + +namespace mori { +namespace cco { + +struct ccoSdma { + ccoDevComm const& comm; + + __device__ inline ccoSdma(ccoDevComm const& c) : comm(c) {} + + // put: copy local srcWin[srcOffset] -> peer dstWin[dstOffset] (bytes). + // Coop=ccoCoopThread — a single thread drives one queue (queueId). + // Coop=ccoCoopWarp — a warp drives all queues, one lane per queue; bytes + // are split across sdmaNumQueue and queueId is ignored. + template + __device__ inline void put(int peer, ccoWindow_t dstWin, size_t dstOffset, ccoWindow_t srcWin, + size_t srcOffset, size_t bytes, int queueId = 0) { + static_assert(std::is_same_v || std::is_same_v, + "ccoSdma::put supports ccoCoopThread or ccoCoopWarp"); + const ccoSdmaContext& s = comm.sdma; + const uint32_t n = s.sdmaNumQueue; + void* dst = ccoGetLsaPeerPtr(dstWin, peer, dstOffset); + void* src = ccoGetLocalPtr(srcWin, srcOffset); + if constexpr (std::is_same_v) { + // Appends COPY + ATOMIC(signal++) on the (peer,queueId) queue, bumps expect. + core::SdmaPutThread(src, dst, bytes, s.deviceHandles + peer * n, s.signalBuf + peer * n, + s.expectSignals + peer * n, n, queueId); + } else { + core::SdmaPutWarp(src, dst, bytes, s.deviceHandles + peer * n, s.signalBuf + peer * n, + s.expectSignals + peer * n, n); + } + } + + // get: copy peer srcWin[srcOffset] -> local dstWin[dstOffset] (bytes). Coop + // selects thread (single queue) vs warp (all queues) scope, as in put(). + template + __device__ inline void get(int peer, ccoWindow_t dstWin, size_t dstOffset, ccoWindow_t srcWin, + size_t srcOffset, size_t bytes, int queueId = 0) { + static_assert(std::is_same_v || std::is_same_v, + "ccoSdma::get supports ccoCoopThread or ccoCoopWarp"); + const ccoSdmaContext& s = comm.sdma; + const uint32_t n = s.sdmaNumQueue; + void* dst = ccoGetLocalPtr(dstWin, dstOffset); + void* src = ccoGetLsaPeerPtr(srcWin, peer, srcOffset); + if constexpr (std::is_same_v) { + core::SdmaPutThread(src, dst, bytes, s.deviceHandles + peer * n, s.signalBuf + peer * n, + s.expectSignals + peer * n, n, queueId); + } else { + core::SdmaPutWarp(src, dst, bytes, s.deviceHandles + peer * n, s.signalBuf + peer * n, + s.expectSignals + peer * n, n); + } + } + + // quiet: block until outstanding ops have landed. Coop=ccoCoopThread waits on + // the single (peer,queueId) queue; Coop=ccoCoopWarp waits on every queue (one + // lane per queue), pairing with a warp-scope put/get. + template + __device__ inline void quiet(int peer, int queueId = 0) { + static_assert(std::is_same_v || std::is_same_v, + "ccoSdma::quiet supports ccoCoopThread or ccoCoopWarp"); + const ccoSdmaContext& s = comm.sdma; + const uint32_t n = s.sdmaNumQueue; + if constexpr (std::is_same_v) { + anvil::waitForSignal(s.signalBuf + peer * n + queueId, s.expectSignals[peer * n + queueId]); + } else { + core::SdmaQueitWarp(s.signalBuf + peer * n, s.expectSignals + peer * n, n); + } + } +}; + +} // namespace cco +} // namespace mori +#endif // __HIPCC__ || __CUDACC__ diff --git a/include/mori/core/transport/sdma/device_primitives.hpp b/include/mori/core/transport/sdma/device_primitives.hpp index f3b670625..8ac986bdb 100644 --- a/include/mori/core/transport/sdma/device_primitives.hpp +++ b/include/mori/core/transport/sdma/device_primitives.hpp @@ -68,31 +68,29 @@ inline __device__ void SdmaPutWarp(void* srcBuf, void* dstBuf, size_t copy_size, uint64_t base = 0; uint64_t pendingWptr = 0; uint64_t startBase = 0; - size_t perq_send_size = 0; uint64_t offset = 0; const int laneId = threadIdx.x % warpSize; if (laneId >= queNum) return; int queueId = laneId; - const size_t rand_size = copy_size / queNum; // per queue rand data + const size_t rand_size = copy_size / queNum; // per queue slice size - char* srcPtr = reinterpret_cast(srcBuf); - char* dstPtr = reinterpret_cast(dstBuf); + // Each queue owns the contiguous slice [queueId*rand_size, ...); the last + // queue absorbs the remainder so uneven sizes are fully covered. + const size_t perq_send_size = + (queueId < (queNum - 1)) ? rand_size : (copy_size - (queNum - 1) * rand_size); + const size_t byteOffset = static_cast(queueId) * rand_size; + + char* srcPtr = reinterpret_cast(srcBuf) + byteOffset; + char* dstPtr = reinterpret_cast(dstBuf) + byteOffset; anvil::SdmaQueueDeviceHandle handle = **(deviceHandles + queueId); base = handle.ReserveQueueSpace(sizeof(SDMA_PKT_COPY_LINEAR), offset); pendingWptr = base; startBase = base; - if (queueId < (queNum - 1)) - perq_send_size = rand_size; - else - perq_send_size = copy_size - (queNum - 1) * rand_size; - auto packet_d = anvil::CreateCopyPacket(srcPtr, dstPtr, perq_send_size); handle.template placePacket(packet_d, pendingWptr, offset); - srcPtr += perq_send_size; - dstPtr += perq_send_size; base = handle.ReserveQueueSpace(sizeof(SDMA_PKT_ATOMIC), offset); pendingWptr = base; diff --git a/tests/cpp/cco/test_sdma_get.cpp b/tests/cpp/cco/test_sdma_get.cpp new file mode 100644 index 000000000..a7a3a6242 --- /dev/null +++ b/tests/cpp/cco/test_sdma_get.cpp @@ -0,0 +1,186 @@ +// 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. +// +// test: cco sdma get — intra-node copy-engine (ccoSdma::get), single node. +// +// alltoall via SDMA get. Layout (owner, dest-slot, idx): +// sendBuf[p*COUNT + i] = myRank*1000 + p*100 + i +// Each rank pulls peer p's slice destined for it into recv slot p: +// peer p's sendBuf[myRank*COUNT+i] = p*1000 + myRank*100 + i +// recv[p*COUNT+i] == p*1000 + myRank*100 + i +// +// One thread per peer (distinct SDMA queue per peer); quiet() drains completion. +// Requires MORI_ENABLE_SDMA=1; otherwise the comm has no SDMA queues and the +// test SKIPs. + +#include "cco_test_harness.hpp" + +static const size_t PER_RANK_VMM_SIZE = 256ULL * 1024 * 1024; +static const size_t COUNT = 256; // elements per rank-pair + +// thread scope: thread p pulls peer p's slice destined for us into recv slot p, +// over peer p's queue 0. +__global__ void SdmaGetKernel(mori::cco::ccoWindowDevice* sendWin, + mori::cco::ccoWindowDevice* recvWin, size_t count, + mori::cco::ccoDevComm devComm) { + using namespace mori::cco; + ccoSdma sdma{devComm}; + int myRank = devComm.rank; + int nRanks = devComm.lsaSize; + size_t perPair = count * sizeof(float); + + int p = threadIdx.x; + if (p >= nRanks || p == myRank) return; + sdma.get(p, reinterpret_cast(recvWin), p * perPair, + reinterpret_cast(sendWin), myRank * perPair, perPair); + sdma.quiet(p); +} + +// warp scope: one warp (block) per peer p pulls peer p's slice into recv slot p, +// with the transfer split across all of peer p's SDMA queues internally. +__global__ void SdmaGetWarpKernel(mori::cco::ccoWindowDevice* sendWin, + mori::cco::ccoWindowDevice* recvWin, size_t count, + mori::cco::ccoDevComm devComm) { + using namespace mori::cco; + ccoSdma sdma{devComm}; + int myRank = devComm.rank; + int nRanks = devComm.lsaSize; + size_t perPair = count * sizeof(float); + + int p = blockIdx.x; // one warp per peer + if (p >= nRanks || p == myRank) return; + sdma.get(p, reinterpret_cast(recvWin), p * perPair, + reinterpret_cast(sendWin), myRank * perPair, perPair); + sdma.quiet(p); +} + +int run_test(int rank, int nranks, const mori::cco::ccoUniqueId& uid) { + g_rank = rank; + + int numDevices = 0; + HIP_CHECK(hipGetDeviceCount(&numDevices)); + int dev = rank % numDevices; + HIP_CHECK(hipSetDevice(dev)); + + printf("[rank %d/%d] pid=%d GPU=%d\n", rank, nranks, getpid(), dev); + + mori::cco::ccoComm* comm = nullptr; + if (mori::cco::ccoCommCreate(uid, nranks, rank, PER_RANK_VMM_SIZE, &comm) != 0) { + fprintf(stderr, "[rank %d] CommCreate failed\n", rank); + return 1; + } + + size_t bufSize = COUNT * nranks * sizeof(float); + void* sendBuf = nullptr; + void* recvBuf = nullptr; + if (mori::cco::ccoMemAlloc(comm, bufSize, &sendBuf) != 0 || + mori::cco::ccoMemAlloc(comm, bufSize, &recvBuf) != 0) { + fprintf(stderr, "[rank %d] MemAlloc failed\n", rank); + return 1; + } + + std::vector hostSend(COUNT * nranks); + for (int p = 0; p < nranks; p++) + for (size_t i = 0; i < COUNT; i++) + hostSend[p * COUNT + i] = static_cast(rank * 1000 + p * 100 + i); + HIP_CHECK(hipMemcpy(sendBuf, hostSend.data(), bufSize, hipMemcpyHostToDevice)); + HIP_CHECK(hipMemset(recvBuf, 0xff, bufSize)); + + mori::cco::ccoWindow_t sendWin = nullptr, recvWin = nullptr; + if (mori::cco::ccoWindowRegister(comm, sendBuf, bufSize, &sendWin) != 0 || + mori::cco::ccoWindowRegister(comm, recvBuf, bufSize, &recvWin) != 0) { + fprintf(stderr, "[rank %d] WindowRegister failed\n", rank); + return 1; + } + + // SDMA needs no GDA connectivity; the signal pool is materialized whenever the + // comm has SDMA queues (set up in ccoCommCreate for canSDMA peers). + mori::cco::ccoDevCommRequirements reqs = CCO_DEV_COMM_REQUIREMENTS_INITIALIZER; + reqs.gdaConnectionType = mori::cco::CCO_GDA_CONNECTION_NONE; + reqs.gdaContextCount = 0; + reqs.gdaSignalCount = 0; + reqs.gdaCounterCount = 0; + mori::cco::ccoDevComm devComm{}; + if (mori::cco::ccoDevCommCreate(comm, &reqs, &devComm) != 0) { + fprintf(stderr, "[rank %d] DevCommCreate failed\n", rank); + return 1; + } + + bool ok = true; + if (devComm.sdma.sdmaNumQueue == 0) { + printf("[rank %d] SKIP — no SDMA queues (set MORI_ENABLE_SDMA=1)\n", rank); + } else { + hipStream_t stream; + HIP_CHECK(hipStreamCreate(&stream)); + + // We pull peers' slices into our recv; validate what landed locally. + auto verify = [&](const char* scope) { + std::vector host(COUNT * nranks); + HIP_CHECK(hipMemcpy(host.data(), recvBuf, bufSize, hipMemcpyDeviceToHost)); + for (int p = 0; p < nranks; p++) { + if (p == rank) continue; + for (size_t i = 0; i < COUNT; i++) { + float expected = static_cast(p * 1000 + rank * 100 + i); + if (host[p * COUNT + i] != expected) { + fprintf(stderr, "[rank %d] GET(%s) mismatch [peer=%d][%zu]: got %.0f expected %.0f\n", + rank, scope, p, i, host[p * COUNT + i], expected); + return false; + } + } + } + return true; + }; + + // peers' sendBuf must be initialized before any get is issued. + // thread scope: one thread per peer, peer's queue 0. + mori::cco::ccoBarrierAll(comm); + SdmaGetKernel<<<1, 64, 0, stream>>>(sendWin, recvWin, COUNT, devComm); + HIP_CHECK(hipStreamSynchronize(stream)); + mori::cco::ccoBarrierAll(comm); + const bool ok_thread = verify("thread"); + + // warp scope: one warp per peer, transfer split across all of the peer's queues. + HIP_CHECK(hipMemset(recvBuf, 0xff, bufSize)); + mori::cco::ccoBarrierAll(comm); + SdmaGetWarpKernel<<>>(sendWin, recvWin, COUNT, devComm); + HIP_CHECK(hipStreamSynchronize(stream)); + mori::cco::ccoBarrierAll(comm); + const bool ok_warp = verify("warp"); + + ok = ok_thread && ok_warp; + HIP_CHECK(hipStreamDestroy(stream)); + printf("[rank %d] thread=%s warp=%s %s\n", rank, ok_thread ? "PASS" : "FAIL", + ok_warp ? "PASS" : "FAIL", ok ? "PASSED" : "FAILED"); + } + + mori::cco::ccoDevCommDestroy(comm, &devComm); + mori::cco::ccoWindowDeregister(comm, recvWin); + mori::cco::ccoWindowDeregister(comm, sendWin); + mori::cco::ccoMemFree(comm, recvBuf); + mori::cco::ccoMemFree(comm, sendBuf); + mori::cco::ccoCommDestroy(comm); + return ok ? 0 : 1; +} + +int main(int argc, char** argv) { + return ccoTestMain(argc, argv, "CCO SDMA get", "/tmp/cco_sdma_get_uid", 19891); +} diff --git a/tests/cpp/cco/test_sdma_put.cpp b/tests/cpp/cco/test_sdma_put.cpp new file mode 100644 index 000000000..4cc13998c --- /dev/null +++ b/tests/cpp/cco/test_sdma_put.cpp @@ -0,0 +1,184 @@ +// 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. +// +// test: cco sdma put — intra-node copy-engine (ccoSdma::put), single node. +// +// alltoall via SDMA put. Layout (owner, dest-slot, idx): +// sendBuf[p*COUNT + i] = myRank*1000 + p*100 + i +// Each rank copies its slice for peer p into peer p's recv slot `myRank`, so +// recv[s*COUNT+i] == s*1000 + myRank*100 + i (written by rank s) +// +// One thread per peer (distinct SDMA queue per peer); quiet() drains completion. +// Requires MORI_ENABLE_SDMA=1; otherwise the comm has no SDMA queues and the +// test SKIPs. + +#include "cco_test_harness.hpp" + +static const size_t PER_RANK_VMM_SIZE = 256ULL * 1024 * 1024; +static const size_t COUNT = 256; // elements per rank-pair + +// thread scope: thread p writes our slice into peer p's recv slot indexed by +// myRank, over peer p's queue 0. +__global__ void SdmaPutKernel(mori::cco::ccoWindowDevice* sendWin, + mori::cco::ccoWindowDevice* recvWin, size_t count, + mori::cco::ccoDevComm devComm) { + using namespace mori::cco; + ccoSdma sdma{devComm}; + int myRank = devComm.rank; + int nRanks = devComm.lsaSize; + size_t perPair = count * sizeof(float); + + int p = threadIdx.x; + if (p >= nRanks || p == myRank) return; + sdma.put(p, reinterpret_cast(recvWin), myRank * perPair, + reinterpret_cast(sendWin), p * perPair, perPair); + sdma.quiet(p); +} + +// warp scope: one warp (block) per peer p writes our slice into peer p's recv +// slot, with the transfer split across all of peer p's SDMA queues internally. +__global__ void SdmaPutWarpKernel(mori::cco::ccoWindowDevice* sendWin, + mori::cco::ccoWindowDevice* recvWin, size_t count, + mori::cco::ccoDevComm devComm) { + using namespace mori::cco; + ccoSdma sdma{devComm}; + int myRank = devComm.rank; + int nRanks = devComm.lsaSize; + size_t perPair = count * sizeof(float); + + int p = blockIdx.x; // one warp per peer + if (p >= nRanks || p == myRank) return; + sdma.put(p, reinterpret_cast(recvWin), myRank * perPair, + reinterpret_cast(sendWin), p * perPair, perPair); + sdma.quiet(p); +} + +int run_test(int rank, int nranks, const mori::cco::ccoUniqueId& uid) { + g_rank = rank; + + int numDevices = 0; + HIP_CHECK(hipGetDeviceCount(&numDevices)); + int dev = rank % numDevices; + HIP_CHECK(hipSetDevice(dev)); + + printf("[rank %d/%d] pid=%d GPU=%d\n", rank, nranks, getpid(), dev); + + mori::cco::ccoComm* comm = nullptr; + if (mori::cco::ccoCommCreate(uid, nranks, rank, PER_RANK_VMM_SIZE, &comm) != 0) { + fprintf(stderr, "[rank %d] CommCreate failed\n", rank); + return 1; + } + + size_t bufSize = COUNT * nranks * sizeof(float); + void* sendBuf = nullptr; + void* recvBuf = nullptr; + if (mori::cco::ccoMemAlloc(comm, bufSize, &sendBuf) != 0 || + mori::cco::ccoMemAlloc(comm, bufSize, &recvBuf) != 0) { + fprintf(stderr, "[rank %d] MemAlloc failed\n", rank); + return 1; + } + + std::vector hostSend(COUNT * nranks); + for (int p = 0; p < nranks; p++) + for (size_t i = 0; i < COUNT; i++) + hostSend[p * COUNT + i] = static_cast(rank * 1000 + p * 100 + i); + HIP_CHECK(hipMemcpy(sendBuf, hostSend.data(), bufSize, hipMemcpyHostToDevice)); + HIP_CHECK(hipMemset(recvBuf, 0xff, bufSize)); + + mori::cco::ccoWindow_t sendWin = nullptr, recvWin = nullptr; + if (mori::cco::ccoWindowRegister(comm, sendBuf, bufSize, &sendWin) != 0 || + mori::cco::ccoWindowRegister(comm, recvBuf, bufSize, &recvWin) != 0) { + fprintf(stderr, "[rank %d] WindowRegister failed\n", rank); + return 1; + } + + // SDMA needs no GDA connectivity; the signal pool is materialized whenever the + // comm has SDMA queues (set up in ccoCommCreate for canSDMA peers). + mori::cco::ccoDevCommRequirements reqs = CCO_DEV_COMM_REQUIREMENTS_INITIALIZER; + reqs.gdaConnectionType = mori::cco::CCO_GDA_CONNECTION_NONE; + reqs.gdaContextCount = 0; + reqs.gdaSignalCount = 0; + reqs.gdaCounterCount = 0; + mori::cco::ccoDevComm devComm{}; + if (mori::cco::ccoDevCommCreate(comm, &reqs, &devComm) != 0) { + fprintf(stderr, "[rank %d] DevCommCreate failed\n", rank); + return 1; + } + + bool ok = true; + if (devComm.sdma.sdmaNumQueue == 0) { + printf("[rank %d] SKIP — no SDMA queues (set MORI_ENABLE_SDMA=1)\n", rank); + } else { + hipStream_t stream; + HIP_CHECK(hipStreamCreate(&stream)); + + // recv is filled by peers writing into our window; validate the alltoall. + auto verify = [&](const char* scope) { + std::vector host(COUNT * nranks); + HIP_CHECK(hipMemcpy(host.data(), recvBuf, bufSize, hipMemcpyDeviceToHost)); + for (int s = 0; s < nranks; s++) { + if (s == rank) continue; + for (size_t i = 0; i < COUNT; i++) { + float expected = static_cast(s * 1000 + rank * 100 + i); + if (host[s * COUNT + i] != expected) { + fprintf(stderr, "[rank %d] PUT(%s) mismatch [src=%d][%zu]: got %.0f expected %.0f\n", + rank, scope, s, i, host[s * COUNT + i], expected); + return false; + } + } + } + return true; + }; + + // thread scope: one thread per peer, peer's queue 0. + mori::cco::ccoBarrierAll(comm); + SdmaPutKernel<<<1, 64, 0, stream>>>(sendWin, recvWin, COUNT, devComm); + HIP_CHECK(hipStreamSynchronize(stream)); + mori::cco::ccoBarrierAll(comm); + const bool ok_thread = verify("thread"); + + // warp scope: one warp per peer, transfer split across all of the peer's queues. + HIP_CHECK(hipMemset(recvBuf, 0xff, bufSize)); + mori::cco::ccoBarrierAll(comm); + SdmaPutWarpKernel<<>>(sendWin, recvWin, COUNT, devComm); + HIP_CHECK(hipStreamSynchronize(stream)); + mori::cco::ccoBarrierAll(comm); + const bool ok_warp = verify("warp"); + + ok = ok_thread && ok_warp; + HIP_CHECK(hipStreamDestroy(stream)); + printf("[rank %d] thread=%s warp=%s %s\n", rank, ok_thread ? "PASS" : "FAIL", + ok_warp ? "PASS" : "FAIL", ok ? "PASSED" : "FAILED"); + } + + mori::cco::ccoDevCommDestroy(comm, &devComm); + mori::cco::ccoWindowDeregister(comm, recvWin); + mori::cco::ccoWindowDeregister(comm, sendWin); + mori::cco::ccoMemFree(comm, recvBuf); + mori::cco::ccoMemFree(comm, sendBuf); + mori::cco::ccoCommDestroy(comm); + return ok ? 0 : 1; +} + +int main(int argc, char** argv) { + return ccoTestMain(argc, argv, "CCO SDMA put", "/tmp/cco_sdma_put_uid", 19890); +} From 0c5602489ca2e089c7061cffa122062fcb1abf34 Mon Sep 17 00:00:00 2001 From: qizzhang Date: Mon, 6 Jul 2026 01:18:53 -0500 Subject: [PATCH 02/10] modify gitignore --- .gitignore | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 26dff7510..25a37aa5b 100644 --- a/.gitignore +++ b/.gitignore @@ -18,7 +18,9 @@ python/mori/cco/cco.cpp # Cython-generated from cco.pyx (build artifact python/mori/umbp_master python/mori/spdk_proxy python/mori/examples/ # BUILD_EXAMPLES=ON copies example binaries here -python/mori/benchmarks/ # BUILD_BENCHMARK=ON copies benchmark binaries here +python/mori/benchmarks/* +python/mori/cco/cco.cpp + # local cco experiments (not for commit) examples/cco/**/_dump_*.py From 4eeaa59e9956eb51ed19c092e1a4a0d24abed2d9 Mon Sep 17 00:00:00 2001 From: qizzhang Date: Tue, 7 Jul 2026 21:50:18 +0800 Subject: [PATCH 03/10] fmt --- benchmark/cco/p2p_get_bw.cpp | 5 +++-- benchmark/cco/p2p_put_bw.cpp | 5 +++-- benchmark/cco/util.cpp | 3 +-- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/benchmark/cco/p2p_get_bw.cpp b/benchmark/cco/p2p_get_bw.cpp index 533ae162d..6d9eaa045 100644 --- a/benchmark/cco/p2p_get_bw.cpp +++ b/benchmark/cco/p2p_get_bw.cpp @@ -128,8 +128,9 @@ __global__ void sdma_get_bw_warp(ccoWindowDevice* sendWin, ccoWindowDevice* recv sdma.quiet(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) { +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, diff --git a/benchmark/cco/p2p_put_bw.cpp b/benchmark/cco/p2p_put_bw.cpp index 1d1ceaffd..592dfe1c8 100644 --- a/benchmark/cco/p2p_put_bw.cpp +++ b/benchmark/cco/p2p_put_bw.cpp @@ -134,8 +134,9 @@ __global__ void sdma_put_bw_warp(ccoWindowDevice* sendWin, ccoWindowDevice* recv sdma.quiet(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) { +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, diff --git a/benchmark/cco/util.cpp b/benchmark/cco/util.cpp index 8d3af5ae8..916848c94 100644 --- a/benchmark/cco/util.cpp +++ b/benchmark/cco/util.cpp @@ -386,8 +386,7 @@ int PerfInit(int argc, char** argv, PerfContext* ctx) { if (args.transport == Transport::kLsa || args.transport == Transport::kSdma) { if (ctx->devComm.lsaSize < 2) { if (ctx->my_pe == 0) { - std::fprintf(stderr, - "%s transport requires both PEs on the same node (lsaSize=%d).\n", + std::fprintf(stderr, "%s transport requires both PEs on the same node (lsaSize=%d).\n", TransportToChar(args.transport), ctx->devComm.lsaSize); } PerfFinalize(ctx); From 2b56d8e96deaa33e49b6d612c630e00412485e92 Mon Sep 17 00:00:00 2001 From: qizzhang Date: Tue, 7 Jul 2026 21:40:51 -0500 Subject: [PATCH 04/10] add sdma tests to ci --- tests/cpp/cco/test_sdma_put_mt.cpp | 257 +++++++++++++++++++++++++++++ tools/run_cco_tests.sh | 14 +- 2 files changed, 269 insertions(+), 2 deletions(-) create mode 100644 tests/cpp/cco/test_sdma_put_mt.cpp diff --git a/tests/cpp/cco/test_sdma_put_mt.cpp b/tests/cpp/cco/test_sdma_put_mt.cpp new file mode 100644 index 000000000..f49e73ee6 --- /dev/null +++ b/tests/cpp/cco/test_sdma_put_mt.cpp @@ -0,0 +1,257 @@ +// 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. +// +// test: cco sdma put — single process, N host threads, one GPU per thread. +// +// Same alltoall-via-SDMA-put semantics as test_sdma_put.cpp, but instead of +// forking one process per rank, this launches nRanks std::threads inside ONE +// process. Each thread binds its own GPU (hipSetDevice), builds its own ccoComm +// against the shared uniqueId, and runs put/quiet. Every comm reserves its own +// flat VA and maps all peers relative to its own flatBase, so the per-thread +// comms are self-consistent — nothing about the LSA peer addressing assumes a +// one-process-per-rank layout. +// +// Layout (owner, dest-slot, idx): +// sendBuf[p*COUNT + i] = myRank*1000 + p*100 + i +// Each rank copies its slice for peer p into peer p's recv slot `myRank`, so +// recv[s*COUNT+i] == s*1000 + myRank*100 + i (written by rank s) +// +// Requires MORI_ENABLE_SDMA=1; otherwise the comms have no SDMA queues and the +// test SKIPs. + +#include + +#include +#include +#include +#include + +#include "hip/hip_runtime.h" +#include "mori/cco/cco.hpp" + +static const size_t PER_RANK_VMM_SIZE = 256ULL * 1024 * 1024; +static const size_t COUNT = 256; // elements per rank-pair + +// Thread-scoped HIP check: on a hard HIP error, abort the whole process (a +// half-initialized rank would otherwise deadlock its peers at the next barrier). +#define HIP_CHECK_MT(rank, cmd) \ + do { \ + hipError_t e = (cmd); \ + if (e != hipSuccess) { \ + fprintf(stderr, "[rank %d] HIP error %d (%s) at %s:%d\n", (rank), e, hipGetErrorString(e), \ + __FILE__, __LINE__); \ + fflush(stderr); \ + _exit(1); \ + } \ + } while (0) + +// thread scope: thread p writes our slice into peer p's recv slot indexed by +// myRank, over peer p's queue 0. +__global__ void SdmaPutMtKernel(mori::cco::ccoWindowDevice* sendWin, + mori::cco::ccoWindowDevice* recvWin, size_t count, + mori::cco::ccoDevComm devComm) { + using namespace mori::cco; + ccoSdma sdma{devComm}; + int myRank = devComm.rank; + int nRanks = devComm.lsaSize; + size_t perPair = count * sizeof(float); + + int p = threadIdx.x; + if (p >= nRanks || p == myRank) return; + sdma.put(p, reinterpret_cast(recvWin), myRank * perPair, + reinterpret_cast(sendWin), p * perPair, perPair); + sdma.quiet(p); +} + +// warp scope: one warp (block) per peer p writes our slice into peer p's recv +// slot, with the transfer split across all of peer p's SDMA queues internally. +__global__ void SdmaPutMtWarpKernel(mori::cco::ccoWindowDevice* sendWin, + mori::cco::ccoWindowDevice* recvWin, size_t count, + mori::cco::ccoDevComm devComm) { + using namespace mori::cco; + ccoSdma sdma{devComm}; + int myRank = devComm.rank; + int nRanks = devComm.lsaSize; + size_t perPair = count * sizeof(float); + + int p = blockIdx.x; // one warp per peer + if (p >= nRanks || p == myRank) return; + sdma.put(p, reinterpret_cast(recvWin), myRank * perPair, + reinterpret_cast(sendWin), p * perPair, perPair); + sdma.quiet(p); +} + +struct ThreadResult { + bool skipped = false; + bool ok = false; +}; + +static void worker(int rank, int nranks, int dev, const mori::cco::ccoUniqueId& uid, + ThreadResult* res) { + HIP_CHECK_MT(rank, hipSetDevice(dev)); + printf("[rank %d/%d] tid=%ld GPU=%d\n", rank, nranks, (long)gettid(), dev); + fflush(stdout); + + mori::cco::ccoComm* comm = nullptr; + if (mori::cco::ccoCommCreate(uid, nranks, rank, PER_RANK_VMM_SIZE, &comm) != 0) { + fprintf(stderr, "[rank %d] CommCreate failed\n", rank); + _exit(1); + } + + size_t bufSize = COUNT * nranks * sizeof(float); + void* sendBuf = nullptr; + void* recvBuf = nullptr; + if (mori::cco::ccoMemAlloc(comm, bufSize, &sendBuf) != 0 || + mori::cco::ccoMemAlloc(comm, bufSize, &recvBuf) != 0) { + fprintf(stderr, "[rank %d] MemAlloc failed\n", rank); + _exit(1); + } + + std::vector hostSend(COUNT * nranks); + for (int p = 0; p < nranks; p++) + for (size_t i = 0; i < COUNT; i++) + hostSend[p * COUNT + i] = static_cast(rank * 1000 + p * 100 + i); + HIP_CHECK_MT(rank, hipMemcpy(sendBuf, hostSend.data(), bufSize, hipMemcpyHostToDevice)); + HIP_CHECK_MT(rank, hipMemset(recvBuf, 0xff, bufSize)); + + mori::cco::ccoWindow_t sendWin = nullptr, recvWin = nullptr; + if (mori::cco::ccoWindowRegister(comm, sendBuf, bufSize, &sendWin) != 0 || + mori::cco::ccoWindowRegister(comm, recvBuf, bufSize, &recvWin) != 0) { + fprintf(stderr, "[rank %d] WindowRegister failed\n", rank); + _exit(1); + } + + mori::cco::ccoDevCommRequirements reqs = CCO_DEV_COMM_REQUIREMENTS_INITIALIZER; + reqs.gdaConnectionType = mori::cco::CCO_GDA_CONNECTION_NONE; + reqs.gdaContextCount = 0; + reqs.gdaSignalCount = 0; + reqs.gdaCounterCount = 0; + mori::cco::ccoDevComm devComm{}; + if (mori::cco::ccoDevCommCreate(comm, &reqs, &devComm) != 0) { + fprintf(stderr, "[rank %d] DevCommCreate failed\n", rank); + _exit(1); + } + + bool ok = true; + if (devComm.sdma.sdmaNumQueue == 0) { + printf("[rank %d] SKIP — no SDMA queues (set MORI_ENABLE_SDMA=1)\n", rank); + res->skipped = true; + } else { + hipStream_t stream; + HIP_CHECK_MT(rank, hipStreamCreate(&stream)); + + // recv is filled by peers writing into our window; validate the alltoall. + auto verify = [&](const char* scope) { + std::vector host(COUNT * nranks); + HIP_CHECK_MT(rank, hipMemcpy(host.data(), recvBuf, bufSize, hipMemcpyDeviceToHost)); + for (int s = 0; s < nranks; s++) { + if (s == rank) continue; + for (size_t i = 0; i < COUNT; i++) { + float expected = static_cast(s * 1000 + rank * 100 + i); + if (host[s * COUNT + i] != expected) { + fprintf(stderr, "[rank %d] PUT(%s) mismatch [src=%d][%zu]: got %.0f expected %.0f\n", + rank, scope, s, i, host[s * COUNT + i], expected); + return false; + } + } + } + return true; + }; + + // thread scope: one thread per peer, peer's queue 0. + mori::cco::ccoBarrierAll(comm); + SdmaPutMtKernel<<<1, 64, 0, stream>>>(sendWin, recvWin, COUNT, devComm); + HIP_CHECK_MT(rank, hipStreamSynchronize(stream)); + mori::cco::ccoBarrierAll(comm); + const bool ok_thread = verify("thread"); + + // warp scope: one warp per peer, transfer split across all of the peer's queues. + HIP_CHECK_MT(rank, hipMemset(recvBuf, 0xff, bufSize)); + mori::cco::ccoBarrierAll(comm); + SdmaPutMtWarpKernel<<>>(sendWin, recvWin, COUNT, devComm); + HIP_CHECK_MT(rank, hipStreamSynchronize(stream)); + mori::cco::ccoBarrierAll(comm); + const bool ok_warp = verify("warp"); + + ok = ok_thread && ok_warp; + HIP_CHECK_MT(rank, hipStreamDestroy(stream)); + printf("[rank %d] thread=%s warp=%s %s\n", rank, ok_thread ? "PASS" : "FAIL", + ok_warp ? "PASS" : "FAIL", ok ? "PASSED" : "FAILED"); + res->ok = ok; + } + + mori::cco::ccoDevCommDestroy(comm, &devComm); + mori::cco::ccoWindowDeregister(comm, recvWin); + mori::cco::ccoWindowDeregister(comm, sendWin); + mori::cco::ccoMemFree(comm, recvBuf); + mori::cco::ccoMemFree(comm, sendBuf); + mori::cco::ccoCommDestroy(comm); +} + +int main(int argc, char** argv) { + int numDevices = 0; + if (hipGetDeviceCount(&numDevices) != hipSuccess || numDevices < 2) { + printf("Need at least 2 GPUs.\n"); + return 1; + } + int nranks = numDevices; + if (argc > 1) { + int req = atoi(argv[1]); + if (req >= 2 && req < nranks) nranks = req; + } + + printf("=== CCO SDMA put Test (threads, %d ranks / %d GPUs) ===\n", nranks, numDevices); + fflush(stdout); + + // One process => one uniqueId, shared by every thread (rank 0's socket + // rendezvous; created here on the main thread, connected during CommCreate). + mori::cco::ccoUniqueId uid; + if (mori::cco::ccoGetUniqueId(&uid) != 0) { + fprintf(stderr, "ccoGetUniqueId failed (set MORI_SOCKET_IFNAME=)\n"); + return 1; + } + + std::vector results(nranks); + std::vector threads; + threads.reserve(nranks); + for (int r = 0; r < nranks; r++) { + threads.emplace_back(worker, r, nranks, r % numDevices, std::cref(uid), &results[r]); + } + for (auto& t : threads) t.join(); + + bool anySkipped = false; + int passed = 0; + for (int r = 0; r < nranks; r++) { + if (results[r].skipped) { + anySkipped = true; + } else if (results[r].ok) { + passed++; + } + } + + if (anySkipped) { + printf("\n=== SKIPPED (no SDMA queues) ===\n"); + return 0; + } + printf("\n=== %d/%d PASSED ===\n", passed, nranks); + return passed == nranks ? 0 : 1; +} diff --git a/tools/run_cco_tests.sh b/tools/run_cco_tests.sh index a1da7ec98..f096e9d4d 100755 --- a/tools/run_cco_tests.sh +++ b/tools/run_cco_tests.sh @@ -10,10 +10,20 @@ cd "$build_dir" failed=0 for bin in tests/cpp/cco/test_*; do [ -x "$bin" ] || continue - case "$(basename "$bin")" in + name="$(basename "$bin")" + case "$name" in test_lsa_memcheck|test_gda_barrier|test_gda_counter|test_gda_multi_context|test_gda_signal_ut|test_gda_thread_aggregate|test_gda_put|test_gda_get) continue ;; + # SDMA transport tests only exercise real work when SDMA queues exist; run + # them with MORI_ENABLE_SDMA=1 so they don't silently SKIP. + test_sdma_*) continue ;; esac - echo "=== $(basename "$bin") ===" + echo "=== $name ===" timeout -k 10 120 "./$bin" "$nranks" || failed=1 done + +for bin in tests/cpp/cco/test_sdma_*; do + [ -x "$bin" ] || continue + echo "=== $(basename "$bin") (MORI_ENABLE_SDMA=1) ===" + MORI_ENABLE_SDMA=1 timeout -k 10 120 "./$bin" "$nranks" || failed=1 +done exit $failed From 882e5597a2a3b63fa0c16708021e2ea0d8541ffa Mon Sep 17 00:00:00 2001 From: qizzhang Date: Wed, 8 Jul 2026 01:31:31 -0500 Subject: [PATCH 05/10] refine quiet, tested on mi300x --- benchmark/cco/p2p_get_bw.cpp | 2 +- benchmark/cco/p2p_get_latency.cpp | 2 +- benchmark/cco/p2p_put_bw.cpp | 2 +- benchmark/cco/p2p_put_latency.cpp | 2 +- include/mori/cco/cco.hpp | 36 +++++++++++++------------------ 5 files changed, 19 insertions(+), 25 deletions(-) diff --git a/benchmark/cco/p2p_get_bw.cpp b/benchmark/cco/p2p_get_bw.cpp index 6d9eaa045..ce9925f1d 100644 --- a/benchmark/cco/p2p_get_bw.cpp +++ b/benchmark/cco/p2p_get_bw.cpp @@ -113,7 +113,7 @@ __global__ void sdma_get_bw(ccoWindowDevice* sendWin, ccoWindowDevice* recvWin, for (int i = 0; i < iter; i++) sdma.get(peerLsa, reinterpret_cast(recvWin), off, reinterpret_cast(sendWin), off, bytes, q); - sdma.quiet(peerLsa, q); + sdma.quietQueue(peerLsa, q); } // SDMA warp scope: one warp drives all SDMA queues via a single warp-scope get diff --git a/benchmark/cco/p2p_get_latency.cpp b/benchmark/cco/p2p_get_latency.cpp index 6076cd9c0..cb9c93919 100644 --- a/benchmark/cco/p2p_get_latency.cpp +++ b/benchmark/cco/p2p_get_latency.cpp @@ -80,7 +80,7 @@ __global__ void sdma_get_lat(ccoWindowDevice* sendWin, ccoWindowDevice* recvWin, for (int i = 0; i < iter; i++) { sdma.get(peerLsa, reinterpret_cast(recvWin), 0, reinterpret_cast(sendWin), 0, bytes, 0); - sdma.quiet(peerLsa, 0); + sdma.quietQueue(peerLsa, 0); } } diff --git a/benchmark/cco/p2p_put_bw.cpp b/benchmark/cco/p2p_put_bw.cpp index 592dfe1c8..468408e74 100644 --- a/benchmark/cco/p2p_put_bw.cpp +++ b/benchmark/cco/p2p_put_bw.cpp @@ -119,7 +119,7 @@ __global__ void sdma_put_bw(ccoWindowDevice* sendWin, ccoWindowDevice* recvWin, for (int i = 0; i < iter; i++) sdma.put(peerLsa, reinterpret_cast(recvWin), off, reinterpret_cast(sendWin), off, bytes, q); - sdma.quiet(peerLsa, q); + sdma.quietQueue(peerLsa, q); } // SDMA warp scope: one warp drives all SDMA queues via a single warp-scope put diff --git a/benchmark/cco/p2p_put_latency.cpp b/benchmark/cco/p2p_put_latency.cpp index fd8cce162..d7992e0ac 100644 --- a/benchmark/cco/p2p_put_latency.cpp +++ b/benchmark/cco/p2p_put_latency.cpp @@ -83,7 +83,7 @@ __global__ void sdma_put_lat(ccoWindowDevice* sendWin, ccoWindowDevice* recvWin, for (int i = 0; i < iter; i++) { sdma.put(peerLsa, reinterpret_cast(recvWin), 0, reinterpret_cast(sendWin), 0, bytes, 0); - sdma.quiet(peerLsa, 0); + sdma.quietQueue(peerLsa, 0); } } diff --git a/include/mori/cco/cco.hpp b/include/mori/cco/cco.hpp index 4c2b1a7ba..719a5eba7 100644 --- a/include/mori/cco/cco.hpp +++ b/include/mori/cco/cco.hpp @@ -878,15 +878,8 @@ int ccoBarrierAll(ccoComm* comm); } // namespace cco } // namespace mori -/* ════════════════════════════════════════════════════════════════════════════ - * SDMA (intra-node copy-engine) device session. - * - * Moves data via the SDMA queues on ccoDevComm::sdma; peers are addressed by - * LSA rank. put/get are non-blocking (single thread per (peer, queueId)); - * completion is recorded in the local signal pool and awaited with quiet(). - * Device-only — pulls in the SDMA core primitives, so it lives behind the same - * __HIPCC__ guard and is skipped entirely by host translation units. - * ════════════════════════════════════════════════════════════════════════════ */ +// SDMA (intra-node copy-engine) session. Non-blocking put/get over the SDMA +// queues; peers by LSA rank; completion awaited by quiet. Device-only. #if defined(__HIPCC__) || defined(__CUDACC__) #include "mori/core/transport/sdma/device_primitives.hpp" @@ -898,10 +891,8 @@ struct ccoSdma { __device__ inline ccoSdma(ccoDevComm const& c) : comm(c) {} - // put: copy local srcWin[srcOffset] -> peer dstWin[dstOffset] (bytes). - // Coop=ccoCoopThread — a single thread drives one queue (queueId). - // Coop=ccoCoopWarp — a warp drives all queues, one lane per queue; bytes - // are split across sdmaNumQueue and queueId is ignored. + // put: local src -> peer dst. Thread scope = one queue (queueId); warp scope = + // split across all queues (queueId ignored). template __device__ inline void put(int peer, ccoWindow_t dstWin, size_t dstOffset, ccoWindow_t srcWin, size_t srcOffset, size_t bytes, int queueId = 0) { @@ -912,7 +903,6 @@ struct ccoSdma { void* dst = ccoGetLsaPeerPtr(dstWin, peer, dstOffset); void* src = ccoGetLocalPtr(srcWin, srcOffset); if constexpr (std::is_same_v) { - // Appends COPY + ATOMIC(signal++) on the (peer,queueId) queue, bumps expect. core::SdmaPutThread(src, dst, bytes, s.deviceHandles + peer * n, s.signalBuf + peer * n, s.expectSignals + peer * n, n, queueId); } else { @@ -921,8 +911,7 @@ struct ccoSdma { } } - // get: copy peer srcWin[srcOffset] -> local dstWin[dstOffset] (bytes). Coop - // selects thread (single queue) vs warp (all queues) scope, as in put(). + // get: peer src -> local dst. Same scope rules as put(). template __device__ inline void get(int peer, ccoWindow_t dstWin, size_t dstOffset, ccoWindow_t srcWin, size_t srcOffset, size_t bytes, int queueId = 0) { @@ -941,21 +930,26 @@ struct ccoSdma { } } - // quiet: block until outstanding ops have landed. Coop=ccoCoopThread waits on - // the single (peer,queueId) queue; Coop=ccoCoopWarp waits on every queue (one - // lane per queue), pairing with a warp-scope put/get. + // quiet: wait for all outstanding ops to `peer` across every queue. template - __device__ inline void quiet(int peer, int queueId = 0) { + __device__ inline void quiet(int peer) { static_assert(std::is_same_v || std::is_same_v, "ccoSdma::quiet supports ccoCoopThread or ccoCoopWarp"); const ccoSdmaContext& s = comm.sdma; const uint32_t n = s.sdmaNumQueue; if constexpr (std::is_same_v) { - anvil::waitForSignal(s.signalBuf + peer * n + queueId, s.expectSignals[peer * n + queueId]); + core::SdmaQueitThread(s.signalBuf + peer * n, s.expectSignals + peer * n, n); } else { core::SdmaQueitWarp(s.signalBuf + peer * n, s.expectSignals + peer * n, n); } } + + // quietQueue: wait on a single (peer, queueId) queue only. + __device__ inline void quietQueue(int peer, int queueId) { + const ccoSdmaContext& s = comm.sdma; + const uint32_t n = s.sdmaNumQueue; + anvil::waitForSignal(s.signalBuf + peer * n + queueId, s.expectSignals[peer * n + queueId]); + } }; } // namespace cco From 9ecbab1c9418a86cf49a589e0b785c4ec6c24ec6 Mon Sep 17 00:00:00 2001 From: qizzhang Date: Thu, 9 Jul 2026 21:14:19 -0500 Subject: [PATCH 06/10] ci --- include/mori/cco/cco.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/mori/cco/cco.hpp b/include/mori/cco/cco.hpp index 719a5eba7..3a01c007b 100644 --- a/include/mori/cco/cco.hpp +++ b/include/mori/cco/cco.hpp @@ -944,7 +944,7 @@ struct ccoSdma { } } - // quietQueue: wait on a single (peer, queueId) queue only. + // quietQueue: wait on a single (peer, queueId) queue only __device__ inline void quietQueue(int peer, int queueId) { const ccoSdmaContext& s = comm.sdma; const uint32_t n = s.sdmaNumQueue; From 9acc4faf7ea579b286ca36ffa445d9d6543eb696 Mon Sep 17 00:00:00 2001 From: Qizhou Zhang Date: Mon, 13 Jul 2026 03:00:45 +0000 Subject: [PATCH 07/10] add sdma python interface --- examples/cco/python/07_flydsl_sdma/main.py | 234 +++++++++++++++++++++ python/mori/cco/device/flydsl/__init__.py | 3 +- python/mori/cco/device/flydsl/_bindings.py | 17 ++ python/mori/cco/device/flydsl/handles.py | 69 ++++++ src/cco/device/cco_device_wrapper.cpp | 32 +++ 5 files changed, 354 insertions(+), 1 deletion(-) create mode 100644 examples/cco/python/07_flydsl_sdma/main.py diff --git a/examples/cco/python/07_flydsl_sdma/main.py b/examples/cco/python/07_flydsl_sdma/main.py new file mode 100644 index 000000000..7e5179162 --- /dev/null +++ b/examples/cco/python/07_flydsl_sdma/main.py @@ -0,0 +1,234 @@ +#!/usr/bin/env python3 +# 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. +"""CCO Example 07 — FlyDSL SDMA all-reduce (copy-engine all-gather + local reduce) + +A custom all-reduce built on cco's SDMA copy-engine API (``DevComm.sdma()``), +which moves data between intra-node peers via the DMA engines rather than the +compute units: + + * **Scatter (all-gather):** each rank SDMA-``put``s its input region into a + dedicated per-source slot of every *other* peer's gather region, then + ``quiet``s each peer to wait for the DMA copies + completion signals to land. + SDMA is a one-sided copy engine — the put only guarantees *my* writes reached + the peers, so a host ``barrier`` follows to ensure every peer's writes have + also reached *me* before I read them. + * **Reduce:** every rank sums its own input plus the gathered peer slots + (f32, vectorized 16 B loads) into its output region. All ranks run it, so all + end up with the full sum (= all-reduce). + +SDMA addresses peers through the same flat symmetric-window VA as the LSA model +(``Window.lsa_ptr``); ``put``/``quiet`` just drive the copy through cco's +per-DevComm SDMA queue pool (enabled by ``sdma_queue_count``). + +Single node, ``world_size`` ranks (one GPU each): + + mpirun -n 2 python main.py +""" + +import os +import sys + +try: + from mpi4py import MPI +except ImportError: + print("ERROR: mpi4py required. pip install mpi4py") + sys.exit(1) + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import buffer_ops, range_constexpr +from flydsl.expr.typing import Int64, T + +from mori.cco import Communicator, CCODevCommRequirements, GDA_CONNECTION_NONE +import mori.cco.device.flydsl as cco + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from cco_example_common import set_device, sync, fill, zero, read, F32 # noqa: E402 + +WS = 2 # world size (ranks, one GPU each) +THREADS = 256 +NUM_ELEMS = 256 * 1024 # f32 elements to all-reduce (1 MiB) +NBYTES = NUM_ELEMS * 4 +ELEMS_PER_PACK = 4 # vector<4xf32> = 16 B load unit +NUM_PACKS = NUM_ELEMS // ELEMS_PER_PACK +SDMA_QUEUES = 8 # per-DevComm SDMA queue pool + +# Window layout: [ input | gather (WS slots) | output ]. +# gather slot p receives peer p's input; my own slot is unused (own input is +# read straight from the input region during the reduce). +IN_OFF = 0 +GATHER_OFF = NBYTES +OUT_OFF = GATHER_OFF + WS * NBYTES +WIN_BYTES = OUT_OFF + NBYTES + + +def _make_scatter_kernel(my_rank): + """Scatter kernel specialized to this rank (peer loop constant-folds).""" + dst_off = GATHER_OFF + my_rank * NBYTES # my slot in every peer's gather region + + @flyc.kernel(known_block_size=[THREADS, 1, 1]) + def scatter_kernel(dev_comm: Int64, win: Int64): + sdma = cco.DevComm(dev_comm).sdma() + if fx.thread_idx.x == 0: + # issue one copy per peer, then wait on each. + for p in range(WS): + if p != my_rank: + sdma.put( + p, + win, + dst_off, + win, + IN_OFF, + NBYTES, + 0, + coop=cco.CoopScope.THREAD, + ) + for p in range(WS): + if p != my_rank: + sdma.quiet(p, coop=cco.CoopScope.THREAD) + + @flyc.jit + def run(dev_comm: Int64, win: Int64, stream=fx.Stream(None)): + scatter_kernel(dev_comm, win).launch( + grid=(1, 1, 1), block=[THREADS, 1, 1], stream=stream + ) + + return run + + +def _make_reduce_kernel(my_rank): + """Reduce kernel: sum own input + gathered peer slots into output.""" + # Byte offsets (in my own window) of every contribution to the sum. + src_offs = [IN_OFF] + [GATHER_OFF + p * NBYTES for p in range(WS) if p != my_rank] + + @flyc.kernel(known_block_size=[THREADS, 1, 1]) + def reduce_kernel(win: Int64): + tid = fx.thread_idx.x + w = cco.Window(win) + # All reads are local memory — my own symmetric-window VA. + src_rsrc = [ + buffer_ops.create_buffer_resource_from_addr( + fx.Int64(w.lsa_ptr(my_rank, off)) + ) + for off in src_offs + ] + out_rsrc = buffer_ops.create_buffer_resource_from_addr( + fx.Int64(w.lsa_ptr(my_rank, OUT_OFF)) + ) + for pk in range(tid, NUM_PACKS, THREADS): + elem_off = pk * ELEMS_PER_PACK + acc = None + for s in range_constexpr(len(src_rsrc)): + v = fx.Vector( + buffer_ops.buffer_load( + src_rsrc[s], elem_off, vec_width=4, dtype=T.f32 + ) + ) + acc = v if acc is None else acc + v + buffer_ops.buffer_store(acc, out_rsrc, elem_off) + + @flyc.jit + def run(win: Int64, stream=fx.Stream(None)): + reduce_kernel(win).launch(grid=(1, 1, 1), block=[THREADS, 1, 1], stream=stream) + + return run + + +def main() -> int: + mpi = MPI.COMM_WORLD + rank, nranks = mpi.Get_rank(), mpi.Get_size() + if nranks != WS: + if rank == 0: + print(f"This example is built for world_size={WS} (mpirun -n {WS}).") + return 1 + # SDMA is gated by MORI_ENABLE_SDMA; without it the DevComm builds no SDMA + # queue pool (sdmaNumQueue=0, null device handles) and the copy-engine put + # dereferences a null pointer -> opaque GPU memory-access fault. + if os.environ.get("MORI_ENABLE_SDMA") not in ("1", "true", "TRUE", "on", "ON"): + if rank == 0: + print( + "This example needs SDMA enabled. Re-run with:\n" + f" MORI_ENABLE_SDMA=1 mpirun -x MORI_ENABLE_SDMA -n {WS} python main.py", + flush=True, + ) + return 1 + set_device(rank) + uid = Communicator.get_unique_id() if rank == 0 else None + uid = mpi.bcast(uid, root=0) + + scatter = _make_scatter_kernel(rank) + reduce = _make_reduce_kernel(rank) + errors = 0 + + with Communicator.init(nranks, rank, uid, per_rank_vmm=256 * 1024 * 1024) as comm: + mem = comm.alloc_mem(WIN_BYTES) + win = comm.register_window(mem.ptr, mem.size) + + # Zero the window, then fill my input: input[i] = (rank + 1) * (i + 1) + # -> allreduce sum = S * (i + 1), S = sum_{r=1..WS} r = WS*(WS+1)/2. + zero(win.local_ptr, WIN_BYTES) + fill( + win.local_ptr + IN_OFF, + [(rank + 1) * (i + 1) for i in range(NUM_ELEMS)], + F32, + ) + + reqs = CCODevCommRequirements() + reqs.gda_connection_type = GDA_CONNECTION_NONE + reqs.gda_signal_count = 0 + reqs.gda_counter_count = 0 + reqs.sdma_queue_count = SDMA_QUEUES + dc = comm.create_dev_comm(reqs) + + comm.barrier() # all inputs visible + gather regions zeroed + scatter(dc.ptr, win.handle) # push my input into every peer's slot + sync() + comm.barrier() # every peer's slot now populated in my window + reduce(win.handle) # local sum -> output + sync() + + S = WS * (WS + 1) // 2 + host = read(win.local_ptr + OUT_OFF, NUM_ELEMS, F32) + for i in (0, 1, NUM_ELEMS // 2, NUM_ELEMS - 1): + exp = float(S * (i + 1)) + if abs(host[i] - exp) > 1e-3 * max(1.0, exp): + print( + f"[rank {rank}] MISMATCH [{i}]: got {host[i]} expected {exp}", + flush=True, + ) + errors += 1 + print( + f"[rank {rank}] {'allreduce verified' if errors == 0 else 'FAILED'} " + f"(SDMA all-gather + reduce over {WS} ranks of {NUM_ELEMS} f32), " + f"out[0,1,-1]={[host[0], host[1], host[NUM_ELEMS-1]]}", + flush=True, + ) + + all_err = mpi.allreduce(errors, op=MPI.SUM) + if rank == 0: + print("SUCCESS" if all_err == 0 else "FAILED", flush=True) + return 0 if all_err == 0 else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/python/mori/cco/device/flydsl/__init__.py b/python/mori/cco/device/flydsl/__init__.py index 48624975b..5f7372fe6 100644 --- a/python/mori/cco/device/flydsl/__init__.py +++ b/python/mori/cco/device/flydsl/__init__.py @@ -42,7 +42,7 @@ """ from mori.cco.device.bitcode import find_cco_bitcode, get_bitcode_path -from .handles import DevComm, Window, Gda, CoopScope, SignalOp, ThreadMode +from .handles import DevComm, Window, Gda, CoopScope, SignalOp, ThreadMode, Sdma from . import _bindings __all__ = [ @@ -55,4 +55,5 @@ "find_cco_bitcode", "get_bitcode_path", "_bindings", + "Sdma", ] diff --git a/python/mori/cco/device/flydsl/_bindings.py b/python/mori/cco/device/flydsl/_bindings.py index 862d2db6e..09eecb02c 100644 --- a/python/mori/cco/device/flydsl/_bindings.py +++ b/python/mori/cco/device/flydsl/_bindings.py @@ -86,6 +86,23 @@ for c in ("warp", "block") } +# ── SDMA ── + +_SDMA_XFER_ARGS = [_U64, _I32, _U64, _U64, _U64, _U64, _U64, _I32] + +SDMA_XFER = { + f"{op}__{s}": _ffi(f"cco_sdma_{op}__{s}", _SDMA_XFER_ARGS, "void") + for op in ("put", "get") + for s in ("thread", "warp") +} + +# quiet: (devComm, peer) +SDMA_QUIET = { + s: _ffi(f"cco_sdma_quiet__{s}", [_U64, _I32], "void") for s in ("thread", "warp") +} +# quiet_queue: (devComm, peer, queueId) +cco_sdma_quiet_queue = _ffi("cco_sdma_quiet_queue", [_U64, _I32, _I32], "void") + # ── axis-free symbols ── # cco_lsa_ptr(window, peerLsaRank, offset) -> peer's load/store-accessible VA. cco_lsa_ptr = _ffi("cco_lsa_ptr", [_U64, _I32, _U64], _U64, pure=True) diff --git a/python/mori/cco/device/flydsl/handles.py b/python/mori/cco/device/flydsl/handles.py index cba6611f9..bb0e7193b 100644 --- a/python/mori/cco/device/flydsl/handles.py +++ b/python/mori/cco/device/flydsl/handles.py @@ -275,6 +275,71 @@ def flush_peer(self, peer, *, coop=CoopScope.WARP): raw.FLUSH_PEER[_flush_tag(coop)](self.dev_comm, self.ctx, peer) +@cco_struct +class Sdma: + dev_comm: fx.Int64 + + def put( + self, + peer, + dst_win, + dst_off, + src_win, + src_off, + nbytes, + qid, + *, + coop=CoopScope.THREAD, + ): + if _const(coop, "coop") == CoopScope.BLOCK: + raise ValueError("cco: SDMA put has no block-level API yet") + raw.SDMA_XFER[f"put__{_COOP_TAG[coop]}"]( + self.dev_comm, + peer, + _win(dst_win), + dst_off, + _win(src_win), + src_off, + nbytes, + qid, + ) + + def get( + self, + peer, + dst_win, + dst_off, + src_win, + src_off, + nbytes, + qid, + *, + coop=CoopScope.THREAD, + ): + if _const(coop, "coop") == CoopScope.BLOCK: + raise ValueError("cco: SDMA get has no block-level API yet") + raw.SDMA_XFER[f"get__{_COOP_TAG[coop]}"]( + self.dev_comm, + peer, + _win(dst_win), + dst_off, + _win(src_win), + src_off, + nbytes, + qid, + ) + + def quiet(self, peer, *, coop=CoopScope.THREAD): + """Wait for all outstanding SDMA ops to `peer` across every queue.""" + if _const(coop, "coop") == CoopScope.BLOCK: + raise ValueError("cco: SDMA quiet has no block-level API yet") + raw.SDMA_QUIET[_COOP_TAG[coop]](self.dev_comm, peer) + + def quiet_queue(self, peer, qid): + """Wait on a single (peer, queueId) queue only.""" + raw.cco_sdma_quiet_queue(self.dev_comm, peer, qid) + + @cco_struct class DevComm: """Handle for a device-resident ``ccoDevComm``.""" @@ -300,3 +365,7 @@ def lsa_size(self): def gda(self, ctx=0) -> Gda: """Build a GDA handle on this devComm for the given (compile-time) context index.""" return Gda(dev_comm=self.ptr, ctx=ctx) + + def sdma(self) -> Sdma: + """Build an SDMA handle on this devComm (LSA copy-engine put/get).""" + return Sdma(dev_comm=self.ptr) diff --git a/src/cco/device/cco_device_wrapper.cpp b/src/cco/device/cco_device_wrapper.cpp index 39cc3dee5..4286e3787 100644 --- a/src/cco/device/cco_device_wrapper.cpp +++ b/src/cco/device/cco_device_wrapper.cpp @@ -53,6 +53,7 @@ namespace { using namespace mori::cco; using Gda = ccoGda; +using Sdma = ccoSdma; inline __device__ const ccoDevComm* AsDevComm(uint64_t h) { return reinterpret_cast(h); @@ -96,6 +97,37 @@ CCO_DEV uint64_t cco_lsa_ptr(uint64_t window, int peer, uint64_t offset) { return reinterpret_cast(w->winBase) + static_cast(peer) * stride + offset; } +// Expose SDMA C API +#define CCO_DEF_SDMA_XFER(OP, TAG, COOP) \ + CCO_DEV void cco_sdma_##OP##__##TAG(uint64_t dc, int peer, uint64_t dW, uint64_t dO, \ + uint64_t sW, uint64_t sO, uint64_t n, int qid) { \ + Sdma sdma{*AsDevComm(dc)}; \ + sdma.OP(peer, AsWindow(dW), dO, AsWindow(sW), sO, n, qid); \ + } + +CCO_DEF_SDMA_XFER(put, thread, ccoCoopThread) +CCO_DEF_SDMA_XFER(put, warp, ccoCoopWarp) +CCO_DEF_SDMA_XFER(get, thread, ccoCoopThread) +CCO_DEF_SDMA_XFER(get, warp, ccoCoopWarp) + +#undef CCO_DEF_SDMA_XFER + +// ── SDMA quiet: cco_sdma_quiet__ (wait for outstanding ops to peer) ── +#define CCO_DEF_SDMA_QUIET(TAG, COOP) \ + CCO_DEV void cco_sdma_quiet__##TAG(uint64_t dc, int peer) { \ + Sdma sdma{*AsDevComm(dc)}; \ + sdma.quiet(peer); \ + } +CCO_DEF_SDMA_QUIET(thread, ccoCoopThread) +CCO_DEF_SDMA_QUIET(warp, ccoCoopWarp) +#undef CCO_DEF_SDMA_QUIET + +// quiet a single (peer, queueId) queue only. +CCO_DEV void cco_sdma_quiet_queue(uint64_t dc, int peer, int qid) { + Sdma sdma{*AsDevComm(dc)}; + sdma.quietQueue(peer, qid); +} + // ── ccoDevComm field accessors ── CCO_DEV int cco_devcomm_rank(uint64_t dc) { return AsDevComm(dc)->rank; } CCO_DEV int cco_devcomm_world_size(uint64_t dc) { return AsDevComm(dc)->worldSize; } From 0732f5fc6014c04f6619faeac0c6b3a45db72e9c Mon Sep 17 00:00:00 2001 From: Qizhou Zhang Date: Mon, 13 Jul 2026 03:22:14 +0000 Subject: [PATCH 08/10] fix compile --- src/umbp/tests/CMakeLists.txt | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/umbp/tests/CMakeLists.txt b/src/umbp/tests/CMakeLists.txt index 40e431c15..338755038 100644 --- a/src/umbp/tests/CMakeLists.txt +++ b/src/umbp/tests/CMakeLists.txt @@ -21,6 +21,11 @@ enable_testing() include(GoogleTest) +# Discover tests at ctest time: POST_BUILD runs each binary under the parallel +# build, and the mori/HSA/grpc init blows past the 5s timeout -> empty JSON that +# CMake 4.4's strict parser rejects. +set(CMAKE_GTEST_DISCOVER_TESTS_DISCOVERY_MODE PRE_TEST) + # --------------------------------------------------------------------------- # test_external_kv_block_index # --------------------------------------------------------------------------- From 82c1a2732d674f097de840c4b23fbeffa2f95c1c Mon Sep 17 00:00:00 2001 From: Qizhou Zhang Date: Mon, 13 Jul 2026 03:24:37 +0000 Subject: [PATCH 09/10] fix ignore --- .gitignore | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 25a37aa5b..830344ff9 100644 --- a/.gitignore +++ b/.gitignore @@ -17,7 +17,7 @@ python/mori/cco/*.so python/mori/cco/cco.cpp # Cython-generated from cco.pyx (build artifact) python/mori/umbp_master python/mori/spdk_proxy -python/mori/examples/ # BUILD_EXAMPLES=ON copies example binaries here +python/mori/examples/* python/mori/benchmarks/* python/mori/cco/cco.cpp From ba97872387170b2cde46ed7eafff014d095dae89 Mon Sep 17 00:00:00 2001 From: Qizhou Zhang Date: Mon, 13 Jul 2026 07:03:18 +0000 Subject: [PATCH 10/10] add python sdma examples --- .github/workflows/ci_cco.yml | 6 + tests/python/cco/helper.py | 65 ++++++ tests/python/cco/test_sdma_api.py | 347 ++++++++++++++++++++++++++++++ 3 files changed, 418 insertions(+) create mode 100644 tests/python/cco/helper.py create mode 100644 tests/python/cco/test_sdma_api.py diff --git a/.github/workflows/ci_cco.yml b/.github/workflows/ci_cco.yml index 7ce4a6fdf..d9693aed0 100644 --- a/.github/workflows/ci_cco.yml +++ b/.github/workflows/ci_cco.yml @@ -83,6 +83,12 @@ jobs: # mpirun --allow-run-as-root -np 2 python 05_flydsl_lsa_allreduce/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: Cleanup if: always() run: | diff --git a/tests/python/cco/helper.py b/tests/python/cco/helper.py new file mode 100644 index 000000000..0d8f6e53d --- /dev/null +++ b/tests/python/cco/helper.py @@ -0,0 +1,65 @@ +# 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. +"""Host<->device memcpy helpers for the cco python tests. + +cco windows are raw device pointers in cco's own VMM (``win.local_ptr``), so +filling/reading a window goes through hipMemcpy via ``mori.jit.hip_driver`` +(ctypes over libamdhip64) — dependency-free and AMD-portable. + +Owned by the tests so the suite does not depend on examples/ (which is demo code +free to change); mirrors the fill/zero/read/sync helpers used by the cco +examples. +""" + +import ctypes + +from mori.jit.hip_driver import _get_hip_lib, _check + +_H2D, _D2H = 1, 2 + + +def sync() -> None: + _check(_get_hip_lib().hipDeviceSynchronize(), "hipDeviceSynchronize") + + +def _copy(dst, src, nbytes, kind) -> None: + _check( + _get_hip_lib().hipMemcpy(dst, src, ctypes.c_size_t(nbytes), ctypes.c_int(kind)), + "hipMemcpy", + ) + + +def fill(dev_ptr: int, values, ctype=ctypes.c_uint64) -> None: + """Host -> window: write `values` (a sequence) into device memory at dev_ptr.""" + arr = (ctype * len(values))(*values) + _copy(ctypes.c_void_p(dev_ptr), arr, ctypes.sizeof(arr), _H2D) + + +def zero(dev_ptr: int, nbytes: int) -> None: + _copy(ctypes.c_void_p(dev_ptr), (ctypes.c_uint8 * nbytes)(), nbytes, _H2D) + + +def read(dev_ptr: int, count: int, ctype=ctypes.c_uint64): + """Window -> host: read `count` elements of `ctype` from device memory.""" + arr = (ctype * count)() + _copy(arr, ctypes.c_void_p(dev_ptr), ctypes.sizeof(arr), _D2H) + return arr diff --git a/tests/python/cco/test_sdma_api.py b/tests/python/cco/test_sdma_api.py new file mode 100644 index 000000000..479606dee --- /dev/null +++ b/tests/python/cco/test_sdma_api.py @@ -0,0 +1,347 @@ +# 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. +"""Functional tests for the cco SDMA FlyDSL API — all-to-all via put, drained +with quiet (all queues) and quiet_queue (single queue). + +These drive the copy engine end-to-end on real GPUs. Each test spawns +``world_size`` processes (one GPU each) via ``torch.multiprocessing``; rank 0 +mints a cco UniqueId and shares it through a file, so no MPI launcher is needed. + +SDMA is gated by ``MORI_ENABLE_SDMA``; without it the DevComm builds no queue +pool and the copy-engine ops dereference null. The spawn helper sets it for the +children, and the whole module skips if fewer than 2 GPUs are visible. + +Data model (mirrors tests/cpp/cco/test_sdma_put.cpp): each rank fills its input +region with ``rank*1000 + i`` and, after the collective, every rank checks it +received the exact bytes the source rank produced. + +Host<->device memcpy goes through the test-owned ``helper`` module (this dir), +so the suite does not depend on examples/. +""" + +import os +import tempfile +import traceback + +import pytest +import torch + +flydsl = pytest.importorskip("flydsl", reason="cco FlyDSL bindings require flydsl") + +NUM_ELEMS = 64 # uint32 payload per transfer +NBYTES = NUM_ELEMS * 4 +THREADS = 64 + + +def _num_gpus() -> int: + return torch.cuda.device_count() + + +def _require_gpus(world_size: int) -> None: + n = _num_gpus() + if n < world_size: + pytest.skip(f"Need {world_size} GPUs, only {n} available") + + +# ── per-rank worker bodies (run in a spawned process) ────────────────────── +# The workers run in freshly spawned child interpreters, so their imports MUST +# stay inside the function bodies (via _worker_imports): flydsl/cco initialize a +# GPU runtime that must come up only after MORI_ENABLE_SDMA + set_device. + + +def _worker_imports(): + """The full import set a spawned worker needs (see module docstring).""" + import ctypes + + import flydsl.compiler as flyc + import flydsl.expr as fx + from flydsl.expr.typing import Int64 + from mori.cco import CCODevCommRequirements, GDA_CONNECTION_NONE + import mori.cco.device.flydsl as cco + + from helper import sync, fill, zero, read + + return dict( + ctypes=ctypes, + flyc=flyc, + fx=fx, + Int64=Int64, + CCODevCommRequirements=CCODevCommRequirements, + GDA_CONNECTION_NONE=GDA_CONNECTION_NONE, + cco=cco, + sync=sync, + fill=fill, + zero=zero, + read=read, + ) + + +def _init_comm(rank, world_size, uid_path): + """Common setup: set device, load shared UniqueId, create Communicator.""" + from mori.cco import Communicator, UniqueId + + torch.cuda.set_device(rank) + with open(uid_path, "rb") as f: + uid = UniqueId.from_bytes(f.read()) + comm = Communicator.init(world_size, rank, uid, per_rank_vmm=128 * 1024 * 1024) + return comm + + +def _worker_put(rank, world_size, uid_path, coop_name, results_path): + """All-to-all via SDMA put: rank writes its slice into every peer's slot.""" + err = None + try: + mods = _worker_imports() + flyc, fx, Int64, cco = mods["flyc"], mods["fx"], mods["Int64"], mods["cco"] + sync, fill, zero, read = mods["sync"], mods["fill"], mods["zero"], mods["read"] + u32 = mods["ctypes"].c_uint32 + CCODevCommRequirements = mods["CCODevCommRequirements"] + GDA_CONNECTION_NONE = mods["GDA_CONNECTION_NONE"] + + coop = getattr(cco.CoopScope, coop_name) + comm = _init_comm(rank, world_size, uid_path) + + # Window layout: [ send | recv (world_size slots) ] + send_off = 0 + recv_off = NBYTES + win_bytes = recv_off + world_size * NBYTES + + mem = comm.alloc_mem(win_bytes) + win = comm.register_window(mem.ptr, mem.size) + zero(win.local_ptr, win_bytes) + # send[i] = rank*1000 + i + fill(win.local_ptr + send_off, [rank * 1000 + i for i in range(NUM_ELEMS)], u32) + + # Before a2a: show my local send buffer (first 8 elems). + pre = read(win.local_ptr + send_off, NUM_ELEMS, u32) + print(f"[rank {rank}] BEFORE a2a send[:8] = {list(pre[:8])}", flush=True) + + reqs = CCODevCommRequirements() + reqs.gda_connection_type = GDA_CONNECTION_NONE + reqs.gda_signal_count = 0 + reqs.gda_counter_count = 0 + reqs.sdma_queue_count = 8 + dc = comm.create_dev_comm(reqs) + + my_rank = rank + # THREAD coop = one logical op from lane 0; WARP coop needs the whole + # wavefront to enter together (each lane drives one SDMA queue). + gate_single = coop == cco.CoopScope.THREAD + + @flyc.kernel(known_block_size=[THREADS, 1, 1]) + def put_kernel(dev_comm: Int64, w: Int64): + sdma = cco.DevComm(dev_comm).sdma() + active = fx.thread_idx.x == 0 if gate_single else fx.thread_idx.x < 64 + if active: + for p in range(world_size): + if p != my_rank: + # into peer p's recv slot reserved for me (my_rank) + sdma.put( + p, + w, + recv_off + my_rank * NBYTES, + w, + send_off, + NBYTES, + 0, + coop=coop, + ) + for p in range(world_size): + if p != my_rank: + sdma.quiet(p, coop=coop) + + @flyc.jit + def run(dev_comm: Int64, w: Int64, stream=fx.Stream(None)): + put_kernel(dev_comm, w).launch( + grid=(1, 1, 1), block=[THREADS, 1, 1], stream=stream + ) + + comm.barrier() + run(dc.ptr, win.handle) + sync() + comm.barrier() + + # recv slot s must hold rank s's send buffer: s*1000 + i + for s in range(world_size): + if s == my_rank: + continue + got = read(win.local_ptr + recv_off + s * NBYTES, NUM_ELEMS, u32) + print( + f"[rank {my_rank}] AFTER a2a recv[from {s}][:8] = " + f"{list(got[:8])}", + flush=True, + ) + for i in (0, NUM_ELEMS // 2, NUM_ELEMS - 1): + exp = s * 1000 + i + if got[i] != exp: + raise AssertionError( + f"rank {my_rank} recv[{s}][{i}]={got[i]} expected {exp}" + ) + comm.destroy() + except Exception: + err = traceback.format_exc() + with open(results_path % rank, "w") as f: + f.write(err or "") + + +def _worker_put_quiet_queue(rank, world_size, uid_path, coop_name, results_path): + """Same all-to-all put, but drained per queue with quiet_queue(peer, 0). + + THREAD-coop put uses queue 0 only, so a single quiet_queue on queue 0 fully + drains it — this exercises the single-(peer,queue) completion path distinct + from the all-queue quiet() used by _worker_put.""" + err = None + try: + mods = _worker_imports() + flyc, fx, Int64, cco = mods["flyc"], mods["fx"], mods["Int64"], mods["cco"] + sync, fill, zero, read = mods["sync"], mods["fill"], mods["zero"], mods["read"] + u32 = mods["ctypes"].c_uint32 + CCODevCommRequirements = mods["CCODevCommRequirements"] + GDA_CONNECTION_NONE = mods["GDA_CONNECTION_NONE"] + + comm = _init_comm(rank, world_size, uid_path) + send_off = 0 + recv_off = NBYTES + win_bytes = recv_off + world_size * NBYTES + + mem = comm.alloc_mem(win_bytes) + win = comm.register_window(mem.ptr, mem.size) + zero(win.local_ptr, win_bytes) + fill(win.local_ptr + send_off, [rank * 1000 + i for i in range(NUM_ELEMS)], u32) + + reqs = CCODevCommRequirements() + reqs.gda_connection_type = GDA_CONNECTION_NONE + reqs.gda_signal_count = 0 + reqs.gda_counter_count = 0 + reqs.sdma_queue_count = 8 + dc = comm.create_dev_comm(reqs) + + my_rank = rank + + @flyc.kernel(known_block_size=[THREADS, 1, 1]) + def put_kernel(dev_comm: Int64, w: Int64): + sdma = cco.DevComm(dev_comm).sdma() + if fx.thread_idx.x == 0: + for p in range(world_size): + if p != my_rank: + sdma.put( + p, + w, + recv_off + my_rank * NBYTES, + w, + send_off, + NBYTES, + 0, + coop=cco.CoopScope.THREAD, + ) + for p in range(world_size): + if p != my_rank: + sdma.quiet_queue(p, 0) + + @flyc.jit + def run(dev_comm: Int64, w: Int64, stream=fx.Stream(None)): + put_kernel(dev_comm, w).launch( + grid=(1, 1, 1), block=[THREADS, 1, 1], stream=stream + ) + + comm.barrier() + run(dc.ptr, win.handle) + sync() + comm.barrier() + + for s in range(world_size): + if s == my_rank: + continue + got = read(win.local_ptr + recv_off + s * NBYTES, NUM_ELEMS, u32) + for i in (0, NUM_ELEMS // 2, NUM_ELEMS - 1): + exp = s * 1000 + i + if got[i] != exp: + raise AssertionError( + f"rank {my_rank} recv[{s}][{i}]={got[i]} expected {exp}" + ) + comm.destroy() + except Exception: + err = traceback.format_exc() + with open(results_path % rank, "w") as f: + f.write(err or "") + + +# ── spawn harness ────────────────────────────────────────────────────────── + + +def _run(worker, world_size, coop_name): + _require_gpus(world_size) + import mori.cco.device.flydsl # noqa: F401 (fail early if flydsl broken) + from mori.cco import Communicator + + os.environ["MORI_ENABLE_SDMA"] = "1" + os.environ.setdefault("MORI_SOCKET_IFNAME", "lo") + + tmp = tempfile.mkdtemp(prefix="cco_sdma_test_") + uid_path = os.path.join(tmp, "uid.bin") + results_path = os.path.join(tmp, "result_%d.txt") + + uid = Communicator.get_unique_id() + with open(uid_path, "wb") as f: + f.write(bytes(uid)) + + # This dir holds helper.py; spawned children need it on sys.path to import it. + this_dir = os.path.dirname(os.path.abspath(__file__)) + + torch.multiprocessing.spawn( + _spawn_entry, + args=(worker, world_size, uid_path, coop_name, results_path, this_dir), + nprocs=world_size, + join=True, + ) + + errs = [] + for rank in range(world_size): + with open(results_path % rank) as f: + e = f.read() + if e: + errs.append(f"rank {rank}:\n{e}") + assert not errs, "\n\n".join(errs) + + +def _spawn_entry(rank, worker, world_size, uid_path, coop_name, results_path, this_dir): + import sys + + sys.path.insert(0, this_dir) # so the child can import helper + worker(rank, world_size, uid_path, coop_name, results_path) + + +# ── tests ────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("world_size", [2, 4, 8]) +@pytest.mark.parametrize("coop_name", ["THREAD", "WARP"]) +def test_sdma_put_alltoall(coop_name, world_size): + """SDMA put moves the exact source bytes into each peer's slot; quiet drains. + + Scales across 2/4/8 GPUs; cases needing more GPUs than present are skipped.""" + _run(_worker_put, world_size=world_size, coop_name=coop_name) + + +def test_sdma_quiet_queue(): + """Single-queue completion (quiet_queue) drains a thread-coop put correctly.""" + _run(_worker_put_quiet_queue, world_size=2, coop_name="THREAD")