diff --git a/benchmark/nixlbench/src/kernels/nixlbench_device_launch.cu b/benchmark/nixlbench/src/kernels/nixlbench_device_launch.cu index 161d51d28c..0578f738b4 100644 --- a/benchmark/nixlbench/src/kernels/nixlbench_device_launch.cu +++ b/benchmark/nixlbench/src/kernels/nixlbench_device_launch.cu @@ -23,6 +23,13 @@ namespace { constexpr unsigned kWarpSize = 32; // Assumed equal to device warpSize (CUDA guarantee); constexpr unsigned kMaxGroups = 32; // 32 threads or max 1024 / 32 warps per block +__device__ __forceinline__ uint64_t +nixlbenchGetTimeNs() { + uint64_t global_timer; + asm volatile("mov.u64 %0, %globaltimer;" : "=l"(global_timer)); + return global_timer; +} + template __device__ nixl_status_t nixlbenchPollXferStatus(nixl_status_t status, nixlGpuXferStatusH &xfer_status) { @@ -36,22 +43,24 @@ template __device__ nixl_status_t nixlbenchPostPut(const nixlbenchDeviceXferParams ¶ms, size_t region_idx, + unsigned channel_id, nixlGpuXferStatusH &xfer_status) { const nixlMemViewElem src{params.localMvh, region_idx, 0}; const nixlMemViewElem dst{params.remoteMvh, region_idx, 0}; - return nixlPut(src, dst, params.regionSize, 0, 0, &xfer_status); + return nixlPut(src, dst, params.regionSize, channel_id, 0, &xfer_status); } +template __device__ nixl_status_t nixlbenchSignalCounter(const nixlbenchDeviceXferParams ¶ms, size_t counter_offset, uint64_t value, + unsigned channel_id, + nixlGpuXferStatusH &xfer_status, const char *counter_name) { - const nixlMemViewElem counter{params.remoteMvh, params.numRegions, counter_offset}; - nixlGpuXferStatusH xfer_status; - nixl_status_t status = - nixlAtomicAdd(value, counter, 0, 0, &xfer_status); - status = nixlbenchPollXferStatus(status, xfer_status); + const nixlMemViewElem counter{params.remoteMvh, params.counterIndex, counter_offset}; + nixl_status_t status = nixlAtomicAdd(value, counter, channel_id, 0, &xfer_status); + status = nixlbenchPollXferStatus(status, xfer_status); if (status != NIXL_SUCCESS) { printf("[nixlbenchSignalCounter] nixlAtomicAdd(%s) did not complete: final_status=%d\n", @@ -61,19 +70,37 @@ nixlbenchSignalCounter(const nixlbenchDeviceXferParams ¶ms, return status; } +template __device__ nixl_status_t -nixlbenchSignalCompletion(nixlbenchDeviceXferParams params) { - return nixlbenchSignalCounter(params, params.completionCounterOffsetBytes, 1ull, "completion"); +nixlbenchSignalCompletion(const nixlbenchDeviceXferParams ¶ms, + uint64_t num_iterations, + unsigned channel_id, + nixlGpuXferStatusH &xfer_status) { + return nixlbenchSignalCounter(params, + params.completionCounterOffsetBytes, + num_iterations, + channel_id, + xfer_status, + "completion"); } +template __device__ nixl_status_t -nixlbenchSignalError(nixlbenchDeviceXferParams params) { - return nixlbenchSignalCounter(params, params.errorCounterOffsetBytes, 1ull, "error"); +nixlbenchSignalError(const nixlbenchDeviceXferParams ¶ms, + unsigned channel_id, + nixlGpuXferStatusH &xfer_status) { + return nixlbenchSignalCounter( + params, params.errorCounterOffsetBytes, 1ull, channel_id, xfer_status, "error"); } /** * Performs device-initiated NIXL PUT transfers and reports completion or errors * through remote counters. + * + * Every group runs @c numIterations iterations of the complete region list, + * so the block as a whole performs @c numIterations * num_groups transfers of the list. + * Groups keep independent transfer status and timing samples, + * signal its own share of the completion counter, so no group synchronization is needed. */ template __global__ void @@ -87,38 +114,57 @@ nixlbenchPutKernel(nixlbenchDeviceXferParams params) { group_id = threadIdx.x / warpSize; num_groups = (blockDim.x + warpSize - 1) / warpSize; } - nixlGpuXferStatusH &xfer_status = xfer_statuses[group_id]; - nixl_status_t put_status = NIXL_SUCCESS; - for (size_t region_idx = group_id; region_idx < params.numRegions; region_idx += num_groups) { - put_status = nixlbenchPostPut(params, region_idx, xfer_status); - if (put_status != NIXL_IN_PROG) { - break; + const unsigned channel_id = group_id % params.channelNum; + const bool group_leader = Level == nixl_gpu_level_t::THREAD || threadIdx.x % warpSize == 0; + const size_t region_base = group_id * params.numRegions; + bool group_failed = false; + + for (uint64_t iter = 0; iter < params.numIterations && !group_failed; ++iter) { + const uint64_t post_start_ns = nixlbenchGetTimeNs(); + + nixl_status_t put_status = NIXL_SUCCESS; + for (size_t region_idx = 0; region_idx < params.numRegions; ++region_idx) { + put_status = + nixlbenchPostPut(params, region_base + region_idx, channel_id, xfer_status); + if (put_status != NIXL_IN_PROG) { + break; + } } - } - if (put_status == NIXL_IN_PROG) { - put_status = nixlbenchPollXferStatus(put_status, xfer_status); - } - if (put_status != NIXL_SUCCESS && - (Level == nixl_gpu_level_t::THREAD || threadIdx.x % warpSize == 0)) { - printf("[nixlbenchPutKernel] transfer did not complete: " - "threadIdx.x=%u blockIdx.x=%u blockDim.x=%u final_status=%d\n", - threadIdx.x, - blockIdx.x, - blockDim.x, - static_cast(put_status)); - } + const uint64_t post_end_ns = nixlbenchGetTimeNs(); - const bool any_put_failed = __syncthreads_or(put_status != NIXL_SUCCESS); - if (threadIdx.x == 0) { - if (any_put_failed) { - (void)nixlbenchSignalError(params); - return; + if (put_status == NIXL_IN_PROG) { + put_status = nixlbenchPollXferStatus(put_status, xfer_status); + } + if constexpr (Level == nixl_gpu_level_t::WARP) { + __syncwarp(); } + const uint64_t xfer_end_ns = nixlbenchGetTimeNs(); - if (nixlbenchSignalCompletion(params) != NIXL_SUCCESS) { - (void)nixlbenchSignalError(params); + if (group_leader) { + const size_t sample_idx = iter * num_groups + group_id; + params.postDurationNs[sample_idx] = post_end_ns - post_start_ns; + params.xferDurationNs[sample_idx] = xfer_end_ns - post_end_ns; } + + if (put_status != NIXL_SUCCESS) { + if (group_leader) { + printf("[nixlbenchPutKernel] transfer did not complete: " + "threadIdx.x=%u blockIdx.x=%u blockDim.x=%u final_status=%d\n", + threadIdx.x, + blockIdx.x, + blockDim.x, + static_cast(put_status)); + } + group_failed = true; + } + } + + if (group_failed) { + (void)nixlbenchSignalError(params, channel_id, xfer_status); + } else if (nixlbenchSignalCompletion( + params, params.numIterations, channel_id, xfer_status) != NIXL_SUCCESS) { + (void)nixlbenchSignalError(params, channel_id, xfer_status); } } @@ -132,21 +178,32 @@ nixlbenchLaunchDevicePut(const nixlbenchDeviceXferParams ¶ms, unsigned block return NIXL_ERR_INVALID_PARAM; } + if (params.postDurationNs == nullptr || params.xferDurationNs == nullptr) { + std::cerr << "nixlbench: nixlbenchLaunchDevicePut: duration output buffers are required " + "(postDurationNs and xferDurationNs must hold numIterations * num_groups " + "entries)\n"; + return NIXL_ERR_INVALID_PARAM; + } + if (params.channelNum == 0) { + std::cerr << "nixlbench: nixlbenchLaunchDevicePut: channelNum must be greater than zero\n"; + return NIXL_ERR_INVALID_PARAM; + } + if (block_threads == 0 || block_threads > 1024u) { std::cerr << "nixlbench: nixlbenchLaunchDevicePut: invalid block_threads=" << block_threads << " (must be 1..1024)\n"; return NIXL_ERR_INVALID_PARAM; } + if (block_threads > kWarpSize && block_threads % kWarpSize != 0) { + std::cerr << "nixlbench: nixlbenchLaunchDevicePut: block_threads (" << block_threads + << ") must be a multiple of " << kWarpSize + << " (WARP-level nixlPut requires full warps)\n"; + return NIXL_ERR_INVALID_PARAM; + } if (block_threads <= kWarpSize) { nixlbenchPutKernel<<<1, block_threads, 0, nullptr>>>(params); } else { - if (block_threads % kWarpSize != 0) { - std::cerr << "nixlbench: nixlbenchLaunchDevicePut: block_threads (" << block_threads - << ") must be a multiple of " << kWarpSize - << " (WARP-level nixlPut requires full warps)\n"; - return NIXL_ERR_INVALID_PARAM; - } nixlbenchPutKernel<<<1, block_threads, 0, nullptr>>>(params); } diff --git a/benchmark/nixlbench/src/kernels/nixlbench_device_launch.cuh b/benchmark/nixlbench/src/kernels/nixlbench_device_launch.cuh index 0c26349748..e92604084f 100644 --- a/benchmark/nixlbench/src/kernels/nixlbench_device_launch.cuh +++ b/benchmark/nixlbench/src/kernels/nixlbench_device_launch.cuh @@ -8,27 +8,38 @@ #include #include +#include /** * @brief Parameters for @ref nixlbenchPutKernel (passed by value to the device). * * @a localMvh and @a remoteMvh must come from nixlAgent::prepMemView using the same flattening - * order as xferBenchNixlWorker::prepareGPULocalView / prepareGPURemoteView (outer vector = thread - * lists, inner vector = IOVs for that thread). + * order as xferBenchNixlWorker::prepareGPULocalView / prepareGPURemoteView (outer vector = group + * lists, inner vector = IOVs for that group). * - * @a numRegions is the **data** region count (put loop uses indices @c 0 .. @a numRegions-1). - * The host must append a counter buffer as the last remote descriptor, so the view has - * @a numRegions + 1 regions. The counter buffer stores: + * @a numRegions is data region count: group @c g owns indices @c g*numRegions .. @c + * (g+1)*numRegions-1. The host must append a counter buffer after all data descriptors, at index @a + * counterIndex. The counter buffer stores: * - done counter at byte offset @a completionCounterOffsetBytes * - error counter at byte offset @a errorCounterOffsetBytes * - * Kernel uses @c nixlAtomicAdd on @c { remoteMvh, numRegions, offset }. + * Every group in the launched block transfers @a numIterations times. + * The block as a whole performs @c numIterations * num_groups list transfers. + * + * Each group signals independently using @c nixlAtomicAdd on @c { remoteMvh, counterIndex, offset } + * over channel @c group_id%channelNum to add @a numIterations to the done counter. + * Duration outputs contain @c numIterations * num_groups entries in iteration-major order. */ struct nixlbenchDeviceXferParams { nixlMemViewH localMvh; ///< Local memory view from prepMemView nixlMemViewH remoteMvh; ///< Remote memory view from prepMemView - size_t numRegions; ///< Data region count (puts); completion index when signaling + size_t numRegions; ///< Data region count (puts) + size_t counterIndex; ///< Index of counter buffer (= numRegions * num_groups) size_t regionSize; ///< Bytes per region for this transfer pattern + uint64_t numIterations; ///< Per-group number of complete region-list transfers + unsigned channelNum; ///< Logical channels shared by groups using group_id % channelNum + uint64_t *postDurationNs; ///< Per-iteration, per-group PUT posting duration output + uint64_t *xferDurationNs; ///< Per-iteration, per-group completion polling duration output size_t completionCounterOffsetBytes; ///< Done counter offset in the counter region size_t errorCounterOffsetBytes; ///< Error counter offset in the counter region }; @@ -37,10 +48,9 @@ struct nixlbenchDeviceXferParams { * @brief Launches @ref nixlbenchPutKernel with a 1-D block of @a block_threads threads. * * If @a block_threads is less than or equal to the GPU warp size (32), - * @c nixl_gpu_level_t::THREAD is used; - * otherwise @c nixl_gpu_level_t::WARP is used (each warp strides over regions and all lanes in - * the warp participate in each device API call). Typical - * @a block_threads matches nixlbench @c --num_threads. + * @c nixl_gpu_level_t::THREAD is used (one group per thread); + * otherwise @c nixl_gpu_level_t::WARP is used (one group per warp). + * Typical @a block_threads matches nixlbench @c --num_threads. * * Requires NIXL UCX GPU Device API support. @a block_threads must be in [1, 1024]; * values greater than 32 must be a multiple of 32. diff --git a/benchmark/nixlbench/src/main.cpp b/benchmark/nixlbench/src/main.cpp index 7266bf7f23..cf3d86d758 100644 --- a/benchmark/nixlbench/src/main.cpp +++ b/benchmark/nixlbench/src/main.cpp @@ -186,7 +186,7 @@ createWorker() { static int runBenchmark() { int ret = 0; - int num_threads = xferBenchConfig::num_threads; + int num_workers = xferBenchConfig::workerNum(); // Create the appropriate worker based on worker configuration std::unique_ptr worker_ptr = createWorker(); @@ -204,7 +204,7 @@ runBenchmark() { return EXIT_FAILURE; } - std::vector> iov_lists = worker_ptr->allocateMemory(num_threads); + std::vector> iov_lists = worker_ptr->allocateMemory(num_workers); auto mem_guard = make_scope_guard ([&] { worker_ptr->deallocateMemory(iov_lists); }); @@ -225,7 +225,7 @@ runBenchmark() { !worker_ptr->signaled() && block_size <= xferBenchConfig::max_block_size; block_size *= 2) { - ret = processBatchSizes(*worker_ptr, iov_lists, block_size, num_threads); + ret = processBatchSizes(*worker_ptr, iov_lists, block_size, num_workers); if (0 != ret) { return EXIT_FAILURE; } diff --git a/benchmark/nixlbench/src/utils/utils.cpp b/benchmark/nixlbench/src/utils/utils.cpp index 977557d2ac..1ba445d271 100644 --- a/benchmark/nixlbench/src/utils/utils.cpp +++ b/benchmark/nixlbench/src/utils/utils.cpp @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -248,6 +249,11 @@ NB_ARG_BOOL(use_device_api, "When enabled, --num_threads is repurposed as the CUDA kernel " "block size (num_threads <= 32 -> THREAD level; > 32 -> WARP level, must be a " "multiple of 32), and the internal CPU thread count is forced to 1."); +NB_ARG_INT32(device_channel_num, + 1, + "Number of logical UCX Device API channels, default to 1. " + "0 means one channel per execution group. " + "Only used when --use_device_api is enabled."); #undef NB_ARG_INT32 #undef NB_ARG_UINT32 @@ -331,6 +337,7 @@ bool xferBenchConfig::gusli_try_use_uring = false; std::optional xferBenchConfig::plugin_parameters = std::nullopt; bool xferBenchConfig::use_device_api = false; int xferBenchConfig::block_threads = 1; +int xferBenchConfig::device_channel_num = 0; static bool validateDeviceAPIConfig() { @@ -364,6 +371,10 @@ validateDeviceAPIConfig() { if (xferBenchConfig::pipeline_depth != 1) { return reject("pipeline_depth must be 1"); } + if (std::getenv("UCX_RC_GDA_NUM_CHANNELS") != nullptr) { + return reject("UCX_RC_GDA_NUM_CHANNELS must not be set; " + "use --device_channel_num to configure Device API channels"); + } return true; #else return reject("UCX GPU Device API support is not enabled in this build. " @@ -371,6 +382,56 @@ validateDeviceAPIConfig() { #endif } +static bool +setupDeviceAPIConfig() { + const int num_threads = xferBenchConfig::num_threads; + + if ((num_threads < 1) || (num_threads > 1024)) { + std::cerr << "Invalid value for --num_threads: " << num_threads + << ". Device API requires a GPU kernel block thread count in [1, 1024]" + << std::endl; + return false; + } + if ((num_threads > XFERBENCH_DEVICE_WARP_SIZE) && + (num_threads % XFERBENCH_DEVICE_WARP_SIZE != 0)) { + std::cerr << "Invalid value for --num_threads: " << num_threads + << ". Device API requires block_threads > 32 must be a multiple of 32" + << std::endl; + return false; + } + if (xferBenchConfig::device_channel_num < 0) { + std::cerr << "Invalid value for --device_channel_num: " + << xferBenchConfig::device_channel_num + << ". Device API channel number must be >= 0" << std::endl; + return false; + } + + xferBenchConfig::block_threads = num_threads; + const int group_num = xferBenchConfig::deviceGroupNum(); + + if (xferBenchConfig::device_channel_num == 0) { + xferBenchConfig::device_channel_num = group_num; + } else if (xferBenchConfig::device_channel_num > group_num) { + std::cout << "WARNING: Adjusting --device_channel_num from " + << xferBenchConfig::device_channel_num << " to Device API group number " + << group_num << std::endl; + xferBenchConfig::device_channel_num = group_num; + } + if (xferBenchConfig::num_iter < group_num) { + std::cerr << "Invalid value for --num_iter: " << xferBenchConfig::num_iter + << " , must not be smaller than Device API group number: " << group_num + << std::endl; + return false; + } + + xferBenchConfig::num_threads = 1; + std::cout << "Device API mode: kernel block_threads = " << xferBenchConfig::block_threads + << ", group_num = " << group_num + << ", channel_num = " << xferBenchConfig::device_channel_num + << ", num_threads forced to 1" << std::endl; + return true; +} + int xferBenchConfig::parseConfig(int argc, char *argv[]) { plugin_parameters.reset(); @@ -575,27 +636,10 @@ xferBenchConfig::loadParams(void) { return -1; } use_device_api = NB_ARG(use_device_api); - if (use_device_api && !validateDeviceAPIConfig()) { + device_channel_num = NB_ARG(device_channel_num); + if (use_device_api && (!validateDeviceAPIConfig() || !setupDeviceAPIConfig())) { return -1; } - if (use_device_api) { - if (num_threads < 1 || num_threads > 1024) { - std::cerr << "Invalid value for --num_threads: " << num_threads - << ". Device API requires a GPU kernel block thread count in [1, 1024]" - << std::endl; - return -1; - } - if (num_threads > 32 && num_threads % 32 != 0) { - std::cerr << "Invalid value for --num_threads: " << num_threads - << ". Device API requires block_threads > 32 must be a multiple of 32" - << std::endl; - return -1; - } - block_threads = num_threads; - num_threads = 1; - std::cout << "Device API mode: kernel block_threads=" << block_threads - << ", num_threads forced to 1" << std::endl; - } etcd_endpoints = NB_ARG(etcd_endpoints); asio_address = NB_ARG(asio_address); asio_port = NB_ARG(asio_port); @@ -732,12 +776,15 @@ xferBenchConfig::loadParams(void) { << std::endl; return -1; } - if ((max_block_size * max_batch_size) > (total_buffer_size / num_threads)) { + const int workers = workerNum(); + const char *worker_kind = use_device_api ? "groups" : "threads"; + + if ((max_block_size * max_batch_size) > (total_buffer_size / workers)) { std::cerr << "Incorrect buffer size configuration " << "(max_block_size * max_batch_size) " << "(" << (max_block_size * max_batch_size) << ")" - << " is > (total_buffer_size / num_threads) (" - << (total_buffer_size / num_threads) << ")" << std::endl; + << " is > (total_buffer_size / " << workers << " " << worker_kind << ") (" + << (total_buffer_size / workers) << ")" << std::endl; return -1; } @@ -746,29 +793,31 @@ xferBenchConfig::loadParams(void) { return -1; } - int partition = (num_threads * large_blk_iter_ftr); + int partition = (workers * large_blk_iter_ftr); if (num_iter % partition) { num_iter += partition - (num_iter % partition); std::cout << "WARNING: Adjusting num_iter to " << num_iter - << " to allow equal distribution to " << num_threads << " threads" << std::endl; + << " to allow equal distribution to " << workers << " " << worker_kind + << std::endl; } if (warmup_iter % partition) { warmup_iter += partition - (warmup_iter % partition); std::cout << "WARNING: Adjusting warmup_iter to " << warmup_iter - << " to allow equal distribution to " << num_threads << " threads" << std::endl; + << " to allow equal distribution to " << workers << " " << worker_kind + << std::endl; } - partition = (num_initiator_dev * num_threads); + partition = (num_initiator_dev * workers); if (total_buffer_size % partition) { - std::cerr << "Total_buffer_size must be divisible by the product of num_threads and " - "num_initiator_dev" + std::cerr << "Total_buffer_size must be divisible by the product of " << workers << " " + << worker_kind << " and num_initiator_dev" << ", next such value is " << total_buffer_size + partition - (total_buffer_size % partition) << std::endl; return -1; } - partition = (num_target_dev * num_threads); + partition = (num_target_dev * workers); if (total_buffer_size % partition) { - std::cerr << "Total_buffer_size must be divisible by the product of num_threads and " - "num_target_dev" + std::cerr << "Total_buffer_size must be divisible by the product of " << workers << " " + << worker_kind << " and num_target_dev" << ", next such value is " << total_buffer_size + partition - (total_buffer_size % partition) << std::endl; return -1; @@ -923,6 +972,8 @@ xferBenchConfig::printConfig() { if (use_device_api) { printOption("Device API Kernel block threads (--num_threads=N)", std::to_string(block_threads)); + printOption("Device API channels (--device_channel_num=N)", + std::to_string(device_channel_num)); } printSeparator('-'); std::cout << std::endl; @@ -955,6 +1006,19 @@ xferBenchConfig::parseDeviceList() { return devices; } +int +xferBenchConfig::deviceGroupNum() { + return xferBenchConfig::block_threads <= XFERBENCH_DEVICE_WARP_SIZE ? + xferBenchConfig::block_threads : + xferBenchConfig::block_threads / XFERBENCH_DEVICE_WARP_SIZE; +} + +int +xferBenchConfig::workerNum() { + return xferBenchConfig::use_device_api ? xferBenchConfig::deviceGroupNum() : + xferBenchConfig::num_threads; +} + bool xferBenchConfig::isStorageBackend() { return (XFERBENCH_BACKEND_GDS == xferBenchConfig::backend || @@ -1344,7 +1408,7 @@ xferBenchUtils::printStats(bool is_target, double totalbw = 0; int total_iter = xferBenchConfig::num_iter; - int per_thread_iter = total_iter / xferBenchConfig::num_threads; + int per_thread_iter = total_iter / xferBenchConfig::workerNum(); if (block_size > LARGE_BLOCK_SIZE) { total_iter /= xferBenchConfig::large_blk_iter_ftr; diff --git a/benchmark/nixlbench/src/utils/utils.h b/benchmark/nixlbench/src/utils/utils.h index 8af75c0219..c6ea26acba 100644 --- a/benchmark/nixlbench/src/utils/utils.h +++ b/benchmark/nixlbench/src/utils/utils.h @@ -81,6 +81,9 @@ #define XFERBENCH_INITIATOR_BUFFER_ELEMENT 0xbb #define XFERBENCH_TARGET_BUFFER_ELEMENT 0xaa +// CUDA warp size, to derive Device API group number from block_threads +#define XFERBENCH_DEVICE_WARP_SIZE 32 + // Runtime types #define XFERBENCH_RT_ETCD "ETCD" #define XFERBENCH_RT_ASIO "ASIO" @@ -232,6 +235,15 @@ class xferBenchConfig { static std::optional plugin_parameters; static bool use_device_api; static int block_threads; + static int device_channel_num; + + /* Number of independent groups Device API kernel runs with. */ + static int + deviceGroupNum(); + + /* Parallel workers split iterations across both CPU and Device API. */ + static int + workerNum(); static int parseConfig(int argc, char *argv[]); diff --git a/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp b/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp index a54f2db702..d6a5d75b25 100644 --- a/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp +++ b/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp @@ -25,6 +25,7 @@ #include #include #include "kernels/nixlbench_device_launch.cuh" +#include #include #include #include @@ -77,6 +78,7 @@ resolveVramSegment() { constexpr size_t kDeviceCounterDoneOffsetBytes = 0; constexpr size_t kDeviceCounterErrorOffsetBytes = sizeof(uint64_t); constexpr size_t kDeviceCounterBytes = 2 * sizeof(uint64_t); +constexpr uint64_t kUnwrittenDeviceDurationNs = std::numeric_limits::max(); // Reuse parser from utils @@ -155,6 +157,13 @@ xferBenchNixlWorker::xferBenchNixlWorker(const std::vector &devices if (0 == xferBenchConfig::backend.compare(XFERBENCH_BACKEND_UCX)) { backend_params["num_threads"] = std::to_string(xferBenchConfig::progress_threads); + if (xferBenchConfig::use_device_api) { + backend_params["ucx_num_device_channels"] = + std::to_string(xferBenchConfig::device_channel_num); + std::cout << "Device API mode: configuring " << xferBenchConfig::device_channel_num + << " UCX device channels for " << xferBenchConfig::deviceGroupNum() + << " execution groups" << std::endl; + } // No need to set device_list if all is specified // fallback to backend preference @@ -1887,50 +1896,124 @@ execTransfer(nixlAgent *agent, return ret; } +// Runs per_group_iter iterations of the full region list on every kernel group static int execDeviceTransfer(nixlMemViewH local_mvh, nixlMemViewH remote_mvh, - const int num_iter, - const int num_threads, + const int per_group_iter, + const int block_threads, size_t num_regions, size_t region_size, + const bool collect_iteration_stats, xferBenchStats &stats, const std::atomic *terminate_ptr = nullptr) { #ifdef HAVE_UCX_GPU_DEVICE_API stats.clear(); + if (per_group_iter < 1) { + std::cerr << "NIXL Device API requires at least one transfer iteration per group" + << std::endl; + return -1; + } + if (__builtin_expect(terminate_ptr && terminate_ptr->load(), 0)) { + return -1; + } + + const size_t num_groups = static_cast(xferBenchConfig::deviceGroupNum()); + const size_t sample_count = static_cast(per_group_iter) * num_groups; + const size_t duration_bytes = sample_count * sizeof(uint64_t); + uint64_t *device_post_duration_ns = nullptr; + uint64_t *device_xfer_duration_ns = nullptr; + auto duration_buffer_guard = + make_scope_guard([&device_post_duration_ns, &device_xfer_duration_ns] { + auto free_buffer = [](uint64_t *buffer) { + if (buffer == nullptr) { + return; + } + const cudaError_t error = cudaFree(buffer); + if (error != cudaSuccess) { + std::cerr << "Failed to free Device API duration buffer: " + << cudaGetErrorString(error) << std::endl; + } + }; + free_buffer(device_post_duration_ns); + free_buffer(device_xfer_duration_ns); + }); + CHECK_CUDA_ERROR(cudaMalloc(&device_post_duration_ns, duration_bytes), + "Failed to allocate Device API post duration buffer"); + CHECK_CUDA_ERROR(cudaMalloc(&device_xfer_duration_ns, duration_bytes), + "Failed to allocate Device API transfer duration buffer"); + CHECK_CUDA_ERROR(cudaMemset(device_post_duration_ns, 0xFF, duration_bytes), + "Failed to initialize Device API post duration buffer"); + CHECK_CUDA_ERROR(cudaMemset(device_xfer_duration_ns, 0xFF, duration_bytes), + "Failed to initialize Device API transfer duration buffer"); + CHECK_CUDA_ERROR(cudaStreamSynchronize(0), + "Failed to synchronize Device API duration buffer initialization"); + nixlbenchDeviceXferParams params; params.localMvh = local_mvh; params.remoteMvh = remote_mvh; params.numRegions = num_regions; + // prepareGPURemoteView() flattens num_groups lists of num_regions descriptors and appends the + // counter buffer, so the counter sits right after the last data descriptor. + params.counterIndex = num_regions * num_groups; params.regionSize = region_size; + params.numIterations = static_cast(per_group_iter); + params.channelNum = static_cast(xferBenchConfig::device_channel_num); + params.postDurationNs = device_post_duration_ns; + params.xferDurationNs = device_xfer_duration_ns; params.completionCounterOffsetBytes = kDeviceCounterDoneOffsetBytes; params.errorCounterOffsetBytes = kDeviceCounterErrorOffsetBytes; xferBenchTimer total_timer; - stats.transfer_duration.reserve(num_iter); - xferBenchTimer timer; - for (int i = 0; i < num_iter; ++i) { - if (__builtin_expect(terminate_ptr && terminate_ptr->load(), 0)) { - stats.total_duration.add(total_timer.lap()); - return -1; - } - nixl_status_t st = nixlbenchLaunchDevicePut(params, static_cast(num_threads)); - if (__builtin_expect(st != NIXL_SUCCESS, 0)) { - std::cerr << "nixlbenchLaunchDevicePut failed: " << nixlEnumStrings::statusStr(st) - << std::endl; - stats.total_duration.add(total_timer.lap()); - return -1; + nixl_status_t st = nixlbenchLaunchDevicePut(params, static_cast(block_threads)); + const nixlTime::us_t total_duration = total_timer.lap(); + stats.total_duration.add(total_duration); + if (__builtin_expect(st != NIXL_SUCCESS, 0)) { + std::cerr << "nixlbenchLaunchDevicePut failed: " << nixlEnumStrings::statusStr(st) + << std::endl; + return -1; + } + if (__builtin_expect(terminate_ptr && terminate_ptr->load(), 0)) { + return -1; + } + if (collect_iteration_stats) { + std::vector post_duration_ns(sample_count); + std::vector xfer_duration_ns(sample_count); + CHECK_CUDA_ERROR(cudaMemcpy(post_duration_ns.data(), + device_post_duration_ns, + duration_bytes, + cudaMemcpyDeviceToHost), + "Failed to copy Device API post duration samples"); + CHECK_CUDA_ERROR(cudaMemcpy(xfer_duration_ns.data(), + device_xfer_duration_ns, + duration_bytes, + cudaMemcpyDeviceToHost), + "Failed to copy Device API transfer duration samples"); + stats.post_duration.reserve(sample_count); + stats.transfer_duration.reserve(sample_count); + for (size_t iter = 0; iter < static_cast(per_group_iter); ++iter) { + for (size_t group_id = 0; group_id < num_groups; ++group_id) { + const size_t sample_idx = iter * num_groups + group_id; + if (post_duration_ns[sample_idx] == kUnwrittenDeviceDurationNs || + xfer_duration_ns[sample_idx] == kUnwrittenDeviceDurationNs) { + std::cerr << "NIXL Device API produced an incomplete duration sample" + << std::endl; + return -1; + } + stats.post_duration.add(static_cast(post_duration_ns[sample_idx]) / 1000.0); + stats.transfer_duration.add(static_cast(xfer_duration_ns[sample_idx]) / + 1000.0); + } } - stats.transfer_duration.add(timer.lap()); } - stats.total_duration.add(total_timer.lap()); return 0; #else (void)local_mvh; (void)remote_mvh; - (void)num_iter; - (void)num_threads; + (void)per_group_iter; + (void)block_threads; (void)num_regions; (void)region_size; + (void)collect_iteration_stats; (void)stats; (void)terminate_ptr; std::cerr << "NIXL Device API support is not enabled in this build" << std::endl; @@ -2005,8 +2088,9 @@ std::variant xferBenchNixlWorker::transfer(size_t block_size, const std::vector> &local_iovs, const std::vector> &remote_iovs) { - int num_iter = xferBenchConfig::num_iter / xferBenchConfig::num_threads; - int skip = xferBenchConfig::warmup_iter / xferBenchConfig::num_threads; + const int workers = xferBenchConfig::workerNum(); + int num_iter = xferBenchConfig::num_iter / workers; + int skip = xferBenchConfig::warmup_iter / workers; xferBenchStats stats; int ret = 0; nixl_xfer_op_t xfer_op = XFERBENCH_OP_READ == xferBenchConfig::op_type ? NIXL_READ : NIXL_WRITE; @@ -2029,12 +2113,15 @@ xferBenchNixlWorker::transfer(size_t block_size, releaseMemView(local_mvh); }); + // Regions owned by a single execution group. The memory views flatten one IOV list per group + // in order, so group g owns indices [g * num_regions, (g + 1) * num_regions). size_t num_regions = 0; if (xferBenchConfig::use_device_api) { - if (local_iovs.size() != 1 || remote_iovs.size() != 1) { - std::cerr << "NIXL Device API requires exactly one local and one remote IOV list: " - << "local=" << local_iovs.size() << ", remote=" << remote_iovs.size() - << std::endl; + const size_t num_groups = static_cast(xferBenchConfig::deviceGroupNum()); + if (local_iovs.size() != num_groups || remote_iovs.size() != num_groups) { + std::cerr << "NIXL Device API requires one local and one remote IOV list per group " + << "(" << num_groups << "): local=" << local_iovs.size() + << ", remote=" << remote_iovs.size() << std::endl; return std::variant(-1); } const size_t local_regions = local_iovs.front().size(); @@ -2057,6 +2144,7 @@ xferBenchNixlWorker::transfer(size_t block_size, xferBenchConfig::block_threads, num_regions, block_size, + false, // collect_iteration_stats stats, &terminate); } else { @@ -2085,6 +2173,7 @@ xferBenchNixlWorker::transfer(size_t block_size, xferBenchConfig::block_threads, num_regions, block_size, + isMasterRank(), // collect_iteration_stats stats, &terminate); } else {