From c2f88f5b9de6dfc4b0a71796aad3dc6306b60618 Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Sun, 17 May 2026 10:59:30 +0300 Subject: [PATCH 01/18] test: fix device API max threads constant Signed-off-by: Tomer Davidor --- test/gtest/device_api/utils.cuh | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/gtest/device_api/utils.cuh b/test/gtest/device_api/utils.cuh index 9f3489980d..3d9ebc36e7 100644 --- a/test/gtest/device_api/utils.cuh +++ b/test/gtest/device_api/utils.cuh @@ -31,7 +31,9 @@ #include #include -#define MAX_THREADS 1024 +// 512 * sizeof(nixlGpuXferStatusH) (64B) = 32 KiB, under the 48 KiB static +// __shared__ limit (same on all supported arches). 1024 would overflow. +#define MAX_THREADS 512 #define UCS_NSEC_PER_SEC 1000000000ul #define NS_TO_SEC(ns) ((ns) * 1.0 / (UCS_NSEC_PER_SEC)) From aa9a28cfbf65f7732c815af36bc0c5c6e49d3f0d Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Wed, 29 Apr 2026 18:11:40 +0300 Subject: [PATCH 02/18] gpu: device_api: restructure GPU Device API into common/ucx/proxy layers Factor the existing UCX-only GPU device headers into a backend-agnostic common layer (nixl_device_types, nixl_device_api, nixl_device_wrappers) and move the UCX implementation detail into ucx/nixl_device_impl.cuh. Add a proxy backend variant under src/api/gpu/proxy/ that compiles into the nixl_gpu_proxy_lib static library. The proxy headers depend on proxy_protocol.h (the shared GPU/CPU ring-buffer protocol), which is introduced here alongside the GPU layer it serves. The top-level src/api/gpu/nixl_device.cuh dispatches to the selected backend implementation via the NIXL_GPU_DEVICE_BACKEND_PROXY macro. Signed-off-by: Tomer Davidor --- meson.build | 1 + src/api/gpu/common/nixl_device_api.cuh | 74 ++++++++ src/api/gpu/common/nixl_device_types.cuh | 41 ++++ src/api/gpu/common/nixl_device_wrappers.cuh | 61 ++++++ src/api/gpu/meson.build | 28 +++ src/api/gpu/nixl_device.cuh | 32 ++++ src/api/gpu/proxy/nixl_device.cuh | 34 ++++ src/api/gpu/proxy/nixl_device_impl.cuh | 131 +++++++++++++ src/api/gpu/proxy/nixl_device_proxy.cu | 22 +++ src/api/gpu/proxy/nixl_device_proxy.cuh | 195 ++++++++++++++++++++ src/api/gpu/ucx/nixl_device.cuh | 189 ++----------------- src/api/gpu/ucx/nixl_device_impl.cuh | 153 +++++++++++++++ src/api/meson.build | 3 + src/core/device_proxy/proxy_protocol.h | 86 +++++++++ 14 files changed, 873 insertions(+), 177 deletions(-) create mode 100644 src/api/gpu/common/nixl_device_api.cuh create mode 100644 src/api/gpu/common/nixl_device_types.cuh create mode 100644 src/api/gpu/common/nixl_device_wrappers.cuh create mode 100644 src/api/gpu/meson.build create mode 100644 src/api/gpu/nixl_device.cuh create mode 100644 src/api/gpu/proxy/nixl_device.cuh create mode 100644 src/api/gpu/proxy/nixl_device_impl.cuh create mode 100644 src/api/gpu/proxy/nixl_device_proxy.cu create mode 100644 src/api/gpu/proxy/nixl_device_proxy.cuh create mode 100644 src/api/gpu/ucx/nixl_device_impl.cuh create mode 100644 src/core/device_proxy/proxy_protocol.h diff --git a/meson.build b/meson.build index 6e1f7ad7ca..fb5af2ef18 100644 --- a/meson.build +++ b/meson.build @@ -510,6 +510,7 @@ endif nixl_inc_dirs = include_directories('src/api/cpp', 'src/api/cpp/backend', 'src/infra', 'src/core', 'src/core/telemetry') nixl_gpu_inc_dirs = include_directories('src/api/gpu/ucx') +nixl_gpu_proxy_inc_dirs = include_directories('src/api/gpu/proxy') plugins_inc_dirs = include_directories('src/plugins') utils_inc_dirs = include_directories('src/utils') diff --git a/src/api/gpu/common/nixl_device_api.cuh b/src/api/gpu/common/nixl_device_api.cuh new file mode 100644 index 0000000000..d84276af37 --- /dev/null +++ b/src/api/gpu/common/nixl_device_api.cuh @@ -0,0 +1,74 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef NIXL_SRC_API_GPU_COMMON_NIXL_DEVICE_API_CUH +#define NIXL_SRC_API_GPU_COMMON_NIXL_DEVICE_API_CUH + +#include "nixl_device_types.cuh" + +#if defined(NIXL_GPU_DEVICE_BACKEND_PROXY) +#include "../proxy/nixl_device_impl.cuh" + +namespace nixl::gpu { +namespace selected_impl = proxy_impl; +} +#elif defined(NIXL_GPU_DEVICE_BACKEND_UCX) +#include "../ucx/nixl_device_impl.cuh" + +namespace nixl::gpu { +namespace selected_impl = ucx_impl; +} +#else +#error "No GPU device backend implementation selected" +#endif + +namespace nixl::gpu::api { + +template +__device__ inline nixl_status_t +get_xfer_status(nixlGpuXferStatusH &xfer_status) { + return selected_impl::get_xfer_status(xfer_status); +} + +template +__device__ inline nixl_status_t +put(const nixlMemViewElem &src, + const nixlMemViewElem &dst, + size_t size, + unsigned channel_id = 0, + uint64_t flags = 0, + nixlGpuXferStatusH *xfer_status = nullptr) { + return selected_impl::put(src, dst, size, channel_id, flags, xfer_status); +} + +template +__device__ inline nixl_status_t +atomic_add(uint64_t value, + const nixlMemViewElem &counter, + unsigned channel_id = 0, + uint64_t flags = 0, + nixlGpuXferStatusH *xfer_status = nullptr) { + return selected_impl::atomic_add(value, counter, channel_id, flags, xfer_status); +} + +__device__ inline void * +get_ptr(nixlMemViewH mvh, size_t index) { + return selected_impl::get_ptr(mvh, index); +} + +} // namespace nixl::gpu::api + +#endif // NIXL_SRC_API_GPU_COMMON_NIXL_DEVICE_API_CUH diff --git a/src/api/gpu/common/nixl_device_types.cuh b/src/api/gpu/common/nixl_device_types.cuh new file mode 100644 index 0000000000..3a3962f736 --- /dev/null +++ b/src/api/gpu/common/nixl_device_types.cuh @@ -0,0 +1,41 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef NIXL_SRC_API_GPU_COMMON_NIXL_DEVICE_TYPES_CUH +#define NIXL_SRC_API_GPU_COMMON_NIXL_DEVICE_TYPES_CUH + +#include +#include + +#include + +struct nixlGpuXferStatusH { + alignas(16) unsigned char storage[64] = {}; +}; + +enum class nixl_gpu_level_t : uint64_t { THREAD = 0, WARP = 1, BLOCK = 2, GRID = 3 }; + +namespace nixl_gpu_flags { +constexpr uint64_t defer = 1; +} // namespace nixl_gpu_flags + +struct nixlMemViewElem { + nixlMemViewH mvh; + size_t index; /**< Index in the memory view */ + size_t offset; /**< Offset within the buffer */ +}; + +#endif // NIXL_SRC_API_GPU_COMMON_NIXL_DEVICE_TYPES_CUH diff --git a/src/api/gpu/common/nixl_device_wrappers.cuh b/src/api/gpu/common/nixl_device_wrappers.cuh new file mode 100644 index 0000000000..8764fb27b6 --- /dev/null +++ b/src/api/gpu/common/nixl_device_wrappers.cuh @@ -0,0 +1,61 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Public __device__ wrapper functions for the NIXL GPU transfer API. +// +// This file must be included AFTER nixl::gpu::selected_impl has been aliased +// to the desired backend namespace (ucx_impl, proxy_impl, etc.). It is +// included automatically by each backend-specific nixl_device.cuh facade and +// by the generic GPU API header; callers should not include it directly. + +#ifndef NIXL_SRC_API_GPU_COMMON_NIXL_DEVICE_WRAPPERS_CUH +#define NIXL_SRC_API_GPU_COMMON_NIXL_DEVICE_WRAPPERS_CUH + +template +__device__ inline nixl_status_t +nixlGpuGetXferStatus(nixlGpuXferStatusH &xfer_status) { + return nixl::gpu::selected_impl::get_xfer_status(xfer_status); +} + +template +__device__ inline nixl_status_t +nixlPut(const nixlMemViewElem &src, + const nixlMemViewElem &dst, + size_t size, + unsigned channel_id = 0, + uint64_t flags = 0, + nixlGpuXferStatusH *xfer_status = nullptr) { + return nixl::gpu::selected_impl::put(src, dst, size, channel_id, flags, xfer_status); +} + +template +__device__ inline nixl_status_t +nixlAtomicAdd(uint64_t value, + const nixlMemViewElem &counter, + unsigned channel_id = 0, + uint64_t flags = 0, + nixlGpuXferStatusH *xfer_status = nullptr) { + return nixl::gpu::selected_impl::atomic_add( + value, counter, channel_id, flags, xfer_status); +} + +__device__ inline void * +nixlGetPtr(nixlMemViewH mvh, size_t index) { + return nixl::gpu::selected_impl::get_ptr(mvh, index); +} + +#endif // NIXL_SRC_API_GPU_COMMON_NIXL_DEVICE_WRAPPERS_CUH diff --git a/src/api/gpu/meson.build b/src/api/gpu/meson.build new file mode 100644 index 0000000000..c8f502ae57 --- /dev/null +++ b/src/api/gpu/meson.build @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# proxy/nixl_device_proxy.cu defines g_nixl_proxy_ctx, the __device__ global +# that ProxyRuntime publishes via cudaMemcpyToSymbol and device kernels read +# through load_proxy_context(). It must be compiled as a separate CUDA TU so +# that the symbol exists at device link time. +if cuda_dep.found() + nixl_gpu_proxy_lib = static_library('nixl_gpu_proxy', + 'proxy/nixl_device_proxy.cu', + dependencies: [cuda_dep], + include_directories: [nixl_inc_dirs], + install: false) + + nixl_gpu_proxy_dep = declare_dependency(link_with: nixl_gpu_proxy_lib) +endif diff --git a/src/api/gpu/nixl_device.cuh b/src/api/gpu/nixl_device.cuh new file mode 100644 index 0000000000..92ecc8abd1 --- /dev/null +++ b/src/api/gpu/nixl_device.cuh @@ -0,0 +1,32 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generic backend-agnostic entry point. Selects an implementation at +// compile time via NIXL_GPU_DEVICE_BACKEND_UCX or +// NIXL_GPU_DEVICE_BACKEND_PROXY; exactly one must be defined. +// +// When the backend is already known at the include site prefer the +// backend-specific facade (ucx/nixl_device.cuh or proxy/nixl_device.cuh) +// which requires no macro. + +#ifndef NIXL_SRC_API_GPU_NIXL_DEVICE_CUH +#define NIXL_SRC_API_GPU_NIXL_DEVICE_CUH + +#include "common/nixl_device_api.cuh" +#include "common/nixl_device_wrappers.cuh" + +#endif // NIXL_SRC_API_GPU_NIXL_DEVICE_CUH diff --git a/src/api/gpu/proxy/nixl_device.cuh b/src/api/gpu/proxy/nixl_device.cuh new file mode 100644 index 0000000000..3f95a16dab --- /dev/null +++ b/src/api/gpu/proxy/nixl_device.cuh @@ -0,0 +1,34 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Proxy-backend device API facade. Self-contained: no backend-selection +// macro required. Include this header (or place src/api/gpu/proxy on the +// include path so that resolves here) when building CUDA +// translation units that submit work through the CPU proxy runtime. + +#ifndef NIXL_SRC_API_GPU_PROXY_NIXL_DEVICE_CUH +#define NIXL_SRC_API_GPU_PROXY_NIXL_DEVICE_CUH + +#include "nixl_device_impl.cuh" + +namespace nixl::gpu { +namespace selected_impl = proxy_impl; +} + +#include "../common/nixl_device_wrappers.cuh" + +#endif // NIXL_SRC_API_GPU_PROXY_NIXL_DEVICE_CUH diff --git a/src/api/gpu/proxy/nixl_device_impl.cuh b/src/api/gpu/proxy/nixl_device_impl.cuh new file mode 100644 index 0000000000..54d8c0ccbf --- /dev/null +++ b/src/api/gpu/proxy/nixl_device_impl.cuh @@ -0,0 +1,131 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef NIXL_SRC_API_GPU_PROXY_NIXL_DEVICE_IMPL_CUH +#define NIXL_SRC_API_GPU_PROXY_NIXL_DEVICE_IMPL_CUH + +#include "nixl_device_proxy.cuh" +#include "nixl_types.h" + +namespace nixl::gpu::proxy_impl { + +template +__device__ inline nixl_status_t +get_xfer_status(nixlGpuXferStatusH &xfer_status) { + uint32_t lane_id; + nixlProxyExecInit(lane_id); + + ProxyDeviceContext *ctx = load_proxy_context(); + + nixl_status_t status = NIXL_IN_PROG; + if (lane_id == 0) { + if (ctx == nullptr) { + status = NIXL_ERR_BACKEND; + } else { + status = ctx->pollXferStatus(xfer_status); + } + } + + if constexpr (level == nixl_gpu_level_t::WARP) { + status = static_cast( + __shfl_sync(0xffffffff, static_cast(status), 0)); + } else if constexpr (level == nixl_gpu_level_t::BLOCK) { + __shared__ nixl_status_t s_status; + if (threadIdx.x == 0) { + s_status = status; + } + __syncthreads(); + status = s_status; + __syncthreads(); + } + + return status; +} + +template +__device__ inline nixl_status_t +put(const nixlMemViewElem &src, + const nixlMemViewElem &dst, + size_t size, + unsigned channel_id = 0, + uint64_t flags = 0, + nixlGpuXferStatusH *xfer_status = nullptr) { + + uint32_t lane_id; + nixlProxyExecInit(lane_id); + nixl_status_t status = NIXL_IN_PROG; + if (lane_id == 0) { + ProxyDeviceContext *ctx = load_proxy_context(); + if (ctx == nullptr) { + status = NIXL_ERR_BACKEND; + } else { + nixlProxySubmission submission{}; + submission.opcode = nixl_proxy_opcode_t::PUT; + submission.channel_id = static_cast(channel_id); + submission.flags = flags; + submission.src_proxy_memview_id = proxyMemViewIdFromHandle(src.mvh); + submission.src_index = src.index; + submission.src_offset = src.offset; + submission.dst_proxy_memview_id = proxyMemViewIdFromHandle(dst.mvh); + submission.dst_index = dst.index; + submission.dst_offset = dst.offset; + submission.size = size; + status = ctx->enqueue(submission, xfer_status); + } + } + nixlProxySync(); + return status; +} + +template +__device__ inline nixl_status_t +atomic_add(uint64_t value, + const nixlMemViewElem &counter, + unsigned channel_id = 0, + uint64_t flags = 0, + nixlGpuXferStatusH *xfer_status = nullptr) { + uint32_t lane_id; + nixlProxyExecInit(lane_id); + nixl_status_t status = NIXL_IN_PROG; + if (lane_id == 0) { + ProxyDeviceContext *ctx = load_proxy_context(); + if (ctx == nullptr) { + status = NIXL_ERR_BACKEND; + } else { + nixlProxySubmission submission{}; + submission.opcode = nixl_proxy_opcode_t::ATOMIC_ADD; + submission.channel_id = static_cast(channel_id); + submission.flags = flags; + submission.dst_proxy_memview_id = proxyMemViewIdFromHandle(counter.mvh); + submission.dst_index = counter.index; + submission.dst_offset = counter.offset; + submission.value = value; + status = ctx->enqueue(submission, xfer_status); + } + } + nixlProxySync(); + return status; +} + +__device__ inline void * +get_ptr(nixlMemViewH, size_t) { + // TODO: Implement support for NVLink fast-path over proxy - NIX-1342 + return nullptr; +} + +} // namespace nixl::gpu::proxy_impl + +#endif // NIXL_SRC_API_GPU_PROXY_NIXL_DEVICE_IMPL_CUH diff --git a/src/api/gpu/proxy/nixl_device_proxy.cu b/src/api/gpu/proxy/nixl_device_proxy.cu new file mode 100644 index 0000000000..0e87cb2e5d --- /dev/null +++ b/src/api/gpu/proxy/nixl_device_proxy.cu @@ -0,0 +1,22 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Defines the device-visible proxy context pointer. Device kernels read it +// through load_proxy_context(). +#include "nixl_device_proxy.cuh" + +__device__ ProxyDeviceContext *g_nixl_proxy_ctx = nullptr; diff --git a/src/api/gpu/proxy/nixl_device_proxy.cuh b/src/api/gpu/proxy/nixl_device_proxy.cuh new file mode 100644 index 0000000000..7a8d90885f --- /dev/null +++ b/src/api/gpu/proxy/nixl_device_proxy.cuh @@ -0,0 +1,195 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef NIXL_SRC_API_GPU_PROXY_NIXL_DEVICE_PROXY_CUH +#define NIXL_SRC_API_GPU_PROXY_NIXL_DEVICE_PROXY_CUH + +#include +#include + +#include "../common/nixl_device_types.cuh" +#include "../../../core/device_proxy/proxy_protocol.h" + +struct ProxyDeviceContext; + +// Overlay struct written into nixlGpuXferStatusH::storage by enqueue() +// and read back by pollXferStatus(). Must fit within the 64-byte opaque blob. +struct ProxyXferStatus { + nixlProxyCompletionSlot *slot; // device pointer to the channel's nixlProxyCompletionSlot + uint64_t op_idx; +}; +static_assert(sizeof(ProxyXferStatus) <= sizeof(nixlGpuXferStatusH), + "ProxyXferStatus must fit in nixlGpuXferStatusH::storage"); + +// Defined in nixl_device_proxy.cu and read by device kernels through +// load_proxy_context(). +extern __device__ ProxyDeviceContext *g_nixl_proxy_ctx; + +// Host-callable helpers. Keeping these inline in CUDA translation units avoids +// cross-DSO symbol ownership issues for g_nixl_proxy_ctx. +__host__ inline cudaError_t +nixlProxyPublishContext(nixlProxyDeviceContextData *ctx) { + ProxyDeviceContext *device_ctx = reinterpret_cast(ctx); + cudaError_t err = cudaMemcpyToSymbol(g_nixl_proxy_ctx, &device_ctx, sizeof(ProxyDeviceContext *)); + if (err != cudaSuccess) { + fprintf(stderr, + "nixlProxyPublishContext: cudaMemcpyToSymbol failed: code=%d msg=%s\n", + static_cast(err), + cudaGetErrorString(err)); + } + return err; +} + +__host__ inline cudaError_t +nixlProxyClearContext() { + ProxyDeviceContext *null_ctx = nullptr; + cudaError_t err = cudaMemcpyToSymbol(g_nixl_proxy_ctx, &null_ctx, sizeof(ProxyDeviceContext *)); + if (err != cudaSuccess) { + fprintf(stderr, + "nixlProxyClearContext: cudaMemcpyToSymbol failed: code=%d msg=%s\n", + static_cast(err), + cudaGetErrorString(err)); + } + return err; +} + +__device__ inline uint64_t +proxyMemViewIdFromHandle(nixlMemViewH mvh) { + return static_cast(reinterpret_cast(mvh)); +} + +__device__ inline ProxyDeviceContext * +load_proxy_context() { + return g_nixl_proxy_ctx; +} + +static_assert(sizeof(nixlProxyWorkRing::running_op_idx) == 8, + "running_op_idx must be 64-bit to avoid wrap-around false completions"); +static_assert(sizeof(nixlProxyCompletionSlot::completed_idx) == 8, + "completed_idx must be 64-bit to match running_op_idx"); + +template +__device__ inline void nixlProxyExecInit(uint32_t &lane_id) { + static_assert(level != nixl_gpu_level_t::GRID, + "Proxy GPU backend does not support GRID-level operations"); + + if constexpr (level == nixl_gpu_level_t::THREAD) { + lane_id = 0; + } else if constexpr (level == nixl_gpu_level_t::WARP) { + lane_id = threadIdx.x % warpSize; + } else if constexpr (level == nixl_gpu_level_t::BLOCK) { + lane_id = threadIdx.x; + } +} + +template +__device__ inline void nixlProxySync() { + static_assert(level != nixl_gpu_level_t::GRID, + "Proxy GPU backend does not support GRID-level operations"); + + if constexpr (level == nixl_gpu_level_t::WARP) { + __syncwarp(); + } else if constexpr (level == nixl_gpu_level_t::BLOCK) { + __syncthreads(); + } +} + +struct ProxyDeviceContext : nixlProxyDeviceContextData { + + // Enqueue a transfer submission into the MPSC work ring for the selected + // channel, spinning if the ring is full. Optionally records a completion + // token in *xfer_status for later polling via pollXferStatus(). + // + // producer_idx lives in HBM; consumer_idx lives in pinned host memory + // (accessible from device via UVA mapped pointer). Both are accessed with + // system-scope atomics so the CPU proxy worker sees the update coherently. + __device__ inline nixl_status_t + enqueue(nixlProxySubmission submission, nixlGpuXferStatusH *xfer_status = nullptr) { + if (submission.channel_id >= num_channels) { + return NIXL_ERR_INVALID_PARAM; + } + + nixlProxyChannelView &channel_view = channels[submission.channel_id]; + nixlProxyWorkRing *ring = channel_view.work_ring; + + cuda::atomic_ref prod(*ring->producer_idx); + cuda::atomic_ref cons(*ring->consumer_idx); + cuda::atomic_ref shut(*shutdown_word); + + // Atomically claim a unique slot in the ring. + uint32_t my_slot = prod.fetch_add(1, cuda::memory_order_relaxed); + + // Spin until the claimed slot has space (consumer has freed it). + while (my_slot - cons.load(cuda::memory_order_acquire) >= ring->depth) { + if (shut.load(cuda::memory_order_acquire) + == static_cast(nixl_proxy_control_state_t::SHUTDOWN)) { + return NIXL_ERR_BACKEND; + } + } + + cuda::atomic_ref op_idx(ring->running_op_idx); + submission.op_idx = op_idx.fetch_add(1, cuda::memory_order_relaxed); + ring->records[my_slot % ring->depth] = submission; + + // Signal this slot is ready for the consumer. The release + // guarantees the record write above is visible before the + // consumer reads it via an acquire load on ready_flag. + cuda::atomic_ref ready( + ring->records[my_slot % ring->depth].ready_flag); + ready.store(1, cuda::memory_order_release); + + if (xfer_status != nullptr) { + ProxyXferStatus pxs{channel_view.completion_slot, submission.op_idx}; + memcpy(xfer_status->storage, &pxs, sizeof(ProxyXferStatus)); + } + + return NIXL_IN_PROG; + } + + // Poll the completion slot recorded by enqueue(). + // + // The completion slot implements collapsed-CQ semantics: + // - completed_idx > op_idx => this op completed earlier, so it succeeded + // - completed_idx == op_idx => next_status is this op's terminal status + // - completed_idx < op_idx => this op is still pending, unless an earlier + // completion published a terminal error and + // latched the channel + __device__ inline nixl_status_t + pollXferStatus(const nixlGpuXferStatusH &xfer_status) const { + const ProxyXferStatus *pxs = + reinterpret_cast(xfer_status.storage); + + cuda::atomic_ref comp_idx( + pxs->slot->completed_idx); + + const uint64_t completed_idx = comp_idx.load(cuda::memory_order_acquire); + const nixl_status_t current_status = pxs->slot->next_status; + + if (completed_idx > pxs->op_idx) { + return NIXL_SUCCESS; + } + if (completed_idx == pxs->op_idx) { + return current_status; + } + if (current_status < 0) { + return current_status; + } + + return NIXL_IN_PROG; + } +}; + +#endif // NIXL_SRC_API_GPU_PROXY_NIXL_DEVICE_PROXY_CUH diff --git a/src/api/gpu/ucx/nixl_device.cuh b/src/api/gpu/ucx/nixl_device.cuh index 6e4f5ef215..e98baaecf0 100644 --- a/src/api/gpu/ucx/nixl_device.cuh +++ b/src/api/gpu/ucx/nixl_device.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,186 +14,21 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#ifndef _NIXL_DEVICE_CUH -#define _NIXL_DEVICE_CUH -#include -#include +// UCX-specific device API facade. Self-contained: no backend-selection macro +// required. Include this header (or place src/api/gpu/ucx on the include +// path so that resolves here) when building CUDA +// translation units that use the UCX GPU device API directly. -#include +#ifndef NIXL_SRC_API_GPU_UCX_NIXL_DEVICE_CUH +#define NIXL_SRC_API_GPU_UCX_NIXL_DEVICE_CUH -struct nixlGpuXferStatusH { - ucp_device_request_t device_request; -}; +#include "nixl_device_impl.cuh" -/** - * @enum nixl_gpu_level_t - * @brief An enumeration of different levels for GPU transfer requests. - */ -enum class nixl_gpu_level_t : uint64_t { - THREAD = UCS_DEVICE_LEVEL_THREAD, - WARP = UCS_DEVICE_LEVEL_WARP, - BLOCK = UCS_DEVICE_LEVEL_BLOCK, - GRID = UCS_DEVICE_LEVEL_GRID -}; - -namespace nixl_gpu_flags { -constexpr uint64_t defer = 1; - -__device__ inline uint64_t -to_ucp_flags(uint64_t nixl_flags) noexcept { - constexpr uint64_t all_known_nixl_flags{defer}; - assert((nixl_flags & ~all_known_nixl_flags) == 0); - - uint64_t ucp_flags{UCP_DEVICE_FLAG_NODELAY}; - if (nixl_flags & defer) { - ucp_flags &= ~UCP_DEVICE_FLAG_NODELAY; - } - return ucp_flags; +namespace nixl::gpu { +namespace selected_impl = ucx_impl; } -} // namespace nixl_gpu_flags - -struct nixlMemViewElem { - nixlMemViewH mvh; - size_t index; /**< Index in the memory view */ - size_t offset; /**< Offset within the buffer */ -}; -/** - * @brief Convert UCS status to NIXL status. - * - * @param status [in] UCS status code. - * - * @return NIXL_IN_PROG If the UCS status is not an error. - * @return NIXL_ERR_BACKEND If the UCS status is an error. - */ -__device__ inline nixl_status_t -nixlGpuConvertUcsStatus(ucs_status_t status) { - if (!UCS_STATUS_IS_ERR(status)) { - return NIXL_IN_PROG; - } - printf("UCX returned error: %d\n", status); - return NIXL_ERR_BACKEND; -} - -/** - * @brief Get the status of the transfer request. - * - * @param xfer_status [in] Status of the transfer. - * - * @return NIXL_SUCCESS The request has completed, no more operations are in progress. - * @return NIXL_IN_PROG One or more operations in the request have not completed. - * @return NIXL_ERR_BACKEND An error occurred in UCX backend. - */ -template -__device__ nixl_status_t -nixlGpuGetXferStatus(nixlGpuXferStatusH &xfer_status) { - const auto status = ucp_device_progress_req(level)>( - &xfer_status.device_request); - - switch (status) { - case UCS_OK: - return NIXL_SUCCESS; - case UCS_INPROGRESS: - return NIXL_IN_PROG; - default: - return NIXL_ERR_BACKEND; - } -} - -/** - * @brief Post a single-region memory transfer from local to remote GPU. - * - * This function creates and posts a transfer request using memory view elements @a src and @a dst. - * - * @param src [in] Source memory view element - * @param dst [in] Destination memory view element - * @param size [in] Size in bytes to transfer - * @param channel_id [in] Channel ID to use for the transfer - * @param flags [in] Transfer flags - * @param xfer_status [in,out] Optional status handle (use @ref nixlGpuGetXferStatus) - * - * @return NIXL_IN_PROG Transfer posted successfully. - * @return NIXL_ERR_BACKEND An error occurred in UCX backend. - */ -template -__device__ nixl_status_t -nixlPut(const nixlMemViewElem &src, - const nixlMemViewElem &dst, - size_t size, - unsigned channel_id = 0, - uint64_t flags = 0, - nixlGpuXferStatusH *xfer_status = nullptr) { - auto src_mem_list = static_cast(src.mvh); - auto dst_mem_list = static_cast(dst.mvh); - ucp_device_request_t *ucp_request{xfer_status ? &xfer_status->device_request : nullptr}; - const auto status = - ucp_device_put(level)>(src_mem_list, - src.index, - src.offset, - dst_mem_list, - dst.index, - dst.offset, - size, - channel_id, - nixl_gpu_flags::to_ucp_flags(flags), - ucp_request); - return nixlGpuConvertUcsStatus(status); -} - -/** - * @brief Atomic add to remote GPU memory. - * - * This function performs an atomic increment on a remote counter. - * The increment is visible only after previous writes complete. - * - * @param value [in] Value to add to the counter - * @param counter [in] Counter memory view element - * @param channel_id [in] Channel ID to use for the transfer - * @param flags [in] Transfer flags - * @param xfer_status [in,out] Optional status handle (use @ref nixlGpuGetXferStatus) - * - * @return NIXL_IN_PROG Atomic add posted successfully. - * @return NIXL_ERR_BACKEND An error occurred in UCX backend. - */ -template -__device__ nixl_status_t -nixlAtomicAdd(uint64_t value, - const nixlMemViewElem &counter, - unsigned channel_id = 0, - uint64_t flags = 0, - nixlGpuXferStatusH *xfer_status = nullptr) { - auto mem_list = static_cast(counter.mvh); - ucp_device_request_t *ucp_request{xfer_status ? &xfer_status->device_request : nullptr}; - const auto status = ucp_device_counter_inc(level)>( - value, - mem_list, - counter.index, - counter.offset, - channel_id, - nixl_gpu_flags::to_ucp_flags(flags), - ucp_request); - return nixlGpuConvertUcsStatus(status); -} - -/** - * @brief Get a local pointer to remote memory. - * - * This function returns a local pointer to the mapped memory of the - * remote memory view handle at the given index. - * The memory view must be prepared on the host using @ref nixlAgent::prepMemView. - * - * @param mem_view [in] Memory view handle (remote buffers) - * @param index [in] Index in the memory view - - * @return Pointer to the mapped memory, or nullptr if not available. - */ -__device__ inline void * -nixlGetPtr(nixlMemViewH mvh, size_t index) { - auto mem_list = static_cast(mvh); - void *ptr = nullptr; - ucp_device_get_ptr(mem_list, index, &ptr); - return ptr; -} +#include "../common/nixl_device_wrappers.cuh" -#endif // _NIXL_DEVICE_CUH +#endif // NIXL_SRC_API_GPU_UCX_NIXL_DEVICE_CUH diff --git a/src/api/gpu/ucx/nixl_device_impl.cuh b/src/api/gpu/ucx/nixl_device_impl.cuh new file mode 100644 index 0000000000..db8d02e44a --- /dev/null +++ b/src/api/gpu/ucx/nixl_device_impl.cuh @@ -0,0 +1,153 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef NIXL_SRC_API_GPU_UCX_NIXL_DEVICE_IMPL_CUH +#define NIXL_SRC_API_GPU_UCX_NIXL_DEVICE_IMPL_CUH + +#include "../common/nixl_device_types.cuh" + +#include + +#include +#include + +namespace nixl::gpu::ucx_impl { + +template struct UcpDeviceLevel; + +template<> struct UcpDeviceLevel { + static constexpr ucs_device_level_t value = UCS_DEVICE_LEVEL_THREAD; +}; + +template<> struct UcpDeviceLevel { + static constexpr ucs_device_level_t value = UCS_DEVICE_LEVEL_WARP; +}; + +template<> struct UcpDeviceLevel { + static constexpr ucs_device_level_t value = UCS_DEVICE_LEVEL_BLOCK; +}; + +template<> struct UcpDeviceLevel { + static constexpr ucs_device_level_t value = UCS_DEVICE_LEVEL_GRID; +}; + +__device__ inline uint64_t +to_ucp_flags(uint64_t nixl_flags) noexcept { + constexpr uint64_t all_known_nixl_flags{nixl_gpu_flags::defer}; + assert((nixl_flags & ~all_known_nixl_flags) == 0); + + uint64_t ucp_flags{UCP_DEVICE_FLAG_NODELAY}; + if (nixl_flags & nixl_gpu_flags::defer) { + ucp_flags &= ~UCP_DEVICE_FLAG_NODELAY; + } + return ucp_flags; +} + +__device__ inline nixl_status_t +convert_status(ucs_status_t status) { + if (!UCS_STATUS_IS_ERR(status)) { + return NIXL_IN_PROG; + } + printf("UCX returned error: %d\n", status); + return NIXL_ERR_BACKEND; +} + +__device__ inline ucp_device_request_t * +request_ptr(nixlGpuXferStatusH *xfer_status) { + static_assert(sizeof(ucp_device_request_t) <= sizeof(nixlGpuXferStatusH{}.storage), + "nixlGpuXferStatusH storage is too small for UCX device request"); + return xfer_status ? reinterpret_cast(xfer_status->storage) : nullptr; +} + +__device__ inline ucp_device_local_mem_list_h +local_mem_list(nixlMemViewH mvh) { + return static_cast(mvh); +} + +__device__ inline ucp_device_remote_mem_list_h +remote_mem_list(nixlMemViewH mvh) { + return static_cast(mvh); +} + +template +__device__ inline nixl_status_t +get_xfer_status(nixlGpuXferStatusH &xfer_status) { + const auto status = + ucp_device_progress_req::value>(request_ptr(&xfer_status)); + + switch (status) { + case UCS_OK: + return NIXL_SUCCESS; + case UCS_INPROGRESS: + return NIXL_IN_PROG; + default: + return NIXL_ERR_BACKEND; + } +} + +template +__device__ inline nixl_status_t +put(const nixlMemViewElem &src, + const nixlMemViewElem &dst, + size_t size, + unsigned channel_id = 0, + uint64_t flags = 0, + nixlGpuXferStatusH *xfer_status = nullptr) { + auto src_mem_list = local_mem_list(src.mvh); + auto dst_mem_list = remote_mem_list(dst.mvh); + const auto status = ucp_device_put::value>(src_mem_list, + src.index, + src.offset, + dst_mem_list, + dst.index, + dst.offset, + size, + channel_id, + to_ucp_flags(flags), + request_ptr(xfer_status)); + return convert_status(status); +} + +template +__device__ inline nixl_status_t +atomic_add(uint64_t value, + const nixlMemViewElem &counter, + unsigned channel_id = 0, + uint64_t flags = 0, + nixlGpuXferStatusH *xfer_status = nullptr) { + auto mem_list = remote_mem_list(counter.mvh); + const auto status = + ucp_device_counter_inc::value>(value, + mem_list, + counter.index, + counter.offset, + channel_id, + to_ucp_flags(flags), + request_ptr(xfer_status)); + return convert_status(status); +} + +__device__ inline void * +get_ptr(nixlMemViewH mvh, size_t index) { + auto mem_list = remote_mem_list(mvh); + void *ptr = nullptr; + ucp_device_get_ptr(mem_list, index, &ptr); + return ptr; +} + +} // namespace nixl::gpu::ucx_impl + +#endif // NIXL_SRC_API_GPU_UCX_NIXL_DEVICE_IMPL_CUH diff --git a/src/api/meson.build b/src/api/meson.build index 662e15fbd8..c59466b809 100644 --- a/src/api/meson.build +++ b/src/api/meson.build @@ -15,3 +15,6 @@ subdir('cpp') subdir('python') +if cuda_dep.found() + subdir('gpu') +endif diff --git a/src/core/device_proxy/proxy_protocol.h b/src/core/device_proxy/proxy_protocol.h new file mode 100644 index 0000000000..7f0c36c35f --- /dev/null +++ b/src/core/device_proxy/proxy_protocol.h @@ -0,0 +1,86 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef NIXL_SRC_CORE_DEVICE_PROXY_PROXY_PROTOCOL_H +#define NIXL_SRC_CORE_DEVICE_PROXY_PROXY_PROTOCOL_H + +#include +#include + +#include + +enum class nixl_proxy_opcode_t : uint32_t { + PUT = 0, + ATOMIC_ADD = 1, +}; + +enum class nixl_proxy_control_state_t : uint32_t { + RUNNING = 0, + SHUTDOWN = 1, +}; + +struct nixlProxySubmission { + uint64_t op_idx = 0; + nixl_proxy_opcode_t opcode = nixl_proxy_opcode_t::PUT; + uint32_t channel_id = 0; + uint64_t flags = 0; + + uint64_t src_proxy_memview_id = 0; + size_t src_index = 0; + size_t src_offset = 0; + + uint64_t dst_proxy_memview_id = 0; + size_t dst_index = 0; + size_t dst_offset = 0; + + uint32_t ready_flag = 0; + size_t size = 0; + uint64_t value = 0; +}; + +struct nixlProxyWorkRing { + /** Host-accessible (e.g. cudaMallocHost); GPU may read via mapped pointer if needed. */ + nixlProxySubmission *records = nullptr; + /** Mapped pinned producer; GPU advances with CUDA atomics; host reads via __atomic_*. */ + uint32_t *producer_idx = nullptr; + /** Mapped pinned consumer; host proxy uses __atomic_* on host alias (nixlProxyChannelState). */ + uint32_t *consumer_idx = nullptr; + /** The depth of the work ring. */ + uint32_t depth = 0; + /** Monotonic 64-bit counter; starts at 1 so completed_idx==0 means + * "no operation completed yet" and the first op_idx is never 0. */ + uint64_t running_op_idx = 1; +}; + +struct alignas(16) nixlProxyCompletionSlot { + uint64_t completed_idx = 0; + nixl_status_t next_status = NIXL_IN_PROG; +}; + +struct nixlProxyChannelView { + nixlProxyWorkRing *work_ring = nullptr; + /** Mapped pinned host memory (device alias); host writes via host pointer with atomics. */ + nixlProxyCompletionSlot *completion_slot = nullptr; + uint32_t channel_id = 0; +}; + +struct nixlProxyDeviceContextData { + nixlProxyChannelView *channels = nullptr; + uint32_t num_channels = 0; + uint32_t *shutdown_word = nullptr; +}; + +#endif // NIXL_SRC_CORE_DEVICE_PROXY_PROXY_PROTOCOL_H From 9c944591d8cf318a72268f6b315293e8653d4378 Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Wed, 29 Apr 2026 18:13:43 +0300 Subject: [PATCH 03/18] device_proxy: add proxy runtime, worker, and unit tests Introduce the CPU-side proxy machinery under src/core/device_proxy/: - proxy_runtime.h/cpp: ProxyRuntime orchestrates channel state and the ProxyMemViewRegistry that maps GPU-visible proxy handles to backend memory-view descriptors. - proxy_worker.h/cpp: ProxyWorker is a CPU thread that drains GPU submission rings and dispatches transfers through a backend adapter. - backend_adapter.h: pure-virtual interface for proxy transport backends. - backend_provider.h: factory interface for backend engines that can supply a proxy adapter. Both new .cpp files are compiled into nixl_lib (src/core/meson.build). Unit tests for ProxyMemViewRegistry and ProxyRuntime are added under test/gtest/unit/. They depend only on internal core headers (no public nixl.h API) and run against a stub backend. Signed-off-by: Tomer Davidor --- src/api/cpp/backend/backend_engine.h | 14 + src/core/agent_data.h | 5 + src/core/device_proxy/backend_adapter.h | 78 ++ src/core/device_proxy/proxy_runtime.cpp | 827 ++++++++++++++++++ src/core/device_proxy/proxy_runtime.h | 301 +++++++ src/core/device_proxy/proxy_worker.cpp | 182 ++++ src/core/device_proxy/proxy_worker.h | 68 ++ src/core/meson.build | 25 +- src/core/nixl_agent.cpp | 13 + src/plugins/ucx/ucx_backend.h | 5 + test/gtest/unit/meson.build | 6 + .../unit/proxy_memview_registry/meson.build | 21 + .../proxy_memview_registry.cpp | 449 ++++++++++ test/gtest/unit/proxy_runtime/meson.build | 27 + .../unit/proxy_runtime/proxy_runtime.cpp | 418 +++++++++ 15 files changed, 2437 insertions(+), 2 deletions(-) create mode 100644 src/core/device_proxy/backend_adapter.h create mode 100644 src/core/device_proxy/proxy_runtime.cpp create mode 100644 src/core/device_proxy/proxy_runtime.h create mode 100644 src/core/device_proxy/proxy_worker.cpp create mode 100644 src/core/device_proxy/proxy_worker.h create mode 100644 test/gtest/unit/proxy_memview_registry/meson.build create mode 100644 test/gtest/unit/proxy_memview_registry/proxy_memview_registry.cpp create mode 100644 test/gtest/unit/proxy_runtime/meson.build create mode 100644 test/gtest/unit/proxy_runtime/proxy_runtime.cpp diff --git a/src/api/cpp/backend/backend_engine.h b/src/api/cpp/backend/backend_engine.h index c1ff43a24e..6d956c70e0 100644 --- a/src/api/cpp/backend/backend_engine.h +++ b/src/api/cpp/backend/backend_engine.h @@ -27,6 +27,8 @@ #include "backend_aux.h" #include "telemetry_event.h" +class nixlDeviceProxyBackendAdapter; + constexpr size_t MAX_TELEMETRY_QUEUE_SIZE = 1000; // Base backend engine class for different backend implementations @@ -111,6 +113,12 @@ class nixlBackendEngine { // pure virtual, and return errors, as parent shouldn't call if supportsNotif is false. virtual bool supportsNotif() const = 0; + // Determines if a backend supports device proxy runtime creation. + virtual bool + supportsProxy() const { + return false; + } + virtual nixl_mem_list_t getSupportedMems() const = 0; // TODO: Return by const-reference and mark noexcept? @@ -175,6 +183,12 @@ class nixlBackendEngine { virtual void releaseMemView(nixlMemViewH) const {} + virtual nixl_status_t + createDeviceProxyBackendAdapter(const nixlBackendInitParams &, + std::unique_ptr &) { + return NIXL_ERR_NOT_SUPPORTED; + } + // *** Needs to be implemented if supportsRemote() is true *** // // Gets serialized form of public metadata diff --git a/src/core/agent_data.h b/src/core/agent_data.h index 3949a50ca8..d58dfabda3 100644 --- a/src/core/agent_data.h +++ b/src/core/agent_data.h @@ -174,6 +174,11 @@ class nixlBackendH { return engine->supportsNotif(); } + bool + supportsProxy() const { + return engine->supportsProxy(); + } + friend class nixlAgentData; friend class nixlAgent; }; diff --git a/src/core/device_proxy/backend_adapter.h b/src/core/device_proxy/backend_adapter.h new file mode 100644 index 0000000000..03bc77113e --- /dev/null +++ b/src/core/device_proxy/backend_adapter.h @@ -0,0 +1,78 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef NIXL_SRC_CORE_DEVICE_PROXY_BACKEND_ADAPTER_H +#define NIXL_SRC_CORE_DEVICE_PROXY_BACKEND_ADAPTER_H + +#include +#include +#include + +#include +#include "backend_aux.h" +#include "proxy_protocol.h" + +struct nixlBackendProxyXferDesc { + nixl_mem_t mem_type = VRAM_SEG; + nixlMetaDesc desc{}; +}; + +struct nixlBackendProxySubmission { + uint64_t op_idx = 0; + nixl_proxy_opcode_t opcode = nixl_proxy_opcode_t::PUT; + uint32_t channel_id = 0; + uint64_t flags = 0; + + nixlBackendProxyXferDesc local{}; + nixlBackendProxyXferDesc remote{}; + std::string remote_agent; + + size_t size = 0; + uint64_t value = 0; +}; + +class nixlDeviceProxyBackendAdapter { +public: + virtual ~nixlDeviceProxyBackendAdapter() = default; + + virtual nixl_status_t + init(uint32_t, uint32_t) { + return NIXL_ERR_NOT_SUPPORTED; + } + + virtual nixl_status_t + loadRemoteConnInfo(const std::string &, const nixl_blob_t &) { + return NIXL_ERR_NOT_SUPPORTED; + } + + virtual nixl_status_t + submit(const nixlBackendProxySubmission &submission, uint64_t &request_token) = 0; + + virtual nixl_status_t + checkCompletion(uint64_t request_token) = 0; + + virtual nixl_status_t + progress() { + return NIXL_ERR_NOT_SUPPORTED; + } + + virtual nixl_status_t + shutdown() { + return NIXL_ERR_NOT_SUPPORTED; + } +}; + +#endif // NIXL_SRC_CORE_DEVICE_PROXY_BACKEND_ADAPTER_H diff --git a/src/core/device_proxy/proxy_runtime.cpp b/src/core/device_proxy/proxy_runtime.cpp new file mode 100644 index 0000000000..531caff04f --- /dev/null +++ b/src/core/device_proxy/proxy_runtime.cpp @@ -0,0 +1,827 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "proxy_runtime.h" +#include "backend_adapter.h" +#include "nixl_types.h" +#include "proxy_worker.h" +#include "nixl_log.h" +#include +#include +#include + +nixl_status_t +nixlProxyMemViewRegistry::registerProxyMemView(nixlMemViewH backend_memview, + nixlMemViewH *proxy_memview) { + if (proxy_memview == nullptr) { + return NIXL_ERR_INVALID_PARAM; + } + + RegistryEntry entry; + entry.proxy_memview_id = next_proxy_memview_id_; + entry.backend_memview = backend_memview; + entries_.push_back(entry); + + *proxy_memview = reinterpret_cast(entry.proxy_memview_id); + ++next_proxy_memview_id_; + NIXL_DEBUG << "nixlProxyMemViewRegistry::register: backend_mvh=" + << backend_memview << " -> proxy_id=" + << (next_proxy_memview_id_ - 1); + return NIXL_SUCCESS; +} + +nixl_status_t +nixlProxyMemViewRegistry::prepMemView(const nixl_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview) { + return prepMemView(nullptr, dlist, proxy_memview); +} + +nixl_status_t +nixlProxyMemViewRegistry::prepMemView( + const nixl_remote_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview) { + return prepMemView(nullptr, dlist, proxy_memview); +} + +nixl_status_t +nixlProxyMemViewRegistry::prepMemView(nixlMemViewH backend_memview, + const nixl_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview) { + if (proxy_memview == nullptr) { + return NIXL_ERR_INVALID_PARAM; + } + + nixlMemViewH registered_proxy_memview = nullptr; + nixl_status_t status = registerProxyMemView(backend_memview, ®istered_proxy_memview); + if (status != NIXL_SUCCESS) { + return status; + } + + status = storeMetadata(registered_proxy_memview, dlist); + if (status != NIXL_SUCCESS) { + unregisterProxyMemView(registered_proxy_memview); + return status; + } + + *proxy_memview = registered_proxy_memview; + return NIXL_SUCCESS; +} + +nixl_status_t +nixlProxyMemViewRegistry::prepMemView( + nixlMemViewH backend_memview, + const nixl_remote_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview) { + if (proxy_memview == nullptr) { + return NIXL_ERR_INVALID_PARAM; + } + + nixlMemViewH registered_proxy_memview = nullptr; + nixl_status_t status = registerProxyMemView(backend_memview, ®istered_proxy_memview); + if (status != NIXL_SUCCESS) { + return status; + } + + status = storeMetadata(registered_proxy_memview, dlist); + if (status != NIXL_SUCCESS) { + unregisterProxyMemView(registered_proxy_memview); + return status; + } + + *proxy_memview = registered_proxy_memview; + return NIXL_SUCCESS; +} + +nixl_status_t +nixlProxyMemViewRegistry::unregisterProxyMemView(nixlMemViewH proxy_memview) { + RegistryEntry *entry = getEntryForHandle(proxy_memview); + if (entry == nullptr) { + return NIXL_ERR_INVALID_PARAM; + } + entry->state = ProxyMemViewRegEntryState::ENTRY_RETIRED; + NIXL_DEBUG << "nixlProxyMemViewRegistry::unregister: proxy_id=" + << entry->proxy_memview_id; + return NIXL_SUCCESS; +} + +bool +nixlProxyMemViewRegistry::resolveProxyMemView(nixlMemViewH proxy_memview, + nixlMemViewH &backend_memview) const { + const RegistryEntry *entry = getEntryForHandle(proxy_memview); + if (entry == nullptr || entry->state == ProxyMemViewRegEntryState::ENTRY_RETIRED) { + return false; + } + backend_memview = entry->backend_memview; + return true; +} + +bool +nixlProxyMemViewRegistry::resolveProxyMemViewId(uint64_t proxy_memview_id, + nixlMemViewH &backend_memview) const { + const RegistryEntry *entry = getEntryForId(proxy_memview_id); + if (entry == nullptr || entry->state == ProxyMemViewRegEntryState::ENTRY_RETIRED) { + return false; + } + backend_memview = entry->backend_memview; + return true; +} + +nixl_status_t +nixlProxyMemViewRegistry::storeMetadata(nixlMemViewH proxy_memview, + const nixl_meta_dlist_t &dlist) { + RegistryEntry *entry = getEntryForHandle(proxy_memview); + if (entry == nullptr || entry->state == ProxyMemViewRegEntryState::ENTRY_RETIRED) { + return NIXL_ERR_NOT_FOUND; + } + + fillLocalMetadata(dlist, entry->local_metadata); + entry->remote_metadata = RemoteMetadata{}; + entry->metadata_kind = ProxyMemViewRegMetadataKind::METADATA_KIND_LOCAL; + entry->state = ProxyMemViewRegEntryState::ENTRY_READY; + + NIXL_DEBUG << "nixlProxyMemViewRegistry::storeMetadata(local): proxy_id=" + << entry->proxy_memview_id << " entries=" << dlist.descCount(); + return NIXL_SUCCESS; +} + +nixl_status_t +nixlProxyMemViewRegistry::storeMetadata(nixlMemViewH proxy_memview, + const nixl_remote_meta_dlist_t &dlist) { + RegistryEntry *entry = getEntryForHandle(proxy_memview); + if (entry == nullptr || entry->state == ProxyMemViewRegEntryState::ENTRY_RETIRED) { + return NIXL_ERR_NOT_FOUND; + } + + fillRemoteMetadata(dlist, entry->remote_metadata); + entry->local_metadata = LocalMetadata{}; + entry->metadata_kind = ProxyMemViewRegMetadataKind::METADATA_KIND_REMOTE; + entry->state = ProxyMemViewRegEntryState::ENTRY_READY; + + NIXL_DEBUG << "nixlProxyMemViewRegistry::storeMetadata(remote): proxy_id=" + << entry->proxy_memview_id << " entries=" << dlist.descCount(); + return NIXL_SUCCESS; +} + +nixl_status_t +nixlProxyMemViewRegistry::prepareSubmission(const nixlProxySubmission &submission, + nixlBackendProxySubmission &prepared_submission) const { + bool needs_source = false; + size_t transfer_size = 0; + switch (submission.opcode) { + case nixl_proxy_opcode_t::PUT: + needs_source = true; + transfer_size = submission.size; + break; + case nixl_proxy_opcode_t::ATOMIC_ADD: + transfer_size = sizeof(uint64_t); + break; + default: + NIXL_ERROR << "nixlProxyMemViewRegistry::prepareSubmission: unsupported opcode: " + << static_cast(submission.opcode); + return NIXL_ERR_NOT_SUPPORTED; + } + + const RemoteMetadata *remote_metadata = nullptr; + const ProxyMemViewRegStoredEntry *dst_metadata = nullptr; + nixl_status_t status = getRemoteEntryForSubmission( + submission.dst_proxy_memview_id, + submission.dst_index, + submission.dst_offset, + transfer_size, + remote_metadata, + dst_metadata); + if (status != NIXL_SUCCESS) { + return status; + } + + nixlBackendProxySubmission prepared{}; + prepared.op_idx = submission.op_idx; + prepared.opcode = submission.opcode; + prepared.channel_id = submission.channel_id; + prepared.flags = submission.flags; + prepared.size = transfer_size; + prepared.value = submission.value; + prepared.remote_agent = remote_metadata->remote_agent; + prepared.remote.mem_type = remote_metadata->mem_type; + prepared.remote.desc = nixlMetaDesc( + dst_metadata->base_addr + submission.dst_offset, + transfer_size, + dst_metadata->dev_id, + dst_metadata->metadata); + + if (needs_source) { + const LocalMetadata *local_metadata = nullptr; + const ProxyMemViewRegStoredEntry *src_metadata = nullptr; + status = getLocalEntryForSubmission( + submission.src_proxy_memview_id, + submission.src_index, + submission.src_offset, + transfer_size, + local_metadata, + src_metadata); + if (status != NIXL_SUCCESS) { + return status; + } + + prepared.local.mem_type = local_metadata->mem_type; + prepared.local.desc = nixlMetaDesc( + src_metadata->base_addr + submission.src_offset, + transfer_size, + src_metadata->dev_id, + src_metadata->metadata); + } + + prepared_submission = prepared; + return NIXL_SUCCESS; +} + +void +nixlProxyMemViewRegistry::clear() noexcept { + for (auto &entry : entries_) { + entry.state = ProxyMemViewRegEntryState::ENTRY_RETIRED; + } +} + +nixlProxyMemViewRegistry::RegistryEntry * +nixlProxyMemViewRegistry::getEntryForHandle(nixlMemViewH proxy_memview) { + return getEntryForId(reinterpret_cast(proxy_memview)); +} + +const nixlProxyMemViewRegistry::RegistryEntry * +nixlProxyMemViewRegistry::getEntryForHandle(nixlMemViewH proxy_memview) const { + return getEntryForId(reinterpret_cast(proxy_memview)); +} + +nixlProxyMemViewRegistry::RegistryEntry * +nixlProxyMemViewRegistry::getEntryForId(uint64_t proxy_memview_id) { + if (proxy_memview_id < 1 + || proxy_memview_id >= next_proxy_memview_id_ + || proxy_memview_id > entries_.size()) { + return nullptr; + } + return &entries_[proxy_memview_id - 1]; +} + +const nixlProxyMemViewRegistry::RegistryEntry * +nixlProxyMemViewRegistry::getEntryForId(uint64_t proxy_memview_id) const { + if (proxy_memview_id < 1 + || proxy_memview_id >= next_proxy_memview_id_ + || proxy_memview_id > entries_.size()) { + return nullptr; + } + return &entries_[proxy_memview_id - 1]; +} + +nixl_status_t +nixlProxyMemViewRegistry::getRemoteEntryForSubmission(uint64_t proxy_memview_id, + size_t index, + size_t offset, + size_t size, + const RemoteMetadata *&metadata, + const ProxyMemViewRegStoredEntry *&entry) const { + metadata = nullptr; + entry = nullptr; + + const RegistryEntry *registry_entry = getEntryForId(proxy_memview_id); + if (registry_entry == nullptr || registry_entry->state != ProxyMemViewRegEntryState::ENTRY_READY) { + NIXL_DEBUG << "nixlProxyMemViewRegistry::prepareSubmission: dst not ready" + << " dst_proxy_id=" << proxy_memview_id; + return NIXL_ERR_NOT_FOUND; + } + if (registry_entry->metadata_kind != ProxyMemViewRegMetadataKind::METADATA_KIND_REMOTE) { + NIXL_DEBUG << "nixlProxyMemViewRegistry::prepareSubmission: dst metadata kind invalid" + << " dst_proxy_id=" << proxy_memview_id; + return NIXL_ERR_INVALID_PARAM; + } + + const auto &remote_metadata = registry_entry->remote_metadata; + if (index >= remote_metadata.entries.size()) { + return NIXL_ERR_INVALID_PARAM; + } + + const ProxyMemViewRegStoredEntry &remote_entry = remote_metadata.entries[index]; + if (!rangeFits(remote_entry, offset, size)) { + return NIXL_ERR_INVALID_PARAM; + } + + metadata = &remote_metadata; + entry = &remote_entry; + return NIXL_SUCCESS; +} + +nixl_status_t +nixlProxyMemViewRegistry::getLocalEntryForSubmission(uint64_t proxy_memview_id, + size_t index, + size_t offset, + size_t size, + const LocalMetadata *&metadata, + const ProxyMemViewRegStoredEntry *&entry) const { + metadata = nullptr; + entry = nullptr; + + const RegistryEntry *registry_entry = getEntryForId(proxy_memview_id); + if (registry_entry == nullptr || registry_entry->state != ProxyMemViewRegEntryState::ENTRY_READY) { + NIXL_DEBUG << "nixlProxyMemViewRegistry::prepareSubmission: src not ready" + << " src_proxy_id=" << proxy_memview_id; + return NIXL_ERR_NOT_FOUND; + } + if (registry_entry->metadata_kind != ProxyMemViewRegMetadataKind::METADATA_KIND_LOCAL) { + NIXL_DEBUG << "nixlProxyMemViewRegistry::prepareSubmission: src metadata kind invalid" + << " src_proxy_id=" << proxy_memview_id; + return NIXL_ERR_INVALID_PARAM; + } + + const auto &local_metadata = registry_entry->local_metadata; + if (index >= local_metadata.entries.size()) { + return NIXL_ERR_INVALID_PARAM; + } + + const ProxyMemViewRegStoredEntry &local_entry = local_metadata.entries[index]; + if (!rangeFits(local_entry, offset, size)) { + return NIXL_ERR_INVALID_PARAM; + } + + metadata = &local_metadata; + entry = &local_entry; + return NIXL_SUCCESS; +} + +bool +nixlProxyMemViewRegistry::rangeFits(const ProxyMemViewRegStoredEntry &entry, size_t offset, size_t size) { + return offset <= entry.len && size <= entry.len - offset; +} + +void +nixlProxyMemViewRegistry::fillLocalMetadata(const nixl_meta_dlist_t &dlist, + LocalMetadata &out) { + out = LocalMetadata{}; + out.mem_type = dlist.getType(); + out.entries.reserve(dlist.descCount()); + for (const auto &desc : dlist) { + out.entries.push_back(ProxyMemViewRegStoredEntry{desc.addr, desc.len, desc.devId, desc.metadataP}); + } +} + +void +nixlProxyMemViewRegistry::fillRemoteMetadata(const nixl_remote_meta_dlist_t &dlist, + RemoteMetadata &out) { + out = RemoteMetadata{}; + out.mem_type = dlist.getType(); + out.entries.reserve(dlist.descCount()); + for (const auto &desc : dlist) { + if (out.remote_agent.empty() && desc.remoteAgent != nixl_null_agent) { + out.remote_agent = desc.remoteAgent; + } + out.entries.push_back(ProxyMemViewRegStoredEntry{desc.addr, desc.len, desc.devId, desc.metadataP}); + } +} + +nixl_status_t +nixlProxyChannelState::allocate(uint32_t channel_id, uint32_t depth) { + NIXL_INFO << "nixlProxyChannelState::allocate: channel_id=" << channel_id + << " depth=" << depth; + if (cudaMallocHost(&work_ring_, sizeof(nixlProxyWorkRing)) != cudaSuccess + || cudaMallocHost(&records_, sizeof(nixlProxySubmission) * depth) != cudaSuccess + || cudaMallocHost(reinterpret_cast(&consumer_idx_host_), + sizeof(uint32_t)) != cudaSuccess + || cudaMallocHost(reinterpret_cast(&producer_idx_host_), + sizeof(uint32_t)) != cudaSuccess + || cudaMallocHost(&completion_slot_host_, sizeof(nixlProxyCompletionSlot)) != cudaSuccess) { + NIXL_ERROR << "nixlProxyChannelState::allocate: CUDA allocation failed for channel " + << channel_id; + deallocate(); + return NIXL_ERR_BACKEND; + } + + void *consumer_dev = nullptr; + if (cudaHostGetDevicePointer(&consumer_dev, consumer_idx_host_, 0) != cudaSuccess) { + deallocate(); + return NIXL_ERR_BACKEND; + } + consumer_idx_dev_ = static_cast(consumer_dev); + + void *producer_dev = nullptr; + if (cudaHostGetDevicePointer(&producer_dev, producer_idx_host_, 0) != cudaSuccess) { + deallocate(); + return NIXL_ERR_BACKEND; + } + producer_idx_dev_ = static_cast(producer_dev); + + void *completion_dev = nullptr; + if (cudaHostGetDevicePointer(&completion_dev, completion_slot_host_, 0) != cudaSuccess) { + deallocate(); + return NIXL_ERR_BACKEND; + } + completion_slot_dev_ = static_cast(completion_dev); + + for (uint32_t i = 0; i < depth; ++i) { + records_[i] = nixlProxySubmission{}; + } + __atomic_store_n(producer_idx_host_, 0, __ATOMIC_RELEASE); + __atomic_store_n(consumer_idx_host_, 0, __ATOMIC_RELEASE); + completion_slot_host_->next_status = NIXL_IN_PROG; + __atomic_store_n(&completion_slot_host_->completed_idx, + uint64_t{0}, __ATOMIC_RELEASE); + *work_ring_ = nixlProxyWorkRing{ + records_, + producer_idx_dev_, + consumer_idx_dev_, + depth, + }; + device_view = nixlProxyChannelView{ work_ring_, completion_slot_dev_, channel_id }; + + inflight_requests.clear(); + NIXL_INFO << "nixlProxyChannelState::allocate: channel " << channel_id << " ready" + << " work_ring=" << work_ring_ + << " records=" << records_ + << " producer_idx(host)=" << producer_idx_host_ + << " producer_idx(dev)=" << producer_idx_dev_ + << " consumer_idx(host)=" << consumer_idx_host_ + << " consumer_idx(dev)=" << consumer_idx_dev_ + << " completion_slot(host)=" << completion_slot_host_ + << " completion_slot(dev)=" << completion_slot_dev_; + return NIXL_SUCCESS; +} + +void +nixlProxyChannelState::deallocate() noexcept { + if (completion_slot_host_) { + cudaFreeHost(completion_slot_host_); + completion_slot_host_ = nullptr; + completion_slot_dev_ = nullptr; + } + if (producer_idx_host_) { + cudaFreeHost(producer_idx_host_); + producer_idx_host_ = nullptr; + producer_idx_dev_ = nullptr; + } + if (consumer_idx_host_) { + cudaFreeHost(consumer_idx_host_); + consumer_idx_host_ = nullptr; + consumer_idx_dev_ = nullptr; + } + if (records_) { cudaFreeHost(records_); records_ = nullptr; } + if (work_ring_) { cudaFreeHost(work_ring_); work_ring_ = nullptr; } + device_view = nixlProxyChannelView{}; +} + +nixlProxyChannelState::~nixlProxyChannelState() { + deallocate(); +} + +nixlProxyChannelState::nixlProxyChannelState(nixlProxyChannelState &&other) noexcept + : device_view(other.device_view), + inflight_requests(std::move(other.inflight_requests)), + work_ring_(other.work_ring_), + records_(other.records_), + producer_idx_host_(other.producer_idx_host_), + producer_idx_dev_(other.producer_idx_dev_), + consumer_idx_host_(other.consumer_idx_host_), + consumer_idx_dev_(other.consumer_idx_dev_), + completion_slot_host_(other.completion_slot_host_), + completion_slot_dev_(other.completion_slot_dev_) { + other.work_ring_ = nullptr; + other.records_ = nullptr; + other.producer_idx_host_ = nullptr; + other.producer_idx_dev_ = nullptr; + other.consumer_idx_host_ = nullptr; + other.consumer_idx_dev_ = nullptr; + other.completion_slot_host_ = nullptr; + other.completion_slot_dev_ = nullptr; + other.device_view = nixlProxyChannelView{}; +} + +nixlProxyChannelState & +nixlProxyChannelState::operator=(nixlProxyChannelState &&other) noexcept { + if (this != &other) { + deallocate(); + device_view = other.device_view; + inflight_requests = std::move(other.inflight_requests); + work_ring_ = other.work_ring_; + records_ = other.records_; + producer_idx_host_ = other.producer_idx_host_; + producer_idx_dev_ = other.producer_idx_dev_; + consumer_idx_host_ = other.consumer_idx_host_; + consumer_idx_dev_ = other.consumer_idx_dev_; + completion_slot_host_ = other.completion_slot_host_; + completion_slot_dev_ = other.completion_slot_dev_; + other.work_ring_ = nullptr; + other.records_ = nullptr; + other.producer_idx_host_ = nullptr; + other.producer_idx_dev_ = nullptr; + other.consumer_idx_host_ = nullptr; + other.consumer_idx_dev_ = nullptr; + other.completion_slot_host_ = nullptr; + other.completion_slot_dev_ = nullptr; + other.device_view = nixlProxyChannelView{}; + } + return *this; +} + +nixlProxyRuntime::nixlProxyRuntime() = default; + +nixlProxyRuntime::~nixlProxyRuntime() { + if (backend_) { + shutdown(); + } +} + +nixl_status_t +nixlProxyRuntime::init(std::unique_ptr backend, + uint32_t channel_count, + uint32_t worker_count, + uint64_t pthr_delay_us) { + NIXL_INFO << "ProxyRuntime::init: channel_count=" << channel_count + << " worker_count=" << worker_count + << " pthr_delay_us=" << pthr_delay_us + << " backend=" << backend.get(); + if (backend == nullptr || channel_count == 0 || worker_count == 0) { + NIXL_ERROR << "ProxyRuntime::init: invalid params"; + return NIXL_ERR_INVALID_PARAM; + } + + backend_ = std::move(backend); + memview_registry_.clear(); + + if (cudaMallocHost(reinterpret_cast(&shutdown_word_host_), + sizeof(uint32_t)) != cudaSuccess) { + NIXL_ERROR << "ProxyRuntime::init: failed to allocate shutdown_word"; + shutdown_word_host_ = nullptr; + backend_.reset(); + return NIXL_ERR_BACKEND; + } + void *shutdown_dev = nullptr; + if (cudaHostGetDevicePointer(&shutdown_dev, shutdown_word_host_, 0) != cudaSuccess) { + cudaFreeHost(shutdown_word_host_); + shutdown_word_host_ = nullptr; + backend_.reset(); + return NIXL_ERR_BACKEND; + } + shutdown_word_dev_ = static_cast(shutdown_dev); + __atomic_store_n(shutdown_word_host_, + static_cast(nixl_proxy_control_state_t::RUNNING), + __ATOMIC_RELEASE); + + worker_count = std::min(worker_count, channel_count); + NIXL_INFO << "ProxyRuntime::init: effective worker_count=" << worker_count + << " (clamped to channel_count)"; + + nixl_status_t rc = backend_->init(worker_count, channel_count); + if ((rc != NIXL_SUCCESS) && (rc != NIXL_ERR_NOT_SUPPORTED)) { + NIXL_ERROR << "ProxyRuntime::init: backend init failed: " << rc; + cudaFreeHost(shutdown_word_host_); + shutdown_word_host_ = nullptr; + shutdown_word_dev_ = nullptr; + backend_.reset(); + return rc; + } + if (rc == NIXL_ERR_NOT_SUPPORTED) { + NIXL_INFO << "ProxyRuntime::init: backend init hook not supported; continuing"; + } + + channels_.resize(channel_count); + for (uint32_t channel_id = 0; channel_id < channel_count; ++channel_id) { + rc = channels_[channel_id].allocate(channel_id, ring_depth_); + if (rc != NIXL_SUCCESS) { + channels_.clear(); + backend_->shutdown(); + cudaFreeHost(shutdown_word_host_); + shutdown_word_host_ = nullptr; + shutdown_word_dev_ = nullptr; + backend_.reset(); + return rc; + } + } + + if (cudaMallocHost(&device_channel_views_, + sizeof(nixlProxyChannelView) * channel_count) != cudaSuccess) { + channels_.clear(); + backend_->shutdown(); + cudaFreeHost(shutdown_word_host_); + shutdown_word_host_ = nullptr; + shutdown_word_dev_ = nullptr; + backend_.reset(); + return NIXL_ERR_BACKEND; + } + for (uint32_t channel_id = 0; channel_id < channel_count; ++channel_id) { + device_channel_views_[channel_id] = channels_[channel_id].device_view; + } + + if (cudaMallocHost(&device_context_, + sizeof(nixlProxyDeviceContextData)) != cudaSuccess) { + cudaFreeHost(device_channel_views_); + device_channel_views_ = nullptr; + channels_.clear(); + backend_->shutdown(); + cudaFreeHost(shutdown_word_host_); + shutdown_word_host_ = nullptr; + shutdown_word_dev_ = nullptr; + backend_.reset(); + return NIXL_ERR_BACKEND; + } + *device_context_ = nixlProxyDeviceContextData{ + device_channel_views_, + channel_count, + shutdown_word_dev_ + }; + + workers_.clear(); + workers_.reserve(worker_count); + workers_started_ = false; + + for (uint32_t w = 0; w < worker_count; ++w) { + uint32_t first_ch = (w * channel_count) / worker_count; + uint32_t end_ch = ((w + 1) * channel_count) / worker_count; + uint32_t n_ch = end_ch - first_ch; + + NIXL_INFO << "ProxyRuntime::init: worker " << w + << " assigned channels [" << first_ch << ", " << end_ch << ")"; + workers_.push_back(std::make_unique( + backend_.get(), + &memview_registry_, + shutdown_word_host_, + &channels_[first_ch], + n_ch, + pthr_delay_us)); + } + + NIXL_INFO << "ProxyRuntime::init: complete — " + << channel_count << " channels, " + << worker_count << " workers, " + << "device_context=" << device_context_; + return NIXL_SUCCESS; +} + +nixl_status_t +nixlProxyRuntime::loadRemoteConnInfo(const std::string &remote_name, + const nixl_blob_t &conn_info) { + NIXL_INFO << "ProxyRuntime::loadRemoteConnInfo: remote='" << remote_name + << "' conn_info_size=" << conn_info.size(); + if (backend_ == nullptr) { + NIXL_ERROR << "ProxyRuntime::loadRemoteConnInfo: no backend"; + return NIXL_ERR_NOT_SUPPORTED; + } + nixl_status_t rc = backend_->loadRemoteConnInfo(remote_name, conn_info); + NIXL_INFO << "ProxyRuntime::loadRemoteConnInfo: result=" << rc; + return rc; +} + +nixl_status_t +nixlProxyRuntime::registerProxyMemView(nixlMemViewH backend_memview, + nixlMemViewH *proxy_memview) { + return memview_registry_.registerProxyMemView(backend_memview, proxy_memview); +} + +nixl_status_t +nixlProxyRuntime::prepMemView(const nixl_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview) { + return memview_registry_.prepMemView(dlist, proxy_memview); +} + +nixl_status_t +nixlProxyRuntime::prepMemView(const nixl_remote_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview) { + return memview_registry_.prepMemView(dlist, proxy_memview); +} + +nixl_status_t +nixlProxyRuntime::prepMemView(nixlMemViewH backend_memview, + const nixl_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview) { + return memview_registry_.prepMemView(backend_memview, dlist, proxy_memview); +} + +nixl_status_t +nixlProxyRuntime::prepMemView( + nixlMemViewH backend_memview, + const nixl_remote_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview) { + return memview_registry_.prepMemView(backend_memview, dlist, proxy_memview); +} + +nixl_status_t +nixlProxyRuntime::unregisterProxyMemView(nixlMemViewH proxy_memview) { + return memview_registry_.unregisterProxyMemView(proxy_memview); +} + +nixl_status_t +nixlProxyRuntime::storeMetadata(nixlMemViewH proxy_memview, + const nixl_meta_dlist_t &dlist) { + return memview_registry_.storeMetadata(proxy_memview, dlist); +} + +nixl_status_t +nixlProxyRuntime::storeMetadata(nixlMemViewH proxy_memview, + const nixl_remote_meta_dlist_t &dlist) { + return memview_registry_.storeMetadata(proxy_memview, dlist); +} + +bool +nixlProxyRuntime::resolveProxyMemView(nixlMemViewH proxy_memview, + nixlMemViewH &backend_memview) const { + return memview_registry_.resolveProxyMemView(proxy_memview, backend_memview); +} + +bool +nixlProxyRuntime::resolveProxyMemViewId(uint64_t proxy_memview_id, + nixlMemViewH &backend_memview) const { + return memview_registry_.resolveProxyMemViewId(proxy_memview_id, backend_memview); +} + +nixl_status_t +nixlProxyRuntime::startWorkers() { + NIXL_INFO << "ProxyRuntime::startWorkers: launching " + << workers_.size() << " worker thread(s)"; + if (shutdown_word_host_ == nullptr) { + NIXL_ERROR << "ProxyRuntime::startWorkers: runtime not initialized"; + return NIXL_ERR_NOT_SUPPORTED; + } + __atomic_store_n(shutdown_word_host_, + static_cast(nixl_proxy_control_state_t::SHUTDOWN), + __ATOMIC_RELEASE); + joinWorkerThreads(); + for (auto &channel : channels_) { + channel.inflight_requests.clear(); + channel.error_latched = false; + } + + __atomic_store_n(shutdown_word_host_, + static_cast(nixl_proxy_control_state_t::RUNNING), + __ATOMIC_RELEASE); + + uint32_t idx = 0; + for (auto &worker : workers_) { + worker->start(idx); + ++idx; + } + workers_started_ = true; + + NIXL_INFO << "ProxyRuntime::startWorkers: all threads launched"; + return NIXL_SUCCESS; +} + +void +nixlProxyRuntime::joinWorkerThreads() noexcept { + for (auto &worker : workers_) { + worker->join(); + } +} + +nixl_status_t +nixlProxyRuntime::shutdown() { + NIXL_INFO << "ProxyRuntime::shutdown: signalling workers to stop"; + if (shutdown_word_host_ != nullptr) { + __atomic_store_n(shutdown_word_host_, + static_cast(nixl_proxy_control_state_t::SHUTDOWN), + __ATOMIC_RELEASE); + } + + joinWorkerThreads(); + workers_started_ = false; + NIXL_INFO << "ProxyRuntime::shutdown: all worker threads joined"; + + nixl_status_t backend_status = NIXL_SUCCESS; + if (backend_ != nullptr) { + NIXL_INFO << "ProxyRuntime::shutdown: shutting down backend"; + backend_status = backend_->shutdown(); + NIXL_INFO << "ProxyRuntime::shutdown: backend shutdown status=" << backend_status; + if (backend_status == NIXL_ERR_NOT_SUPPORTED) { + backend_status = NIXL_SUCCESS; + } + } + + workers_.clear(); + memview_registry_.clear(); + + if (device_context_) { + cudaFreeHost(device_context_); + device_context_ = nullptr; + } + if (shutdown_word_host_) { + cudaFreeHost(shutdown_word_host_); + shutdown_word_host_ = nullptr; + shutdown_word_dev_ = nullptr; + } + if (device_channel_views_) { + cudaFreeHost(device_channel_views_); + device_channel_views_ = nullptr; + } + + channels_.clear(); + backend_.reset(); + NIXL_INFO << "ProxyRuntime::shutdown: complete"; + return backend_status; +} diff --git a/src/core/device_proxy/proxy_runtime.h b/src/core/device_proxy/proxy_runtime.h new file mode 100644 index 0000000000..cdbe9e3b09 --- /dev/null +++ b/src/core/device_proxy/proxy_runtime.h @@ -0,0 +1,301 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef NIXL_SRC_CORE_DEVICE_PROXY_PROXY_RUNTIME_H +#define NIXL_SRC_CORE_DEVICE_PROXY_PROXY_RUNTIME_H + +#include +#include +#include +#include +#include +#include + +#include "backend_aux.h" +#include "proxy_protocol.h" +#include "backend_adapter.h" + +class ProxyWorker; + +static constexpr uint32_t kDefaultProxyRingDepth = 256; + +struct nixlProxyRequestState { + uint64_t op_idx = 0; + uint64_t backend_req_token = 0; + nixl_status_t status = NIXL_IN_PROG; +}; + +struct alignas(64) nixlProxyChannelState { + nixlProxyChannelView device_view{}; + std::deque inflight_requests; + bool error_latched = false; + + nixlProxyWorkRing *work_ring_ = nullptr; + nixlProxySubmission *records_ = nullptr; + /** Mapped pinned host memory; host proxy uses __atomic_* on host alias. */ + uint32_t *producer_idx_host_ = nullptr; + /** Device-mapped alias of producer_idx_host_ for nixlProxyWorkRing (GPU-writable). */ + uint32_t *producer_idx_dev_ = nullptr; + /** Consumer count: host pinned; proxy uses __atomic_* on consumer_idx_host_. */ + uint32_t *consumer_idx_host_ = nullptr; + /** Same word as consumer_idx_host_, for nixlProxyWorkRing::consumer_idx (GPU-readable). */ + uint32_t *consumer_idx_dev_ = nullptr; + /** Mapped pinned host memory; proxy worker writes directly via host alias. */ + nixlProxyCompletionSlot *completion_slot_host_ = nullptr; + /** Device-mapped alias of completion_slot_host_ for nixlProxyChannelView. */ + nixlProxyCompletionSlot *completion_slot_dev_ = nullptr; + + nixlProxyChannelState() = default; + ~nixlProxyChannelState(); + nixlProxyChannelState(nixlProxyChannelState &&) noexcept; + nixlProxyChannelState &operator=(nixlProxyChannelState &&) noexcept; + nixlProxyChannelState(const nixlProxyChannelState &) = delete; + nixlProxyChannelState &operator=(const nixlProxyChannelState &) = delete; + + nixl_status_t + allocate(uint32_t channel_id, uint32_t depth); + + void + deallocate() noexcept; +}; + +class nixlProxyMemViewRegistry { + public: + nixl_status_t + registerProxyMemView(nixlMemViewH backend_memview, + nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(const nixl_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(const nixl_remote_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(nixlMemViewH backend_memview, + const nixl_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(nixlMemViewH backend_memview, + const nixl_remote_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview); + + nixl_status_t + unregisterProxyMemView(nixlMemViewH proxy_memview); + + nixl_status_t + storeMetadata(nixlMemViewH proxy_memview, + const nixl_meta_dlist_t &dlist); + + nixl_status_t + storeMetadata(nixlMemViewH proxy_memview, + const nixl_remote_meta_dlist_t &dlist); + + bool + resolveProxyMemView(nixlMemViewH proxy_memview, + nixlMemViewH &backend_memview) const; + + bool + resolveProxyMemViewId(uint64_t proxy_memview_id, + nixlMemViewH &backend_memview) const; + + nixl_status_t + prepareSubmission(const nixlProxySubmission &submission, + nixlBackendProxySubmission &prepared_submission) const; + + void + clear() noexcept; + + private: + struct ProxyMemViewRegStoredEntry { + uintptr_t base_addr = 0; + size_t len = 0; + uint64_t dev_id = 0; + nixlBackendMD *metadata = nullptr; + }; + + struct LocalMetadata { + nixl_mem_t mem_type = DRAM_SEG; + std::vector entries; + }; + + struct RemoteMetadata { + nixl_mem_t mem_type = DRAM_SEG; + std::string remote_agent; + std::vector entries; + }; + + enum class ProxyMemViewRegEntryState : uint8_t { + ENTRY_ALLOCATED, + ENTRY_READY, + ENTRY_RETIRED, + }; + + enum class ProxyMemViewRegMetadataKind : uint8_t { + METADATA_KIND_NONE, + METADATA_KIND_LOCAL, + METADATA_KIND_REMOTE, + }; + + struct RegistryEntry { + uint64_t proxy_memview_id = 0; + nixlMemViewH backend_memview = nullptr; + ProxyMemViewRegEntryState state = ProxyMemViewRegEntryState::ENTRY_ALLOCATED; + ProxyMemViewRegMetadataKind metadata_kind = ProxyMemViewRegMetadataKind::METADATA_KIND_NONE; + LocalMetadata local_metadata{}; + RemoteMetadata remote_metadata{}; + }; + + RegistryEntry * + getEntryForHandle(nixlMemViewH proxy_memview); + + const RegistryEntry * + getEntryForHandle(nixlMemViewH proxy_memview) const; + + RegistryEntry * + getEntryForId(uint64_t proxy_memview_id); + + const RegistryEntry * + getEntryForId(uint64_t proxy_memview_id) const; + + nixl_status_t + getRemoteEntryForSubmission(uint64_t proxy_memview_id, + size_t index, + size_t offset, + size_t size, + const RemoteMetadata *&metadata, + const ProxyMemViewRegStoredEntry *&entry) const; + + nixl_status_t + getLocalEntryForSubmission(uint64_t proxy_memview_id, + size_t index, + size_t offset, + size_t size, + const LocalMetadata *&metadata, + const ProxyMemViewRegStoredEntry *&entry) const; + + static bool + rangeFits(const ProxyMemViewRegStoredEntry &entry, size_t offset, size_t size); + + static void + fillLocalMetadata(const nixl_meta_dlist_t &dlist, LocalMetadata &out); + + static void + fillRemoteMetadata(const nixl_remote_meta_dlist_t &dlist, RemoteMetadata &out); + + std::vector entries_; + uint64_t next_proxy_memview_id_ = 1; +}; + +class nixlProxyRuntime { + public: + nixlProxyRuntime(); + ~nixlProxyRuntime(); + + nixlProxyRuntime(nixlProxyRuntime &&) = delete; + nixlProxyRuntime(const nixlProxyRuntime &) = delete; + nixlProxyRuntime& operator=(nixlProxyRuntime &&) = delete; + nixlProxyRuntime& operator=(const nixlProxyRuntime &) = delete; + + nixl_status_t + init(std::unique_ptr backend, + uint32_t channel_count, + uint32_t worker_count, + uint64_t pthr_delay_us = 0); + + nixl_status_t + loadRemoteConnInfo(const std::string &remote_name, + const nixl_blob_t &conn_info); + + nixl_status_t + registerProxyMemView(nixlMemViewH backend_memview, + nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(const nixl_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(const nixl_remote_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(nixlMemViewH backend_memview, + const nixl_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(nixlMemViewH backend_memview, + const nixl_remote_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview); + + nixl_status_t + unregisterProxyMemView(nixlMemViewH proxy_memview); + + nixl_status_t + storeMetadata(nixlMemViewH proxy_memview, + const nixl_meta_dlist_t &dlist); + + nixl_status_t + storeMetadata(nixlMemViewH proxy_memview, + const nixl_remote_meta_dlist_t &dlist); + + bool + resolveProxyMemView(nixlMemViewH proxy_memview, + nixlMemViewH &backend_memview) const; + + bool + resolveProxyMemViewId(uint64_t proxy_memview_id, + nixlMemViewH &backend_memview) const; + + nixl_status_t + startWorkers(); + + nixl_status_t + shutdown(); + + const nixlProxyMemViewRegistry & + memviewRegistry() const { return memview_registry_; } + + uint32_t + channelCount() const { return static_cast(channels_.size()); } + + const nixlProxyChannelView * + deviceChannelViews() const { return device_channel_views_; } + + nixlProxyDeviceContextData * + deviceContext() const { return device_context_; } + + private: + void + joinWorkerThreads() noexcept; + + std::vector channels_; + nixlProxyChannelView *device_channel_views_ = nullptr; + nixlProxyDeviceContextData *device_context_ = nullptr; + std::vector> workers_; + nixlProxyMemViewRegistry memview_registry_; + std::unique_ptr backend_; + uint32_t *shutdown_word_host_ = nullptr; + uint32_t *shutdown_word_dev_ = nullptr; + uint32_t ring_depth_ = kDefaultProxyRingDepth; + bool workers_started_ = false; +}; + +#endif // NIXL_SRC_CORE_DEVICE_PROXY_PROXY_RUNTIME_H diff --git a/src/core/device_proxy/proxy_worker.cpp b/src/core/device_proxy/proxy_worker.cpp new file mode 100644 index 0000000000..83e32a6f3d --- /dev/null +++ b/src/core/device_proxy/proxy_worker.cpp @@ -0,0 +1,182 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "proxy_worker.h" +#include "proxy_runtime.h" +#include "backend_adapter.h" +#include "nixl_log.h" +#include +#include + +ProxyWorker::ProxyWorker(nixlDeviceProxyBackendAdapter *backend, + const nixlProxyMemViewRegistry *proxy_memview_registry, + uint32_t *shutdown_word, + nixlProxyChannelState *assigned_channels, + uint32_t assigned_channel_count, + uint64_t pthr_delay_us) noexcept + : backend_(backend), + proxy_memview_registry_(proxy_memview_registry), + shutdown_word_(shutdown_word), + assigned_channels_(assigned_channels), + assigned_channel_count_(assigned_channel_count), + pthr_delay_us_(pthr_delay_us) {} + +ProxyWorker::~ProxyWorker() { + join(); +} + +void +ProxyWorker::start(uint32_t worker_idx) { + thread_ = std::thread([this, worker_idx]() { + NIXL_INFO << "ProxyWorker thread " << worker_idx << " started"; + while (__atomic_load_n(shutdown_word_, __ATOMIC_ACQUIRE) + == static_cast(nixl_proxy_control_state_t::RUNNING)) { + runOnce(); + if (pthr_delay_us_ > 0) { + std::this_thread::sleep_for(std::chrono::microseconds(pthr_delay_us_)); + } + } + NIXL_INFO << "ProxyWorker thread " << worker_idx << " exiting"; + }); +} + +void +ProxyWorker::join() noexcept { + if (thread_.joinable()) { + thread_.join(); + } +} + +void +ProxyWorker::runOnce() { + for (uint32_t i = 0; i < assigned_channel_count_; i++) { + nixlProxyChannelState &channel = assigned_channels_[i]; + nixlProxySubmission submission; + while (tryDequeue(channel, submission)) { + submitToBackend(channel, submission); + } + } + driveBackendProgress(); + for (uint32_t i = 0; i < assigned_channel_count_; i++) { + nixlProxyChannelState &channel = assigned_channels_[i]; + publishCompletions(channel); + } +} + +bool +ProxyWorker::tryDequeue(nixlProxyChannelState &channel, nixlProxySubmission &submission) { + nixlProxyWorkRing *ring = channel.work_ring_; + // Sole writer of consumer_idx on host — relaxed load is sufficient. + uint32_t local_consumer_idx = + __atomic_load_n(channel.consumer_idx_host_, __ATOMIC_RELAXED); + uint32_t slot = local_consumer_idx % ring->depth; + // ready_flag is the GPU-to-CPU signal that the record is written + // (pairs with release store in device enqueue). No producer_idx + // read on host — it is GPU-internal for slot allocation. + if (!__atomic_load_n(&ring->records[slot].ready_flag, __ATOMIC_ACQUIRE)) { + return false; + } + submission = ring->records[slot]; + __atomic_store_n(&ring->records[slot].ready_flag, 0, __ATOMIC_RELAXED); + __atomic_store_n(channel.consumer_idx_host_, + local_consumer_idx + 1, + __ATOMIC_RELEASE); + NIXL_DEBUG << "ProxyWorker::tryDequeue: channel=" << channel.device_view.channel_id + << " consumer=" << local_consumer_idx + << " opcode=" << static_cast(submission.opcode) + << " op_idx=" << submission.op_idx + << " size=" << submission.size; + return true; +} + +void +ProxyWorker::submitToBackend(nixlProxyChannelState &channel, const nixlProxySubmission &submission) { + nixlBackendProxySubmission prepared_submission; + nixl_status_t status = + proxy_memview_registry_->prepareSubmission(submission, prepared_submission); + if (status != NIXL_SUCCESS) { + NIXL_DEBUG << "ProxyWorker::submitToBackend: submission preparation failed" + << " op_idx=" << submission.op_idx + << " status=" << status; + channel.inflight_requests.push_back( + {submission.op_idx, 0, status}); + // The terminal error is queued for publishCompletions(); the worker handled it. + return; + } + + NIXL_DEBUG << "ProxyWorker::submitToBackend: op_idx=" << submission.op_idx + << " opcode=" << static_cast(submission.opcode) + << " channel=" << submission.channel_id + << " local_addr=0x" << std::hex << prepared_submission.local.desc.addr + << " remote_addr=0x" << prepared_submission.remote.desc.addr << std::dec + << " size=" << submission.size + << " remote_agent='" << prepared_submission.remote_agent << "'"; + + uint64_t request_token = 0; + nixlProxyRequestState inflight{}; + inflight.op_idx = submission.op_idx; + status = backend_->submit(prepared_submission, request_token); + inflight.backend_req_token = request_token; + if (status != NIXL_SUCCESS) { + // backend submit failed, so status is already terminal and can be + // published without polling the backend. + NIXL_ERROR << "ProxyWorker::submitToBackend: backend submit failed" + << " status=" << status << " op_idx=" << submission.op_idx + << " request_token=" << request_token; + inflight.status = status; + } + + NIXL_DEBUG << "ProxyWorker::submitToBackend: submitted op_idx=" << submission.op_idx + << " request_token=" << request_token << " status=" << status; + channel.inflight_requests.push_back(inflight); +} + +void +ProxyWorker::driveBackendProgress() { + backend_->progress(); +} + +void +ProxyWorker::publishCompletions(nixlProxyChannelState &channel) { + if (channel.error_latched) { + return; + } + while (!channel.inflight_requests.empty()) { + nixlProxyRequestState &front = channel.inflight_requests.front(); + nixl_status_t st; + if (front.status != NIXL_IN_PROG) { + st = front.status; + } else { + st = backend_->checkCompletion(front.backend_req_token); + if (st == NIXL_IN_PROG) { + break; + } + } + NIXL_DEBUG << "ProxyWorker::publishCompletions: channel=" + << channel.device_view.channel_id + << " op_idx=" << front.op_idx + << " status=" << st + << " token=" << front.backend_req_token; + channel.completion_slot_host_->next_status = st; + __atomic_store_n(&channel.completion_slot_host_->completed_idx, + front.op_idx, __ATOMIC_RELEASE); + channel.inflight_requests.pop_front(); + if (st != NIXL_SUCCESS) { + channel.error_latched = true; + break; + } + } +} diff --git a/src/core/device_proxy/proxy_worker.h b/src/core/device_proxy/proxy_worker.h new file mode 100644 index 0000000000..3eea79a534 --- /dev/null +++ b/src/core/device_proxy/proxy_worker.h @@ -0,0 +1,68 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef NIXL_SRC_CORE_DEVICE_PROXY_PROXY_WORKER_H +#define NIXL_SRC_CORE_DEVICE_PROXY_PROXY_WORKER_H + +#include +#include +#include "proxy_protocol.h" + +class nixlDeviceProxyBackendAdapter; +class nixlProxyMemViewRegistry; +struct nixlProxyChannelState; + +class ProxyWorker { +public: + ProxyWorker(nixlDeviceProxyBackendAdapter *backend, + const nixlProxyMemViewRegistry *proxy_memview_registry, + uint32_t *shutdown_word, + nixlProxyChannelState *assigned_channels, + uint32_t assigned_channel_count, + uint64_t pthr_delay_us) noexcept; + ~ProxyWorker(); + + void + start(uint32_t worker_idx); + void + join() noexcept; + + void + runOnce(); + +private: + bool + tryDequeue(nixlProxyChannelState &channel, nixlProxySubmission &submission); + + void + submitToBackend(nixlProxyChannelState &channel, const nixlProxySubmission &submission); + + void + driveBackendProgress(); + + void + publishCompletions(nixlProxyChannelState &channel); + + nixlDeviceProxyBackendAdapter *backend_ = nullptr; + const nixlProxyMemViewRegistry *proxy_memview_registry_ = nullptr; + uint32_t *shutdown_word_ = nullptr; + nixlProxyChannelState *assigned_channels_ = nullptr; + uint32_t assigned_channel_count_ = 0; + uint64_t pthr_delay_us_ = 0; + std::thread thread_; +}; + +#endif // NIXL_SRC_CORE_DEVICE_PROXY_PROXY_WORKER_H diff --git a/src/core/meson.build b/src/core/meson.build index ee542bf94d..a94b5ee108 100644 --- a/src/core/meson.build +++ b/src/core/meson.build @@ -17,6 +17,10 @@ nixl_lib_deps = [nixl_infra, serdes_interface, stream_interface, dl_dep, nixl_common_dep, thread_dep] +if cuda_dep.found() + nixl_lib_deps += [ cuda_dep ] +endif + if etcd_dep.found() nixl_lib_deps += [ etcd_dep ] endif @@ -58,6 +62,20 @@ if libtransfer_engine.found() and not disable_mooncake_backend and 'Mooncake' in nixl_lib_deps += [ mooncake_backend_interface, cuda_dep ] endif +nixl_lib_link_args = ['-lstdc++fs'] +nixl_lib_build_rpath = '' +nixl_lib_install_rpath = join_paths(get_option('prefix'), get_option('libdir')) +if 'UCX' in static_plugins and ucx_path != '' + nixl_lib_link_args += [ + '-L' + ucx_path + '/lib', + '-Wl,--no-as-needed', + '-lucp', '-lucs', '-luct', + '-Wl,--as-needed', + ] + nixl_lib_build_rpath = ucx_path + '/lib' + nixl_lib_install_rpath += ':' + ucx_path + '/lib' +endif + # Tracing (nixl::trace). Only the facade (tracer.cpp) and the plugin loader are # compiled into libnixl; each backend ships as a separate .so plugin under # src/plugins/tracing, loaded on demand by nixlPluginManager. @@ -71,14 +89,17 @@ nixl_lib_sources = [ 'telemetry/buffer_plugin.cpp', 'telemetry/nop_plugin.cpp', 'tracing/tracer.cpp', + 'device_proxy/proxy_runtime.cpp', + 'device_proxy/proxy_worker.cpp', ] nixl_lib = library('nixl', nixl_lib_sources, include_directories: [ nixl_inc_dirs, utils_inc_dirs, include_directories('tracing') ], - link_args: ['-lstdc++fs'], + link_args: nixl_lib_link_args, dependencies: nixl_lib_deps, install: true, - install_rpath: join_paths(get_option('prefix'), get_option('libdir'))) + build_rpath: nixl_lib_build_rpath, + install_rpath: nixl_lib_install_rpath) nixl_dep = declare_dependency(link_with: nixl_lib, include_directories: nixl_inc_dirs) diff --git a/src/core/nixl_agent.cpp b/src/core/nixl_agent.cpp index 5e123ce896..a08a9831a7 100644 --- a/src/core/nixl_agent.cpp +++ b/src/core/nixl_agent.cpp @@ -448,6 +448,19 @@ nixlAgent::createBackend(const nixl_backend_t &type, data->connMd_[type] = conn_info; } + if (data->proxyModeEnabled()) { + if (backend->supportsProxy()) { + const nixl_status_t ret = data->createProxyRuntime(backend.get(), type, init_params); + if (ret != NIXL_SUCCESS) { + NIXL_ERROR_FUNC << "Failed to initialize proxy runtime on backend '" << type + << "' with status " << ret; + return ret; + } + } else { + NIXL_WARN << "Proxy runtime is enabled but backend '" << type << "' does not support it"; + } + } + // TODO: Simplify, e.g. by making nixlBackendH's c'tor public? std::unique_ptr bknd_temp(new nixlBackendH(backend.get())); const auto [it, inserted] = data->backendHandles_.try_emplace(type, std::move(bknd_temp)); diff --git a/src/plugins/ucx/ucx_backend.h b/src/plugins/ucx/ucx_backend.h index a2a9e2dbc2..51b4b4fce1 100644 --- a/src/plugins/ucx/ucx_backend.h +++ b/src/plugins/ucx/ucx_backend.h @@ -114,6 +114,11 @@ class nixlUcxEngine : public nixlBackendEngine { return true; } + bool + supportsProxy() const override { + return true; + } + nixl_mem_list_t getSupportedMems() const override; diff --git a/test/gtest/unit/meson.build b/test/gtest/unit/meson.build index f684362329..9c8b3799a5 100644 --- a/test/gtest/unit/meson.build +++ b/test/gtest/unit/meson.build @@ -32,6 +32,12 @@ unit_test_deps += [descriptors_unit_test_dep] subdir('tracing') unit_test_deps += [tracing_unit_test_dep] +subdir('proxy_memview_registry') +unit_test_deps += [proxy_memview_registry_unit_test_dep] + +subdir('proxy_runtime') +unit_test_deps += [proxy_runtime_unit_test_dep] + cpp = meson.get_compiler('cpp') azure_storage_blobs = cpp.find_library('azure-storage-blobs', required: false) if enabled_plugins.get('AZURE_BLOB') and azure_storage_blobs.found() diff --git a/test/gtest/unit/proxy_memview_registry/meson.build b/test/gtest/unit/proxy_memview_registry/meson.build new file mode 100644 index 0000000000..031e4340b5 --- /dev/null +++ b/test/gtest/unit/proxy_memview_registry/meson.build @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +proxy_memview_registry_unit_test_dep = declare_dependency( + sources: [ + 'proxy_memview_registry.cpp', + ], + include_directories: [nixl_inc_dirs, gtest_inc_dirs], +) diff --git a/test/gtest/unit/proxy_memview_registry/proxy_memview_registry.cpp b/test/gtest/unit/proxy_memview_registry/proxy_memview_registry.cpp new file mode 100644 index 0000000000..aed0565ae1 --- /dev/null +++ b/test/gtest/unit/proxy_memview_registry/proxy_memview_registry.cpp @@ -0,0 +1,449 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include + +#include + +#include "device_proxy/proxy_runtime.h" + +namespace gtest { +namespace proxy_memview_registry { + + class ProxyMemViewRegistryTest : public testing::Test { + protected: + class DummyBackendMD : public nixlBackendMD { + public: + DummyBackendMD() : nixlBackendMD(false) {} + }; + + nixlProxyMemViewRegistry registry_; + DummyBackendMD local_md_; + DummyBackendMD remote_md_; + + nixlMemViewH + makeFakeBackendHandle(uint64_t id) { + return reinterpret_cast(id); + } + + nixl_meta_dlist_t + makeLocalMetadata(uintptr_t base_addr, uint64_t dev_id = 0) { + nixl_meta_dlist_t dlist(DRAM_SEG); + dlist.addDesc(nixlMetaDesc(base_addr, 64, dev_id, &local_md_)); + return dlist; + } + + nixl_remote_meta_dlist_t + makeRemoteMetadata(uintptr_t base_addr, + const std::string &remote_agent = "peer", + uint64_t dev_id = 0) { + nixl_remote_meta_dlist_t dlist(DRAM_SEG); + nixlRemoteMetaDesc desc(remote_agent); + desc.addr = base_addr; + desc.len = 64; + desc.devId = dev_id; + desc.metadataP = &remote_md_; + dlist.addDesc(desc); + return dlist; + } + }; + + TEST_F(ProxyMemViewRegistryTest, RegisterSingle) { + nixlMemViewH proxy_handle = nullptr; + EXPECT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(100), &proxy_handle), + NIXL_SUCCESS); + EXPECT_NE(proxy_handle, nullptr); + } + + TEST_F(ProxyMemViewRegistryTest, RegisterNullOutputReturnsError) { + EXPECT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(100), nullptr), + NIXL_ERR_INVALID_PARAM); + } + + TEST_F(ProxyMemViewRegistryTest, RegisterMultipleAssignsUniqueIds) { + nixlMemViewH h1 = nullptr, h2 = nullptr, h3 = nullptr; + EXPECT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &h1), NIXL_SUCCESS); + EXPECT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &h2), NIXL_SUCCESS); + EXPECT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(30), &h3), NIXL_SUCCESS); + + EXPECT_NE(h1, h2); + EXPECT_NE(h2, h3); + EXPECT_NE(h1, h3); + } + + TEST_F(ProxyMemViewRegistryTest, ResolveByHandle) { + auto backend = makeFakeBackendHandle(42); + nixlMemViewH proxy_handle = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(backend, &proxy_handle), NIXL_SUCCESS); + + nixlMemViewH resolved = nullptr; + EXPECT_TRUE(registry_.resolveProxyMemView(proxy_handle, resolved)); + EXPECT_EQ(resolved, backend); + } + + TEST_F(ProxyMemViewRegistryTest, ResolveById) { + auto backend = makeFakeBackendHandle(42); + nixlMemViewH proxy_handle = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(backend, &proxy_handle), NIXL_SUCCESS); + + auto proxy_id = reinterpret_cast(proxy_handle); + nixlMemViewH resolved = nullptr; + EXPECT_TRUE(registry_.resolveProxyMemViewId(proxy_id, resolved)); + EXPECT_EQ(resolved, backend); + } + + TEST_F(ProxyMemViewRegistryTest, ResolveMultiple) { + auto b1 = makeFakeBackendHandle(10), b2 = makeFakeBackendHandle(20); + nixlMemViewH h1 = nullptr, h2 = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(b1, &h1), NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(b2, &h2), NIXL_SUCCESS); + + nixlMemViewH r1 = nullptr, r2 = nullptr; + EXPECT_TRUE(registry_.resolveProxyMemView(h1, r1)); + EXPECT_TRUE(registry_.resolveProxyMemView(h2, r2)); + EXPECT_EQ(r1, b1); + EXPECT_EQ(r2, b2); + } + + TEST_F(ProxyMemViewRegistryTest, AllocatedEntryIsResolvableBeforeMetadataPublish) { + nixlMemViewH proxy_handle = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(42), &proxy_handle), + NIXL_SUCCESS); + + nixlMemViewH resolved = nullptr; + EXPECT_TRUE(registry_.resolveProxyMemView(proxy_handle, resolved)); + EXPECT_EQ(resolved, makeFakeBackendHandle(42)); + } + + TEST_F(ProxyMemViewRegistryTest, PrepareSubmissionRequiresReadyEntries) { + nixlMemViewH src_proxy = nullptr; + nixlMemViewH dst_proxy = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), NIXL_SUCCESS); + + nixlProxySubmission submission{}; + submission.opcode = nixl_proxy_opcode_t::PUT; + submission.src_proxy_memview_id = reinterpret_cast(src_proxy); + submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + submission.size = 16; + + nixlBackendProxySubmission prepared_submission; + EXPECT_EQ(registry_.prepareSubmission(submission, prepared_submission), NIXL_ERR_NOT_FOUND); + } + + TEST_F(ProxyMemViewRegistryTest, ReadyEntriesProducePreparedTransportDescriptors) { + nixlMemViewH src_proxy = nullptr; + nixlMemViewH dst_proxy = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(src_proxy, makeLocalMetadata(0x1000)), NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(dst_proxy, makeRemoteMetadata(0x2000, "remote-agent")), + NIXL_SUCCESS); + + nixlProxySubmission submission{}; + submission.opcode = nixl_proxy_opcode_t::PUT; + submission.op_idx = 7; + submission.channel_id = 3; + submission.src_proxy_memview_id = reinterpret_cast(src_proxy); + submission.src_offset = 5; + submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + submission.dst_offset = 9; + submission.size = 16; + + nixlBackendProxySubmission prepared_submission; + ASSERT_EQ(registry_.prepareSubmission(submission, prepared_submission), NIXL_SUCCESS); + EXPECT_EQ(prepared_submission.op_idx, 7u); + EXPECT_EQ(prepared_submission.channel_id, 3u); + EXPECT_EQ(prepared_submission.local.mem_type, DRAM_SEG); + EXPECT_EQ(prepared_submission.local.desc.addr, 0x1005u); + EXPECT_EQ(prepared_submission.local.desc.len, 16u); + EXPECT_EQ(prepared_submission.local.desc.metadataP, &local_md_); + EXPECT_EQ(prepared_submission.remote.mem_type, DRAM_SEG); + EXPECT_EQ(prepared_submission.remote.desc.addr, 0x2009u); + EXPECT_EQ(prepared_submission.remote.desc.len, 16u); + EXPECT_EQ(prepared_submission.remote.desc.metadataP, &remote_md_); + EXPECT_EQ(prepared_submission.remote_agent, "remote-agent"); + } + + TEST_F(ProxyMemViewRegistryTest, PrepMemViewProducesReadyEntries) { + nixlMemViewH src_proxy = nullptr; + nixlMemViewH dst_proxy = nullptr; + ASSERT_EQ(registry_.prepMemView(makeLocalMetadata(0x1000), &src_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.prepMemView(makeRemoteMetadata(0x2000), &dst_proxy), + NIXL_SUCCESS); + + nixlMemViewH resolved = makeFakeBackendHandle(42); + EXPECT_TRUE(registry_.resolveProxyMemView(src_proxy, resolved)); + EXPECT_EQ(resolved, nullptr); + + nixlProxySubmission submission{}; + submission.opcode = nixl_proxy_opcode_t::PUT; + submission.src_proxy_memview_id = reinterpret_cast(src_proxy); + submission.src_offset = 4; + submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + submission.dst_offset = 8; + submission.size = 16; + + nixlBackendProxySubmission prepared_submission; + ASSERT_EQ(registry_.prepareSubmission(submission, prepared_submission), NIXL_SUCCESS); + EXPECT_EQ(prepared_submission.local.desc.addr, 0x1004u); + EXPECT_EQ(prepared_submission.local.desc.len, 16u); + EXPECT_EQ(prepared_submission.local.desc.metadataP, &local_md_); + EXPECT_EQ(prepared_submission.remote.desc.addr, 0x2008u); + EXPECT_EQ(prepared_submission.remote.desc.len, 16u); + EXPECT_EQ(prepared_submission.remote.desc.metadataP, &remote_md_); + } + + TEST_F(ProxyMemViewRegistryTest, PrepareSubmissionAllowsRangesEndingAtDescriptorBoundary) { + nixlMemViewH src_proxy = nullptr; + nixlMemViewH dst_proxy = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(src_proxy, makeLocalMetadata(0x1000)), NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(dst_proxy, makeRemoteMetadata(0x2000)), NIXL_SUCCESS); + + nixlProxySubmission submission{}; + submission.opcode = nixl_proxy_opcode_t::PUT; + submission.src_proxy_memview_id = reinterpret_cast(src_proxy); + submission.src_offset = 48; + submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + submission.dst_offset = 48; + submission.size = 16; + + nixlBackendProxySubmission prepared_submission; + ASSERT_EQ(registry_.prepareSubmission(submission, prepared_submission), NIXL_SUCCESS); + EXPECT_EQ(prepared_submission.local.desc.addr, 0x1030u); + EXPECT_EQ(prepared_submission.local.desc.len, 16u); + EXPECT_EQ(prepared_submission.remote.desc.addr, 0x2030u); + EXPECT_EQ(prepared_submission.remote.desc.len, 16u); + } + + TEST_F(ProxyMemViewRegistryTest, PrepareSubmissionRejectsSourceRangeOutsideDescriptor) { + nixlMemViewH src_proxy = nullptr; + nixlMemViewH dst_proxy = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(src_proxy, makeLocalMetadata(0x1000)), NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(dst_proxy, makeRemoteMetadata(0x2000)), NIXL_SUCCESS); + + nixlProxySubmission submission{}; + submission.opcode = nixl_proxy_opcode_t::PUT; + submission.src_proxy_memview_id = reinterpret_cast(src_proxy); + submission.src_offset = 60; + submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + submission.size = 8; + + nixlBackendProxySubmission prepared_submission; + prepared_submission.op_idx = 123; + EXPECT_EQ(registry_.prepareSubmission(submission, prepared_submission), + NIXL_ERR_INVALID_PARAM); + EXPECT_EQ(prepared_submission.op_idx, 123u); + } + + TEST_F(ProxyMemViewRegistryTest, PrepareSubmissionRejectsDestinationRangeOutsideDescriptor) { + nixlMemViewH src_proxy = nullptr; + nixlMemViewH dst_proxy = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(src_proxy, makeLocalMetadata(0x1000)), NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(dst_proxy, makeRemoteMetadata(0x2000)), NIXL_SUCCESS); + + nixlProxySubmission submission{}; + submission.opcode = nixl_proxy_opcode_t::PUT; + submission.src_proxy_memview_id = reinterpret_cast(src_proxy); + submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + submission.dst_offset = 60; + submission.size = 8; + + nixlBackendProxySubmission prepared_submission; + EXPECT_EQ(registry_.prepareSubmission(submission, prepared_submission), + NIXL_ERR_INVALID_PARAM); + } + + TEST_F(ProxyMemViewRegistryTest, PrepareSubmissionRejectsOverflowingRange) { + nixlMemViewH src_proxy = nullptr; + nixlMemViewH dst_proxy = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(src_proxy, makeLocalMetadata(0x1000)), NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(dst_proxy, makeRemoteMetadata(0x2000)), NIXL_SUCCESS); + + nixlProxySubmission submission{}; + submission.opcode = nixl_proxy_opcode_t::PUT; + submission.src_proxy_memview_id = reinterpret_cast(src_proxy); + submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + submission.dst_offset = std::numeric_limits::max(); + submission.size = 1; + + nixlBackendProxySubmission prepared_submission; + EXPECT_EQ(registry_.prepareSubmission(submission, prepared_submission), + NIXL_ERR_INVALID_PARAM); + } + + TEST_F(ProxyMemViewRegistryTest, PrepareSubmissionRejectsUnsupportedOpcode) { + nixlMemViewH dst_proxy = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(dst_proxy, makeRemoteMetadata(0x2000)), NIXL_SUCCESS); + + nixlProxySubmission submission{}; + submission.opcode = static_cast(99); + submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + + nixlBackendProxySubmission prepared_submission; + prepared_submission.op_idx = 123; + EXPECT_EQ(registry_.prepareSubmission(submission, prepared_submission), + NIXL_ERR_NOT_SUPPORTED); + EXPECT_EQ(prepared_submission.op_idx, 123u); + } + + TEST_F(ProxyMemViewRegistryTest, PreparedDescriptorsPreserveDeviceIds) { + nixlMemViewH src_proxy = nullptr; + nixlMemViewH dst_proxy = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(src_proxy, makeLocalMetadata(0x1000, 7)), + NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(dst_proxy, makeRemoteMetadata(0x2000, "peer", 11)), + NIXL_SUCCESS); + + nixlProxySubmission submission{}; + submission.opcode = nixl_proxy_opcode_t::PUT; + submission.src_proxy_memview_id = reinterpret_cast(src_proxy); + submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + submission.size = 8; + + nixlBackendProxySubmission prepared_submission; + ASSERT_EQ(registry_.prepareSubmission(submission, prepared_submission), NIXL_SUCCESS); + EXPECT_EQ(prepared_submission.local.desc.devId, 7u); + EXPECT_EQ(prepared_submission.remote.desc.devId, 11u); + } + + TEST_F(ProxyMemViewRegistryTest, AtomicAddUsesCounterSizeForDestinationBounds) { + nixlMemViewH dst_proxy = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(dst_proxy, makeRemoteMetadata(0x2000)), NIXL_SUCCESS); + + nixlProxySubmission submission{}; + submission.opcode = nixl_proxy_opcode_t::ATOMIC_ADD; + submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + submission.dst_offset = 56; + + nixlBackendProxySubmission prepared_submission; + ASSERT_EQ(registry_.prepareSubmission(submission, prepared_submission), NIXL_SUCCESS); + EXPECT_EQ(prepared_submission.size, sizeof(uint64_t)); + EXPECT_EQ(prepared_submission.remote.desc.addr, 0x2038u); + EXPECT_EQ(prepared_submission.remote.desc.len, sizeof(uint64_t)); + + submission.dst_offset = 60; + EXPECT_EQ(registry_.prepareSubmission(submission, prepared_submission), + NIXL_ERR_INVALID_PARAM); + } + + TEST_F(ProxyMemViewRegistryTest, MetadataKindMustMatchSubmissionRole) { + nixlMemViewH src_proxy = nullptr; + nixlMemViewH dst_proxy = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(src_proxy, makeRemoteMetadata(0x1000)), NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(dst_proxy, makeLocalMetadata(0x2000)), NIXL_SUCCESS); + + nixlProxySubmission submission{}; + submission.opcode = nixl_proxy_opcode_t::PUT; + submission.src_proxy_memview_id = reinterpret_cast(src_proxy); + submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + submission.size = 16; + + nixlBackendProxySubmission prepared_submission; + EXPECT_EQ(registry_.prepareSubmission(submission, prepared_submission), + NIXL_ERR_INVALID_PARAM); + } + + TEST_F(ProxyMemViewRegistryTest, RetiredEntriesStopFutureDispatchButKeepOtherEntriesUsable) { + nixlMemViewH src_proxy = nullptr; + nixlMemViewH dst_proxy = nullptr; + nixlMemViewH other_proxy = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(30), &other_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(src_proxy, makeLocalMetadata(0x1000)), NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(dst_proxy, makeRemoteMetadata(0x2000)), NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(other_proxy, makeRemoteMetadata(0x3000)), NIXL_SUCCESS); + + ASSERT_EQ(registry_.unregisterProxyMemView(dst_proxy), NIXL_SUCCESS); + EXPECT_EQ(registry_.unregisterProxyMemView(dst_proxy), NIXL_SUCCESS); + + nixlProxySubmission retired_submission{}; + retired_submission.opcode = nixl_proxy_opcode_t::PUT; + retired_submission.src_proxy_memview_id = reinterpret_cast(src_proxy); + retired_submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + retired_submission.size = 8; + + nixlBackendProxySubmission prepared_submission; + EXPECT_EQ(registry_.prepareSubmission(retired_submission, prepared_submission), + NIXL_ERR_NOT_FOUND); + + nixlProxySubmission live_submission{}; + live_submission.opcode = nixl_proxy_opcode_t::PUT; + live_submission.src_proxy_memview_id = reinterpret_cast(src_proxy); + live_submission.dst_proxy_memview_id = reinterpret_cast(other_proxy); + live_submission.size = 8; + + EXPECT_EQ(registry_.prepareSubmission(live_submission, prepared_submission), NIXL_SUCCESS); + } + + TEST_F(ProxyMemViewRegistryTest, ClearRetiresExistingEntriesAndPreservesFreshIds) { + nixlMemViewH old_proxy = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &old_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(old_proxy, makeLocalMetadata(0x1000)), NIXL_SUCCESS); + + registry_.clear(); + + nixlMemViewH resolved = nullptr; + EXPECT_FALSE(registry_.resolveProxyMemView(old_proxy, resolved)); + + nixlMemViewH new_proxy = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &new_proxy), NIXL_SUCCESS); + EXPECT_NE(old_proxy, new_proxy); + EXPECT_TRUE(registry_.resolveProxyMemView(new_proxy, resolved)); + EXPECT_EQ(resolved, makeFakeBackendHandle(20)); + } + + TEST_F(ProxyMemViewRegistryTest, StoreMetadataRejectsRetiredEntries) { + nixlMemViewH proxy_handle = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &proxy_handle), NIXL_SUCCESS); + ASSERT_EQ(registry_.unregisterProxyMemView(proxy_handle), NIXL_SUCCESS); + EXPECT_EQ(registry_.storeMetadata(proxy_handle, makeLocalMetadata(0x1000)), + NIXL_ERR_NOT_FOUND); + } + +} // namespace proxy_memview_registry +} // namespace gtest diff --git a/test/gtest/unit/proxy_runtime/meson.build b/test/gtest/unit/proxy_runtime/meson.build new file mode 100644 index 0000000000..4f0d39ae53 --- /dev/null +++ b/test/gtest/unit/proxy_runtime/meson.build @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +proxy_runtime_test_deps = [] +if cuda_dep.found() + proxy_runtime_test_deps += [cuda_dep] +endif + +proxy_runtime_unit_test_dep = declare_dependency( + sources: [ + 'proxy_runtime.cpp', + ], + include_directories: [nixl_inc_dirs, gtest_inc_dirs], + dependencies: proxy_runtime_test_deps, +) diff --git a/test/gtest/unit/proxy_runtime/proxy_runtime.cpp b/test/gtest/unit/proxy_runtime/proxy_runtime.cpp new file mode 100644 index 0000000000..4442b14ce7 --- /dev/null +++ b/test/gtest/unit/proxy_runtime/proxy_runtime.cpp @@ -0,0 +1,418 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "device_proxy/backend_adapter.h" +#include "device_proxy/proxy_runtime.h" + +namespace gtest { +namespace proxy_runtime { + +class DummyBackendMD : public nixlBackendMD { + public: + DummyBackendMD() : nixlBackendMD(false) {} +}; + +class StubBackend : public nixlDeviceProxyBackendAdapter { + public: + nixl_status_t + init(uint32_t worker_count, uint32_t channel_count) override { + init_called_ = true; + init_worker_count_ = worker_count; + init_channel_count_ = channel_count; + return init_rc_; + } + + nixl_status_t + loadRemoteConnInfo(const std::string &, const nixl_blob_t &) override { + return NIXL_SUCCESS; + } + + nixl_status_t + submit(const nixlBackendProxySubmission &submission, uint64_t &request_token) override { + { + std::lock_guard lock(submit_mutex_); + submissions_.push_back(submission); + } + request_token = ++next_request_token_; + return NIXL_SUCCESS; + } + + nixl_status_t + checkCompletion(uint64_t) override { + return NIXL_SUCCESS; + } + + nixl_status_t + progress() override { + ++progress_calls_; + return NIXL_SUCCESS; + } + + nixl_status_t + shutdown() override { + return NIXL_SUCCESS; + } + + bool init_called_ = false; + uint32_t init_worker_count_ = 0; + uint32_t init_channel_count_ = 0; + nixl_status_t init_rc_ = NIXL_SUCCESS; + std::atomic progress_calls_{0}; + mutable std::mutex submit_mutex_; + std::vector submissions_; + uint64_t next_request_token_ = 0; +}; + +class ProxyRuntimeTest : public testing::Test { + protected: + nixl_status_t + initRuntime(uint32_t channel_count, + uint32_t worker_count, + nixl_status_t init_rc = NIXL_SUCCESS) { + auto backend = std::make_unique(); + backend_ = backend.get(); + backend_->init_rc_ = init_rc; + return runtime_.init(std::move(backend), channel_count, worker_count); + } + + void TearDown() override { + runtime_.shutdown(); + } + + StubBackend *backend_ = nullptr; + nixlProxyRuntime runtime_; +}; + +TEST_F(ProxyRuntimeTest, InitCallsBackendInit) { + ASSERT_EQ(initRuntime(4, 2), NIXL_SUCCESS); + EXPECT_TRUE(backend_->init_called_); + EXPECT_EQ(backend_->init_worker_count_, 2u); + EXPECT_EQ(backend_->init_channel_count_, 4u); +} + +TEST_F(ProxyRuntimeTest, InitRejectsNullBackend) { + EXPECT_EQ(runtime_.init(nullptr, 4, 2), NIXL_ERR_INVALID_PARAM); +} + +TEST_F(ProxyRuntimeTest, InitRejectsZeroChannels) { + EXPECT_EQ(initRuntime(0, 2), NIXL_ERR_INVALID_PARAM); +} + +TEST_F(ProxyRuntimeTest, InitRejectsZeroWorkers) { + EXPECT_EQ(initRuntime(4, 0), NIXL_ERR_INVALID_PARAM); +} + +TEST_F(ProxyRuntimeTest, InitPropagatesBackendFailure) { + EXPECT_EQ(initRuntime(4, 2, NIXL_ERR_BACKEND), NIXL_ERR_BACKEND); +} + +TEST_F(ProxyRuntimeTest, InitSetsChannelCount) { + ASSERT_EQ(initRuntime(4, 2), NIXL_SUCCESS); + EXPECT_EQ(runtime_.channelCount(), 4u); +} + +TEST_F(ProxyRuntimeTest, DeviceChannelViewsPopulated) { + ASSERT_EQ(initRuntime(3, 1), NIXL_SUCCESS); + const nixlProxyChannelView *views = runtime_.deviceChannelViews(); + ASSERT_NE(views, nullptr); + for (uint32_t i = 0; i < 3; ++i) { + EXPECT_EQ(views[i].channel_id, i); + EXPECT_NE(views[i].work_ring, nullptr); + EXPECT_NE(views[i].work_ring->producer_idx, nullptr); + EXPECT_NE(views[i].work_ring->consumer_idx, nullptr); + EXPECT_NE(views[i].completion_slot, nullptr); + EXPECT_EQ(views[i].work_ring->depth, kDefaultProxyRingDepth); + } +} + +TEST_F(ProxyRuntimeTest, WorkRingIndicesStartAtZero) { + ASSERT_EQ(initRuntime(2, 1), NIXL_SUCCESS); + const nixlProxyChannelView *views = runtime_.deviceChannelViews(); + for (uint32_t i = 0; i < 2; ++i) { + uint32_t producer = 0; + uint32_t consumer = 0; + ASSERT_EQ(cudaMemcpy(&producer, + views[i].work_ring->producer_idx, + sizeof(producer), + cudaMemcpyDeviceToHost), + cudaSuccess); + ASSERT_EQ(cudaMemcpy(&consumer, + views[i].work_ring->consumer_idx, + sizeof(consumer), + cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_EQ(producer, 0u); + EXPECT_EQ(consumer, 0u); + } +} + +TEST_F(ProxyRuntimeTest, CompletionSlotsInitialized) { + ASSERT_EQ(initRuntime(2, 1), NIXL_SUCCESS); + const nixlProxyChannelView *views = runtime_.deviceChannelViews(); + for (uint32_t i = 0; i < 2; ++i) { + nixlProxyCompletionSlot slot{}; + ASSERT_EQ(cudaMemcpy(&slot, + views[i].completion_slot, + sizeof(nixlProxyCompletionSlot), + cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_EQ(slot.completed_idx, 0u); + EXPECT_EQ(slot.next_status, NIXL_IN_PROG); + } +} + +TEST_F(ProxyRuntimeTest, WorkerCountClampedToChannels) { + ASSERT_EQ(initRuntime(2, 8), NIXL_SUCCESS); + EXPECT_EQ(runtime_.channelCount(), 2u); + EXPECT_EQ(backend_->init_worker_count_, 2u); + EXPECT_EQ(backend_->init_channel_count_, 2u); +} + +TEST_F(ProxyRuntimeTest, DeviceContextPopulated) { + ASSERT_EQ(initRuntime(3, 1), NIXL_SUCCESS); + auto *ctx = runtime_.deviceContext(); + ASSERT_NE(ctx, nullptr); + EXPECT_EQ(ctx->num_channels, 3u); + EXPECT_EQ(ctx->channels, runtime_.deviceChannelViews()); + EXPECT_NE(ctx->shutdown_word, nullptr); +} + +TEST_F(ProxyRuntimeTest, DeviceContextNullAfterShutdown) { + ASSERT_EQ(initRuntime(2, 1), NIXL_SUCCESS); + ASSERT_NE(runtime_.deviceContext(), nullptr); + ASSERT_EQ(runtime_.shutdown(), NIXL_SUCCESS); + EXPECT_EQ(runtime_.deviceContext(), nullptr); +} + +TEST_F(ProxyRuntimeTest, StartWorkersAndShutdown) { + ASSERT_EQ(initRuntime(2, 2), NIXL_SUCCESS); + ASSERT_EQ(runtime_.startWorkers(), NIXL_SUCCESS); + + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + + ASSERT_EQ(runtime_.shutdown(), NIXL_SUCCESS); +} + +TEST_F(ProxyRuntimeTest, RepeatedStartWorkersIsRejected) { + ASSERT_EQ(initRuntime(2, 2), NIXL_SUCCESS); + ASSERT_EQ(runtime_.startWorkers(), NIXL_SUCCESS); + + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + + EXPECT_EQ(runtime_.startWorkers(), NIXL_ERR_INVALID_PARAM); + + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + + ASSERT_EQ(runtime_.shutdown(), NIXL_SUCCESS); +} + +TEST_F(ProxyRuntimeTest, ShutdownWithoutStartIsHarmless) { + ASSERT_EQ(initRuntime(2, 1), NIXL_SUCCESS); + EXPECT_EQ(runtime_.shutdown(), NIXL_SUCCESS); +} + +TEST_F(ProxyRuntimeTest, ShutdownBeforeInitIsHarmless) { + EXPECT_EQ(runtime_.shutdown(), NIXL_SUCCESS); +} + +TEST_F(ProxyRuntimeTest, DoubleShutdownIsHarmless) { + ASSERT_EQ(initRuntime(2, 1), NIXL_SUCCESS); + ASSERT_EQ(runtime_.startWorkers(), NIXL_SUCCESS); + EXPECT_EQ(runtime_.shutdown(), NIXL_SUCCESS); + EXPECT_EQ(runtime_.shutdown(), NIXL_SUCCESS); +} + +TEST_F(ProxyRuntimeTest, InitAfterShutdownWorks) { + ASSERT_EQ(initRuntime(2, 1), NIXL_SUCCESS); + ASSERT_EQ(runtime_.startWorkers(), NIXL_SUCCESS); + ASSERT_EQ(runtime_.shutdown(), NIXL_SUCCESS); + + ASSERT_EQ(initRuntime(4, 2), NIXL_SUCCESS); + EXPECT_EQ(runtime_.channelCount(), 4u); + ASSERT_EQ(runtime_.startWorkers(), NIXL_SUCCESS); + EXPECT_EQ(runtime_.shutdown(), NIXL_SUCCESS); +} + +TEST_F(ProxyRuntimeTest, SingleChannelSingleWorker) { + ASSERT_EQ(initRuntime(1, 1), NIXL_SUCCESS); + ASSERT_EQ(runtime_.startWorkers(), NIXL_SUCCESS); + + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + + ASSERT_EQ(runtime_.shutdown(), NIXL_SUCCESS); + EXPECT_EQ(runtime_.channelCount(), 0u); +} + +TEST_F(ProxyRuntimeTest, ManyChannelsManyWorkers) { + ASSERT_EQ(initRuntime(16, 4), NIXL_SUCCESS); + ASSERT_EQ(runtime_.startWorkers(), NIXL_SUCCESS); + + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + + ASSERT_EQ(runtime_.shutdown(), NIXL_SUCCESS); +} + +TEST_F(ProxyRuntimeTest, PrepMemViewProducesReadyEntries) { + DummyBackendMD local_md; + DummyBackendMD remote_md; + const auto local_backend = reinterpret_cast(uintptr_t{0x10}); + const auto remote_backend = reinterpret_cast(uintptr_t{0x20}); + + nixl_meta_dlist_t local_dlist(DRAM_SEG); + local_dlist.addDesc(nixlMetaDesc(0x1000, 64, 0, &local_md)); + + nixl_remote_meta_dlist_t remote_dlist(DRAM_SEG); + nixlRemoteMetaDesc remote_desc("peer"); + remote_desc.addr = 0x2000; + remote_desc.len = 64; + remote_desc.devId = 0; + remote_desc.metadataP = &remote_md; + remote_dlist.addDesc(remote_desc); + + nixlMemViewH src_proxy = nullptr; + nixlMemViewH dst_proxy = nullptr; + ASSERT_EQ(runtime_.prepMemView(local_backend, local_dlist, &src_proxy), + NIXL_SUCCESS); + ASSERT_EQ(runtime_.prepMemView(remote_backend, remote_dlist, &dst_proxy), + NIXL_SUCCESS); + + nixlMemViewH resolved = nullptr; + EXPECT_TRUE(runtime_.resolveProxyMemView(src_proxy, resolved)); + EXPECT_EQ(resolved, local_backend); + EXPECT_TRUE(runtime_.resolveProxyMemView(dst_proxy, resolved)); + EXPECT_EQ(resolved, remote_backend); + + nixlProxySubmission submission{}; + submission.opcode = nixl_proxy_opcode_t::PUT; + submission.src_proxy_memview_id = reinterpret_cast(src_proxy); + submission.src_offset = 4; + submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + submission.dst_offset = 8; + submission.size = 32; + + nixlBackendProxySubmission prepared_submission; + ASSERT_EQ(runtime_.memviewRegistry().prepareSubmission(submission, prepared_submission), + NIXL_SUCCESS); + EXPECT_EQ(prepared_submission.local.desc.addr, 0x1004u); + EXPECT_EQ(prepared_submission.local.desc.len, 32u); + EXPECT_EQ(prepared_submission.local.desc.metadataP, &local_md); + EXPECT_EQ(prepared_submission.remote.desc.addr, 0x2008u); + EXPECT_EQ(prepared_submission.remote.desc.len, 32u); + EXPECT_EQ(prepared_submission.remote.desc.metadataP, &remote_md); + EXPECT_EQ(prepared_submission.remote_agent, "peer"); +} + +TEST_F(ProxyRuntimeTest, PrepMemViewRejectsNullOutput) { + DummyBackendMD local_md; + nixl_meta_dlist_t local_dlist(DRAM_SEG); + local_dlist.addDesc(nixlMetaDesc(0x1000, 64, 0, &local_md)); + + EXPECT_EQ(runtime_.prepMemView(local_dlist, nullptr), + NIXL_ERR_INVALID_PARAM); +} + +TEST_F(ProxyRuntimeTest, WorkerSubmitsPreparedTransportDescriptors) { + DummyBackendMD local_md; + DummyBackendMD remote_md; + + ASSERT_EQ(initRuntime(1, 1), NIXL_SUCCESS); + + nixlMemViewH src_proxy = nullptr; + nixlMemViewH dst_proxy = nullptr; + ASSERT_EQ(runtime_.registerProxyMemView(reinterpret_cast(uintptr_t{0x10}), + &src_proxy), + NIXL_SUCCESS); + ASSERT_EQ(runtime_.registerProxyMemView(reinterpret_cast(uintptr_t{0x20}), + &dst_proxy), + NIXL_SUCCESS); + + nixl_meta_dlist_t local_dlist(DRAM_SEG); + local_dlist.addDesc(nixlMetaDesc(0x1000, 64, 0, &local_md)); + ASSERT_EQ(runtime_.storeMetadata(src_proxy, local_dlist), NIXL_SUCCESS); + + nixl_remote_meta_dlist_t remote_dlist(DRAM_SEG); + nixlRemoteMetaDesc remote_desc("peer"); + remote_desc.addr = 0x2000; + remote_desc.len = 64; + remote_desc.devId = 0; + remote_desc.metadataP = &remote_md; + remote_dlist.addDesc(remote_desc); + ASSERT_EQ(runtime_.storeMetadata(dst_proxy, remote_dlist), NIXL_SUCCESS); + + ASSERT_EQ(runtime_.startWorkers(), NIXL_SUCCESS); + + nixlProxySubmission submission{}; + submission.op_idx = 11; + submission.opcode = nixl_proxy_opcode_t::PUT; + submission.channel_id = 0; + submission.src_proxy_memview_id = reinterpret_cast(src_proxy); + submission.src_offset = 4; + submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + submission.dst_offset = 8; + submission.size = 32; + + auto *ring = runtime_.deviceChannelViews()[0].work_ring; + ring->records[0] = submission; + __atomic_store_n(&ring->records[0].ready_flag, 1u, __ATOMIC_RELEASE); + + const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(250); + while (std::chrono::steady_clock::now() < deadline) { + { + std::lock_guard lock(backend_->submit_mutex_); + if (!backend_->submissions_.empty()) { + break; + } + } + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + + std::vector submissions; + { + std::lock_guard lock(backend_->submit_mutex_); + submissions = backend_->submissions_; + } + + ASSERT_EQ(runtime_.shutdown(), NIXL_SUCCESS); + + ASSERT_EQ(submissions.size(), 1u); + const auto &prepared = submissions.front(); + EXPECT_EQ(prepared.op_idx, 11u); + EXPECT_EQ(prepared.channel_id, 0u); + EXPECT_EQ(prepared.local.mem_type, DRAM_SEG); + EXPECT_EQ(prepared.local.desc.addr, 0x1004u); + EXPECT_EQ(prepared.local.desc.len, 32u); + EXPECT_EQ(prepared.local.desc.metadataP, &local_md); + EXPECT_EQ(prepared.remote.mem_type, DRAM_SEG); + EXPECT_EQ(prepared.remote.desc.addr, 0x2008u); + EXPECT_EQ(prepared.remote.desc.len, 32u); + EXPECT_EQ(prepared.remote.desc.metadataP, &remote_md); + EXPECT_EQ(prepared.remote_agent, "peer"); +} + +} // namespace proxy_runtime +} // namespace gtest From 8b2940d0ffc16162ee4334403e9cfbc0e7817dcd Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Wed, 29 Apr 2026 18:15:23 +0300 Subject: [PATCH 04/18] plugins/ucx: wire UCX engine as a proxy backend provider Make nixlUcxEngine inherit from DeviceProxyBackendProvider and implement createDeviceProxyBackendAdapter(), which returns a nixlUcxProxyBackend. nixlUcxProxyBackend (src/plugins/ucx/device_proxy/) implements the DeviceProxyBackendAdapter interface: it translates PreparedProxyTransferDesc objects into UCX postXfer calls, using the engine's existing UCX endpoint and memory-registration infrastructure. The UCX plugin meson.build is updated to compile ucx_proxy_backend.cpp into the plugin DSO. Signed-off-by: Tomer Davidor --- .../ucx/device_proxy/ucx_proxy_backend.cpp | 143 ++++++++++++++++++ .../ucx/device_proxy/ucx_proxy_backend.h | 68 +++++++++ src/plugins/ucx/meson.build | 8 +- src/plugins/ucx/ucx_backend.cpp | 9 ++ src/plugins/ucx/ucx_backend.h | 5 + 5 files changed, 231 insertions(+), 2 deletions(-) create mode 100644 src/plugins/ucx/device_proxy/ucx_proxy_backend.cpp create mode 100644 src/plugins/ucx/device_proxy/ucx_proxy_backend.h diff --git a/src/plugins/ucx/device_proxy/ucx_proxy_backend.cpp b/src/plugins/ucx/device_proxy/ucx_proxy_backend.cpp new file mode 100644 index 0000000000..b001427bd6 --- /dev/null +++ b/src/plugins/ucx/device_proxy/ucx_proxy_backend.cpp @@ -0,0 +1,143 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "ucx_proxy_backend.h" +#include "../ucx_backend.h" +#include "nixl_log.h" +#include "nixl_types.h" + +namespace { +constexpr uint64_t kInvalidToken = 0; +} + +nixl_status_t +nixlUcxProxyBackendAdapter::submit(const nixlBackendProxySubmission &submission, + uint64_t &request_token) { + request_token = kInvalidToken; + if (engine_ == nullptr) { + return NIXL_ERR_INVALID_PARAM; + } + + switch (submission.opcode) { + case nixl_proxy_opcode_t::PUT: + return submitPut(submission, request_token); + case nixl_proxy_opcode_t::ATOMIC_ADD: + return submitAtomicAdd(submission, request_token); + default: + return NIXL_ERR_NOT_SUPPORTED; + } +} + +nixl_status_t +nixlUcxProxyBackendAdapter::submitPut(const nixlBackendProxySubmission &submission, + uint64_t &request_token) { + nixl_meta_dlist_t local_list(submission.local.mem_type); + local_list.addDesc(submission.local.desc); + + nixl_meta_dlist_t remote_list(submission.remote.mem_type); + remote_list.addDesc(submission.remote.desc); + + nixlBackendReqH *handle = nullptr; + nixl_status_t status = engine_->prepXfer( + NIXL_WRITE, local_list, remote_list, submission.remote_agent, handle); + if (status != NIXL_SUCCESS) { + NIXL_DEBUG << "nixlUcxProxyBackendAdapter::submitPut: prepXfer failed status=" << status; + return status; + } + + status = engine_->postXfer( + NIXL_WRITE, local_list, remote_list, submission.remote_agent, handle); + if (status != NIXL_SUCCESS && status != NIXL_IN_PROG) { + NIXL_DEBUG << "nixlUcxProxyBackendAdapter::submitPut: postXfer failed status=" << status; + engine_->releaseReqH(handle); + return status; + } + + request_token = trackRequest(handle); + NIXL_DEBUG << "nixlUcxProxyBackendAdapter::submitPut: posted RDMA write" + << " src_addr=0x" << std::hex + << submission.local.desc.addr << std::dec + << " dst_addr=0x" << std::hex + << submission.remote.desc.addr << std::dec + << " size=" << submission.size + << " remote_agent='" << submission.remote_agent << "'" + << " token=" << request_token; + return NIXL_SUCCESS; +} + +nixl_status_t +nixlUcxProxyBackendAdapter::submitAtomicAdd(const nixlBackendProxySubmission &, + uint64_t &) { + return NIXL_ERR_NOT_SUPPORTED; +} + +nixl_status_t +nixlUcxProxyBackendAdapter::checkCompletion(uint64_t request_token) { + if (engine_ == nullptr) { + return NIXL_ERR_INVALID_PARAM; + } + + std::lock_guard lock(request_mutex_); + const auto it = tracked_requests_.find(request_token); + if (it == tracked_requests_.end()) { + return NIXL_ERR_NOT_FOUND; + } + + nixlBackendReqH *handle = it->second; + const nixl_status_t status = engine_->checkXfer(handle); + if (status == NIXL_IN_PROG) { + return NIXL_IN_PROG; + } + + NIXL_DEBUG << "nixlUcxProxyBackendAdapter::checkCompletion: token=" << request_token + << " status=" << status; + engine_->releaseReqH(handle); + tracked_requests_.erase(it); + return status; +} + +nixl_status_t +nixlUcxProxyBackendAdapter::progress() { + if (engine_ != nullptr && !progress_thread_enabled_) { + engine_->progress(); + } + + return NIXL_SUCCESS; +} + +nixl_status_t +nixlUcxProxyBackendAdapter::shutdown() { + NIXL_INFO << "nixlUcxProxyBackendAdapter::shutdown: releasing " + << tracked_requests_.size() << " tracked request(s)"; + { + std::lock_guard lock(request_mutex_); + if (engine_ != nullptr) { + for (auto &entry : tracked_requests_) { + engine_->releaseReqH(entry.second); + } + } + tracked_requests_.clear(); + } + return NIXL_SUCCESS; +} + +uint64_t +nixlUcxProxyBackendAdapter::trackRequest(nixlBackendReqH *handle) { + std::lock_guard lock(request_mutex_); + const uint64_t token = next_request_token_++; + tracked_requests_.emplace(token, handle); + return token; +} diff --git a/src/plugins/ucx/device_proxy/ucx_proxy_backend.h b/src/plugins/ucx/device_proxy/ucx_proxy_backend.h new file mode 100644 index 0000000000..91b7e0d483 --- /dev/null +++ b/src/plugins/ucx/device_proxy/ucx_proxy_backend.h @@ -0,0 +1,68 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef NIXL_SRC_PLUGINS_UCX_DEVICE_PROXY_UCX_PROXY_BACKEND_H +#define NIXL_SRC_PLUGINS_UCX_DEVICE_PROXY_UCX_PROXY_BACKEND_H + +#include +#include +#include +#include + +#include "backend/backend_aux.h" +#include "../../../core/device_proxy/backend_adapter.h" + +class nixlUcxEngine; + +class nixlUcxProxyBackendAdapter : public nixlDeviceProxyBackendAdapter { +public: + explicit nixlUcxProxyBackendAdapter(nixlUcxEngine *engine = nullptr, + bool progress_thread_enabled = false) noexcept + : engine_(engine), + progress_thread_enabled_(progress_thread_enabled) {} + + ~nixlUcxProxyBackendAdapter() override = default; + + nixl_status_t + submit(const nixlBackendProxySubmission &submission, uint64_t &request_token) override; + + nixl_status_t + checkCompletion(uint64_t request_token) override; + + nixl_status_t + progress() override; + + nixl_status_t + shutdown() override; + +private: + nixl_status_t + submitPut(const nixlBackendProxySubmission &submission, uint64_t &request_token); + + nixl_status_t + submitAtomicAdd(const nixlBackendProxySubmission &submission, uint64_t &request_token); + + uint64_t + trackRequest(nixlBackendReqH *handle); + + nixlUcxEngine *engine_ = nullptr; + bool progress_thread_enabled_ = false; + std::mutex request_mutex_; + std::unordered_map tracked_requests_; + uint64_t next_request_token_ = 1; +}; + +#endif // NIXL_SRC_PLUGINS_UCX_DEVICE_PROXY_UCX_PROXY_BACKEND_H diff --git a/src/plugins/ucx/meson.build b/src/plugins/ucx/meson.build index 6e15a1720d..2d323a5f0f 100644 --- a/src/plugins/ucx/meson.build +++ b/src/plugins/ucx/meson.build @@ -24,7 +24,8 @@ ucx_backend_sources = ['config.cpp', 'ucx_backend.cpp', 'ucx_enums.cpp', 'ucx_plugin.cpp', - 'ucx_utils.cpp'] + 'ucx_utils.cpp', + 'device_proxy/ucx_proxy_backend.cpp'] ucx_backend_dependencies = [asio_dep, nixl_common_dep, @@ -64,4 +65,7 @@ else endif endif -ucx_backend_interface = declare_dependency(link_with: ucx_backend_lib, include_directories: ucx_backend_includes) +ucx_backend_interface = declare_dependency( + link_with: ucx_backend_lib, + include_directories: ucx_backend_includes, + dependencies: [ucx_dep]) diff --git a/src/plugins/ucx/ucx_backend.cpp b/src/plugins/ucx/ucx_backend.cpp index 8f8835c32a..6f10791d99 100644 --- a/src/plugins/ucx/ucx_backend.cpp +++ b/src/plugins/ucx/ucx_backend.cpp @@ -16,6 +16,7 @@ */ #include "ucx_backend.h" +#include "device_proxy/ucx_proxy_backend.h" #include "common/nixl_log.h" #include "serdes/serdes.h" #include "common/backend.h" @@ -870,6 +871,14 @@ nixl_status_t nixlUcxEngine::checkConn(const std::string &remote_agent) { return remoteConnMap.count(remote_agent) ? NIXL_SUCCESS : NIXL_ERR_NOT_FOUND; } +nixl_status_t +nixlUcxEngine::createDeviceProxyBackendAdapter( + const nixlBackendInitParams &init_params, + std::unique_ptr &adapter) { + adapter = std::make_unique(this, init_params.enableProgTh); + return NIXL_SUCCESS; +} + nixl_status_t nixlUcxEngine::getConnInfo(std::string &str) const { str = workerAddr; return NIXL_SUCCESS; diff --git a/src/plugins/ucx/ucx_backend.h b/src/plugins/ucx/ucx_backend.h index 51b4b4fce1..a2f94598cd 100644 --- a/src/plugins/ucx/ucx_backend.h +++ b/src/plugins/ucx/ucx_backend.h @@ -200,6 +200,11 @@ class nixlUcxEngine : public nixlBackendEngine { nixl_status_t checkConn(const std::string &remote_agent); + nixl_status_t + createDeviceProxyBackendAdapter( + const nixlBackendInitParams &init_params, + std::unique_ptr &adapter) override; + nixl_status_t prepMemView(const nixl_remote_meta_dlist_t &, nixlMemViewH &, From e4f166d36b858fbfb0140893b82cfad7115ac449 Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Wed, 29 Apr 2026 18:21:58 +0300 Subject: [PATCH 05/18] core: expose proxy runtime through nixlAgent API Add three fields to nixlAgentConfig: enableDeviceProxy, proxyWorkerCount, and proxyChannelCount, all defaulting to disabled/1. Add nixlAgent::getProxyDeviceContext() to the public API, returning the ProxyDeviceContextData pointer that GPU kernels need to enqueue transfers. Wire proxy lifecycle into nixlAgentData: - createProxyRuntime(): called when a backend that implements DeviceProxyBackendProvider is registered and enableDeviceProxy is set. Initialises ProxyRuntime and starts worker threads. - shutdownProxyRuntime(): called from the agent destructor. nixl_listener.cpp: propagate remote connection info into the proxy runtime (proxyRuntime->loadRemoteConnInfo) so proxy workers can post to remote peers after metadata exchange. Signed-off-by: Tomer Davidor --- src/api/cpp/nixl.h | 9 +++ src/api/cpp/nixl_params.h | 24 ++++++++ src/core/agent_data.h | 31 +++++++++- src/core/nixl_agent.cpp | 120 +++++++++++++++++++++++++++++++++---- src/core/nixl_listener.cpp | 7 +++ 5 files changed, 178 insertions(+), 13 deletions(-) diff --git a/src/api/cpp/nixl.h b/src/api/cpp/nixl.h index 9d9080606e..364f5efaa5 100644 --- a/src/api/cpp/nixl.h +++ b/src/api/cpp/nixl.h @@ -561,6 +561,15 @@ class nixlAgent { checkRemoteMD (const std::string remote_name, const nixl_xfer_dlist_t &descs) const; + /** + * @brief Return the proxy runtime's device context data pointer, + * or nullptr when proxy mode is not active. + * + * The caller can pass this to nixlProxyPublishContext() from a CUDA TU + * to make the context visible to GPU kernels. + */ + void * + getProxyDeviceContext() const; }; #endif diff --git a/src/api/cpp/nixl_params.h b/src/api/cpp/nixl_params.h index 8de369308f..08e9c0f214 100644 --- a/src/api/cpp/nixl_params.h +++ b/src/api/cpp/nixl_params.h @@ -33,6 +33,9 @@ struct nixlAgentConfig { static constexpr nixl_thread_sync_t kDefaultSyncMode = nixl_thread_sync_t::NIXL_THREAD_SYNC_DEFAULT; static constexpr bool kDefaultCaptureTelemetry = false; + static constexpr bool kDefaultEnableDeviceProxy = false; + static constexpr uint32_t kDefaultProxyWorkerCount = 1; + static constexpr uint32_t kDefaultProxyChannelCount = 1; static constexpr uint64_t kDefaultPthrDelayUs = 0; static constexpr uint64_t kDefaultLthrDelayUs = 100000; static constexpr std::chrono::microseconds kDefaultEtcdWatchTimeout = @@ -48,6 +51,21 @@ struct nixlAgentConfig { nixl_thread_sync_t syncMode = kDefaultSyncMode; /** @var Capture telemetry info regardless of environment variables*/ bool captureTelemetry = kDefaultCaptureTelemetry; + /** @var Enable the device proxy orchestration skeleton. */ + bool enableDeviceProxy = kDefaultEnableDeviceProxy; + /** + * @var Desired number of proxy workers per proxy runtime. + * The UCX proxy backend currently shares a nixlUcxEngine across proxy + * workers; keep this at 1 unless the backend has been validated for + * concurrent postXfer calls. + */ + uint32_t proxyWorkerCount = kDefaultProxyWorkerCount; + /** + * @var Desired number of proxy channels per proxy runtime. + * Channels may exceed workers; this is the supported way to expose + * multiple GPU submission queues while using one proxy worker. + */ + uint32_t proxyChannelCount = kDefaultProxyChannelCount; /** * @var Progress thread event waiting timeout. @@ -94,12 +112,18 @@ struct nixlAgentConfig { uint64_t pthr_delay_us = kDefaultPthrDelayUs, uint64_t lthr_delay_us = kDefaultLthrDelayUs, bool capture_telemetry = kDefaultCaptureTelemetry, + bool enable_device_proxy = kDefaultEnableDeviceProxy, + uint32_t proxy_worker_count = kDefaultProxyWorkerCount, + uint32_t proxy_channel_count = kDefaultProxyChannelCount, std::chrono::microseconds etcd_watch_timeout = kDefaultEtcdWatchTimeout) noexcept : useProgThread(use_prog_thread), useListenThread(use_listen_thread), listenPort(port), syncMode(sync_mode), captureTelemetry(capture_telemetry), + enableDeviceProxy(enable_device_proxy), + proxyWorkerCount(proxy_worker_count), + proxyChannelCount(proxy_channel_count), pthrDelay(pthr_delay_us), lthrDelay(lthr_delay_us), etcdWatchTimeout(etcd_watch_timeout) {} diff --git a/src/core/agent_data.h b/src/core/agent_data.h index d58dfabda3..945e3a39a4 100644 --- a/src/core/agent_data.h +++ b/src/core/agent_data.h @@ -21,9 +21,11 @@ #include "telemetry.h" #include "tracing/trace.h" #include "stream/metadata_stream.h" +#include "device_proxy/proxy_runtime.h" #include "sync.h" #include +#include #include #if HAVE_ETCD @@ -36,8 +38,21 @@ class SyncClient; #define NIXL_ETCD_NAMESPACE_DEFAULT "/nixl/agents/" #endif // HAVE_ETCD +class nixlBackendEngine; +class nixlBackendInitParams; +class nixlProxyRuntime; + using backend_list_t = std::vector; +enum class ProxyOrchestrationPhase : uint8_t { + Disabled = 0, + Registered, + MetadataPending, + ReadyToBootstrap, + Active, + ShuttingDown, +}; + //Internal typedef to define metadata communication request types //To be extended with ETCD operations enum nixl_comm_t { @@ -79,7 +94,7 @@ class nixlAgentData { backend_list_t notifEngines; std::array memToBackend; - // Bookkeeping from memory view handles to backend engines + // Bookkeeping from public memory view handles to backend engines std::unordered_map mvhToEngine; std::unordered_map> @@ -100,6 +115,8 @@ class nixlAgentData { std::unordered_map> backendHandles_; std::unordered_map connMd_; backend_map_t backendEngines_; + std::unique_ptr proxyRuntime; + nixlBackendEngine *proxyTransportEngine = nullptr; std::unordered_map remoteSections_; std::unique_ptr telemetry_; // Composite tracer (fans out to every enabled backend); null when no @@ -125,6 +142,16 @@ class nixlAgentData { getBackends(const nixl_opt_args_t *opt_args, const nixlMemSection §ion, nixl_mem_t mem_type); + [[nodiscard]] bool + proxyModeEnabled() const; + [[nodiscard]] bool + hasProxyRuntime() const; + nixl_status_t + createProxyRuntime(nixlBackendEngine *engine, + const nixl_backend_t &backend, + const nixlBackendInitParams &init_params); + void + shutdownProxyRuntime(); void warnAboutEfaHardwareMismatch(); @@ -141,8 +168,6 @@ class nixlAgentData { friend class nixlAgent; }; -class nixlBackendEngine; - // This class hides away the nixlBackendEngine from user of the Agent API class nixlBackendH { private: diff --git a/src/core/nixl_agent.cpp b/src/core/nixl_agent.cpp index a08a9831a7..288581e282 100644 --- a/src/core/nixl_agent.cpp +++ b/src/core/nixl_agent.cpp @@ -18,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -239,6 +240,8 @@ nixlAgent::nixlAgent(const std::string &name, const nixlAgentConfig &cfg) : } nixlAgent::~nixlAgent() { + data->shutdownProxyRuntime(); + if (data->needsCommThread_) { data->agentShutdown = true; // commQueue is guarded by commLock (see enqueueCommWork/getCommWork); @@ -353,6 +356,75 @@ nixlAgentData::warnAboutEfaHardwareMismatch() { } } +bool +nixlAgentData::proxyModeEnabled() const { + return config_.enableDeviceProxy; +} + +bool +nixlAgentData::hasProxyRuntime() const { + return proxyRuntime != nullptr; +} + +nixl_status_t +nixlAgentData::createProxyRuntime(nixlBackendEngine *engine, + const nixl_backend_t &backend, + const nixlBackendInitParams &init_params) { + if (hasProxyRuntime()) { + return NIXL_SUCCESS; + } + + std::unique_ptr proxy_adapter; + nixl_status_t status = engine->createDeviceProxyBackendAdapter(init_params, proxy_adapter); + if (status != NIXL_SUCCESS) { + return status; + } + if (!proxy_adapter) { + return NIXL_ERR_BACKEND; + } + + proxyRuntime = std::make_unique(); + + status = proxyRuntime->init(std::move(proxy_adapter), + config_.proxyChannelCount, + config_.proxyWorkerCount, + config_.pthrDelay); + if (status != NIXL_SUCCESS) { + proxyRuntime.reset(); + return status; + } + + status = proxyRuntime->startWorkers(); + if (status != NIXL_SUCCESS) { + proxyRuntime->shutdown(); + proxyRuntime.reset(); + return status; + } + + proxyTransportEngine = engine; + NIXL_INFO << "Enabled device proxy runtime for backend '" << backend << "' with " + << config_.proxyWorkerCount << " worker(s) and " << config_.proxyChannelCount + << " channel(s)"; + return NIXL_SUCCESS; +} + +void +nixlAgentData::shutdownProxyRuntime() { + if (proxyRuntime) { + proxyRuntime->shutdown(); + proxyRuntime.reset(); + } + proxyTransportEngine = nullptr; +} + +void * +nixlAgent::getProxyDeviceContext() const { + if (data->proxyRuntime) { + return data->proxyRuntime->deviceContext(); + } + return nullptr; +} + nixl_status_t nixlAgent::createBackend(const nixl_backend_t &type, const nixl_b_params_t ¶ms, @@ -457,7 +529,8 @@ nixlAgent::createBackend(const nixl_backend_t &type, return ret; } } else { - NIXL_WARN << "Proxy runtime is enabled but backend '" << type << "' does not support it"; + NIXL_WARN << "Proxy runtime is enabled but backend '" << type + << "' does not support it"; } } @@ -1946,12 +2019,20 @@ nixlAgent::prepMemView(const nixl_remote_dlist_t &dlist, return NIXL_ERR_NOT_FOUND; } - const auto status = engine->prepMemView(remote_meta_dlist, mvh, &opt_args); - if (status == NIXL_SUCCESS) { - data->mvhToEngine.emplace(mvh, *engine); + if (data->hasProxyRuntime() && (data->proxyTransportEngine == engine)) { + const auto status = data->proxyRuntime->prepMemView(remote_meta_dlist, &mvh); + if (status != NIXL_SUCCESS) { + return status; + } + } else { + const auto status = engine->prepMemView(remote_meta_dlist, mvh, &opt_args); + if (status != NIXL_SUCCESS) { + return status; + } } - return status; + data->mvhToEngine.emplace(mvh, *engine); + return NIXL_SUCCESS; } nixl_status_t @@ -1988,12 +2069,20 @@ nixlAgent::prepMemView(const nixl_local_dlist_t &dlist, return NIXL_ERR_NOT_FOUND; } - const auto status = engine->prepMemView(meta_dlist, mvh, &opt_args); - if (status == NIXL_SUCCESS) { - data->mvhToEngine.emplace(mvh, *engine); + if (data->hasProxyRuntime() && (data->proxyTransportEngine == engine)) { + const auto status = data->proxyRuntime->prepMemView(meta_dlist, &mvh); + if (status != NIXL_SUCCESS) { + return status; + } + } else { + const auto status = engine->prepMemView(meta_dlist, mvh, &opt_args); + if (status != NIXL_SUCCESS) { + return status; + } } - return status; + data->mvhToEngine.emplace(mvh, *engine); + return NIXL_SUCCESS; } void @@ -2008,6 +2097,17 @@ nixlAgent::releaseMemView(nixlMemViewH mvh) const { return; } - it->second.releaseMemView(mvh); + nixlMemViewH backend_mvh = mvh; + if (data->hasProxyRuntime()) { + nixlMemViewH resolved = nullptr; + if (data->proxyRuntime->resolveProxyMemView(mvh, resolved)) { + backend_mvh = resolved; + data->proxyRuntime->unregisterProxyMemView(mvh); + } + } + + if (backend_mvh != nullptr) { + it->second.releaseMemView(backend_mvh); + } data->mvhToEngine.erase(it); } diff --git a/src/core/nixl_listener.cpp b/src/core/nixl_listener.cpp index 8a8efa22b9..c0d95bec0a 100644 --- a/src/core/nixl_listener.cpp +++ b/src/core/nixl_listener.cpp @@ -773,6 +773,13 @@ nixlAgentData::loadConnInfo(const std::string &remote_name, return ret; } + if (hasProxyRuntime() && (proxyTransportEngine == eng)) { + const nixl_status_t proxy_ret = proxyRuntime->loadRemoteConnInfo(remote_name, conn_info); + if ((proxy_ret != NIXL_SUCCESS) && (proxy_ret != NIXL_ERR_NOT_SUPPORTED)) { + return proxy_ret; + } + } + remoteBackends_[remote_name].emplace(backend, conn_info); return NIXL_SUCCESS; } From 839b224d4ad6248fc35feb7281cbe17842f6694a Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Wed, 29 Apr 2026 18:24:13 +0300 Subject: [PATCH 06/18] test: add proxy device write coverage Add proxy_write_test.cu: an end-to-end device-API test that stands up a ProxyRuntime with a stub backend, publishes the device context to the GPU via nixlProxyPublishContext(), enqueues PUT operations from a CUDA kernel, and polls for completion. Extend single_write_test.cu with conditional proxy-backend support: when compiled with -DNIXL_GPU_DEVICE_BACKEND_PROXY the test uses the proxy code path (nixlProxyPublishContext / nixlProxyClearContext) instead of the UCX device API path. Introduce the gtest_proxy executable in test/gtest/meson.build. It links nixl_gpu_proxy_lib directly (so Meson includes nixl_device_proxy.cu.o in the RDC device-link step) and is compiled with -DNIXL_GPU_DEVICE_BACKEND_PROXY. Gate it on cuda_dep.found() rather than the UCX GPU flag used by gtest. Signed-off-by: Tomer Davidor --- test/gtest/device_api/meson.build | 7 + test/gtest/device_api/proxy_write_test.cu | 980 +++++++++++++++++++++ test/gtest/device_api/single_write_test.cu | 50 +- test/gtest/meson.build | 32 +- 4 files changed, 1066 insertions(+), 3 deletions(-) create mode 100644 test/gtest/device_api/proxy_write_test.cu diff --git a/test/gtest/device_api/meson.build b/test/gtest/device_api/meson.build index 0fa19e0f37..1be6c91867 100644 --- a/test/gtest/device_api/meson.build +++ b/test/gtest/device_api/meson.build @@ -21,5 +21,12 @@ device_api_sources = [ 'utils.cu' ] +proxy_device_api_sources = [ + 'proxy_write_test.cu', + 'single_write_test.cu', + 'utils.cu', +] + # Export for parent meson.build device_api_test_sources = files(device_api_sources) +proxy_device_api_test_sources = files(proxy_device_api_sources) diff --git a/test/gtest/device_api/proxy_write_test.cu b/test/gtest/device_api/proxy_write_test.cu new file mode 100644 index 0000000000..863cf5152d --- /dev/null +++ b/test/gtest/device_api/proxy_write_test.cu @@ -0,0 +1,980 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Verifies that the proxy backend compiles, links, and that GPU kernels can +// reach the ProxyDeviceContext published by ProxyRuntime::startWorkers(). +// +// nixl_device.cuh resolves to proxy/nixl_device.cuh via the proxy include +// path supplied by the build system - no backend macro is needed. + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "device_proxy/proxy_runtime.h" +#include "device_proxy/backend_adapter.h" +#include "common.h" + +// --------------------------------------------------------------------------- +// Minimal stub backend — satisfies the pure-virtual interface without doing +// any real I/O. Sufficient for testing the runtime lifecycle and the GPU +// device-context path. +// --------------------------------------------------------------------------- +class StubProxyBackendAdapter : public nixlDeviceProxyBackendAdapter { +public: + nixl_status_t + init(uint32_t, uint32_t) override { return NIXL_SUCCESS; } + + nixl_status_t + loadRemoteConnInfo(const std::string &, const nixl_blob_t &) override + { + return NIXL_SUCCESS; + } + + nixl_status_t + submit(const nixlBackendProxySubmission &, uint64_t &token) override + { + token = 0; + return NIXL_SUCCESS; + } + + nixl_status_t + checkCompletion(uint64_t) override { return NIXL_SUCCESS; } + + nixl_status_t + progress() override { return NIXL_SUCCESS; } + + nixl_status_t + shutdown() override { return NIXL_SUCCESS; } +}; + +// --------------------------------------------------------------------------- +// Controllable stub — lets the test thread decide when each submission +// completes. submit() assigns unique monotonic tokens; checkCompletion() +// returns NIXL_IN_PROG until markComplete() is called for a token. +// --------------------------------------------------------------------------- +class ControllableStubAdapter : public nixlDeviceProxyBackendAdapter { +public: + nixl_status_t + init(uint32_t, uint32_t) override { return NIXL_SUCCESS; } + + nixl_status_t + loadRemoteConnInfo(const std::string &, const nixl_blob_t &) override + { + return NIXL_SUCCESS; + } + + nixl_status_t + submit(const nixlBackendProxySubmission &submission, uint64_t &token) override + { + std::lock_guard lk(mu_); + token = next_token_++; + pending_.insert(token); + token_channel_[token] = submission.channel_id; + return NIXL_SUCCESS; + } + + nixl_status_t + checkCompletion(uint64_t token) override + { + std::lock_guard lk(mu_); + auto it = completed_.find(token); + if (it != completed_.end()) { + nixl_status_t status = it->second; + completed_.erase(it); + return status; + } + return NIXL_IN_PROG; + } + + nixl_status_t + progress() override { return NIXL_SUCCESS; } + + nixl_status_t + shutdown() override { return NIXL_SUCCESS; } + + void + markComplete(uint64_t token) + { + markCompleteWithStatus(token, NIXL_SUCCESS); + } + + void + markCompleteWithStatus(uint64_t token, nixl_status_t status) + { + std::lock_guard lk(mu_); + pending_.erase(token); + completed_[token] = status; + } + + bool + hasPending() const + { + std::lock_guard lk(mu_); + return !pending_.empty(); + } + + size_t + pendingCount() const + { + std::lock_guard lk(mu_); + return pending_.size(); + } + + bool + hasPendingForChannel(uint32_t channel_id) const + { + std::lock_guard lk(mu_); + for (uint64_t token : pending_) { + auto it = token_channel_.find(token); + if (it != token_channel_.end() && it->second == channel_id) { + return true; + } + } + return false; + } + + bool + markFirstPendingForChannel(uint32_t channel_id, uint64_t *token = nullptr) + { + std::lock_guard lk(mu_); + for (uint64_t pending_token : pending_) { + auto it = token_channel_.find(pending_token); + if (it != token_channel_.end() && it->second == channel_id) { + pending_.erase(pending_token); + completed_[pending_token] = NIXL_SUCCESS; + if (token != nullptr) { + *token = pending_token; + } + return true; + } + } + return false; + } + +private: + mutable std::mutex mu_; + uint64_t next_token_ = 1; + std::set pending_; + std::map completed_; + std::map token_channel_; +}; + +// --------------------------------------------------------------------------- +// Error-returning stub — submit succeeds but checkCompletion always returns +// NIXL_ERR_BACKEND, simulating a backend failure. +// --------------------------------------------------------------------------- +class ErrorStubAdapter : public nixlDeviceProxyBackendAdapter { +public: + nixl_status_t + init(uint32_t, uint32_t) override { return NIXL_SUCCESS; } + + nixl_status_t + loadRemoteConnInfo(const std::string &, const nixl_blob_t &) override + { + return NIXL_SUCCESS; + } + + nixl_status_t + submit(const nixlBackendProxySubmission &, uint64_t &token) override + { + token = 0; + return NIXL_SUCCESS; + } + + nixl_status_t + checkCompletion(uint64_t) override { return NIXL_ERR_BACKEND; } + + nixl_status_t + progress() override { return NIXL_SUCCESS; } + + nixl_status_t + shutdown() override { return NIXL_SUCCESS; } +}; + +// --------------------------------------------------------------------------- +// Submit-error stub - submit() fails immediately and should be published back +// to the GPU without going through checkCompletion(). +// --------------------------------------------------------------------------- +class SubmitErrorStubAdapter : public nixlDeviceProxyBackendAdapter { +public: + nixl_status_t + init(uint32_t, uint32_t) override { return NIXL_SUCCESS; } + + nixl_status_t + loadRemoteConnInfo(const std::string &, const nixl_blob_t &) override + { + return NIXL_SUCCESS; + } + + nixl_status_t + submit(const nixlBackendProxySubmission &, uint64_t &) override + { + ++submit_calls_; + return NIXL_ERR_BACKEND; + } + + nixl_status_t + checkCompletion(uint64_t) override + { + ++check_completion_calls_; + return NIXL_SUCCESS; + } + + nixl_status_t + progress() override { return NIXL_SUCCESS; } + + nixl_status_t + shutdown() override { return NIXL_SUCCESS; } + + std::atomic submit_calls_{0}; + std::atomic check_completion_calls_{0}; +}; + +class DummyBackendMD : public nixlBackendMD { +public: + DummyBackendMD() : nixlBackendMD(false) {} +}; + +struct DummyProxyMemViews { + nixlMemViewH src = nullptr; + nixlMemViewH dst = nullptr; +}; + +static DummyProxyMemViews +registerDummyMemViews(nixlProxyRuntime &runtime); + +// --------------------------------------------------------------------------- +// Device kernels +// --------------------------------------------------------------------------- + +// Writes true if load_proxy_context() returns a non-null pointer. +__global__ void +proxyContextKernel(bool *out_has_ctx) +{ + *out_has_ctx = (load_proxy_context() != nullptr); +} + +// Calls nixlPut with zero-initialised operands and records the status. +__global__ void +proxyPutKernel(nixlMemViewH src_mvh, + nixlMemViewH dst_mvh, + nixl_status_t *out_status) +{ + nixlMemViewElem src{src_mvh, 0, 0}, dst{dst_mvh, 0, 0}; + *out_status = nixlPut(src, dst, /*size=*/0); +} + +static void +publishProxyContext(nixlProxyRuntime &runtime) +{ + bool *d_warmup = nullptr; + ASSERT_EQ(cudaMalloc(&d_warmup, sizeof(bool)), cudaSuccess); + proxyContextKernel<<<1, 1>>>(d_warmup); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + ASSERT_EQ(cudaFree(d_warmup), cudaSuccess); + + ASSERT_NE(runtime.deviceContext(), nullptr); + ASSERT_EQ(nixlProxyPublishContext(runtime.deviceContext()), cudaSuccess); +} + +static void +clearProxyContext() +{ + ASSERT_EQ(nixlProxyClearContext(), cudaSuccess); +} + +// --------------------------------------------------------------------------- +// Test fixture +// --------------------------------------------------------------------------- + +class ProxyDeviceApiTest : public ::testing::Test { +protected: + void + SetUp() override + { + if (!gtest::hasCudaGpu()) { + GTEST_SKIP() << "No CUDA-capable GPU, skipping proxy device API test."; + } + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + } + + template + T + deviceGet(T *d_ptr) + { + T val{}; + cudaMemcpy(&val, d_ptr, sizeof(T), cudaMemcpyDeviceToHost); + return val; + } + + template + T * + deviceAlloc() + { + T *ptr = nullptr; + EXPECT_EQ(cudaMalloc(&ptr, sizeof(T)), cudaSuccess); + EXPECT_EQ(cudaMemset(ptr, 0, sizeof(T)), cudaSuccess); + return ptr; + } + + template + bool + waitForCondition(Predicate predicate, + std::chrono::milliseconds timeout = std::chrono::milliseconds(500)) + { + const auto deadline = std::chrono::steady_clock::now() + timeout; + while (std::chrono::steady_clock::now() < deadline) { + if (predicate()) { + return true; + } + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + return predicate(); + } +}; + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +// After startWorkers() the GPU should see a non-null proxy context. +TEST_F(ProxyDeviceApiTest, ContextPublishedAfterStartWorkers) +{ + auto adapter = std::make_unique(); + nixlProxyRuntime runtime; + + ASSERT_EQ(runtime.init(std::move(adapter), /*channel_count=*/1, /*worker_count=*/1), + NIXL_SUCCESS); + ASSERT_EQ(runtime.startWorkers(), NIXL_SUCCESS); + publishProxyContext(runtime); + + bool *d_has_ctx = deviceAlloc(); + proxyContextKernel<<<1, 1>>>(d_has_ctx); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + + EXPECT_TRUE(deviceGet(d_has_ctx)); + cudaFree(d_has_ctx); + + clearProxyContext(); + ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); +} + +// After shutdown() the GPU should no longer see a proxy context. +TEST_F(ProxyDeviceApiTest, ContextClearedAfterShutdown) +{ + auto adapter = std::make_unique(); + nixlProxyRuntime runtime; + + ASSERT_EQ(runtime.init(std::move(adapter), /*channel_count=*/1, /*worker_count=*/1), + NIXL_SUCCESS); + ASSERT_EQ(runtime.startWorkers(), NIXL_SUCCESS); + publishProxyContext(runtime); + ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); + clearProxyContext(); + + bool *d_has_ctx = deviceAlloc(); + // Initialise to true so a no-op kernel would give a false pass. + bool init_val = true; + cudaMemcpy(d_has_ctx, &init_val, sizeof(bool), cudaMemcpyHostToDevice); + + proxyContextKernel<<<1, 1>>>(d_has_ctx); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + + EXPECT_FALSE(deviceGet(d_has_ctx)); + cudaFree(d_has_ctx); +} + +// nixlPut() via the proxy backend should report NIXL_IN_PROG once the +// submission is accepted into the proxy ring. +TEST_F(ProxyDeviceApiTest, PutReturnsInProgWhenEnqueued) +{ + auto adapter = std::make_unique(); + nixlProxyRuntime runtime; + + ASSERT_EQ(runtime.init(std::move(adapter), /*channel_count=*/1, /*worker_count=*/1), + NIXL_SUCCESS); + ASSERT_EQ(runtime.startWorkers(), NIXL_SUCCESS); + publishProxyContext(runtime); + const auto mvhs = registerDummyMemViews(runtime); + + nixl_status_t *d_status = deviceAlloc(); + proxyPutKernel<<<1, 1>>>(mvhs.src, mvhs.dst, d_status); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + + EXPECT_EQ(deviceGet(d_status), NIXL_IN_PROG); + cudaFree(d_status); + + clearProxyContext(); + ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); +} + +// --------------------------------------------------------------------------- +// Completion round-trip kernels +// +// Kernels that spin on pollXferStatus require valid proxy memview handles so +// the worker's dispatch succeeds and a completion is actually published. +// The test registers a dummy memview and passes the proxy handle here. +// --------------------------------------------------------------------------- + +// Enqueues a put and spins until pollXferStatus returns a final status. +__global__ void +proxyPutAndPollKernel(nixlMemViewH src_mvh, nixlMemViewH dst_mvh, + uint32_t channel_id, + nixl_status_t *out_put_status, + nixl_status_t *out_poll_status) +{ + nixlMemViewElem src{src_mvh, 0, 0}, dst{dst_mvh, 0, 0}; + nixlGpuXferStatusH xfer_status{}; + *out_put_status = nixlPut(src, dst, /*size=*/0, channel_id, + /*flags=*/0, &xfer_status); + + nixl_status_t poll; + do { + poll = nixlGpuGetXferStatus(xfer_status); + } while (poll == NIXL_IN_PROG); + *out_poll_status = poll; +} + +// Enqueues a put and immediately returns; saves xfer_status to device memory +// so the test thread can later launch a poll kernel. +__global__ void +proxyPutAsyncKernel(nixlMemViewH src_mvh, nixlMemViewH dst_mvh, + uint32_t channel_id, + nixl_status_t *out_put_status, + nixlGpuXferStatusH *out_xfer_status) +{ + nixlMemViewElem src{src_mvh, 0, 0}, dst{dst_mvh, 0, 0}; + *out_put_status = nixlPut(src, dst, /*size=*/0, channel_id, + /*flags=*/0, out_xfer_status); +} + +// Enqueues op_count puts on one channel and records each immediate enqueue +// status. The final submission may block if the ring is full. +__global__ void +proxyPutBurstKernel(uint32_t op_count, uint32_t channel_id, + nixl_status_t *out_put_statuses) +{ + nixlMemViewElem src{}, dst{}; + for (uint32_t i = 0; i < op_count; ++i) { + out_put_statuses[i] = nixlPut(src, dst, /*size=*/0, channel_id); + } +} + +// Non-blocking single poll: returns current status without spinning. +__global__ void +proxyPollOnceKernel(nixlGpuXferStatusH *xfer_status, + nixl_status_t *out_poll_status) +{ + *out_poll_status = nixlGpuGetXferStatus(*xfer_status); +} + +// --------------------------------------------------------------------------- +// Completion round-trip helpers +// --------------------------------------------------------------------------- + +// Register one local and one remote proxy memview so dispatch can prepare +// transport-ready descriptors before submit(). +static DummyProxyMemViews +registerDummyMemViews(nixlProxyRuntime &runtime) +{ + static DummyBackendMD local_md; + static DummyBackendMD remote_md; + + DummyProxyMemViews handles; + nixlMemViewH dummy_local_backend = reinterpret_cast(uintptr_t{0xBEEF}); + nixlMemViewH dummy_remote_backend = reinterpret_cast(uintptr_t{0xFEED}); + + EXPECT_EQ(runtime.registerProxyMemView(dummy_local_backend, &handles.src), NIXL_SUCCESS); + EXPECT_EQ(runtime.registerProxyMemView(dummy_remote_backend, &handles.dst), NIXL_SUCCESS); + + nixl_meta_dlist_t local_dlist(DRAM_SEG); + local_dlist.addDesc(nixlMetaDesc(0x1000, 64, 0, &local_md)); + EXPECT_EQ(runtime.storeMetadata(handles.src, local_dlist), NIXL_SUCCESS); + + nixl_remote_meta_dlist_t remote_dlist(DRAM_SEG); + nixlRemoteMetaDesc remote_desc("peer"); + remote_desc.addr = 0x2000; + remote_desc.len = 64; + remote_desc.devId = 0; + remote_desc.metadataP = &remote_md; + remote_dlist.addDesc(remote_desc); + EXPECT_EQ(runtime.storeMetadata(handles.dst, remote_dlist), NIXL_SUCCESS); + + return handles; +} + +static void +signalProxyShutdown(nixlProxyRuntime &runtime) +{ + cudaPointerAttributes attrs{}; + ASSERT_EQ(cudaPointerGetAttributes(&attrs, runtime.deviceContext()->shutdown_word), + cudaSuccess); + ASSERT_NE(attrs.hostPointer, nullptr); + auto *shutdown_host = static_cast(attrs.hostPointer); + __atomic_store_n(shutdown_host, + static_cast(nixl_proxy_control_state_t::SHUTDOWN), + __ATOMIC_RELEASE); +} + +// --------------------------------------------------------------------------- +// Completion round-trip tests +// --------------------------------------------------------------------------- + +// Full round-trip: GPU enqueues -> worker dequeues -> backend completes +// (immediately via StubProxyBackendAdapter) -> worker publishes -> GPU polls +// NIXL_SUCCESS. +TEST_F(ProxyDeviceApiTest, PutCompletionRoundTrip) +{ + auto adapter = std::make_unique(); + nixlProxyRuntime runtime; + + ASSERT_EQ(runtime.init(std::move(adapter), /*channel_count=*/1, /*worker_count=*/1), + NIXL_SUCCESS); + ASSERT_EQ(runtime.startWorkers(), NIXL_SUCCESS); + publishProxyContext(runtime); + + const auto mvhs = registerDummyMemViews(runtime); + + nixl_status_t *d_put_status = deviceAlloc(); + nixl_status_t *d_poll_status = deviceAlloc(); + + proxyPutAndPollKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status, d_poll_status); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + + EXPECT_EQ(deviceGet(d_put_status), NIXL_IN_PROG); + EXPECT_EQ(deviceGet(d_poll_status), NIXL_SUCCESS); + + cudaFree(d_put_status); + cudaFree(d_poll_status); + clearProxyContext(); + ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); +} + +// Verifies that the GPU kernel stays spinning until the test thread +// explicitly marks the backend token complete. +TEST_F(ProxyDeviceApiTest, CompletionNotVisibleUntilPublished) +{ + auto adapter_owner = std::make_unique(); + auto *adapter = adapter_owner.get(); + nixlProxyRuntime runtime; + + ASSERT_EQ(runtime.init(std::move(adapter_owner), /*channel_count=*/1, /*worker_count=*/1), + NIXL_SUCCESS); + ASSERT_EQ(runtime.startWorkers(), NIXL_SUCCESS); + publishProxyContext(runtime); + + const auto mvhs = registerDummyMemViews(runtime); + + nixl_status_t *d_put_status = deviceAlloc(); + nixl_status_t *d_poll_status = deviceAlloc(); + + // Launch async — kernel will spin on pollXferStatus. + proxyPutAndPollKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status, d_poll_status); + + // Give the worker time to pick up and submit the request. + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + + // Kernel should still be running (spinning on completion). + ASSERT_EQ(cudaStreamQuery(nullptr), cudaErrorNotReady); + + // Release the completion from the test thread. + adapter->markComplete(1); + + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + + EXPECT_EQ(deviceGet(d_put_status), NIXL_IN_PROG); + EXPECT_EQ(deviceGet(d_poll_status), NIXL_SUCCESS); + + cudaFree(d_put_status); + cudaFree(d_poll_status); + clearProxyContext(); + ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); +} + +// Enqueue 3 operations, complete them in order, and verify the collapsed-CQ +// frontier semantics: each pollXferStatus returns NIXL_SUCCESS only after its +// op_idx has been reached. +TEST_F(ProxyDeviceApiTest, MultipleSubmissionsCompletionFrontier) +{ + auto adapter_owner = std::make_unique(); + auto *adapter = adapter_owner.get(); + nixlProxyRuntime runtime; + + ASSERT_EQ(runtime.init(std::move(adapter_owner), /*channel_count=*/1, /*worker_count=*/1), + NIXL_SUCCESS); + ASSERT_EQ(runtime.startWorkers(), NIXL_SUCCESS); + publishProxyContext(runtime); + + const auto mvhs = registerDummyMemViews(runtime); + + constexpr int kOps = 3; + nixl_status_t *d_put_status[kOps]; + nixlGpuXferStatusH *d_xfer_status[kOps]; + + for (int i = 0; i < kOps; i++) { + d_put_status[i] = deviceAlloc(); + ASSERT_EQ(cudaMalloc(&d_xfer_status[i], sizeof(nixlGpuXferStatusH)), + cudaSuccess); + ASSERT_EQ(cudaMemset(d_xfer_status[i], 0, sizeof(nixlGpuXferStatusH)), + cudaSuccess); + } + + // Enqueue 3 operations sequentially (each kernel returns after enqueue). + for (int i = 0; i < kOps; i++) { + proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status[i], + d_xfer_status[i]); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + EXPECT_EQ(deviceGet(d_put_status[i]), NIXL_IN_PROG); + } + + // All three should still be in-progress. + nixl_status_t *d_poll = deviceAlloc(); + for (int i = 0; i < kOps; i++) { + proxyPollOnceKernel<<<1, 1>>>(d_xfer_status[i], d_poll); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + EXPECT_EQ(deviceGet(d_poll), NIXL_IN_PROG) + << "op " << i << " should be in-progress before any markComplete"; + } + + // Complete them one at a time and verify frontier advances. + for (int i = 0; i < kOps; i++) { + adapter->markComplete(static_cast(i + 1)); + + // Give worker time to publish. + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + + // Poll this op — should now be complete. + proxyPollOnceKernel<<<1, 1>>>(d_xfer_status[i], d_poll); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + EXPECT_EQ(deviceGet(d_poll), NIXL_SUCCESS) + << "op " << i << " should be complete after markComplete"; + } + + cudaFree(d_poll); + for (int i = 0; i < kOps; i++) { + cudaFree(d_put_status[i]); + cudaFree(d_xfer_status[i]); + } + clearProxyContext(); + ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); +} + +// Once a later op publishes an error, an earlier op whose op_idx is already +// behind the completion frontier must still observe NIXL_SUCCESS. +TEST_F(ProxyDeviceApiTest, EarlierCompletionStaysSuccessfulAfterLaterError) +{ + auto adapter_owner = std::make_unique(); + auto *adapter = adapter_owner.get(); + nixlProxyRuntime runtime; + + ASSERT_EQ(runtime.init(std::move(adapter_owner), /*channel_count=*/1, /*worker_count=*/1), + NIXL_SUCCESS); + ASSERT_EQ(runtime.startWorkers(), NIXL_SUCCESS); + publishProxyContext(runtime); + + const auto mvhs = registerDummyMemViews(runtime); + + nixl_status_t *d_put_status[2]; + nixlGpuXferStatusH *d_xfer_status[2]; + for (int i = 0; i < 2; ++i) { + d_put_status[i] = deviceAlloc(); + ASSERT_EQ(cudaMalloc(&d_xfer_status[i], sizeof(nixlGpuXferStatusH)), + cudaSuccess); + ASSERT_EQ(cudaMemset(d_xfer_status[i], 0, sizeof(nixlGpuXferStatusH)), + cudaSuccess); + proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status[i], + d_xfer_status[i]); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + EXPECT_EQ(deviceGet(d_put_status[i]), NIXL_IN_PROG); + } + + ASSERT_TRUE(waitForCondition([&]() { return adapter->pendingCount() == 2; })); + + nixl_status_t *d_poll = deviceAlloc(); + adapter->markCompleteWithStatus(1, NIXL_SUCCESS); + ASSERT_TRUE(waitForCondition([&]() { + proxyPollOnceKernel<<<1, 1>>>(d_xfer_status[0], d_poll); + return cudaDeviceSynchronize() == cudaSuccess && deviceGet(d_poll) == NIXL_SUCCESS; + })); + + adapter->markCompleteWithStatus(2, NIXL_ERR_BACKEND); + ASSERT_TRUE(waitForCondition([&]() { + proxyPollOnceKernel<<<1, 1>>>(d_xfer_status[1], d_poll); + return cudaDeviceSynchronize() == cudaSuccess && deviceGet(d_poll) == NIXL_ERR_BACKEND; + })); + + proxyPollOnceKernel<<<1, 1>>>(d_xfer_status[0], d_poll); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + EXPECT_EQ(deviceGet(d_poll), NIXL_SUCCESS); + + cudaFree(d_poll); + for (int i = 0; i < 2; ++i) { + cudaFree(d_put_status[i]); + cudaFree(d_xfer_status[i]); + } + clearProxyContext(); + ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); +} + +// Once an earlier op publishes a terminal error, later queued ops must also +// observe that error instead of spinning forever. +TEST_F(ProxyDeviceApiTest, EarlierErrorPropagatesToLaterQueuedOp) +{ + auto adapter_owner = std::make_unique(); + auto *adapter = adapter_owner.get(); + nixlProxyRuntime runtime; + + ASSERT_EQ(runtime.init(std::move(adapter_owner), /*channel_count=*/1, /*worker_count=*/1), + NIXL_SUCCESS); + ASSERT_EQ(runtime.startWorkers(), NIXL_SUCCESS); + publishProxyContext(runtime); + + const auto mvhs = registerDummyMemViews(runtime); + + nixl_status_t *d_put_status[2]; + nixlGpuXferStatusH *d_xfer_status[2]; + for (int i = 0; i < 2; ++i) { + d_put_status[i] = deviceAlloc(); + ASSERT_EQ(cudaMalloc(&d_xfer_status[i], sizeof(nixlGpuXferStatusH)), + cudaSuccess); + ASSERT_EQ(cudaMemset(d_xfer_status[i], 0, sizeof(nixlGpuXferStatusH)), + cudaSuccess); + proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status[i], + d_xfer_status[i]); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + EXPECT_EQ(deviceGet(d_put_status[i]), NIXL_IN_PROG); + } + + ASSERT_TRUE(waitForCondition([&]() { return adapter->pendingCount() == 2; })); + + nixl_status_t *d_poll = deviceAlloc(); + adapter->markCompleteWithStatus(1, NIXL_ERR_BACKEND); + ASSERT_TRUE(waitForCondition([&]() { + proxyPollOnceKernel<<<1, 1>>>(d_xfer_status[1], d_poll); + return cudaDeviceSynchronize() == cudaSuccess && deviceGet(d_poll) == NIXL_ERR_BACKEND; + })); + + proxyPollOnceKernel<<<1, 1>>>(d_xfer_status[0], d_poll); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + EXPECT_EQ(deviceGet(d_poll), NIXL_ERR_BACKEND); + + cudaFree(d_poll); + for (int i = 0; i < 2; ++i) { + cudaFree(d_put_status[i]); + cudaFree(d_xfer_status[i]); + } + clearProxyContext(); + ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); +} + +// Backend returns NIXL_ERR_BACKEND on checkCompletion; verify the GPU kernel +// receives the error status through the completion slot. +TEST_F(ProxyDeviceApiTest, CompletionPropagatesErrorStatus) +{ + auto adapter = std::make_unique(); + nixlProxyRuntime runtime; + + ASSERT_EQ(runtime.init(std::move(adapter), /*channel_count=*/1, /*worker_count=*/1), + NIXL_SUCCESS); + ASSERT_EQ(runtime.startWorkers(), NIXL_SUCCESS); + publishProxyContext(runtime); + + const auto mvhs = registerDummyMemViews(runtime); + + nixl_status_t *d_put_status = deviceAlloc(); + nixl_status_t *d_poll_status = deviceAlloc(); + + proxyPutAndPollKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status, d_poll_status); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + + EXPECT_EQ(deviceGet(d_put_status), NIXL_IN_PROG); + EXPECT_EQ(deviceGet(d_poll_status), NIXL_ERR_BACKEND); + + cudaFree(d_put_status); + cudaFree(d_poll_status); + clearProxyContext(); + ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); +} + +// Backend submit() failure should be published to the GPU as the terminal +// transfer status, without going through checkCompletion(). +TEST_F(ProxyDeviceApiTest, SubmitFailurePropagatesErrorStatus) +{ + const gtest::LogIgnoreGuard lig("ProxyWorker::submitToBackend: backend submit failed"); + auto adapter_owner = std::make_unique(); + auto *adapter = adapter_owner.get(); + nixlProxyRuntime runtime; + + ASSERT_EQ(runtime.init(std::move(adapter_owner), /*channel_count=*/1, /*worker_count=*/1), + NIXL_SUCCESS); + ASSERT_EQ(runtime.startWorkers(), NIXL_SUCCESS); + publishProxyContext(runtime); + + const auto mvhs = registerDummyMemViews(runtime); + + nixl_status_t *d_put_status = deviceAlloc(); + nixl_status_t *d_poll_status = deviceAlloc(); + + proxyPutAndPollKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status, d_poll_status); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + + EXPECT_EQ(deviceGet(d_put_status), NIXL_IN_PROG); + EXPECT_EQ(deviceGet(d_poll_status), NIXL_ERR_BACKEND); + EXPECT_EQ(adapter->submit_calls_.load(), 1u); + EXPECT_EQ(adapter->check_completion_calls_.load(), 0u); + + cudaFree(d_put_status); + cudaFree(d_poll_status); + clearProxyContext(); + ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); +} + +// When the ring is full and no worker can drain it, the next enqueue should +// spin until shutdown is signalled and then return NIXL_ERR_BACKEND. +TEST_F(ProxyDeviceApiTest, RingOverflowReturnsBackendErrorOnShutdown) +{ + auto adapter = std::make_unique(); + nixlProxyRuntime runtime; + + ASSERT_EQ(runtime.init(std::move(adapter), /*channel_count=*/1, /*worker_count=*/1), + NIXL_SUCCESS); + publishProxyContext(runtime); + + constexpr uint32_t kBurstOps = kDefaultProxyRingDepth + 1; + nixl_status_t *d_statuses = nullptr; + ASSERT_EQ(cudaMalloc(&d_statuses, sizeof(nixl_status_t) * kBurstOps), + cudaSuccess); + ASSERT_EQ(cudaMemset(d_statuses, 0, sizeof(nixl_status_t) * kBurstOps), + cudaSuccess); + + proxyPutBurstKernel<<<1, 1>>>(kBurstOps, 0, d_statuses); + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + ASSERT_EQ(cudaStreamQuery(nullptr), cudaErrorNotReady); + + signalProxyShutdown(runtime); + + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + + std::vector statuses(kBurstOps); + ASSERT_EQ(cudaMemcpy(statuses.data(), + d_statuses, + sizeof(nixl_status_t) * kBurstOps, + cudaMemcpyDeviceToHost), + cudaSuccess); + for (uint32_t i = 0; i < kDefaultProxyRingDepth; ++i) { + EXPECT_EQ(statuses[i], NIXL_IN_PROG) << "unexpected status at op " << i; + } + EXPECT_EQ(statuses.back(), NIXL_ERR_BACKEND); + + cudaFree(d_statuses); + clearProxyContext(); + ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); +} + +// Completions are tracked per-channel, so publishing one channel should not +// advance an unrelated channel's xfer status. +TEST_F(ProxyDeviceApiTest, ChannelCompletionsAdvanceIndependently) +{ + auto adapter_owner = std::make_unique(); + auto *adapter = adapter_owner.get(); + nixlProxyRuntime runtime; + + ASSERT_EQ(runtime.init(std::move(adapter_owner), /*channel_count=*/2, /*worker_count=*/2), + NIXL_SUCCESS); + ASSERT_EQ(runtime.startWorkers(), NIXL_SUCCESS); + publishProxyContext(runtime); + + const auto mvhs = registerDummyMemViews(runtime); + + nixl_status_t *d_put_status[2]; + nixlGpuXferStatusH *d_xfer_status[2]; + for (int i = 0; i < 2; ++i) { + d_put_status[i] = deviceAlloc(); + ASSERT_EQ(cudaMalloc(&d_xfer_status[i], sizeof(nixlGpuXferStatusH)), + cudaSuccess); + ASSERT_EQ(cudaMemset(d_xfer_status[i], 0, sizeof(nixlGpuXferStatusH)), + cudaSuccess); + } + + proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status[0], d_xfer_status[0]); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 1, d_put_status[1], d_xfer_status[1]); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + + EXPECT_EQ(deviceGet(d_put_status[0]), NIXL_IN_PROG); + EXPECT_EQ(deviceGet(d_put_status[1]), NIXL_IN_PROG); + ASSERT_TRUE(waitForCondition([adapter]() { + return adapter->pendingCount() == 2 + && adapter->hasPendingForChannel(0) + && adapter->hasPendingForChannel(1); + })); + + uint64_t completed_token = 0; + ASSERT_TRUE(adapter->markFirstPendingForChannel(1, &completed_token)); + EXPECT_GT(completed_token, 0u); + ASSERT_TRUE(waitForCondition([adapter]() { + return adapter->hasPendingForChannel(0) && !adapter->hasPendingForChannel(1); + })); + + nixl_status_t *d_poll_status = deviceAlloc(); + proxyPollOnceKernel<<<1, 1>>>(d_xfer_status[0], d_poll_status); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + EXPECT_EQ(deviceGet(d_poll_status), NIXL_IN_PROG); + + proxyPollOnceKernel<<<1, 1>>>(d_xfer_status[1], d_poll_status); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + EXPECT_EQ(deviceGet(d_poll_status), NIXL_SUCCESS); + + ASSERT_TRUE(adapter->markFirstPendingForChannel(0)); + ASSERT_TRUE(waitForCondition([adapter]() { return !adapter->hasPending(); })); + + proxyPollOnceKernel<<<1, 1>>>(d_xfer_status[0], d_poll_status); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + EXPECT_EQ(deviceGet(d_poll_status), NIXL_SUCCESS); + + cudaFree(d_poll_status); + for (int i = 0; i < 2; ++i) { + cudaFree(d_put_status[i]); + cudaFree(d_xfer_status[i]); + } + clearProxyContext(); + ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); +} diff --git a/test/gtest/device_api/single_write_test.cu b/test/gtest/device_api/single_write_test.cu index 7bf90f96c6..56743b572b 100644 --- a/test/gtest/device_api/single_write_test.cu +++ b/test/gtest/device_api/single_write_test.cu @@ -18,9 +18,14 @@ #include "utils.cuh" #include "common.h" +#include #include #include +#ifdef NIXL_GPU_DEVICE_BACKEND_PROXY +#include +#endif + namespace gtest::nixl::gpu::single_write { struct putParams { nixlMemViewElem src; @@ -164,6 +169,12 @@ protected: cfg.useProgThread = true; cfg.syncMode = nixl_thread_sync_t::NIXL_THREAD_SYNC_RW; cfg.pthrDelay = 100000; +#ifdef NIXL_GPU_DEVICE_BACKEND_PROXY + cfg.enableDeviceProxy = true; + cfg.proxyChannelCount = numWorkers; + cfg.proxyWorkerCount = + std::min(static_cast(numWorkers), kProxyPostXferWorkerLimit); +#endif return cfg; } @@ -199,10 +210,20 @@ protected: EXPECT_NE(backend_handle, nullptr); backend_handles.push_back(backend_handle); } + +#ifdef NIXL_GPU_DEVICE_BACKEND_PROXY + auto *ctx = static_cast( + agents[SENDER_AGENT]->getProxyDeviceContext()); + ASSERT_NE(ctx, nullptr) << "Proxy device context not available"; + ASSERT_EQ(nixlProxyPublishContext(ctx), cudaSuccess); +#endif } void TearDown() override { +#ifdef NIXL_GPU_DEVICE_BACKEND_PROXY + nixlProxyClearContext(); +#endif agents.clear(); backend_handles.clear(); } @@ -308,6 +329,7 @@ protected: static constexpr size_t SENDER_AGENT = 0; static constexpr size_t RECEIVER_AGENT = 1; static constexpr size_t numWorkers = 32; + static constexpr uint32_t kProxyPostXferWorkerLimit = 1; private: static constexpr uint64_t DEV_ID = 0; @@ -401,7 +423,7 @@ TEST_P(SingleWriteTest, SingleWorkerPut) { ASSERT_EQ(status, NIXL_SUCCESS); putParams put_params{{src_mvh, 0, 0}, {dst_mvh, 0, 0}, size}; - constexpr size_t num_iters = 1000; + constexpr size_t num_iters = 10; gpuTimer gpu_timer; status = dispatchLaunchPutKernel(GetParam(), put_params, num_iters, &gpu_timer); ASSERT_EQ(status, NIXL_SUCCESS); @@ -422,6 +444,13 @@ TEST_P(SingleWriteTest, SingleWorkerPut) { } TEST_P(SingleWriteTest, MultipleWorkersPut) { +#ifdef NIXL_GPU_DEVICE_BACKEND_PROXY + GTEST_LOG_(WARNING) + << "Proxy backend caps worker threads at " << kProxyPostXferWorkerLimit + << " until the UCX postXfer path is validated for concurrent proxy workers; " + << "this test exercises explicit channel selection across " + << numWorkers << " channels"; +#endif constexpr size_t size = 4 * 1024; constexpr nixl_mem_t mem_type = VRAM_SEG; @@ -472,8 +501,12 @@ TEST_P(SingleWriteTest, MultipleWorkersPut) { for (size_t worker_id = 0; worker_id < numWorkers; worker_id++) { putParams put_params{{src_mvhs[worker_id], 0, 0}, {dst_mvhs[worker_id], 0, 0}, size}; +#ifdef NIXL_GPU_DEVICE_BACKEND_PROXY + put_params.channelId = static_cast(worker_id); +#endif constexpr size_t num_iters = 1; - const auto status = dispatchLaunchPutKernel(GetParam(), put_params, num_iters); + const auto status = launchPutKernel( + put_params, num_iters, nullptr, 1); ASSERT_EQ(status, NIXL_SUCCESS) << "Kernel launch failed for worker " << worker_id; } @@ -500,6 +533,9 @@ TEST_P(SingleWriteTest, MultipleWorkersPut) { } TEST_P(SingleWriteTest, SingleWorkerPutGap) { +#ifdef NIXL_GPU_DEVICE_BACKEND_PROXY + GTEST_SKIP() << "FIXME: get_ptr not implemented for proxy backend"; +#endif std::vector src_buffers, dst_buffers; constexpr size_t size = 4 * 1024; constexpr size_t count = 1; @@ -555,6 +591,15 @@ TEST_P(SingleWriteTest, SingleWorkerPutGap) { using gtest::nixl::gpu::single_write::SingleWriteTest; +#ifdef NIXL_GPU_DEVICE_BACKEND_PROXY +INSTANTIATE_TEST_SUITE_P( + proxyDeviceApi, + SingleWriteTest, + testing::ValuesIn(gtest::gpu::_test_levels), + [](const testing::TestParamInfo &info) { + return std::string("Proxy_") + gtest::gpu::GetGpuXferLevelStr(info.param); + }); +#else INSTANTIATE_TEST_SUITE_P( ucxDeviceApi, SingleWriteTest, @@ -562,3 +607,4 @@ INSTANTIATE_TEST_SUITE_P( [](const testing::TestParamInfo &info) { return std::string("UCX_") + gtest::gpu::GetGpuXferLevelStr(info.param); }); +#endif diff --git a/test/gtest/meson.build b/test/gtest/meson.build index b2be5f7bb1..ae6f95215e 100644 --- a/test/gtest/meson.build +++ b/test/gtest/meson.build @@ -52,7 +52,7 @@ subdir('mocks') subdir('unit') subdir('plugins') -if ucx_gpu_device_api_available +if cuda_dep.found() subdir('device_api') endif @@ -159,6 +159,36 @@ if nixl_trace_nvtx_enabled endif endif +proxy_gtest_available = cuda_dep.found() + +if proxy_gtest_available + proxy_gtest_sources = [ + 'main.cpp', + 'mocks/gmock_engine.cpp', + 'plugin_manager.cpp', + 'error_handling.cpp', + 'common.cpp', + ] + proxy_device_api_test_sources + + proxy_define = ['-DHAVE_CUDA', '-DNIXL_GPU_DEVICE_BACKEND_PROXY'] + + proxy_test_exe = executable('gtest_proxy', + sources : proxy_gtest_sources, + include_directories : [nixl_inc_dirs, utils_inc_dirs, gtest_inc_dirs, + nixl_gpu_proxy_inc_dirs, include_directories('device_api')], + cpp_args : cpp_flags + proxy_define, + cuda_args : proxy_define, + dependencies : [nixl_dep, nixl_common_dep, cuda_dep, + gtest_dep, gmock_dep, absl_strings_dep, absl_time_dep, + file_utils_interface], + # nixl_gpu_proxy_lib must be linked directly so Meson includes + # nixl_device_proxy.cu.o in the RDC device-link step for gtest_proxy. + link_with : [nixl_build_lib, nixl_gpu_proxy_lib], + install : false) + + test('gtest_proxy', proxy_test_exe, args: [plugin_dirs_arg]) +endif + if get_option('b_sanitize').split(',').contains('thread') test_env = environment() test_env.set('TSAN_OPTIONS', 'halt_on_error=1') From 2ffcd736368e1eb72b11f30a8154e485dca9919b Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Mon, 11 May 2026 13:38:20 +0300 Subject: [PATCH 07/18] build: Add meson build option for GPU Device API backend Change implicit behavior of UCX device api selection to explicit by adding a new option `gpu_device_api_backend`: auto -> choose default backend if available ucx -> require UCX GPU Device API proxy -> require CUDA/proxy support none -> disable default GPU Device API backend Also centralize backend-specific compile args, deps, flags and link inputs so tests and examples will use the same resolved GPU Device API selected backend. Signed-off-by: Tomer Davidor --- examples/device/ep/meson.build | 10 +++- meson.build | 68 +++++++++++++++++++--- meson_options.txt | 1 + src/api/gpu/meson.build | 4 ++ test/gtest/device_api/single_write_test.cu | 25 ++++---- test/gtest/meson.build | 25 ++++---- 6 files changed, 98 insertions(+), 35 deletions(-) diff --git a/examples/device/ep/meson.build b/examples/device/ep/meson.build index 678d75e7b5..d3cc82e355 100644 --- a/examples/device/ep/meson.build +++ b/examples/device/ep/meson.build @@ -18,6 +18,10 @@ if not cuda_dep.found() subdir_done() endif +if gpu_device_api_backend_resolved != 'ucx' + error('nixl_ep requires gpu_device_api_backend=ucx because it currently uses the UCX GPU Device API') +endif + py = import('python').find_installation('python3', pure: false) # Check if PyTorch is available @@ -72,7 +76,7 @@ nixl_ep_sources = [ nixl_ep_inc_dirs = [ nixl_inc_dirs, - nixl_gpu_inc_dirs, + nixl_gpu_device_api_inc_dirs, include_directories('csrc'), include_directories('csrc/kernels'), torch_inc_dirs, @@ -96,7 +100,7 @@ nixl_ep_cpp_args = [ '-Wno-sign-compare', '-Wno-reorder', '-Wno-attributes', -] +] + gpu_device_api_cpp_args nixl_ep_cuda_args = [ '-DHAVE_CUDA', @@ -109,7 +113,7 @@ nixl_ep_cuda_args = [ '-Xcompiler', '-Wno-sign-compare', '-Xcompiler', '-Wno-reorder', '-Xcompiler', '-Wno-attributes', -] +] + gpu_device_api_cuda_args topk_idx_bits = meson.get_external_property('topk_idx_bits', '32') nixl_ep_cpp_args += ['-DTOPK_IDX_BITS=' + topk_idx_bits] diff --git a/meson.build b/meson.build index fb5af2ef18..8fb140458a 100644 --- a/meson.build +++ b/meson.build @@ -396,6 +396,8 @@ endif # UCX GPU device API detection nvcc_prog = find_program('nvcc', required: false) ucx_gpu_device_api_available = false +have_gpu_side = false +have_host_side = false if ucx_dep.found() and cuda_dep.found() and nvcc_prog.found() cuda = meson.get_compiler('cuda') # TODO: Expose doca_gpunetio_dep through UCX @@ -412,16 +414,67 @@ if ucx_dep.found() and cuda_dep.found() and nvcc_prog.found() ucx_gpu_device_api_available = true add_project_arguments('-DHAVE_UCX_GPU_DEVICE_API', language: ['cpp', 'cuda']) endif +endif - summary({ - 'UCX GPU Device API' : ucx_gpu_device_api_available, - 'GPU-side compile' : have_gpu_side, - 'Host-side compile' : have_host_side, - 'nvcc available' : nvcc_prog.found(), - 'DOCA GPUNETIO found' : doca_gpunetio_dep.found(), - }, section: 'UCX GPU Device API', bool_yn: true) +gpu_device_api_backend_requested = get_option('gpu_device_api_backend') +gpu_device_api_backend_resolved = gpu_device_api_backend_requested + +if gpu_device_api_backend_requested == 'auto' + if ucx_gpu_device_api_available + gpu_device_api_backend_resolved = 'ucx' + else + gpu_device_api_backend_resolved = 'none' + endif +elif gpu_device_api_backend_requested == 'ucx' + if not ucx_gpu_device_api_available + error('gpu_device_api_backend=ucx requires UCX GPU Device API support') + endif +elif gpu_device_api_backend_requested == 'proxy' + if not cuda_dep.found() + error('gpu_device_api_backend=proxy requires CUDA') + endif endif +gpu_device_api_define = '' +gpu_device_api_supported = gpu_device_api_backend_resolved != 'none' +gpu_device_api_is_ucx = gpu_device_api_backend_resolved == 'ucx' +gpu_device_api_is_proxy = gpu_device_api_backend_resolved == 'proxy' +gpu_device_api_cpp_args = [] +gpu_device_api_cuda_args = [] +gpu_device_api_deps = [] +gpu_device_api_link_with = [] + +if gpu_device_api_is_ucx + gpu_device_api_define = '-DNIXL_GPU_DEVICE_BACKEND_UCX' + gpu_device_api_deps = [ucx_dep, doca_gpunetio_dep] +elif gpu_device_api_is_proxy + gpu_device_api_define = '-DNIXL_GPU_DEVICE_BACKEND_PROXY' + gpu_device_api_deps = [cuda_dep] +endif + +if gpu_device_api_define != '' + gpu_device_api_cpp_args += [gpu_device_api_define] + gpu_device_api_cuda_args += [gpu_device_api_define] +endif + +summary({ + 'Requested backend' : gpu_device_api_backend_requested, + 'Resolved backend' : gpu_device_api_backend_resolved, +}, section: 'GPU Device API') + +summary({ + 'UCX backend available' : ucx_gpu_device_api_available, + 'Proxy backend available' : cuda_dep.found(), +}, section: 'GPU Device API', bool_yn: true) + +summary({ + 'UCX GPU Device API' : ucx_gpu_device_api_available, + 'GPU-side compile' : have_gpu_side, + 'Host-side compile' : have_host_side, + 'nvcc available' : nvcc_prog.found(), + 'DOCA GPUNETIO found' : doca_gpunetio_dep.found(), +}, section: 'UCX GPU Device API', bool_yn: true) + if get_option('disable_gds_backend') add_project_arguments('-DDISABLE_GDS_BACKEND', language: 'cpp') endif @@ -509,6 +562,7 @@ if get_option('buildtype') == 'debug' endif nixl_inc_dirs = include_directories('src/api/cpp', 'src/api/cpp/backend', 'src/infra', 'src/core', 'src/core/telemetry') +nixl_gpu_device_api_inc_dirs = include_directories('src/api/gpu') nixl_gpu_inc_dirs = include_directories('src/api/gpu/ucx') nixl_gpu_proxy_inc_dirs = include_directories('src/api/gpu/proxy') plugins_inc_dirs = include_directories('src/plugins') diff --git a/meson_options.txt b/meson_options.txt index eb875fc36c..cda2f694e0 100644 --- a/meson_options.txt +++ b/meson_options.txt @@ -28,6 +28,7 @@ option('cudapath_lib', type: 'string', value: '', description: 'Library path for option('cudapath_stub', type: 'string', value: '', description: 'Extra Stub path for CUDA') option('nixl_cuda_arch_list', type: 'string', value: 'auto', description: 'Comma-separated CUDA SM targets (e.g. 90,100), or auto to select defaults (sm_80+ normally, sm_90+ when build_nixl_ep is enabled). Use a single target like 100 for faster single-GPU builds.') option('wheel_variant', type: 'string', value: '', description: 'Override wheel variant suffix (e.g. -Dwheel_variant=rocm yields nixl_rocm). Empty = autodetect from CUDA major version.') +option('gpu_device_api_backend', type: 'combo', choices: ['auto', 'ucx', 'proxy', 'none'], value: 'auto', description: 'GPU Device API backend to select for default device-side builds') option('static_plugins', type: 'string', value: '', description: 'Plugins to be built-in, comma-separated') option('enable_plugins', type: 'string', value: '', description: 'Comma-separated list of plugins to build. Default (empty) builds all available plugins.') option('disable_plugins', type: 'string', value: '', description: 'Comma-separated list of plugins to exclude from build. Cannot be used with enable_plugins.') diff --git a/src/api/gpu/meson.build b/src/api/gpu/meson.build index c8f502ae57..76faeeeea9 100644 --- a/src/api/gpu/meson.build +++ b/src/api/gpu/meson.build @@ -25,4 +25,8 @@ if cuda_dep.found() install: false) nixl_gpu_proxy_dep = declare_dependency(link_with: nixl_gpu_proxy_lib) + + if gpu_device_api_is_proxy + gpu_device_api_link_with = [nixl_gpu_proxy_lib] + endif endif diff --git a/test/gtest/device_api/single_write_test.cu b/test/gtest/device_api/single_write_test.cu index 56743b572b..a616313b9b 100644 --- a/test/gtest/device_api/single_write_test.cu +++ b/test/gtest/device_api/single_write_test.cu @@ -22,10 +22,6 @@ #include #include -#ifdef NIXL_GPU_DEVICE_BACKEND_PROXY -#include -#endif - namespace gtest::nixl::gpu::single_write { struct putParams { nixlMemViewElem src; @@ -448,8 +444,7 @@ TEST_P(SingleWriteTest, MultipleWorkersPut) { GTEST_LOG_(WARNING) << "Proxy backend caps worker threads at " << kProxyPostXferWorkerLimit << " until the UCX postXfer path is validated for concurrent proxy workers; " - << "this test exercises explicit channel selection across " - << numWorkers << " channels"; + << "this test exercises explicit channel selection across " << numWorkers << " channels"; #endif constexpr size_t size = 4 * 1024; constexpr nixl_mem_t mem_type = VRAM_SEG; @@ -505,8 +500,8 @@ TEST_P(SingleWriteTest, MultipleWorkersPut) { put_params.channelId = static_cast(worker_id); #endif constexpr size_t num_iters = 1; - const auto status = launchPutKernel( - put_params, num_iters, nullptr, 1); + const auto status = + launchPutKernel(put_params, num_iters, nullptr, 1); ASSERT_EQ(status, NIXL_SUCCESS) << "Kernel launch failed for worker " << worker_id; } @@ -592,13 +587,13 @@ TEST_P(SingleWriteTest, SingleWorkerPutGap) { using gtest::nixl::gpu::single_write::SingleWriteTest; #ifdef NIXL_GPU_DEVICE_BACKEND_PROXY -INSTANTIATE_TEST_SUITE_P( - proxyDeviceApi, - SingleWriteTest, - testing::ValuesIn(gtest::gpu::_test_levels), - [](const testing::TestParamInfo &info) { - return std::string("Proxy_") + gtest::gpu::GetGpuXferLevelStr(info.param); - }); +INSTANTIATE_TEST_SUITE_P(proxyDeviceApi, + SingleWriteTest, + testing::ValuesIn(gtest::gpu::_test_levels), + [](const testing::TestParamInfo &info) { + return std::string("Proxy_") + + gtest::gpu::GetGpuXferLevelStr(info.param); + }); #else INSTANTIATE_TEST_SUITE_P( ucxDeviceApi, diff --git a/test/gtest/meson.build b/test/gtest/meson.build index ae6f95215e..011712cf25 100644 --- a/test/gtest/meson.build +++ b/test/gtest/meson.build @@ -59,9 +59,11 @@ endif plugin_dirs_arg = '--tests_plugin_dirs=' + mocks_dep.get_variable('path') cpp_flags = [] +cuda_flags = [] if cuda_dep.found() cpp_flags += '-DHAVE_CUDA' + cuda_flags += '-DHAVE_CUDA' cuda_dependencies = [cuda_dep] else cuda_dependencies = [] @@ -94,13 +96,15 @@ gtest_sources = [ 'configuration.cpp' ] -if ucx_gpu_device_api_available - gtest_sources += device_api_test_sources - device_api_inc = [nixl_gpu_inc_dirs, include_directories('device_api')] - device_api_dep = [ucx_dep, doca_gpunetio_dep] +if gpu_device_api_supported + if gpu_device_api_is_proxy + gtest_sources += proxy_device_api_test_sources + else + gtest_sources += device_api_test_sources + endif + device_api_inc = [nixl_gpu_device_api_inc_dirs, include_directories('device_api')] else device_api_inc = [] - device_api_dep = [] endif if ucx_dep.found() and is_variable('ucx_backend_interface') @@ -113,9 +117,10 @@ endif test_exe = executable('gtest', sources : gtest_sources, include_directories: [nixl_inc_dirs, utils_inc_dirs, device_api_inc, include_directories('../doca-telemetry')], - cpp_args : cpp_flags, - dependencies : [nixl_dep, nixl_common_dep, cuda_dependencies, device_api_dep, gtest_dep, gmock_dep, absl_strings_dep, absl_time_dep, ucx_hw_warning_dep, file_utils_interface], - link_with: [nixl_build_lib], + cpp_args : cpp_flags + gpu_device_api_cpp_args, + cuda_args : cuda_flags + gpu_device_api_cuda_args, + dependencies : [nixl_dep, nixl_common_dep, cuda_dependencies, gpu_device_api_deps, gtest_dep, gmock_dep, absl_strings_dep, absl_time_dep, ucx_hw_warning_dep, file_utils_interface], + link_with: [nixl_build_lib, gpu_device_api_link_with], install : true ) @@ -170,14 +175,14 @@ if proxy_gtest_available 'common.cpp', ] + proxy_device_api_test_sources - proxy_define = ['-DHAVE_CUDA', '-DNIXL_GPU_DEVICE_BACKEND_PROXY'] + proxy_define = ['-DNIXL_GPU_DEVICE_BACKEND_PROXY'] proxy_test_exe = executable('gtest_proxy', sources : proxy_gtest_sources, include_directories : [nixl_inc_dirs, utils_inc_dirs, gtest_inc_dirs, nixl_gpu_proxy_inc_dirs, include_directories('device_api')], cpp_args : cpp_flags + proxy_define, - cuda_args : proxy_define, + cuda_args : cuda_flags + proxy_define, dependencies : [nixl_dep, nixl_common_dep, cuda_dep, gtest_dep, gmock_dep, absl_strings_dep, absl_time_dep, file_utils_interface], From 9b55f6c709b14f242e5e9dfb7ef130a4bff0827e Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Sun, 7 Jun 2026 13:42:21 +0300 Subject: [PATCH 08/18] proxy_runtime: Remove auto-restart ability for workers Signed-off-by: Tomer Davidor --- src/core/device_proxy/proxy_runtime.cpp | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/core/device_proxy/proxy_runtime.cpp b/src/core/device_proxy/proxy_runtime.cpp index 531caff04f..493070ae37 100644 --- a/src/core/device_proxy/proxy_runtime.cpp +++ b/src/core/device_proxy/proxy_runtime.cpp @@ -749,10 +749,12 @@ nixlProxyRuntime::startWorkers() { NIXL_ERROR << "ProxyRuntime::startWorkers: runtime not initialized"; return NIXL_ERR_NOT_SUPPORTED; } - __atomic_store_n(shutdown_word_host_, - static_cast(nixl_proxy_control_state_t::SHUTDOWN), - __ATOMIC_RELEASE); - joinWorkerThreads(); + + if (workers_started_) { + NIXL_ERROR << "ProxyRuntime::startWorkers: workers already started"; + return NIXL_ERR_INVALID_PARAM; + } + for (auto &channel : channels_) { channel.inflight_requests.clear(); channel.error_latched = false; From 9dfb44e062b754dbeebea78b36e4cf1fba7644b9 Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Tue, 28 Apr 2026 20:00:33 +0300 Subject: [PATCH 09/18] ucx_proxy_backend: add direct UCX put submission Route proxy PUT through a lower-level UCX RMA write helper while preserving the existing request handle and flush-backed completion semantics. Signed-off-by: Tomer Davidor --- .../ucx/device_proxy/ucx_proxy_backend.cpp | 24 ++++------- src/plugins/ucx/ucx_backend.cpp | 41 +++++++++++++++++++ src/plugins/ucx/ucx_backend.h | 15 +++++++ 3 files changed, 63 insertions(+), 17 deletions(-) diff --git a/src/plugins/ucx/device_proxy/ucx_proxy_backend.cpp b/src/plugins/ucx/device_proxy/ucx_proxy_backend.cpp index b001427bd6..5a738972fc 100644 --- a/src/plugins/ucx/device_proxy/ucx_proxy_backend.cpp +++ b/src/plugins/ucx/device_proxy/ucx_proxy_backend.cpp @@ -44,25 +44,15 @@ nixlUcxProxyBackendAdapter::submit(const nixlBackendProxySubmission &submission, nixl_status_t nixlUcxProxyBackendAdapter::submitPut(const nixlBackendProxySubmission &submission, uint64_t &request_token) { - nixl_meta_dlist_t local_list(submission.local.mem_type); - local_list.addDesc(submission.local.desc); - - nixl_meta_dlist_t remote_list(submission.remote.mem_type); - remote_list.addDesc(submission.remote.desc); - nixlBackendReqH *handle = nullptr; - nixl_status_t status = engine_->prepXfer( - NIXL_WRITE, local_list, remote_list, submission.remote_agent, handle); - if (status != NIXL_SUCCESS) { - NIXL_DEBUG << "nixlUcxProxyBackendAdapter::submitPut: prepXfer failed status=" << status; - return status; - } - - status = engine_->postXfer( - NIXL_WRITE, local_list, remote_list, submission.remote_agent, handle); + nixl_status_t status = engine_->submitProxyRmaWrite(submission.local.desc, + submission.remote.desc, + submission.size, + handle); if (status != NIXL_SUCCESS && status != NIXL_IN_PROG) { - NIXL_DEBUG << "nixlUcxProxyBackendAdapter::submitPut: postXfer failed status=" << status; - engine_->releaseReqH(handle); + NIXL_DEBUG << "nixlUcxProxyBackendAdapter::submitPut: submitProxyRmaWrite failed " + "status=" + << status; return status; } diff --git a/src/plugins/ucx/ucx_backend.cpp b/src/plugins/ucx/ucx_backend.cpp index 6f10791d99..8b91ceac24 100644 --- a/src/plugins/ucx/ucx_backend.cpp +++ b/src/plugins/ucx/ucx_backend.cpp @@ -1103,6 +1103,47 @@ nixl_status_t nixlUcxEngine::prepXfer (const nixl_xfer_op_t &operation, return NIXL_SUCCESS; } +nixl_status_t +nixlUcxEngine::submitProxyRmaWrite(const nixlMetaDesc &local, + const nixlMetaDesc &remote, + size_t size, + nixlBackendReqH *&handle) const { + handle = nullptr; + + if (local.len != size || remote.len != size) { + return NIXL_ERR_INVALID_PARAM; + } + + auto *lmd = static_cast(local.metadataP); + auto *rmd = static_cast(remote.metadataP); + if (lmd == nullptr || rmd == nullptr || rmd->conn == nullptr) { + return NIXL_ERR_INVALID_PARAM; + } + + const auto worker_id = getWorkerId(); + auto *ucx_handle = new nixlUcxBackendReqH(getWorker(worker_id).get(), worker_id); + handle = ucx_handle; + ucx_handle->reserve(1); + + auto &ep = rmd->conn->getEp(worker_id); + nixlUcxReq req = nullptr; + const nixl_status_t submit_status = ep->write(reinterpret_cast(local.addr), + lmd->mem, + static_cast(remote.addr), + rmd->getRkey(worker_id), + size, + req); + const nixl_status_t append_status = ucx_handle->append(submit_status, req, rmd->conn); + if (append_status != NIXL_SUCCESS) { + releaseReqH(handle); + handle = nullptr; + return append_status; + } + + return submit_status; +} + + nixl_status_t nixlUcxEngine::estimateXferCost (const nixl_xfer_op_t &operation, const nixl_meta_dlist_t &local, const nixl_meta_dlist_t &remote, diff --git a/src/plugins/ucx/ucx_backend.h b/src/plugins/ucx/ucx_backend.h index a2f94598cd..8a6a0bc174 100644 --- a/src/plugins/ucx/ucx_backend.h +++ b/src/plugins/ucx/ucx_backend.h @@ -52,6 +52,7 @@ class nixlUcxConnection : public nixlBackendConnMD { }; using ucx_connection_ptr_t = std::shared_ptr; +class nixlUcxProxyBackendAdapter; // A private metadata has to implement get, and has all the metadata class nixlUcxPrivateMetadata : public nixlBackendMD { @@ -185,6 +186,12 @@ class nixlUcxEngine : public nixlBackendEngine { nixl_status_t releaseReqH(nixlBackendReqH *handle) const override; + nixl_status_t + submitRmaWrite(const nixlMetaDesc &local, + const nixlMetaDesc &remote, + size_t size, + nixlBackendReqH *&handle) const; + unsigned progress(); @@ -253,6 +260,14 @@ class nixlUcxEngine : public nixlBackendEngine { notif_list_t notifList_; private: + friend class nixlUcxProxyBackendAdapter; + + nixl_status_t + submitProxyRmaWrite(const nixlMetaDesc &local, + const nixlMetaDesc &remote, + size_t size, + nixlBackendReqH *&handle) const; + // Memory management helpers nixl_status_t internalMDHelper(const nixl_blob_t &blob, const std::string &agent, nixlBackendMD *&output); From 0f46034de643757b9faa689d59a95e1211aa38ab Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Tue, 28 Apr 2026 20:05:51 +0300 Subject: [PATCH 10/18] ucx_proxy_backend: Add UCX atomic add support Signed-off-by: Tomer Davidor --- src/api/gpu/proxy/nixl_device_impl.cuh | 1 + .../ucx/device_proxy/ucx_proxy_backend.cpp | 41 ++++++++++++------- src/plugins/ucx/ucx_backend.cpp | 33 +++++++++++++++ src/plugins/ucx/ucx_backend.h | 11 +++-- src/plugins/ucx/ucx_utils.cpp | 22 ++++++++++ src/plugins/ucx/ucx_utils.h | 2 + 6 files changed, 89 insertions(+), 21 deletions(-) diff --git a/src/api/gpu/proxy/nixl_device_impl.cuh b/src/api/gpu/proxy/nixl_device_impl.cuh index 54d8c0ccbf..2db9261d4b 100644 --- a/src/api/gpu/proxy/nixl_device_impl.cuh +++ b/src/api/gpu/proxy/nixl_device_impl.cuh @@ -112,6 +112,7 @@ atomic_add(uint64_t value, submission.dst_proxy_memview_id = proxyMemViewIdFromHandle(counter.mvh); submission.dst_index = counter.index; submission.dst_offset = counter.offset; + submission.size = sizeof(uint64_t); submission.value = value; status = ctx->enqueue(submission, xfer_status); } diff --git a/src/plugins/ucx/device_proxy/ucx_proxy_backend.cpp b/src/plugins/ucx/device_proxy/ucx_proxy_backend.cpp index 5a738972fc..df79d429d6 100644 --- a/src/plugins/ucx/device_proxy/ucx_proxy_backend.cpp +++ b/src/plugins/ucx/device_proxy/ucx_proxy_backend.cpp @@ -45,10 +45,8 @@ nixl_status_t nixlUcxProxyBackendAdapter::submitPut(const nixlBackendProxySubmission &submission, uint64_t &request_token) { nixlBackendReqH *handle = nullptr; - nixl_status_t status = engine_->submitProxyRmaWrite(submission.local.desc, - submission.remote.desc, - submission.size, - handle); + nixl_status_t status = engine_->submitProxyRmaWrite( + submission.local.desc, submission.remote.desc, submission.size, handle); if (status != NIXL_SUCCESS && status != NIXL_IN_PROG) { NIXL_DEBUG << "nixlUcxProxyBackendAdapter::submitPut: submitProxyRmaWrite failed " "status=" @@ -58,20 +56,33 @@ nixlUcxProxyBackendAdapter::submitPut(const nixlBackendProxySubmission &submissi request_token = trackRequest(handle); NIXL_DEBUG << "nixlUcxProxyBackendAdapter::submitPut: posted RDMA write" - << " src_addr=0x" << std::hex - << submission.local.desc.addr << std::dec - << " dst_addr=0x" << std::hex - << submission.remote.desc.addr << std::dec - << " size=" << submission.size - << " remote_agent='" << submission.remote_agent << "'" + << " src_addr=0x" << std::hex << submission.local.desc.addr << std::dec + << " dst_addr=0x" << std::hex << submission.remote.desc.addr << std::dec + << " size=" << submission.size << " remote_agent='" << submission.remote_agent << "'" << " token=" << request_token; return NIXL_SUCCESS; } nixl_status_t -nixlUcxProxyBackendAdapter::submitAtomicAdd(const nixlBackendProxySubmission &, - uint64_t &) { - return NIXL_ERR_NOT_SUPPORTED; +nixlUcxProxyBackendAdapter::submitAtomicAdd(const nixlBackendProxySubmission &submission, + uint64_t &request_token) { + nixlBackendReqH *handle = nullptr; + nixl_status_t status = + engine_->submitProxyAtomicAdd(submission.remote.desc, submission.value, handle); + if (status != NIXL_SUCCESS && status != NIXL_IN_PROG) { + NIXL_DEBUG << "nixlUcxProxyBackendAdapter::submitAtomicAdd: submitProxyAtomicAdd " + "failed status=" + << status; + return status; + } + + request_token = trackRequest(handle); + NIXL_DEBUG << "nixlUcxProxyBackendAdapter::submitAtomicAdd: posted RDMA atomic add" + << " dst_addr=0x" << std::hex << submission.remote.desc.addr << std::dec + << " size=" << submission.size << " value=" << submission.value << " remote_agent='" + << submission.remote_agent << "'" + << " token=" << request_token; + return NIXL_SUCCESS; } nixl_status_t @@ -110,8 +121,8 @@ nixlUcxProxyBackendAdapter::progress() { nixl_status_t nixlUcxProxyBackendAdapter::shutdown() { - NIXL_INFO << "nixlUcxProxyBackendAdapter::shutdown: releasing " - << tracked_requests_.size() << " tracked request(s)"; + NIXL_INFO << "nixlUcxProxyBackendAdapter::shutdown: releasing " << tracked_requests_.size() + << " tracked request(s)"; { std::lock_guard lock(request_mutex_); if (engine_ != nullptr) { diff --git a/src/plugins/ucx/ucx_backend.cpp b/src/plugins/ucx/ucx_backend.cpp index 8b91ceac24..a4e9f9ae02 100644 --- a/src/plugins/ucx/ucx_backend.cpp +++ b/src/plugins/ucx/ucx_backend.cpp @@ -1143,6 +1143,39 @@ nixlUcxEngine::submitProxyRmaWrite(const nixlMetaDesc &local, return submit_status; } +nixl_status_t +nixlUcxEngine::submitProxyAtomicAdd(const nixlMetaDesc &remote, + uint64_t value, + nixlBackendReqH *&handle) const { + handle = nullptr; + + if (remote.len != sizeof(uint64_t)) { + return NIXL_ERR_INVALID_PARAM; + } + + auto *rmd = static_cast(remote.metadataP); + if (rmd == nullptr || rmd->conn == nullptr) { + return NIXL_ERR_INVALID_PARAM; + } + + const auto worker_id = getWorkerId(); + auto *ucx_handle = new nixlUcxBackendReqH(getWorker(worker_id).get(), worker_id); + handle = ucx_handle; + ucx_handle->reserve(1); + + auto &ep = rmd->conn->getEp(worker_id); + nixlUcxReq req = nullptr; + const nixl_status_t submit_status = + ep->atomicAdd(value, static_cast(remote.addr), rmd->getRkey(worker_id), req); + const nixl_status_t append_status = ucx_handle->append(submit_status, req, rmd->conn); + if (append_status != NIXL_SUCCESS) { + releaseReqH(handle); + handle = nullptr; + return append_status; + } + + return submit_status; +} nixl_status_t nixlUcxEngine::estimateXferCost (const nixl_xfer_op_t &operation, const nixl_meta_dlist_t &local, diff --git a/src/plugins/ucx/ucx_backend.h b/src/plugins/ucx/ucx_backend.h index 8a6a0bc174..8a545c38e3 100644 --- a/src/plugins/ucx/ucx_backend.h +++ b/src/plugins/ucx/ucx_backend.h @@ -186,12 +186,6 @@ class nixlUcxEngine : public nixlBackendEngine { nixl_status_t releaseReqH(nixlBackendReqH *handle) const override; - nixl_status_t - submitRmaWrite(const nixlMetaDesc &local, - const nixlMetaDesc &remote, - size_t size, - nixlBackendReqH *&handle) const; - unsigned progress(); @@ -268,6 +262,11 @@ class nixlUcxEngine : public nixlBackendEngine { size_t size, nixlBackendReqH *&handle) const; + nixl_status_t + submitProxyAtomicAdd(const nixlMetaDesc &remote, + uint64_t value, + nixlBackendReqH *&handle) const; + // Memory management helpers nixl_status_t internalMDHelper(const nixl_blob_t &blob, const std::string &agent, nixlBackendMD *&output); diff --git a/src/plugins/ucx/ucx_utils.cpp b/src/plugins/ucx/ucx_utils.cpp index a0ad10bd56..044c075e6c 100644 --- a/src/plugins/ucx/ucx_utils.cpp +++ b/src/plugins/ucx/ucx_utils.cpp @@ -313,6 +313,28 @@ nixlUcxEp::write(void *laddr, return nixl::ucx::ucsToNixlStatus(UCS_PTR_STATUS(request)); } +nixl_status_t +nixlUcxEp::atomicAdd(uint64_t value, uint64_t raddr, const nixl::ucx::rkey &rkey, nixlUcxReq &req) { + nixl_status_t status = checkTxState(); + if (status != NIXL_SUCCESS) { + return status; + } + + ucp_request_param_t param = { + .op_attr_mask = UCP_OP_ATTR_FIELD_DATATYPE, + .datatype = ucp_dt_make_contig(sizeof(uint64_t)), + }; + + const ucs_status_ptr_t request = + ucp_atomic_op_nbx(eph, UCP_ATOMIC_OP_ADD, &value, 1, raddr, rkey.get(), ¶m); + if (UCS_PTR_IS_PTR(request)) { + req = static_cast(request); + return NIXL_IN_PROG; + } + + return nixl::ucx::ucsToNixlStatus(UCS_PTR_STATUS(request)); +} + nixl_status_t nixlUcxEp::estimateCost(size_t size, std::chrono::microseconds &duration, diff --git a/src/plugins/ucx/ucx_utils.h b/src/plugins/ucx/ucx_utils.h index f22fd7b327..fd7cd46b6b 100644 --- a/src/plugins/ucx/ucx_utils.h +++ b/src/plugins/ucx/ucx_utils.h @@ -101,6 +101,8 @@ class nixlUcxEp { const nixl::ucx::rkey &rkey, size_t size, nixlUcxReq &req); + [[nodiscard]] nixl_status_t + atomicAdd(uint64_t value, uint64_t raddr, const nixl::ucx::rkey &rkey, nixlUcxReq &req); nixl_status_t estimateCost(size_t size, std::chrono::microseconds &duration, From 985e56ff1558bbcfba1923079b9e00dd4922706c Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Tue, 28 Apr 2026 20:16:45 +0300 Subject: [PATCH 11/18] device_proxy: Add tests for atomicAdd And make sure it is reflected in proxy memview registry as it only requires remote memview Signed-off-by: Tomer Davidor --- test/gtest/device_api/proxy_write_test.cu | 177 ++++++++++++++++++ .../proxy_memview_registry.cpp | 29 +++ .../unit/proxy_runtime/proxy_runtime.cpp | 66 +++++++ 3 files changed, 272 insertions(+) diff --git a/test/gtest/device_api/proxy_write_test.cu b/test/gtest/device_api/proxy_write_test.cu index 863cf5152d..5fa6e6f6d8 100644 --- a/test/gtest/device_api/proxy_write_test.cu +++ b/test/gtest/device_api/proxy_write_test.cu @@ -95,6 +95,7 @@ public: token = next_token_++; pending_.insert(token); token_channel_[token] = submission.channel_id; + submitted_opcodes_.push_back(submission.opcode); return NIXL_SUCCESS; } @@ -176,12 +177,20 @@ public: return false; } + std::vector + submittedOpcodes() const + { + std::lock_guard lk(mu_); + return submitted_opcodes_; + } + private: mutable std::mutex mu_; uint64_t next_token_ = 1; std::set pending_; std::map completed_; std::map token_channel_; + std::vector submitted_opcodes_; }; // --------------------------------------------------------------------------- @@ -289,6 +298,15 @@ proxyPutKernel(nixlMemViewH src_mvh, *out_status = nixlPut(src, dst, /*size=*/0); } +__global__ void +proxyAtomicAddKernel(nixlMemViewH counter_mvh, + uint64_t value, + nixl_status_t *out_status) +{ + nixlMemViewElem counter{counter_mvh, 0, 0}; + *out_status = nixlAtomicAdd(value, counter); +} + static void publishProxyContext(nixlProxyRuntime &runtime) { @@ -437,6 +455,29 @@ TEST_F(ProxyDeviceApiTest, PutReturnsInProgWhenEnqueued) ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); } +TEST_F(ProxyDeviceApiTest, AtomicAddReturnsInProgWhenEnqueued) +{ + auto adapter = std::make_unique(); + nixlProxyRuntime runtime; + + ASSERT_EQ(runtime.init(std::move(adapter), /*channel_count=*/1, /*worker_count=*/1), + NIXL_SUCCESS); + ASSERT_EQ(runtime.startWorkers(), NIXL_SUCCESS); + publishProxyContext(runtime); + const auto mvhs = registerDummyMemViews(runtime); + + nixl_status_t *d_status = deviceAlloc(); + proxyAtomicAddKernel<<<1, 1>>>(mvhs.dst, 42, d_status); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + + EXPECT_EQ(deviceGet(d_status), NIXL_IN_PROG); + cudaFree(d_status); + + clearProxyContext(); + ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); +} + // --------------------------------------------------------------------------- // Completion round-trip kernels // @@ -464,6 +505,25 @@ proxyPutAndPollKernel(nixlMemViewH src_mvh, nixlMemViewH dst_mvh, *out_poll_status = poll; } +__global__ void +proxyAtomicAddAndPollKernel(nixlMemViewH counter_mvh, + uint64_t value, + uint32_t channel_id, + nixl_status_t *out_atomic_status, + nixl_status_t *out_poll_status) +{ + nixlMemViewElem counter{counter_mvh, 0, 0}; + nixlGpuXferStatusH xfer_status{}; + *out_atomic_status = nixlAtomicAdd(value, counter, channel_id, + /*flags=*/0, &xfer_status); + + nixl_status_t poll; + do { + poll = nixlGpuGetXferStatus(xfer_status); + } while (poll == NIXL_IN_PROG); + *out_poll_status = poll; +} + // Enqueues a put and immediately returns; saves xfer_status to device memory // so the test thread can later launch a poll kernel. __global__ void @@ -477,6 +537,18 @@ proxyPutAsyncKernel(nixlMemViewH src_mvh, nixlMemViewH dst_mvh, /*flags=*/0, out_xfer_status); } +__global__ void +proxyAtomicAddAsyncKernel(nixlMemViewH counter_mvh, + uint64_t value, + uint32_t channel_id, + nixl_status_t *out_atomic_status, + nixlGpuXferStatusH *out_xfer_status) +{ + nixlMemViewElem counter{counter_mvh, 0, 0}; + *out_atomic_status = nixlAtomicAdd(value, counter, channel_id, + /*flags=*/0, out_xfer_status); +} + // Enqueues op_count puts on one channel and records each immediate enqueue // status. The final submission may block if the ring is full. __global__ void @@ -580,6 +652,35 @@ TEST_F(ProxyDeviceApiTest, PutCompletionRoundTrip) ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); } +TEST_F(ProxyDeviceApiTest, AtomicAddCompletionRoundTrip) +{ + auto adapter = std::make_unique(); + nixlProxyRuntime runtime; + + ASSERT_EQ(runtime.init(std::move(adapter), /*channel_count=*/1, /*worker_count=*/1), + NIXL_SUCCESS); + ASSERT_EQ(runtime.startWorkers(), NIXL_SUCCESS); + publishProxyContext(runtime); + + const auto mvhs = registerDummyMemViews(runtime); + + nixl_status_t *d_atomic_status = deviceAlloc(); + nixl_status_t *d_poll_status = deviceAlloc(); + + proxyAtomicAddAndPollKernel<<<1, 1>>>(mvhs.dst, 42, 0, + d_atomic_status, d_poll_status); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + + EXPECT_EQ(deviceGet(d_atomic_status), NIXL_IN_PROG); + EXPECT_EQ(deviceGet(d_poll_status), NIXL_SUCCESS); + + cudaFree(d_atomic_status); + cudaFree(d_poll_status); + clearProxyContext(); + ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); +} + // Verifies that the GPU kernel stays spinning until the test thread // explicitly marks the backend token complete. TEST_F(ProxyDeviceApiTest, CompletionNotVisibleUntilPublished) @@ -691,6 +792,82 @@ TEST_F(ProxyDeviceApiTest, MultipleSubmissionsCompletionFrontier) ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); } +TEST_F(ProxyDeviceApiTest, PutPutAtomicAddCompletionFrontier) +{ + auto adapter_owner = std::make_unique(); + auto *adapter = adapter_owner.get(); + nixlProxyRuntime runtime; + + ASSERT_EQ(runtime.init(std::move(adapter_owner), /*channel_count=*/1, /*worker_count=*/1), + NIXL_SUCCESS); + ASSERT_EQ(runtime.startWorkers(), NIXL_SUCCESS); + publishProxyContext(runtime); + + const auto mvhs = registerDummyMemViews(runtime); + + constexpr int kOps = 3; + nixl_status_t *d_submit_status[kOps]; + nixlGpuXferStatusH *d_xfer_status[kOps]; + + for (int i = 0; i < kOps; i++) { + d_submit_status[i] = deviceAlloc(); + ASSERT_EQ(cudaMalloc(&d_xfer_status[i], sizeof(nixlGpuXferStatusH)), + cudaSuccess); + ASSERT_EQ(cudaMemset(d_xfer_status[i], 0, sizeof(nixlGpuXferStatusH)), + cudaSuccess); + } + + proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_submit_status[0], + d_xfer_status[0]); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + EXPECT_EQ(deviceGet(d_submit_status[0]), NIXL_IN_PROG); + + proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_submit_status[1], + d_xfer_status[1]); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + EXPECT_EQ(deviceGet(d_submit_status[1]), NIXL_IN_PROG); + + proxyAtomicAddAsyncKernel<<<1, 1>>>(mvhs.dst, 42, 0, d_submit_status[2], + d_xfer_status[2]); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + EXPECT_EQ(deviceGet(d_submit_status[2]), NIXL_IN_PROG); + + ASSERT_TRUE(waitForCondition([adapter]() { + return adapter->pendingCount() == kOps; + })); + EXPECT_EQ(adapter->submittedOpcodes(), + std::vector({nixl_proxy_opcode_t::PUT, + nixl_proxy_opcode_t::PUT, + nixl_proxy_opcode_t::ATOMIC_ADD})); + + nixl_status_t *d_poll = deviceAlloc(); + for (int i = 0; i < kOps; i++) { + proxyPollOnceKernel<<<1, 1>>>(d_xfer_status[i], d_poll); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + EXPECT_EQ(deviceGet(d_poll), NIXL_IN_PROG) + << "op " << i << " should be in-progress before any markComplete"; + } + + for (int i = 0; i < kOps; i++) { + adapter->markComplete(static_cast(i + 1)); + ASSERT_TRUE(waitForCondition([&]() { + proxyPollOnceKernel<<<1, 1>>>(d_xfer_status[i], d_poll); + return cudaDeviceSynchronize() == cudaSuccess && deviceGet(d_poll) == NIXL_SUCCESS; + })) << "op " << i << " should complete after markComplete"; + } + + cudaFree(d_poll); + for (int i = 0; i < kOps; i++) { + cudaFree(d_submit_status[i]); + cudaFree(d_xfer_status[i]); + } + clearProxyContext(); + ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); +} + // Once a later op publishes an error, an earlier op whose op_idx is already // behind the completion frontier must still observe NIXL_SUCCESS. TEST_F(ProxyDeviceApiTest, EarlierCompletionStaysSuccessfulAfterLaterError) diff --git a/test/gtest/unit/proxy_memview_registry/proxy_memview_registry.cpp b/test/gtest/unit/proxy_memview_registry/proxy_memview_registry.cpp index aed0565ae1..2ae3953d22 100644 --- a/test/gtest/unit/proxy_memview_registry/proxy_memview_registry.cpp +++ b/test/gtest/unit/proxy_memview_registry/proxy_memview_registry.cpp @@ -368,6 +368,35 @@ namespace proxy_memview_registry { NIXL_ERR_INVALID_PARAM); } + TEST_F(ProxyMemViewRegistryTest, ReadyRemoteEntryProducesAtomicPreparedDescriptor) { + nixlMemViewH dst_proxy = nullptr; + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(dst_proxy, makeRemoteMetadata(0x2000, "remote-agent")), + NIXL_SUCCESS); + + nixlProxySubmission submission{}; + submission.opcode = nixl_proxy_opcode_t::ATOMIC_ADD; + submission.op_idx = 7; + submission.channel_id = 3; + submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + submission.dst_offset = 9; + submission.size = sizeof(uint64_t); + submission.value = 42; + + nixlBackendProxySubmission prepared_submission; + ASSERT_EQ(registry_.prepareSubmission(submission, prepared_submission), NIXL_SUCCESS); + EXPECT_EQ(prepared_submission.opcode, nixl_proxy_opcode_t::ATOMIC_ADD); + EXPECT_EQ(prepared_submission.op_idx, 7u); + EXPECT_EQ(prepared_submission.channel_id, 3u); + EXPECT_EQ(prepared_submission.remote.mem_type, DRAM_SEG); + EXPECT_EQ(prepared_submission.remote.desc.addr, 0x2009u); + EXPECT_EQ(prepared_submission.remote.desc.len, sizeof(uint64_t)); + EXPECT_EQ(prepared_submission.remote.desc.metadataP, &remote_md_); + EXPECT_EQ(prepared_submission.remote_agent, "remote-agent"); + EXPECT_EQ(prepared_submission.value, 42u); + } + TEST_F(ProxyMemViewRegistryTest, MetadataKindMustMatchSubmissionRole) { nixlMemViewH src_proxy = nullptr; nixlMemViewH dst_proxy = nullptr; diff --git a/test/gtest/unit/proxy_runtime/proxy_runtime.cpp b/test/gtest/unit/proxy_runtime/proxy_runtime.cpp index 4442b14ce7..236eb1d408 100644 --- a/test/gtest/unit/proxy_runtime/proxy_runtime.cpp +++ b/test/gtest/unit/proxy_runtime/proxy_runtime.cpp @@ -414,5 +414,71 @@ TEST_F(ProxyRuntimeTest, WorkerSubmitsPreparedTransportDescriptors) { EXPECT_EQ(prepared.remote_agent, "peer"); } +TEST_F(ProxyRuntimeTest, WorkerSubmitsPreparedAtomicAddDescriptor) { + DummyBackendMD remote_md; + + ASSERT_EQ(initRuntime(1, 1), NIXL_SUCCESS); + + nixlMemViewH dst_proxy = nullptr; + ASSERT_EQ(runtime_.registerProxyMemView(reinterpret_cast(uintptr_t{0x20}), + &dst_proxy), + NIXL_SUCCESS); + + nixl_remote_meta_dlist_t remote_dlist(DRAM_SEG); + nixlRemoteMetaDesc remote_desc("peer"); + remote_desc.addr = 0x2000; + remote_desc.len = 64; + remote_desc.devId = 0; + remote_desc.metadataP = &remote_md; + remote_dlist.addDesc(remote_desc); + ASSERT_EQ(runtime_.storeMetadata(dst_proxy, remote_dlist), NIXL_SUCCESS); + + ASSERT_EQ(runtime_.startWorkers(), NIXL_SUCCESS); + + nixlProxySubmission submission{}; + submission.op_idx = 11; + submission.opcode = nixl_proxy_opcode_t::ATOMIC_ADD; + submission.channel_id = 0; + submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); + submission.dst_offset = 8; + submission.size = sizeof(uint64_t); + submission.value = 42; + + auto *ring = runtime_.deviceChannelViews()[0].work_ring; + ring->records[0] = submission; + __atomic_store_n(&ring->records[0].ready_flag, 1u, __ATOMIC_RELEASE); + + const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(250); + while (std::chrono::steady_clock::now() < deadline) { + { + std::lock_guard lock(backend_->submit_mutex_); + if (!backend_->submissions_.empty()) { + break; + } + } + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + + std::vector submissions; + { + std::lock_guard lock(backend_->submit_mutex_); + submissions = backend_->submissions_; + } + + ASSERT_EQ(runtime_.shutdown(), NIXL_SUCCESS); + + ASSERT_EQ(submissions.size(), 1u); + const auto &prepared = submissions.front(); + EXPECT_EQ(prepared.op_idx, 11u); + EXPECT_EQ(prepared.opcode, nixl_proxy_opcode_t::ATOMIC_ADD); + EXPECT_EQ(prepared.channel_id, 0u); + EXPECT_EQ(prepared.remote.mem_type, DRAM_SEG); + EXPECT_EQ(prepared.remote.desc.addr, 0x2008u); + EXPECT_EQ(prepared.remote.desc.len, sizeof(uint64_t)); + EXPECT_EQ(prepared.remote.desc.metadataP, &remote_md); + EXPECT_EQ(prepared.remote_agent, "peer"); + EXPECT_EQ(prepared.value, 42u); +} + } // namespace proxy_runtime } // namespace gtest From 163aab8359a44b3b15be8a341d05ed612708c190 Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Mon, 8 Jun 2026 18:37:35 +0300 Subject: [PATCH 12/18] proxy_protocol: use op_idx as readiness word Publish proxy submissions by writing the payload first and committing op_idx last, then have the CPU worker acquire-poll op_idx instead of ready_flag. Signed-off-by: Tomer Davidor --- src/api/gpu/proxy/nixl_device_proxy.cuh | 22 +++++++++---------- src/core/device_proxy/proxy_protocol.h | 1 - src/core/device_proxy/proxy_worker.cpp | 9 +++++--- .../unit/proxy_runtime/proxy_runtime.cpp | 6 +++-- 4 files changed, 20 insertions(+), 18 deletions(-) diff --git a/src/api/gpu/proxy/nixl_device_proxy.cuh b/src/api/gpu/proxy/nixl_device_proxy.cuh index 7a8d90885f..4561accc8d 100644 --- a/src/api/gpu/proxy/nixl_device_proxy.cuh +++ b/src/api/gpu/proxy/nixl_device_proxy.cuh @@ -112,10 +112,6 @@ struct ProxyDeviceContext : nixlProxyDeviceContextData { // Enqueue a transfer submission into the MPSC work ring for the selected // channel, spinning if the ring is full. Optionally records a completion // token in *xfer_status for later polling via pollXferStatus(). - // - // producer_idx lives in HBM; consumer_idx lives in pinned host memory - // (accessible from device via UVA mapped pointer). Both are accessed with - // system-scope atomics so the CPU proxy worker sees the update coherently. __device__ inline nixl_status_t enqueue(nixlProxySubmission submission, nixlGpuXferStatusH *xfer_status = nullptr) { if (submission.channel_id >= num_channels) { @@ -140,19 +136,21 @@ struct ProxyDeviceContext : nixlProxyDeviceContextData { } } - cuda::atomic_ref op_idx(ring->running_op_idx); - submission.op_idx = op_idx.fetch_add(1, cuda::memory_order_relaxed); - ring->records[my_slot % ring->depth] = submission; + cuda::atomic_ref running_op_idx( + ring->running_op_idx); + const uint64_t submission_op_idx = + running_op_idx.fetch_add(1, cuda::memory_order_relaxed); + const uint32_t slot = my_slot % ring->depth; // Signal this slot is ready for the consumer. The release // guarantees the record write above is visible before the - // consumer reads it via an acquire load on ready_flag. - cuda::atomic_ref ready( - ring->records[my_slot % ring->depth].ready_flag); - ready.store(1, cuda::memory_order_release); + // consumer reads op_idx via an acquire load. op_idx == 0 means empty. + cuda::atomic_ref record_op_idx( + ring->records[slot].op_idx); + record_op_idx.store(submission_op_idx, cuda::memory_order_release); if (xfer_status != nullptr) { - ProxyXferStatus pxs{channel_view.completion_slot, submission.op_idx}; + ProxyXferStatus pxs{channel_view.completion_slot, submission_op_idx}; memcpy(xfer_status->storage, &pxs, sizeof(ProxyXferStatus)); } diff --git a/src/core/device_proxy/proxy_protocol.h b/src/core/device_proxy/proxy_protocol.h index 7f0c36c35f..aae3dff628 100644 --- a/src/core/device_proxy/proxy_protocol.h +++ b/src/core/device_proxy/proxy_protocol.h @@ -46,7 +46,6 @@ struct nixlProxySubmission { size_t dst_index = 0; size_t dst_offset = 0; - uint32_t ready_flag = 0; size_t size = 0; uint64_t value = 0; }; diff --git a/src/core/device_proxy/proxy_worker.cpp b/src/core/device_proxy/proxy_worker.cpp index 83e32a6f3d..fb98110c1e 100644 --- a/src/core/device_proxy/proxy_worker.cpp +++ b/src/core/device_proxy/proxy_worker.cpp @@ -83,14 +83,17 @@ ProxyWorker::tryDequeue(nixlProxyChannelState &channel, nixlProxySubmission &sub uint32_t local_consumer_idx = __atomic_load_n(channel.consumer_idx_host_, __ATOMIC_RELAXED); uint32_t slot = local_consumer_idx % ring->depth; - // ready_flag is the GPU-to-CPU signal that the record is written + // op_idx is the GPU-to-CPU signal that the record is written // (pairs with release store in device enqueue). No producer_idx // read on host — it is GPU-internal for slot allocation. - if (!__atomic_load_n(&ring->records[slot].ready_flag, __ATOMIC_ACQUIRE)) { + const uint64_t op_idx = __atomic_load_n(&ring->records[slot].op_idx, __ATOMIC_ACQUIRE); + if (op_idx == 0) { return false; } + __atomic_thread_fence(__ATOMIC_ACQUIRE); submission = ring->records[slot]; - __atomic_store_n(&ring->records[slot].ready_flag, 0, __ATOMIC_RELAXED); + submission.op_idx = op_idx; + __atomic_store_n(&ring->records[slot].op_idx, 0, __ATOMIC_RELAXED); __atomic_store_n(channel.consumer_idx_host_, local_consumer_idx + 1, __ATOMIC_RELEASE); diff --git a/test/gtest/unit/proxy_runtime/proxy_runtime.cpp b/test/gtest/unit/proxy_runtime/proxy_runtime.cpp index 236eb1d408..ae51ea060d 100644 --- a/test/gtest/unit/proxy_runtime/proxy_runtime.cpp +++ b/test/gtest/unit/proxy_runtime/proxy_runtime.cpp @@ -377,8 +377,9 @@ TEST_F(ProxyRuntimeTest, WorkerSubmitsPreparedTransportDescriptors) { submission.size = 32; auto *ring = runtime_.deviceChannelViews()[0].work_ring; + submission.op_idx = 0; ring->records[0] = submission; - __atomic_store_n(&ring->records[0].ready_flag, 1u, __ATOMIC_RELEASE); + __atomic_store_n(&ring->records[0].op_idx, uint64_t{11}, __ATOMIC_RELEASE); const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(250); while (std::chrono::steady_clock::now() < deadline) { @@ -445,8 +446,9 @@ TEST_F(ProxyRuntimeTest, WorkerSubmitsPreparedAtomicAddDescriptor) { submission.value = 42; auto *ring = runtime_.deviceChannelViews()[0].work_ring; + submission.op_idx = 0; ring->records[0] = submission; - __atomic_store_n(&ring->records[0].ready_flag, 1u, __ATOMIC_RELEASE); + __atomic_store_n(&ring->records[0].op_idx, uint64_t{11}, __ATOMIC_RELEASE); const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(250); while (std::chrono::steady_clock::now() < deadline) { From 0827f91dcac40592ecbc4be62e1fa2db9a984e3b Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Mon, 8 Jun 2026 18:38:21 +0300 Subject: [PATCH 13/18] proxy_protocol: pack submissions into 64-byte records Narrow protocol-only submission fields and assert the 64-byte record layout while keeping op_idx at the first word for readiness publication. Signed-off-by: Tomer Davidor --- src/api/gpu/proxy/nixl_device_impl.cuh | 20 ++++---- src/core/device_proxy/proxy_protocol.h | 21 +++++---- .../proxy_memview_registry.cpp | 47 +++++++++++-------- 3 files changed, 50 insertions(+), 38 deletions(-) diff --git a/src/api/gpu/proxy/nixl_device_impl.cuh b/src/api/gpu/proxy/nixl_device_impl.cuh index 2db9261d4b..1096866a14 100644 --- a/src/api/gpu/proxy/nixl_device_impl.cuh +++ b/src/api/gpu/proxy/nixl_device_impl.cuh @@ -75,14 +75,14 @@ put(const nixlMemViewElem &src, nixlProxySubmission submission{}; submission.opcode = nixl_proxy_opcode_t::PUT; submission.channel_id = static_cast(channel_id); - submission.flags = flags; + submission.flags = static_cast(flags); submission.src_proxy_memview_id = proxyMemViewIdFromHandle(src.mvh); - submission.src_index = src.index; - submission.src_offset = src.offset; + submission.src_index = static_cast(src.index); + submission.src_offset = static_cast(src.offset); submission.dst_proxy_memview_id = proxyMemViewIdFromHandle(dst.mvh); - submission.dst_index = dst.index; - submission.dst_offset = dst.offset; - submission.size = size; + submission.dst_index = static_cast(dst.index); + submission.dst_offset = static_cast(dst.offset); + submission.size = static_cast(size); status = ctx->enqueue(submission, xfer_status); } } @@ -108,11 +108,11 @@ atomic_add(uint64_t value, nixlProxySubmission submission{}; submission.opcode = nixl_proxy_opcode_t::ATOMIC_ADD; submission.channel_id = static_cast(channel_id); - submission.flags = flags; + submission.flags = static_cast(flags); submission.dst_proxy_memview_id = proxyMemViewIdFromHandle(counter.mvh); - submission.dst_index = counter.index; - submission.dst_offset = counter.offset; - submission.size = sizeof(uint64_t); + submission.dst_index = static_cast(counter.index); + submission.dst_offset = static_cast(counter.offset); + submission.size = static_cast(sizeof(uint64_t)); submission.value = value; status = ctx->enqueue(submission, xfer_status); } diff --git a/src/core/device_proxy/proxy_protocol.h b/src/core/device_proxy/proxy_protocol.h index aae3dff628..5435bcf940 100644 --- a/src/core/device_proxy/proxy_protocol.h +++ b/src/core/device_proxy/proxy_protocol.h @@ -32,24 +32,27 @@ enum class nixl_proxy_control_state_t : uint32_t { SHUTDOWN = 1, }; -struct nixlProxySubmission { +struct alignas(64) nixlProxySubmission { uint64_t op_idx = 0; nixl_proxy_opcode_t opcode = nixl_proxy_opcode_t::PUT; uint32_t channel_id = 0; - uint64_t flags = 0; + uint32_t flags = 0; - uint64_t src_proxy_memview_id = 0; - size_t src_index = 0; - size_t src_offset = 0; + uint32_t src_index = 0; + uint32_t src_offset = 0; + uint32_t dst_index = 0; + uint32_t dst_offset = 0; + uint32_t size = 0; + uint64_t src_proxy_memview_id = 0; uint64_t dst_proxy_memview_id = 0; - size_t dst_index = 0; - size_t dst_offset = 0; - - size_t size = 0; uint64_t value = 0; }; +static_assert(sizeof(nixlProxySubmission) == 64, "nixlProxySubmission must be 64 bytes"); +static_assert(offsetof(nixlProxySubmission, op_idx) == 0, + "op_idx must be the first word because it publishes record readiness"); + struct nixlProxyWorkRing { /** Host-accessible (e.g. cudaMallocHost); GPU may read via mapped pointer if needed. */ nixlProxySubmission *records = nullptr; diff --git a/test/gtest/unit/proxy_memview_registry/proxy_memview_registry.cpp b/test/gtest/unit/proxy_memview_registry/proxy_memview_registry.cpp index 2ae3953d22..ca97e94d52 100644 --- a/test/gtest/unit/proxy_memview_registry/proxy_memview_registry.cpp +++ b/test/gtest/unit/proxy_memview_registry/proxy_memview_registry.cpp @@ -134,8 +134,10 @@ namespace proxy_memview_registry { TEST_F(ProxyMemViewRegistryTest, PrepareSubmissionRequiresReadyEntries) { nixlMemViewH src_proxy = nullptr; nixlMemViewH dst_proxy = nullptr; - ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), NIXL_SUCCESS); - ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), + NIXL_SUCCESS); nixlProxySubmission submission{}; submission.opcode = nixl_proxy_opcode_t::PUT; @@ -150,8 +152,10 @@ namespace proxy_memview_registry { TEST_F(ProxyMemViewRegistryTest, ReadyEntriesProducePreparedTransportDescriptors) { nixlMemViewH src_proxy = nullptr; nixlMemViewH dst_proxy = nullptr; - ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), NIXL_SUCCESS); - ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), + NIXL_SUCCESS); ASSERT_EQ(registry_.storeMetadata(src_proxy, makeLocalMetadata(0x1000)), NIXL_SUCCESS); ASSERT_EQ(registry_.storeMetadata(dst_proxy, makeRemoteMetadata(0x2000, "remote-agent")), NIXL_SUCCESS); @@ -184,10 +188,8 @@ namespace proxy_memview_registry { TEST_F(ProxyMemViewRegistryTest, PrepMemViewProducesReadyEntries) { nixlMemViewH src_proxy = nullptr; nixlMemViewH dst_proxy = nullptr; - ASSERT_EQ(registry_.prepMemView(makeLocalMetadata(0x1000), &src_proxy), - NIXL_SUCCESS); - ASSERT_EQ(registry_.prepMemView(makeRemoteMetadata(0x2000), &dst_proxy), - NIXL_SUCCESS); + ASSERT_EQ(registry_.prepMemView(makeLocalMetadata(0x1000), &src_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.prepMemView(makeRemoteMetadata(0x2000), &dst_proxy), NIXL_SUCCESS); nixlMemViewH resolved = makeFakeBackendHandle(42); EXPECT_TRUE(registry_.resolveProxyMemView(src_proxy, resolved)); @@ -297,7 +299,7 @@ namespace proxy_memview_registry { submission.opcode = nixl_proxy_opcode_t::PUT; submission.src_proxy_memview_id = reinterpret_cast(src_proxy); submission.dst_proxy_memview_id = reinterpret_cast(dst_proxy); - submission.dst_offset = std::numeric_limits::max(); + submission.dst_offset = std::numeric_limits::max(); submission.size = 1; nixlBackendProxySubmission prepared_submission; @@ -329,8 +331,7 @@ namespace proxy_memview_registry { NIXL_SUCCESS); ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), NIXL_SUCCESS); - ASSERT_EQ(registry_.storeMetadata(src_proxy, makeLocalMetadata(0x1000, 7)), - NIXL_SUCCESS); + ASSERT_EQ(registry_.storeMetadata(src_proxy, makeLocalMetadata(0x1000, 7)), NIXL_SUCCESS); ASSERT_EQ(registry_.storeMetadata(dst_proxy, makeRemoteMetadata(0x2000, "peer", 11)), NIXL_SUCCESS); @@ -400,8 +401,10 @@ namespace proxy_memview_registry { TEST_F(ProxyMemViewRegistryTest, MetadataKindMustMatchSubmissionRole) { nixlMemViewH src_proxy = nullptr; nixlMemViewH dst_proxy = nullptr; - ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), NIXL_SUCCESS); - ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), + NIXL_SUCCESS); ASSERT_EQ(registry_.storeMetadata(src_proxy, makeRemoteMetadata(0x1000)), NIXL_SUCCESS); ASSERT_EQ(registry_.storeMetadata(dst_proxy, makeLocalMetadata(0x2000)), NIXL_SUCCESS); @@ -420,9 +423,12 @@ namespace proxy_memview_registry { nixlMemViewH src_proxy = nullptr; nixlMemViewH dst_proxy = nullptr; nixlMemViewH other_proxy = nullptr; - ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), NIXL_SUCCESS); - ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), NIXL_SUCCESS); - ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(30), &other_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &src_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &dst_proxy), + NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(30), &other_proxy), + NIXL_SUCCESS); ASSERT_EQ(registry_.storeMetadata(src_proxy, makeLocalMetadata(0x1000)), NIXL_SUCCESS); ASSERT_EQ(registry_.storeMetadata(dst_proxy, makeRemoteMetadata(0x2000)), NIXL_SUCCESS); ASSERT_EQ(registry_.storeMetadata(other_proxy, makeRemoteMetadata(0x3000)), NIXL_SUCCESS); @@ -451,7 +457,8 @@ namespace proxy_memview_registry { TEST_F(ProxyMemViewRegistryTest, ClearRetiresExistingEntriesAndPreservesFreshIds) { nixlMemViewH old_proxy = nullptr; - ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &old_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &old_proxy), + NIXL_SUCCESS); ASSERT_EQ(registry_.storeMetadata(old_proxy, makeLocalMetadata(0x1000)), NIXL_SUCCESS); registry_.clear(); @@ -460,7 +467,8 @@ namespace proxy_memview_registry { EXPECT_FALSE(registry_.resolveProxyMemView(old_proxy, resolved)); nixlMemViewH new_proxy = nullptr; - ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &new_proxy), NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(20), &new_proxy), + NIXL_SUCCESS); EXPECT_NE(old_proxy, new_proxy); EXPECT_TRUE(registry_.resolveProxyMemView(new_proxy, resolved)); EXPECT_EQ(resolved, makeFakeBackendHandle(20)); @@ -468,7 +476,8 @@ namespace proxy_memview_registry { TEST_F(ProxyMemViewRegistryTest, StoreMetadataRejectsRetiredEntries) { nixlMemViewH proxy_handle = nullptr; - ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &proxy_handle), NIXL_SUCCESS); + ASSERT_EQ(registry_.registerProxyMemView(makeFakeBackendHandle(10), &proxy_handle), + NIXL_SUCCESS); ASSERT_EQ(registry_.unregisterProxyMemView(proxy_handle), NIXL_SUCCESS); EXPECT_EQ(registry_.storeMetadata(proxy_handle, makeLocalMetadata(0x1000)), NIXL_ERR_NOT_FOUND); From 85b4e612d6bfe7fa1f132892693093501db3afed Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Thu, 7 May 2026 11:02:44 +0300 Subject: [PATCH 14/18] gpu: proxy: inline device helpers and publish context from constant memory Publish the proxy context through constant memory and force-inline device helper paths used by proxy API calls. Signed-off-by: Tomer Davidor --- src/api/gpu/proxy/nixl_device_impl.cuh | 55 ++++++++++++------------- src/api/gpu/proxy/nixl_device_proxy.cu | 2 +- src/api/gpu/proxy/nixl_device_proxy.cuh | 9 ++-- 3 files changed, 32 insertions(+), 34 deletions(-) diff --git a/src/api/gpu/proxy/nixl_device_impl.cuh b/src/api/gpu/proxy/nixl_device_impl.cuh index 1096866a14..695e9d58d1 100644 --- a/src/api/gpu/proxy/nixl_device_impl.cuh +++ b/src/api/gpu/proxy/nixl_device_impl.cuh @@ -23,7 +23,7 @@ namespace nixl::gpu::proxy_impl { template -__device__ inline nixl_status_t +__device__ __forceinline__ nixl_status_t get_xfer_status(nixlGpuXferStatusH &xfer_status) { uint32_t lane_id; nixlProxyExecInit(lane_id); @@ -40,8 +40,7 @@ get_xfer_status(nixlGpuXferStatusH &xfer_status) { } if constexpr (level == nixl_gpu_level_t::WARP) { - status = static_cast( - __shfl_sync(0xffffffff, static_cast(status), 0)); + status = static_cast(__shfl_sync(0xffffffff, static_cast(status), 0)); } else if constexpr (level == nixl_gpu_level_t::BLOCK) { __shared__ nixl_status_t s_status; if (threadIdx.x == 0) { @@ -56,7 +55,7 @@ get_xfer_status(nixlGpuXferStatusH &xfer_status) { } template -__device__ inline nixl_status_t +__device__ __forceinline__ nixl_status_t put(const nixlMemViewElem &src, const nixlMemViewElem &dst, size_t size, @@ -72,18 +71,18 @@ put(const nixlMemViewElem &src, if (ctx == nullptr) { status = NIXL_ERR_BACKEND; } else { - nixlProxySubmission submission{}; - submission.opcode = nixl_proxy_opcode_t::PUT; - submission.channel_id = static_cast(channel_id); - submission.flags = static_cast(flags); - submission.src_proxy_memview_id = proxyMemViewIdFromHandle(src.mvh); - submission.src_index = static_cast(src.index); - submission.src_offset = static_cast(src.offset); - submission.dst_proxy_memview_id = proxyMemViewIdFromHandle(dst.mvh); - submission.dst_index = static_cast(dst.index); - submission.dst_offset = static_cast(dst.offset); - submission.size = static_cast(size); - status = ctx->enqueue(submission, xfer_status); + status = ctx->enqueue( + nixlProxySubmission{.opcode = nixl_proxy_opcode_t::PUT, + .channel_id = static_cast(channel_id), + .flags = static_cast(flags), + .src_index = static_cast(src.index), + .src_offset = static_cast(src.offset), + .dst_index = static_cast(dst.index), + .dst_offset = static_cast(dst.offset), + .size = static_cast(size), + .src_proxy_memview_id = proxyMemViewIdFromHandle(src.mvh), + .dst_proxy_memview_id = proxyMemViewIdFromHandle(dst.mvh)}, + xfer_status); } } nixlProxySync(); @@ -91,7 +90,7 @@ put(const nixlMemViewElem &src, } template -__device__ inline nixl_status_t +__device__ __forceinline__ nixl_status_t atomic_add(uint64_t value, const nixlMemViewElem &counter, unsigned channel_id = 0, @@ -105,23 +104,23 @@ atomic_add(uint64_t value, if (ctx == nullptr) { status = NIXL_ERR_BACKEND; } else { - nixlProxySubmission submission{}; - submission.opcode = nixl_proxy_opcode_t::ATOMIC_ADD; - submission.channel_id = static_cast(channel_id); - submission.flags = static_cast(flags); - submission.dst_proxy_memview_id = proxyMemViewIdFromHandle(counter.mvh); - submission.dst_index = static_cast(counter.index); - submission.dst_offset = static_cast(counter.offset); - submission.size = static_cast(sizeof(uint64_t)); - submission.value = value; - status = ctx->enqueue(submission, xfer_status); + status = ctx->enqueue( + nixlProxySubmission{.opcode = nixl_proxy_opcode_t::ATOMIC_ADD, + .channel_id = static_cast(channel_id), + .flags = static_cast(flags), + .dst_index = static_cast(counter.index), + .dst_offset = static_cast(counter.offset), + .size = static_cast(sizeof(uint64_t)), + .dst_proxy_memview_id = proxyMemViewIdFromHandle(counter.mvh), + .value = value}, + xfer_status); } } nixlProxySync(); return status; } -__device__ inline void * +__device__ __forceinline__ void * get_ptr(nixlMemViewH, size_t) { // TODO: Implement support for NVLink fast-path over proxy - NIX-1342 return nullptr; diff --git a/src/api/gpu/proxy/nixl_device_proxy.cu b/src/api/gpu/proxy/nixl_device_proxy.cu index 0e87cb2e5d..73e00258a0 100644 --- a/src/api/gpu/proxy/nixl_device_proxy.cu +++ b/src/api/gpu/proxy/nixl_device_proxy.cu @@ -19,4 +19,4 @@ // through load_proxy_context(). #include "nixl_device_proxy.cuh" -__device__ ProxyDeviceContext *g_nixl_proxy_ctx = nullptr; +__device__ __constant__ ProxyDeviceContext *g_nixl_proxy_ctx = nullptr; diff --git a/src/api/gpu/proxy/nixl_device_proxy.cuh b/src/api/gpu/proxy/nixl_device_proxy.cuh index 4561accc8d..76685e0b04 100644 --- a/src/api/gpu/proxy/nixl_device_proxy.cuh +++ b/src/api/gpu/proxy/nixl_device_proxy.cuh @@ -36,7 +36,7 @@ static_assert(sizeof(ProxyXferStatus) <= sizeof(nixlGpuXferStatusH), // Defined in nixl_device_proxy.cu and read by device kernels through // load_proxy_context(). -extern __device__ ProxyDeviceContext *g_nixl_proxy_ctx; +extern __device__ __constant__ ProxyDeviceContext *g_nixl_proxy_ctx; // Host-callable helpers. Keeping these inline in CUDA translation units avoids // cross-DSO symbol ownership issues for g_nixl_proxy_ctx. @@ -66,12 +66,12 @@ nixlProxyClearContext() { return err; } -__device__ inline uint64_t +__device__ __forceinline__ uint64_t proxyMemViewIdFromHandle(nixlMemViewH mvh) { return static_cast(reinterpret_cast(mvh)); } -__device__ inline ProxyDeviceContext * +__device__ __forceinline__ ProxyDeviceContext * load_proxy_context() { return g_nixl_proxy_ctx; } @@ -174,11 +174,10 @@ struct ProxyDeviceContext : nixlProxyDeviceContextData { pxs->slot->completed_idx); const uint64_t completed_idx = comp_idx.load(cuda::memory_order_acquire); - const nixl_status_t current_status = pxs->slot->next_status; - if (completed_idx > pxs->op_idx) { return NIXL_SUCCESS; } + const nixl_status_t current_status = pxs->slot->next_status; if (completed_idx == pxs->op_idx) { return current_status; } From 57e25079f004f0f749bd1f8cc4826ed8f2e96fa9 Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Mon, 4 May 2026 15:37:32 +0300 Subject: [PATCH 15/18] device_proxy: move producer index to device memory Move the GPU-owned producer index to device memory, derive operation ids from it, and keep mapped host aliases local to allocation. Signed-off-by: Tomer Davidor --- src/api/gpu/proxy/nixl_device_proxy.cuh | 28 +- src/core/device_proxy/proxy_protocol.h | 11 +- src/core/device_proxy/proxy_runtime.cpp | 171 ++++---- src/core/device_proxy/proxy_runtime.h | 26 +- src/core/device_proxy/proxy_worker.cpp | 52 +-- test/gtest/device_api/proxy_write_test.cu | 382 +++++++++--------- .../unit/proxy_runtime/proxy_runtime.cpp | 71 ++-- 7 files changed, 383 insertions(+), 358 deletions(-) diff --git a/src/api/gpu/proxy/nixl_device_proxy.cuh b/src/api/gpu/proxy/nixl_device_proxy.cuh index 76685e0b04..492abb116f 100644 --- a/src/api/gpu/proxy/nixl_device_proxy.cuh +++ b/src/api/gpu/proxy/nixl_device_proxy.cuh @@ -76,10 +76,12 @@ load_proxy_context() { return g_nixl_proxy_ctx; } -static_assert(sizeof(nixlProxyWorkRing::running_op_idx) == 8, - "running_op_idx must be 64-bit to avoid wrap-around false completions"); +static_assert(sizeof(*nixlProxyWorkRing{}.producer_idx) == 8, + "producer_idx must be 64-bit to avoid wrap-around false completions"); +static_assert(sizeof(*nixlProxyWorkRing{}.consumer_idx) == 8, + "consumer_idx must be 64-bit to match producer_idx"); static_assert(sizeof(nixlProxyCompletionSlot::completed_idx) == 8, - "completed_idx must be 64-bit to match running_op_idx"); + "completed_idx must be 64-bit to match producer_idx"); template __device__ inline void nixlProxyExecInit(uint32_t &lane_id) { @@ -112,6 +114,10 @@ struct ProxyDeviceContext : nixlProxyDeviceContextData { // Enqueue a transfer submission into the MPSC work ring for the selected // channel, spinning if the ring is full. Optionally records a completion // token in *xfer_status for later polling via pollXferStatus(). + // + // producer_idx lives in device memory and only needs device-scope atomicity. + // consumer_idx lives in pinned host memory (accessible from device via + // UVA mapped pointer), so full-ring polling uses system-scope atomics. __device__ inline nixl_status_t enqueue(nixlProxySubmission submission, nixlGpuXferStatusH *xfer_status = nullptr) { if (submission.channel_id >= num_channels) { @@ -121,26 +127,24 @@ struct ProxyDeviceContext : nixlProxyDeviceContextData { nixlProxyChannelView &channel_view = channels[submission.channel_id]; nixlProxyWorkRing *ring = channel_view.work_ring; - cuda::atomic_ref prod(*ring->producer_idx); - cuda::atomic_ref cons(*ring->consumer_idx); + cuda::atomic_ref producer_idx( + *ring->producer_idx); + cuda::atomic_ref cons(*ring->consumer_idx); cuda::atomic_ref shut(*shutdown_word); // Atomically claim a unique slot in the ring. - uint32_t my_slot = prod.fetch_add(1, cuda::memory_order_relaxed); + const uint64_t ticket = producer_idx.fetch_add(1, cuda::memory_order_relaxed); // Spin until the claimed slot has space (consumer has freed it). - while (my_slot - cons.load(cuda::memory_order_acquire) >= ring->depth) { + while (ticket - cons.load(cuda::memory_order_acquire) >= ring->depth) { if (shut.load(cuda::memory_order_acquire) == static_cast(nixl_proxy_control_state_t::SHUTDOWN)) { return NIXL_ERR_BACKEND; } } - cuda::atomic_ref running_op_idx( - ring->running_op_idx); - const uint64_t submission_op_idx = - running_op_idx.fetch_add(1, cuda::memory_order_relaxed); - const uint32_t slot = my_slot % ring->depth; + const uint64_t submission_op_idx = ticket + 1; + const uint32_t slot = static_cast(ticket % ring->depth); // Signal this slot is ready for the consumer. The release // guarantees the record write above is visible before the diff --git a/src/core/device_proxy/proxy_protocol.h b/src/core/device_proxy/proxy_protocol.h index 5435bcf940..5f3add710d 100644 --- a/src/core/device_proxy/proxy_protocol.h +++ b/src/core/device_proxy/proxy_protocol.h @@ -54,17 +54,14 @@ static_assert(offsetof(nixlProxySubmission, op_idx) == 0, "op_idx must be the first word because it publishes record readiness"); struct nixlProxyWorkRing { - /** Host-accessible (e.g. cudaMallocHost); GPU may read via mapped pointer if needed. */ + /** Mapped host records: GPU writes via device alias; CPU worker reads host alias. */ nixlProxySubmission *records = nullptr; - /** Mapped pinned producer; GPU advances with CUDA atomics; host reads via __atomic_*. */ - uint32_t *producer_idx = nullptr; + /** Device-resident producer index; only the GPU updates it. */ + uint64_t *producer_idx = nullptr; /** Mapped pinned consumer; host proxy uses __atomic_* on host alias (nixlProxyChannelState). */ - uint32_t *consumer_idx = nullptr; + uint64_t *consumer_idx = nullptr; /** The depth of the work ring. */ uint32_t depth = 0; - /** Monotonic 64-bit counter; starts at 1 so completed_idx==0 means - * "no operation completed yet" and the first op_idx is never 0. */ - uint64_t running_op_idx = 1; }; struct alignas(16) nixlProxyCompletionSlot { diff --git a/src/core/device_proxy/proxy_runtime.cpp b/src/core/device_proxy/proxy_runtime.cpp index 493070ae37..bcb0d41b7f 100644 --- a/src/core/device_proxy/proxy_runtime.cpp +++ b/src/core/device_proxy/proxy_runtime.cpp @@ -21,6 +21,7 @@ #include "nixl_log.h" #include #include +#include #include nixl_status_t @@ -393,32 +394,34 @@ nixl_status_t nixlProxyChannelState::allocate(uint32_t channel_id, uint32_t depth) { NIXL_INFO << "nixlProxyChannelState::allocate: channel_id=" << channel_id << " depth=" << depth; - if (cudaMallocHost(&work_ring_, sizeof(nixlProxyWorkRing)) != cudaSuccess - || cudaMallocHost(&records_, sizeof(nixlProxySubmission) * depth) != cudaSuccess + ring_depth_ = depth; + if (cudaMalloc(reinterpret_cast(&work_ring_dev_), + sizeof(nixlProxyWorkRing)) != cudaSuccess + || cudaMalloc(reinterpret_cast(&producer_idx_dev_), + sizeof(uint64_t)) != cudaSuccess + || cudaMallocHost(&records_host_, sizeof(nixlProxySubmission) * depth) != cudaSuccess || cudaMallocHost(reinterpret_cast(&consumer_idx_host_), - sizeof(uint32_t)) != cudaSuccess - || cudaMallocHost(reinterpret_cast(&producer_idx_host_), - sizeof(uint32_t)) != cudaSuccess - || cudaMallocHost(&completion_slot_host_, sizeof(nixlProxyCompletionSlot)) != cudaSuccess) { + sizeof(uint64_t)) != cudaSuccess + || cudaMallocHost(&completion_slot_host_, sizeof(nixlProxyCompletionSlot)) != cudaSuccess) { NIXL_ERROR << "nixlProxyChannelState::allocate: CUDA allocation failed for channel " << channel_id; deallocate(); return NIXL_ERR_BACKEND; } - void *consumer_dev = nullptr; - if (cudaHostGetDevicePointer(&consumer_dev, consumer_idx_host_, 0) != cudaSuccess) { + void *records_dev = nullptr; + if (cudaHostGetDevicePointer(&records_dev, records_host_, 0) != cudaSuccess) { deallocate(); return NIXL_ERR_BACKEND; } - consumer_idx_dev_ = static_cast(consumer_dev); + auto *records_dev_ptr = static_cast(records_dev); - void *producer_dev = nullptr; - if (cudaHostGetDevicePointer(&producer_dev, producer_idx_host_, 0) != cudaSuccess) { + void *consumer_dev = nullptr; + if (cudaHostGetDevicePointer(&consumer_dev, consumer_idx_host_, 0) != cudaSuccess) { deallocate(); return NIXL_ERR_BACKEND; } - producer_idx_dev_ = static_cast(producer_dev); + auto *consumer_idx_dev = static_cast(consumer_dev); void *completion_dev = nullptr; if (cudaHostGetDevicePointer(&completion_dev, completion_slot_host_, 0) != cudaSuccess) { @@ -428,29 +431,39 @@ nixlProxyChannelState::allocate(uint32_t channel_id, uint32_t depth) { completion_slot_dev_ = static_cast(completion_dev); for (uint32_t i = 0; i < depth; ++i) { - records_[i] = nixlProxySubmission{}; + records_host_[i] = nixlProxySubmission{}; } - __atomic_store_n(producer_idx_host_, 0, __ATOMIC_RELEASE); - __atomic_store_n(consumer_idx_host_, 0, __ATOMIC_RELEASE); + if (cudaMemset(producer_idx_dev_, 0, sizeof(*producer_idx_dev_)) != cudaSuccess) { + deallocate(); + return NIXL_ERR_BACKEND; + } + __atomic_store_n(consumer_idx_host_, uint64_t{0}, __ATOMIC_RELEASE); completion_slot_host_->next_status = NIXL_IN_PROG; __atomic_store_n(&completion_slot_host_->completed_idx, uint64_t{0}, __ATOMIC_RELEASE); - *work_ring_ = nixlProxyWorkRing{ - records_, + nixlProxyWorkRing work_ring{ + records_dev_ptr, producer_idx_dev_, - consumer_idx_dev_, + consumer_idx_dev, depth, }; - device_view = nixlProxyChannelView{ work_ring_, completion_slot_dev_, channel_id }; + if (cudaMemcpy(work_ring_dev_, + &work_ring, + sizeof(work_ring), + cudaMemcpyHostToDevice) != cudaSuccess) { + deallocate(); + return NIXL_ERR_BACKEND; + } + device_view = nixlProxyChannelView{ work_ring_dev_, completion_slot_dev_, channel_id }; inflight_requests.clear(); NIXL_INFO << "nixlProxyChannelState::allocate: channel " << channel_id << " ready" - << " work_ring=" << work_ring_ - << " records=" << records_ - << " producer_idx(host)=" << producer_idx_host_ + << " work_ring(dev)=" << work_ring_dev_ + << " records=" << records_host_ + << " records(dev)=" << records_dev_ptr << " producer_idx(dev)=" << producer_idx_dev_ << " consumer_idx(host)=" << consumer_idx_host_ - << " consumer_idx(dev)=" << consumer_idx_dev_ + << " consumer_idx(dev)=" << consumer_idx_dev << " completion_slot(host)=" << completion_slot_host_ << " completion_slot(dev)=" << completion_slot_dev_; return NIXL_SUCCESS; @@ -463,18 +476,23 @@ nixlProxyChannelState::deallocate() noexcept { completion_slot_host_ = nullptr; completion_slot_dev_ = nullptr; } - if (producer_idx_host_) { - cudaFreeHost(producer_idx_host_); - producer_idx_host_ = nullptr; - producer_idx_dev_ = nullptr; + if (producer_idx_dev_) { + cudaFree(producer_idx_dev_); + producer_idx_dev_ = nullptr; } if (consumer_idx_host_) { cudaFreeHost(consumer_idx_host_); consumer_idx_host_ = nullptr; - consumer_idx_dev_ = nullptr; } - if (records_) { cudaFreeHost(records_); records_ = nullptr; } - if (work_ring_) { cudaFreeHost(work_ring_); work_ring_ = nullptr; } + if (records_host_) { + cudaFreeHost(records_host_); + records_host_ = nullptr; + } + if (work_ring_dev_) { + cudaFree(work_ring_dev_); + work_ring_dev_ = nullptr; + } + ring_depth_ = 0; device_view = nixlProxyChannelView{}; } @@ -482,26 +500,8 @@ nixlProxyChannelState::~nixlProxyChannelState() { deallocate(); } -nixlProxyChannelState::nixlProxyChannelState(nixlProxyChannelState &&other) noexcept - : device_view(other.device_view), - inflight_requests(std::move(other.inflight_requests)), - work_ring_(other.work_ring_), - records_(other.records_), - producer_idx_host_(other.producer_idx_host_), - producer_idx_dev_(other.producer_idx_dev_), - consumer_idx_host_(other.consumer_idx_host_), - consumer_idx_dev_(other.consumer_idx_dev_), - completion_slot_host_(other.completion_slot_host_), - completion_slot_dev_(other.completion_slot_dev_) { - other.work_ring_ = nullptr; - other.records_ = nullptr; - other.producer_idx_host_ = nullptr; - other.producer_idx_dev_ = nullptr; - other.consumer_idx_host_ = nullptr; - other.consumer_idx_dev_ = nullptr; - other.completion_slot_host_ = nullptr; - other.completion_slot_dev_ = nullptr; - other.device_view = nixlProxyChannelView{}; +nixlProxyChannelState::nixlProxyChannelState(nixlProxyChannelState &&other) noexcept { + *this = std::move(other); } nixlProxyChannelState & @@ -510,20 +510,18 @@ nixlProxyChannelState::operator=(nixlProxyChannelState &&other) noexcept { deallocate(); device_view = other.device_view; inflight_requests = std::move(other.inflight_requests); - work_ring_ = other.work_ring_; - records_ = other.records_; - producer_idx_host_ = other.producer_idx_host_; - producer_idx_dev_ = other.producer_idx_dev_; + work_ring_dev_ = other.work_ring_dev_; + records_host_ = other.records_host_; + producer_idx_dev_ = other.producer_idx_dev_; consumer_idx_host_ = other.consumer_idx_host_; - consumer_idx_dev_ = other.consumer_idx_dev_; + ring_depth_ = other.ring_depth_; completion_slot_host_ = other.completion_slot_host_; completion_slot_dev_ = other.completion_slot_dev_; - other.work_ring_ = nullptr; - other.records_ = nullptr; - other.producer_idx_host_ = nullptr; - other.producer_idx_dev_ = nullptr; + other.work_ring_dev_ = nullptr; + other.records_host_ = nullptr; + other.producer_idx_dev_ = nullptr; other.consumer_idx_host_ = nullptr; - other.consumer_idx_dev_ = nullptr; + other.ring_depth_ = 0; other.completion_slot_host_ = nullptr; other.completion_slot_dev_ = nullptr; other.device_view = nixlProxyChannelView{}; @@ -606,8 +604,9 @@ nixlProxyRuntime::init(std::unique_ptr backend, } } - if (cudaMallocHost(&device_channel_views_, - sizeof(nixlProxyChannelView) * channel_count) != cudaSuccess) { + device_channel_views_.resize(channel_count); + if (cudaMalloc(reinterpret_cast(&device_channel_views_dev_), + sizeof(nixlProxyChannelView) * channel_count) != cudaSuccess) { channels_.clear(); backend_->shutdown(); cudaFreeHost(shutdown_word_host_); @@ -619,11 +618,13 @@ nixlProxyRuntime::init(std::unique_ptr backend, for (uint32_t channel_id = 0; channel_id < channel_count; ++channel_id) { device_channel_views_[channel_id] = channels_[channel_id].device_view; } - - if (cudaMallocHost(&device_context_, - sizeof(nixlProxyDeviceContextData)) != cudaSuccess) { - cudaFreeHost(device_channel_views_); - device_channel_views_ = nullptr; + if (cudaMemcpy(device_channel_views_dev_, + device_channel_views_.data(), + sizeof(nixlProxyChannelView) * channel_count, + cudaMemcpyHostToDevice) != cudaSuccess) { + cudaFree(device_channel_views_dev_); + device_channel_views_dev_ = nullptr; + device_channel_views_.clear(); channels_.clear(); backend_->shutdown(); cudaFreeHost(shutdown_word_host_); @@ -632,11 +633,32 @@ nixlProxyRuntime::init(std::unique_ptr backend, backend_.reset(); return NIXL_ERR_BACKEND; } - *device_context_ = nixlProxyDeviceContextData{ - device_channel_views_, + nixlProxyDeviceContextData device_context{ + device_channel_views_dev_, channel_count, shutdown_word_dev_ }; + if (cudaMalloc(reinterpret_cast(&device_context_), + sizeof(nixlProxyDeviceContextData)) != cudaSuccess + || cudaMemcpy(device_context_, + &device_context, + sizeof(device_context), + cudaMemcpyHostToDevice) != cudaSuccess) { + if (device_context_) { + cudaFree(device_context_); + device_context_ = nullptr; + } + cudaFree(device_channel_views_dev_); + device_channel_views_dev_ = nullptr; + device_channel_views_.clear(); + channels_.clear(); + backend_->shutdown(); + cudaFreeHost(shutdown_word_host_); + shutdown_word_host_ = nullptr; + shutdown_word_dev_ = nullptr; + backend_.reset(); + return NIXL_ERR_BACKEND; + } workers_.clear(); workers_.reserve(worker_count); @@ -661,7 +683,7 @@ nixlProxyRuntime::init(std::unique_ptr backend, NIXL_INFO << "ProxyRuntime::init: complete — " << channel_count << " channels, " << worker_count << " workers, " - << "device_context=" << device_context_; + << "device_context(dev)=" << device_context_; return NIXL_SUCCESS; } @@ -809,7 +831,7 @@ nixlProxyRuntime::shutdown() { memview_registry_.clear(); if (device_context_) { - cudaFreeHost(device_context_); + cudaFree(device_context_); device_context_ = nullptr; } if (shutdown_word_host_) { @@ -817,10 +839,11 @@ nixlProxyRuntime::shutdown() { shutdown_word_host_ = nullptr; shutdown_word_dev_ = nullptr; } - if (device_channel_views_) { - cudaFreeHost(device_channel_views_); - device_channel_views_ = nullptr; + if (device_channel_views_dev_) { + cudaFree(device_channel_views_dev_); + device_channel_views_dev_ = nullptr; } + device_channel_views_.clear(); channels_.clear(); backend_.reset(); diff --git a/src/core/device_proxy/proxy_runtime.h b/src/core/device_proxy/proxy_runtime.h index cdbe9e3b09..1c9e1eefd6 100644 --- a/src/core/device_proxy/proxy_runtime.h +++ b/src/core/device_proxy/proxy_runtime.h @@ -17,7 +17,6 @@ #ifndef NIXL_SRC_CORE_DEVICE_PROXY_PROXY_RUNTIME_H #define NIXL_SRC_CORE_DEVICE_PROXY_PROXY_RUNTIME_H -#include #include #include #include @@ -43,16 +42,14 @@ struct alignas(64) nixlProxyChannelState { std::deque inflight_requests; bool error_latched = false; - nixlProxyWorkRing *work_ring_ = nullptr; - nixlProxySubmission *records_ = nullptr; - /** Mapped pinned host memory; host proxy uses __atomic_* on host alias. */ - uint32_t *producer_idx_host_ = nullptr; - /** Device-mapped alias of producer_idx_host_ for nixlProxyWorkRing (GPU-writable). */ - uint32_t *producer_idx_dev_ = nullptr; + nixlProxyWorkRing *work_ring_dev_ = nullptr; + nixlProxySubmission *records_host_ = nullptr; + /** Device-resident producer index; only the GPU updates it. */ + uint64_t *producer_idx_dev_ = nullptr; /** Consumer count: host pinned; proxy uses __atomic_* on consumer_idx_host_. */ - uint32_t *consumer_idx_host_ = nullptr; - /** Same word as consumer_idx_host_, for nixlProxyWorkRing::consumer_idx (GPU-readable). */ - uint32_t *consumer_idx_dev_ = nullptr; + uint64_t *consumer_idx_host_ = nullptr; + /** Host-side ring depth for the CPU worker; nixlProxyWorkRing itself is device-only. */ + uint32_t ring_depth_ = 0; /** Mapped pinned host memory; proxy worker writes directly via host alias. */ nixlProxyCompletionSlot *completion_slot_host_ = nullptr; /** Device-mapped alias of completion_slot_host_ for nixlProxyChannelView. */ @@ -277,7 +274,9 @@ class nixlProxyRuntime { channelCount() const { return static_cast(channels_.size()); } const nixlProxyChannelView * - deviceChannelViews() const { return device_channel_views_; } + deviceChannelViews() const { + return device_channel_views_.empty() ? nullptr : device_channel_views_.data(); + } nixlProxyDeviceContextData * deviceContext() const { return device_context_; } @@ -287,8 +286,9 @@ class nixlProxyRuntime { joinWorkerThreads() noexcept; std::vector channels_; - nixlProxyChannelView *device_channel_views_ = nullptr; - nixlProxyDeviceContextData *device_context_ = nullptr; + std::vector device_channel_views_; + nixlProxyChannelView *device_channel_views_dev_ = nullptr; + nixlProxyDeviceContextData *device_context_ = nullptr; std::vector> workers_; nixlProxyMemViewRegistry memview_registry_; std::unique_ptr backend_; diff --git a/src/core/device_proxy/proxy_worker.cpp b/src/core/device_proxy/proxy_worker.cpp index fb98110c1e..f1db3219d6 100644 --- a/src/core/device_proxy/proxy_worker.cpp +++ b/src/core/device_proxy/proxy_worker.cpp @@ -42,8 +42,8 @@ void ProxyWorker::start(uint32_t worker_idx) { thread_ = std::thread([this, worker_idx]() { NIXL_INFO << "ProxyWorker thread " << worker_idx << " started"; - while (__atomic_load_n(shutdown_word_, __ATOMIC_ACQUIRE) - == static_cast(nixl_proxy_control_state_t::RUNNING)) { + while (__atomic_load_n(shutdown_word_, __ATOMIC_ACQUIRE) == + static_cast(nixl_proxy_control_state_t::RUNNING)) { runOnce(); if (pthr_delay_us_ > 0) { std::this_thread::sleep_for(std::chrono::microseconds(pthr_delay_us_)); @@ -78,54 +78,46 @@ ProxyWorker::runOnce() { bool ProxyWorker::tryDequeue(nixlProxyChannelState &channel, nixlProxySubmission &submission) { - nixlProxyWorkRing *ring = channel.work_ring_; // Sole writer of consumer_idx on host — relaxed load is sufficient. - uint32_t local_consumer_idx = - __atomic_load_n(channel.consumer_idx_host_, __ATOMIC_RELAXED); - uint32_t slot = local_consumer_idx % ring->depth; + uint64_t local_consumer_idx = __atomic_load_n(channel.consumer_idx_host_, __ATOMIC_RELAXED); + uint32_t slot = static_cast(local_consumer_idx % channel.ring_depth_); // op_idx is the GPU-to-CPU signal that the record is written - // (pairs with release store in device enqueue). No producer_idx + // (pairs with release store in device enqueue). No producer index // read on host — it is GPU-internal for slot allocation. - const uint64_t op_idx = __atomic_load_n(&ring->records[slot].op_idx, __ATOMIC_ACQUIRE); + const uint64_t op_idx = __atomic_load_n(&channel.records_host_[slot].op_idx, __ATOMIC_ACQUIRE); if (op_idx == 0) { return false; } - __atomic_thread_fence(__ATOMIC_ACQUIRE); - submission = ring->records[slot]; + submission = channel.records_host_[slot]; submission.op_idx = op_idx; - __atomic_store_n(&ring->records[slot].op_idx, 0, __ATOMIC_RELAXED); - __atomic_store_n(channel.consumer_idx_host_, - local_consumer_idx + 1, - __ATOMIC_RELEASE); + __atomic_store_n(&channel.records_host_[slot].op_idx, 0, __ATOMIC_RELAXED); + __atomic_store_n(channel.consumer_idx_host_, local_consumer_idx + 1, __ATOMIC_RELEASE); NIXL_DEBUG << "ProxyWorker::tryDequeue: channel=" << channel.device_view.channel_id << " consumer=" << local_consumer_idx << " opcode=" << static_cast(submission.opcode) - << " op_idx=" << submission.op_idx - << " size=" << submission.size; + << " op_idx=" << submission.op_idx << " size=" << submission.size; return true; } void -ProxyWorker::submitToBackend(nixlProxyChannelState &channel, const nixlProxySubmission &submission) { +ProxyWorker::submitToBackend(nixlProxyChannelState &channel, + const nixlProxySubmission &submission) { nixlBackendProxySubmission prepared_submission; nixl_status_t status = proxy_memview_registry_->prepareSubmission(submission, prepared_submission); if (status != NIXL_SUCCESS) { NIXL_DEBUG << "ProxyWorker::submitToBackend: submission preparation failed" - << " op_idx=" << submission.op_idx - << " status=" << status; - channel.inflight_requests.push_back( - {submission.op_idx, 0, status}); + << " op_idx=" << submission.op_idx << " status=" << status; + channel.inflight_requests.push_back({submission.op_idx, 0, status}); // The terminal error is queued for publishCompletions(); the worker handled it. return; } NIXL_DEBUG << "ProxyWorker::submitToBackend: op_idx=" << submission.op_idx << " opcode=" << static_cast(submission.opcode) - << " channel=" << submission.channel_id - << " local_addr=0x" << std::hex << prepared_submission.local.desc.addr - << " remote_addr=0x" << prepared_submission.remote.desc.addr << std::dec - << " size=" << submission.size + << " channel=" << submission.channel_id << " local_addr=0x" << std::hex + << prepared_submission.local.desc.addr << " remote_addr=0x" + << prepared_submission.remote.desc.addr << std::dec << " size=" << submission.size << " remote_agent='" << prepared_submission.remote_agent << "'"; uint64_t request_token = 0; @@ -168,14 +160,12 @@ ProxyWorker::publishCompletions(nixlProxyChannelState &channel) { break; } } - NIXL_DEBUG << "ProxyWorker::publishCompletions: channel=" - << channel.device_view.channel_id - << " op_idx=" << front.op_idx - << " status=" << st + NIXL_DEBUG << "ProxyWorker::publishCompletions: channel=" << channel.device_view.channel_id + << " op_idx=" << front.op_idx << " status=" << st << " token=" << front.backend_req_token; channel.completion_slot_host_->next_status = st; - __atomic_store_n(&channel.completion_slot_host_->completed_idx, - front.op_idx, __ATOMIC_RELEASE); + __atomic_store_n( + &channel.completion_slot_host_->completed_idx, front.op_idx, __ATOMIC_RELEASE); channel.inflight_requests.pop_front(); if (st != NIXL_SUCCESS) { channel.error_latched = true; diff --git a/test/gtest/device_api/proxy_write_test.cu b/test/gtest/device_api/proxy_write_test.cu index 5fa6e6f6d8..326b929d1a 100644 --- a/test/gtest/device_api/proxy_write_test.cu +++ b/test/gtest/device_api/proxy_write_test.cu @@ -47,29 +47,35 @@ class StubProxyBackendAdapter : public nixlDeviceProxyBackendAdapter { public: nixl_status_t - init(uint32_t, uint32_t) override { return NIXL_SUCCESS; } + init(uint32_t, uint32_t) override { + return NIXL_SUCCESS; + } nixl_status_t - loadRemoteConnInfo(const std::string &, const nixl_blob_t &) override - { + loadRemoteConnInfo(const std::string &, const nixl_blob_t &) override { return NIXL_SUCCESS; } nixl_status_t - submit(const nixlBackendProxySubmission &, uint64_t &token) override - { + submit(const nixlBackendProxySubmission &, uint64_t &token) override { token = 0; return NIXL_SUCCESS; } nixl_status_t - checkCompletion(uint64_t) override { return NIXL_SUCCESS; } + checkCompletion(uint64_t) override { + return NIXL_SUCCESS; + } nixl_status_t - progress() override { return NIXL_SUCCESS; } + progress() override { + return NIXL_SUCCESS; + } nixl_status_t - shutdown() override { return NIXL_SUCCESS; } + shutdown() override { + return NIXL_SUCCESS; + } }; // --------------------------------------------------------------------------- @@ -80,17 +86,17 @@ public: class ControllableStubAdapter : public nixlDeviceProxyBackendAdapter { public: nixl_status_t - init(uint32_t, uint32_t) override { return NIXL_SUCCESS; } + init(uint32_t, uint32_t) override { + return NIXL_SUCCESS; + } nixl_status_t - loadRemoteConnInfo(const std::string &, const nixl_blob_t &) override - { + loadRemoteConnInfo(const std::string &, const nixl_blob_t &) override { return NIXL_SUCCESS; } nixl_status_t - submit(const nixlBackendProxySubmission &submission, uint64_t &token) override - { + submit(const nixlBackendProxySubmission &submission, uint64_t &token) override { std::lock_guard lk(mu_); token = next_token_++; pending_.insert(token); @@ -100,8 +106,7 @@ public: } nixl_status_t - checkCompletion(uint64_t token) override - { + checkCompletion(uint64_t token) override { std::lock_guard lk(mu_); auto it = completed_.find(token); if (it != completed_.end()) { @@ -113,42 +118,41 @@ public: } nixl_status_t - progress() override { return NIXL_SUCCESS; } + progress() override { + return NIXL_SUCCESS; + } nixl_status_t - shutdown() override { return NIXL_SUCCESS; } + shutdown() override { + return NIXL_SUCCESS; + } void - markComplete(uint64_t token) - { + markComplete(uint64_t token) { markCompleteWithStatus(token, NIXL_SUCCESS); } void - markCompleteWithStatus(uint64_t token, nixl_status_t status) - { + markCompleteWithStatus(uint64_t token, nixl_status_t status) { std::lock_guard lk(mu_); pending_.erase(token); completed_[token] = status; } bool - hasPending() const - { + hasPending() const { std::lock_guard lk(mu_); return !pending_.empty(); } size_t - pendingCount() const - { + pendingCount() const { std::lock_guard lk(mu_); return pending_.size(); } bool - hasPendingForChannel(uint32_t channel_id) const - { + hasPendingForChannel(uint32_t channel_id) const { std::lock_guard lk(mu_); for (uint64_t token : pending_) { auto it = token_channel_.find(token); @@ -160,8 +164,7 @@ public: } bool - markFirstPendingForChannel(uint32_t channel_id, uint64_t *token = nullptr) - { + markFirstPendingForChannel(uint32_t channel_id, uint64_t *token = nullptr) { std::lock_guard lk(mu_); for (uint64_t pending_token : pending_) { auto it = token_channel_.find(pending_token); @@ -178,8 +181,7 @@ public: } std::vector - submittedOpcodes() const - { + submittedOpcodes() const { std::lock_guard lk(mu_); return submitted_opcodes_; } @@ -200,29 +202,35 @@ private: class ErrorStubAdapter : public nixlDeviceProxyBackendAdapter { public: nixl_status_t - init(uint32_t, uint32_t) override { return NIXL_SUCCESS; } + init(uint32_t, uint32_t) override { + return NIXL_SUCCESS; + } nixl_status_t - loadRemoteConnInfo(const std::string &, const nixl_blob_t &) override - { + loadRemoteConnInfo(const std::string &, const nixl_blob_t &) override { return NIXL_SUCCESS; } nixl_status_t - submit(const nixlBackendProxySubmission &, uint64_t &token) override - { + submit(const nixlBackendProxySubmission &, uint64_t &token) override { token = 0; return NIXL_SUCCESS; } nixl_status_t - checkCompletion(uint64_t) override { return NIXL_ERR_BACKEND; } + checkCompletion(uint64_t) override { + return NIXL_ERR_BACKEND; + } nixl_status_t - progress() override { return NIXL_SUCCESS; } + progress() override { + return NIXL_SUCCESS; + } nixl_status_t - shutdown() override { return NIXL_SUCCESS; } + shutdown() override { + return NIXL_SUCCESS; + } }; // --------------------------------------------------------------------------- @@ -232,33 +240,36 @@ public: class SubmitErrorStubAdapter : public nixlDeviceProxyBackendAdapter { public: nixl_status_t - init(uint32_t, uint32_t) override { return NIXL_SUCCESS; } + init(uint32_t, uint32_t) override { + return NIXL_SUCCESS; + } nixl_status_t - loadRemoteConnInfo(const std::string &, const nixl_blob_t &) override - { + loadRemoteConnInfo(const std::string &, const nixl_blob_t &) override { return NIXL_SUCCESS; } nixl_status_t - submit(const nixlBackendProxySubmission &, uint64_t &) override - { + submit(const nixlBackendProxySubmission &, uint64_t &) override { ++submit_calls_; return NIXL_ERR_BACKEND; } nixl_status_t - checkCompletion(uint64_t) override - { + checkCompletion(uint64_t) override { ++check_completion_calls_; return NIXL_SUCCESS; } nixl_status_t - progress() override { return NIXL_SUCCESS; } + progress() override { + return NIXL_SUCCESS; + } nixl_status_t - shutdown() override { return NIXL_SUCCESS; } + shutdown() override { + return NIXL_SUCCESS; + } std::atomic submit_calls_{0}; std::atomic check_completion_calls_{0}; @@ -283,33 +294,25 @@ registerDummyMemViews(nixlProxyRuntime &runtime); // Writes true if load_proxy_context() returns a non-null pointer. __global__ void -proxyContextKernel(bool *out_has_ctx) -{ +proxyContextKernel(bool *out_has_ctx) { *out_has_ctx = (load_proxy_context() != nullptr); } // Calls nixlPut with zero-initialised operands and records the status. __global__ void -proxyPutKernel(nixlMemViewH src_mvh, - nixlMemViewH dst_mvh, - nixl_status_t *out_status) -{ +proxyPutKernel(nixlMemViewH src_mvh, nixlMemViewH dst_mvh, nixl_status_t *out_status) { nixlMemViewElem src{src_mvh, 0, 0}, dst{dst_mvh, 0, 0}; *out_status = nixlPut(src, dst, /*size=*/0); } __global__ void -proxyAtomicAddKernel(nixlMemViewH counter_mvh, - uint64_t value, - nixl_status_t *out_status) -{ +proxyAtomicAddKernel(nixlMemViewH counter_mvh, uint64_t value, nixl_status_t *out_status) { nixlMemViewElem counter{counter_mvh, 0, 0}; *out_status = nixlAtomicAdd(value, counter); } static void -publishProxyContext(nixlProxyRuntime &runtime) -{ +publishProxyContext(nixlProxyRuntime &runtime) { bool *d_warmup = nullptr; ASSERT_EQ(cudaMalloc(&d_warmup, sizeof(bool)), cudaSuccess); proxyContextKernel<<<1, 1>>>(d_warmup); @@ -322,8 +325,7 @@ publishProxyContext(nixlProxyRuntime &runtime) } static void -clearProxyContext() -{ +clearProxyContext() { ASSERT_EQ(nixlProxyClearContext(), cudaSuccess); } @@ -334,8 +336,7 @@ clearProxyContext() class ProxyDeviceApiTest : public ::testing::Test { protected: void - SetUp() override - { + SetUp() override { if (!gtest::hasCudaGpu()) { GTEST_SKIP() << "No CUDA-capable GPU, skipping proxy device API test."; } @@ -344,8 +345,7 @@ protected: template T - deviceGet(T *d_ptr) - { + deviceGet(T *d_ptr) { T val{}; cudaMemcpy(&val, d_ptr, sizeof(T), cudaMemcpyDeviceToHost); return val; @@ -353,8 +353,7 @@ protected: template T * - deviceAlloc() - { + deviceAlloc() { T *ptr = nullptr; EXPECT_EQ(cudaMalloc(&ptr, sizeof(T)), cudaSuccess); EXPECT_EQ(cudaMemset(ptr, 0, sizeof(T)), cudaSuccess); @@ -364,8 +363,7 @@ protected: template bool waitForCondition(Predicate predicate, - std::chrono::milliseconds timeout = std::chrono::milliseconds(500)) - { + std::chrono::milliseconds timeout = std::chrono::milliseconds(500)) { const auto deadline = std::chrono::steady_clock::now() + timeout; while (std::chrono::steady_clock::now() < deadline) { if (predicate()) { @@ -382,8 +380,7 @@ protected: // --------------------------------------------------------------------------- // After startWorkers() the GPU should see a non-null proxy context. -TEST_F(ProxyDeviceApiTest, ContextPublishedAfterStartWorkers) -{ +TEST_F(ProxyDeviceApiTest, ContextPublishedAfterStartWorkers) { auto adapter = std::make_unique(); nixlProxyRuntime runtime; @@ -405,8 +402,7 @@ TEST_F(ProxyDeviceApiTest, ContextPublishedAfterStartWorkers) } // After shutdown() the GPU should no longer see a proxy context. -TEST_F(ProxyDeviceApiTest, ContextClearedAfterShutdown) -{ +TEST_F(ProxyDeviceApiTest, ContextClearedAfterShutdown) { auto adapter = std::make_unique(); nixlProxyRuntime runtime; @@ -432,8 +428,7 @@ TEST_F(ProxyDeviceApiTest, ContextClearedAfterShutdown) // nixlPut() via the proxy backend should report NIXL_IN_PROG once the // submission is accepted into the proxy ring. -TEST_F(ProxyDeviceApiTest, PutReturnsInProgWhenEnqueued) -{ +TEST_F(ProxyDeviceApiTest, PutReturnsInProgWhenEnqueued) { auto adapter = std::make_unique(); nixlProxyRuntime runtime; @@ -455,8 +450,7 @@ TEST_F(ProxyDeviceApiTest, PutReturnsInProgWhenEnqueued) ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); } -TEST_F(ProxyDeviceApiTest, AtomicAddReturnsInProgWhenEnqueued) -{ +TEST_F(ProxyDeviceApiTest, AtomicAddReturnsInProgWhenEnqueued) { auto adapter = std::make_unique(); nixlProxyRuntime runtime; @@ -488,15 +482,19 @@ TEST_F(ProxyDeviceApiTest, AtomicAddReturnsInProgWhenEnqueued) // Enqueues a put and spins until pollXferStatus returns a final status. __global__ void -proxyPutAndPollKernel(nixlMemViewH src_mvh, nixlMemViewH dst_mvh, +proxyPutAndPollKernel(nixlMemViewH src_mvh, + nixlMemViewH dst_mvh, uint32_t channel_id, nixl_status_t *out_put_status, - nixl_status_t *out_poll_status) -{ + nixl_status_t *out_poll_status) { nixlMemViewElem src{src_mvh, 0, 0}, dst{dst_mvh, 0, 0}; nixlGpuXferStatusH xfer_status{}; - *out_put_status = nixlPut(src, dst, /*size=*/0, channel_id, - /*flags=*/0, &xfer_status); + *out_put_status = nixlPut(src, + dst, + /*size=*/0, + channel_id, + /*flags=*/0, + &xfer_status); nixl_status_t poll; do { @@ -510,12 +508,14 @@ proxyAtomicAddAndPollKernel(nixlMemViewH counter_mvh, uint64_t value, uint32_t channel_id, nixl_status_t *out_atomic_status, - nixl_status_t *out_poll_status) -{ + nixl_status_t *out_poll_status) { nixlMemViewElem counter{counter_mvh, 0, 0}; nixlGpuXferStatusH xfer_status{}; - *out_atomic_status = nixlAtomicAdd(value, counter, channel_id, - /*flags=*/0, &xfer_status); + *out_atomic_status = nixlAtomicAdd(value, + counter, + channel_id, + /*flags=*/0, + &xfer_status); nixl_status_t poll; do { @@ -527,14 +527,18 @@ proxyAtomicAddAndPollKernel(nixlMemViewH counter_mvh, // Enqueues a put and immediately returns; saves xfer_status to device memory // so the test thread can later launch a poll kernel. __global__ void -proxyPutAsyncKernel(nixlMemViewH src_mvh, nixlMemViewH dst_mvh, +proxyPutAsyncKernel(nixlMemViewH src_mvh, + nixlMemViewH dst_mvh, uint32_t channel_id, nixl_status_t *out_put_status, - nixlGpuXferStatusH *out_xfer_status) -{ + nixlGpuXferStatusH *out_xfer_status) { nixlMemViewElem src{src_mvh, 0, 0}, dst{dst_mvh, 0, 0}; - *out_put_status = nixlPut(src, dst, /*size=*/0, channel_id, - /*flags=*/0, out_xfer_status); + *out_put_status = nixlPut(src, + dst, + /*size=*/0, + channel_id, + /*flags=*/0, + out_xfer_status); } __global__ void @@ -542,19 +546,19 @@ proxyAtomicAddAsyncKernel(nixlMemViewH counter_mvh, uint64_t value, uint32_t channel_id, nixl_status_t *out_atomic_status, - nixlGpuXferStatusH *out_xfer_status) -{ + nixlGpuXferStatusH *out_xfer_status) { nixlMemViewElem counter{counter_mvh, 0, 0}; - *out_atomic_status = nixlAtomicAdd(value, counter, channel_id, - /*flags=*/0, out_xfer_status); + *out_atomic_status = nixlAtomicAdd(value, + counter, + channel_id, + /*flags=*/0, + out_xfer_status); } // Enqueues op_count puts on one channel and records each immediate enqueue // status. The final submission may block if the ring is full. __global__ void -proxyPutBurstKernel(uint32_t op_count, uint32_t channel_id, - nixl_status_t *out_put_statuses) -{ +proxyPutBurstKernel(uint32_t op_count, uint32_t channel_id, nixl_status_t *out_put_statuses) { nixlMemViewElem src{}, dst{}; for (uint32_t i = 0; i < op_count; ++i) { out_put_statuses[i] = nixlPut(src, dst, /*size=*/0, channel_id); @@ -563,9 +567,7 @@ proxyPutBurstKernel(uint32_t op_count, uint32_t channel_id, // Non-blocking single poll: returns current status without spinning. __global__ void -proxyPollOnceKernel(nixlGpuXferStatusH *xfer_status, - nixl_status_t *out_poll_status) -{ +proxyPollOnceKernel(nixlGpuXferStatusH *xfer_status, nixl_status_t *out_poll_status) { *out_poll_status = nixlGpuGetXferStatus(*xfer_status); } @@ -576,8 +578,7 @@ proxyPollOnceKernel(nixlGpuXferStatusH *xfer_status, // Register one local and one remote proxy memview so dispatch can prepare // transport-ready descriptors before submit(). static DummyProxyMemViews -registerDummyMemViews(nixlProxyRuntime &runtime) -{ +registerDummyMemViews(nixlProxyRuntime &runtime) { static DummyBackendMD local_md; static DummyBackendMD remote_md; @@ -604,14 +605,30 @@ registerDummyMemViews(nixlProxyRuntime &runtime) return handles; } -static void -signalProxyShutdown(nixlProxyRuntime &runtime) -{ +static uint32_t * +shutdownWordHostFromRuntime(nixlProxyRuntime &runtime) { + nixlProxyDeviceContextData device_ctx{}; + if (runtime.deviceContext() == nullptr) { + return nullptr; + } + if (cudaMemcpy( + &device_ctx, runtime.deviceContext(), sizeof(device_ctx), cudaMemcpyDeviceToHost) != + cudaSuccess) { + return nullptr; + } + if (device_ctx.shutdown_word == nullptr) { + return nullptr; + } + cudaPointerAttributes attrs{}; - ASSERT_EQ(cudaPointerGetAttributes(&attrs, runtime.deviceContext()->shutdown_word), - cudaSuccess); - ASSERT_NE(attrs.hostPointer, nullptr); - auto *shutdown_host = static_cast(attrs.hostPointer); + if (cudaPointerGetAttributes(&attrs, device_ctx.shutdown_word) != cudaSuccess) { + return nullptr; + } + return static_cast(attrs.hostPointer); +} + +static void +signalProxyShutdown(uint32_t *shutdown_host) { __atomic_store_n(shutdown_host, static_cast(nixl_proxy_control_state_t::SHUTDOWN), __ATOMIC_RELEASE); @@ -624,8 +641,7 @@ signalProxyShutdown(nixlProxyRuntime &runtime) // Full round-trip: GPU enqueues -> worker dequeues -> backend completes // (immediately via StubProxyBackendAdapter) -> worker publishes -> GPU polls // NIXL_SUCCESS. -TEST_F(ProxyDeviceApiTest, PutCompletionRoundTrip) -{ +TEST_F(ProxyDeviceApiTest, PutCompletionRoundTrip) { auto adapter = std::make_unique(); nixlProxyRuntime runtime; @@ -636,7 +652,7 @@ TEST_F(ProxyDeviceApiTest, PutCompletionRoundTrip) const auto mvhs = registerDummyMemViews(runtime); - nixl_status_t *d_put_status = deviceAlloc(); + nixl_status_t *d_put_status = deviceAlloc(); nixl_status_t *d_poll_status = deviceAlloc(); proxyPutAndPollKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status, d_poll_status); @@ -652,8 +668,7 @@ TEST_F(ProxyDeviceApiTest, PutCompletionRoundTrip) ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); } -TEST_F(ProxyDeviceApiTest, AtomicAddCompletionRoundTrip) -{ +TEST_F(ProxyDeviceApiTest, AtomicAddCompletionRoundTrip) { auto adapter = std::make_unique(); nixlProxyRuntime runtime; @@ -665,10 +680,9 @@ TEST_F(ProxyDeviceApiTest, AtomicAddCompletionRoundTrip) const auto mvhs = registerDummyMemViews(runtime); nixl_status_t *d_atomic_status = deviceAlloc(); - nixl_status_t *d_poll_status = deviceAlloc(); + nixl_status_t *d_poll_status = deviceAlloc(); - proxyAtomicAddAndPollKernel<<<1, 1>>>(mvhs.dst, 42, 0, - d_atomic_status, d_poll_status); + proxyAtomicAddAndPollKernel<<<1, 1>>>(mvhs.dst, 42, 0, d_atomic_status, d_poll_status); ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); ASSERT_EQ(cudaGetLastError(), cudaSuccess); @@ -683,8 +697,7 @@ TEST_F(ProxyDeviceApiTest, AtomicAddCompletionRoundTrip) // Verifies that the GPU kernel stays spinning until the test thread // explicitly marks the backend token complete. -TEST_F(ProxyDeviceApiTest, CompletionNotVisibleUntilPublished) -{ +TEST_F(ProxyDeviceApiTest, CompletionNotVisibleUntilPublished) { auto adapter_owner = std::make_unique(); auto *adapter = adapter_owner.get(); nixlProxyRuntime runtime; @@ -696,7 +709,7 @@ TEST_F(ProxyDeviceApiTest, CompletionNotVisibleUntilPublished) const auto mvhs = registerDummyMemViews(runtime); - nixl_status_t *d_put_status = deviceAlloc(); + nixl_status_t *d_put_status = deviceAlloc(); nixl_status_t *d_poll_status = deviceAlloc(); // Launch async — kernel will spin on pollXferStatus. @@ -726,8 +739,7 @@ TEST_F(ProxyDeviceApiTest, CompletionNotVisibleUntilPublished) // Enqueue 3 operations, complete them in order, and verify the collapsed-CQ // frontier semantics: each pollXferStatus returns NIXL_SUCCESS only after its // op_idx has been reached. -TEST_F(ProxyDeviceApiTest, MultipleSubmissionsCompletionFrontier) -{ +TEST_F(ProxyDeviceApiTest, MultipleSubmissionsCompletionFrontier) { auto adapter_owner = std::make_unique(); auto *adapter = adapter_owner.get(); nixlProxyRuntime runtime; @@ -740,21 +752,18 @@ TEST_F(ProxyDeviceApiTest, MultipleSubmissionsCompletionFrontier) const auto mvhs = registerDummyMemViews(runtime); constexpr int kOps = 3; - nixl_status_t *d_put_status[kOps]; + nixl_status_t *d_put_status[kOps]; nixlGpuXferStatusH *d_xfer_status[kOps]; for (int i = 0; i < kOps; i++) { - d_put_status[i] = deviceAlloc(); - ASSERT_EQ(cudaMalloc(&d_xfer_status[i], sizeof(nixlGpuXferStatusH)), - cudaSuccess); - ASSERT_EQ(cudaMemset(d_xfer_status[i], 0, sizeof(nixlGpuXferStatusH)), - cudaSuccess); + d_put_status[i] = deviceAlloc(); + ASSERT_EQ(cudaMalloc(&d_xfer_status[i], sizeof(nixlGpuXferStatusH)), cudaSuccess); + ASSERT_EQ(cudaMemset(d_xfer_status[i], 0, sizeof(nixlGpuXferStatusH)), cudaSuccess); } // Enqueue 3 operations sequentially (each kernel returns after enqueue). for (int i = 0; i < kOps; i++) { - proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status[i], - d_xfer_status[i]); + proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status[i], d_xfer_status[i]); ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); ASSERT_EQ(cudaGetLastError(), cudaSuccess); EXPECT_EQ(deviceGet(d_put_status[i]), NIXL_IN_PROG); @@ -792,8 +801,7 @@ TEST_F(ProxyDeviceApiTest, MultipleSubmissionsCompletionFrontier) ASSERT_EQ(runtime.shutdown(), NIXL_SUCCESS); } -TEST_F(ProxyDeviceApiTest, PutPutAtomicAddCompletionFrontier) -{ +TEST_F(ProxyDeviceApiTest, PutPutAtomicAddCompletionFrontier) { auto adapter_owner = std::make_unique(); auto *adapter = adapter_owner.get(); nixlProxyRuntime runtime; @@ -806,42 +814,35 @@ TEST_F(ProxyDeviceApiTest, PutPutAtomicAddCompletionFrontier) const auto mvhs = registerDummyMemViews(runtime); constexpr int kOps = 3; - nixl_status_t *d_submit_status[kOps]; + nixl_status_t *d_submit_status[kOps]; nixlGpuXferStatusH *d_xfer_status[kOps]; for (int i = 0; i < kOps; i++) { d_submit_status[i] = deviceAlloc(); - ASSERT_EQ(cudaMalloc(&d_xfer_status[i], sizeof(nixlGpuXferStatusH)), - cudaSuccess); - ASSERT_EQ(cudaMemset(d_xfer_status[i], 0, sizeof(nixlGpuXferStatusH)), - cudaSuccess); + ASSERT_EQ(cudaMalloc(&d_xfer_status[i], sizeof(nixlGpuXferStatusH)), cudaSuccess); + ASSERT_EQ(cudaMemset(d_xfer_status[i], 0, sizeof(nixlGpuXferStatusH)), cudaSuccess); } - proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_submit_status[0], - d_xfer_status[0]); + proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_submit_status[0], d_xfer_status[0]); ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); ASSERT_EQ(cudaGetLastError(), cudaSuccess); EXPECT_EQ(deviceGet(d_submit_status[0]), NIXL_IN_PROG); - proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_submit_status[1], - d_xfer_status[1]); + proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_submit_status[1], d_xfer_status[1]); ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); ASSERT_EQ(cudaGetLastError(), cudaSuccess); EXPECT_EQ(deviceGet(d_submit_status[1]), NIXL_IN_PROG); - proxyAtomicAddAsyncKernel<<<1, 1>>>(mvhs.dst, 42, 0, d_submit_status[2], - d_xfer_status[2]); + proxyAtomicAddAsyncKernel<<<1, 1>>>(mvhs.dst, 42, 0, d_submit_status[2], d_xfer_status[2]); ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); ASSERT_EQ(cudaGetLastError(), cudaSuccess); EXPECT_EQ(deviceGet(d_submit_status[2]), NIXL_IN_PROG); - ASSERT_TRUE(waitForCondition([adapter]() { - return adapter->pendingCount() == kOps; - })); - EXPECT_EQ(adapter->submittedOpcodes(), - std::vector({nixl_proxy_opcode_t::PUT, - nixl_proxy_opcode_t::PUT, - nixl_proxy_opcode_t::ATOMIC_ADD})); + ASSERT_TRUE(waitForCondition([adapter]() { return adapter->pendingCount() == kOps; })); + EXPECT_EQ( + adapter->submittedOpcodes(), + std::vector( + {nixl_proxy_opcode_t::PUT, nixl_proxy_opcode_t::PUT, nixl_proxy_opcode_t::ATOMIC_ADD})); nixl_status_t *d_poll = deviceAlloc(); for (int i = 0; i < kOps; i++) { @@ -856,7 +857,8 @@ TEST_F(ProxyDeviceApiTest, PutPutAtomicAddCompletionFrontier) ASSERT_TRUE(waitForCondition([&]() { proxyPollOnceKernel<<<1, 1>>>(d_xfer_status[i], d_poll); return cudaDeviceSynchronize() == cudaSuccess && deviceGet(d_poll) == NIXL_SUCCESS; - })) << "op " << i << " should complete after markComplete"; + })) << "op " + << i << " should complete after markComplete"; } cudaFree(d_poll); @@ -870,8 +872,7 @@ TEST_F(ProxyDeviceApiTest, PutPutAtomicAddCompletionFrontier) // Once a later op publishes an error, an earlier op whose op_idx is already // behind the completion frontier must still observe NIXL_SUCCESS. -TEST_F(ProxyDeviceApiTest, EarlierCompletionStaysSuccessfulAfterLaterError) -{ +TEST_F(ProxyDeviceApiTest, EarlierCompletionStaysSuccessfulAfterLaterError) { auto adapter_owner = std::make_unique(); auto *adapter = adapter_owner.get(); nixlProxyRuntime runtime; @@ -883,16 +884,13 @@ TEST_F(ProxyDeviceApiTest, EarlierCompletionStaysSuccessfulAfterLaterError) const auto mvhs = registerDummyMemViews(runtime); - nixl_status_t *d_put_status[2]; + nixl_status_t *d_put_status[2]; nixlGpuXferStatusH *d_xfer_status[2]; for (int i = 0; i < 2; ++i) { d_put_status[i] = deviceAlloc(); - ASSERT_EQ(cudaMalloc(&d_xfer_status[i], sizeof(nixlGpuXferStatusH)), - cudaSuccess); - ASSERT_EQ(cudaMemset(d_xfer_status[i], 0, sizeof(nixlGpuXferStatusH)), - cudaSuccess); - proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status[i], - d_xfer_status[i]); + ASSERT_EQ(cudaMalloc(&d_xfer_status[i], sizeof(nixlGpuXferStatusH)), cudaSuccess); + ASSERT_EQ(cudaMemset(d_xfer_status[i], 0, sizeof(nixlGpuXferStatusH)), cudaSuccess); + proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status[i], d_xfer_status[i]); ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); ASSERT_EQ(cudaGetLastError(), cudaSuccess); EXPECT_EQ(deviceGet(d_put_status[i]), NIXL_IN_PROG); @@ -928,8 +926,7 @@ TEST_F(ProxyDeviceApiTest, EarlierCompletionStaysSuccessfulAfterLaterError) // Once an earlier op publishes a terminal error, later queued ops must also // observe that error instead of spinning forever. -TEST_F(ProxyDeviceApiTest, EarlierErrorPropagatesToLaterQueuedOp) -{ +TEST_F(ProxyDeviceApiTest, EarlierErrorPropagatesToLaterQueuedOp) { auto adapter_owner = std::make_unique(); auto *adapter = adapter_owner.get(); nixlProxyRuntime runtime; @@ -941,16 +938,13 @@ TEST_F(ProxyDeviceApiTest, EarlierErrorPropagatesToLaterQueuedOp) const auto mvhs = registerDummyMemViews(runtime); - nixl_status_t *d_put_status[2]; + nixl_status_t *d_put_status[2]; nixlGpuXferStatusH *d_xfer_status[2]; for (int i = 0; i < 2; ++i) { d_put_status[i] = deviceAlloc(); - ASSERT_EQ(cudaMalloc(&d_xfer_status[i], sizeof(nixlGpuXferStatusH)), - cudaSuccess); - ASSERT_EQ(cudaMemset(d_xfer_status[i], 0, sizeof(nixlGpuXferStatusH)), - cudaSuccess); - proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status[i], - d_xfer_status[i]); + ASSERT_EQ(cudaMalloc(&d_xfer_status[i], sizeof(nixlGpuXferStatusH)), cudaSuccess); + ASSERT_EQ(cudaMemset(d_xfer_status[i], 0, sizeof(nixlGpuXferStatusH)), cudaSuccess); + proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status[i], d_xfer_status[i]); ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); ASSERT_EQ(cudaGetLastError(), cudaSuccess); EXPECT_EQ(deviceGet(d_put_status[i]), NIXL_IN_PROG); @@ -980,8 +974,7 @@ TEST_F(ProxyDeviceApiTest, EarlierErrorPropagatesToLaterQueuedOp) // Backend returns NIXL_ERR_BACKEND on checkCompletion; verify the GPU kernel // receives the error status through the completion slot. -TEST_F(ProxyDeviceApiTest, CompletionPropagatesErrorStatus) -{ +TEST_F(ProxyDeviceApiTest, CompletionPropagatesErrorStatus) { auto adapter = std::make_unique(); nixlProxyRuntime runtime; @@ -992,7 +985,7 @@ TEST_F(ProxyDeviceApiTest, CompletionPropagatesErrorStatus) const auto mvhs = registerDummyMemViews(runtime); - nixl_status_t *d_put_status = deviceAlloc(); + nixl_status_t *d_put_status = deviceAlloc(); nixl_status_t *d_poll_status = deviceAlloc(); proxyPutAndPollKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status, d_poll_status); @@ -1010,8 +1003,7 @@ TEST_F(ProxyDeviceApiTest, CompletionPropagatesErrorStatus) // Backend submit() failure should be published to the GPU as the terminal // transfer status, without going through checkCompletion(). -TEST_F(ProxyDeviceApiTest, SubmitFailurePropagatesErrorStatus) -{ +TEST_F(ProxyDeviceApiTest, SubmitFailurePropagatesErrorStatus) { const gtest::LogIgnoreGuard lig("ProxyWorker::submitToBackend: backend submit failed"); auto adapter_owner = std::make_unique(); auto *adapter = adapter_owner.get(); @@ -1024,7 +1016,7 @@ TEST_F(ProxyDeviceApiTest, SubmitFailurePropagatesErrorStatus) const auto mvhs = registerDummyMemViews(runtime); - nixl_status_t *d_put_status = deviceAlloc(); + nixl_status_t *d_put_status = deviceAlloc(); nixl_status_t *d_poll_status = deviceAlloc(); proxyPutAndPollKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status, d_poll_status); @@ -1044,37 +1036,35 @@ TEST_F(ProxyDeviceApiTest, SubmitFailurePropagatesErrorStatus) // When the ring is full and no worker can drain it, the next enqueue should // spin until shutdown is signalled and then return NIXL_ERR_BACKEND. -TEST_F(ProxyDeviceApiTest, RingOverflowReturnsBackendErrorOnShutdown) -{ +TEST_F(ProxyDeviceApiTest, RingOverflowReturnsBackendErrorOnShutdown) { auto adapter = std::make_unique(); nixlProxyRuntime runtime; ASSERT_EQ(runtime.init(std::move(adapter), /*channel_count=*/1, /*worker_count=*/1), NIXL_SUCCESS); publishProxyContext(runtime); + uint32_t *shutdown_host = shutdownWordHostFromRuntime(runtime); + ASSERT_NE(shutdown_host, nullptr); constexpr uint32_t kBurstOps = kDefaultProxyRingDepth + 1; nixl_status_t *d_statuses = nullptr; - ASSERT_EQ(cudaMalloc(&d_statuses, sizeof(nixl_status_t) * kBurstOps), - cudaSuccess); - ASSERT_EQ(cudaMemset(d_statuses, 0, sizeof(nixl_status_t) * kBurstOps), - cudaSuccess); + ASSERT_EQ(cudaMalloc(&d_statuses, sizeof(nixl_status_t) * kBurstOps), cudaSuccess); + ASSERT_EQ(cudaMemset(d_statuses, 0, sizeof(nixl_status_t) * kBurstOps), cudaSuccess); proxyPutBurstKernel<<<1, 1>>>(kBurstOps, 0, d_statuses); std::this_thread::sleep_for(std::chrono::milliseconds(50)); ASSERT_EQ(cudaStreamQuery(nullptr), cudaErrorNotReady); - signalProxyShutdown(runtime); + signalProxyShutdown(shutdown_host); ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); ASSERT_EQ(cudaGetLastError(), cudaSuccess); std::vector statuses(kBurstOps); - ASSERT_EQ(cudaMemcpy(statuses.data(), - d_statuses, - sizeof(nixl_status_t) * kBurstOps, - cudaMemcpyDeviceToHost), - cudaSuccess); + ASSERT_EQ( + cudaMemcpy( + statuses.data(), d_statuses, sizeof(nixl_status_t) * kBurstOps, cudaMemcpyDeviceToHost), + cudaSuccess); for (uint32_t i = 0; i < kDefaultProxyRingDepth; ++i) { EXPECT_EQ(statuses[i], NIXL_IN_PROG) << "unexpected status at op " << i; } @@ -1087,8 +1077,7 @@ TEST_F(ProxyDeviceApiTest, RingOverflowReturnsBackendErrorOnShutdown) // Completions are tracked per-channel, so publishing one channel should not // advance an unrelated channel's xfer status. -TEST_F(ProxyDeviceApiTest, ChannelCompletionsAdvanceIndependently) -{ +TEST_F(ProxyDeviceApiTest, ChannelCompletionsAdvanceIndependently) { auto adapter_owner = std::make_unique(); auto *adapter = adapter_owner.get(); nixlProxyRuntime runtime; @@ -1104,10 +1093,8 @@ TEST_F(ProxyDeviceApiTest, ChannelCompletionsAdvanceIndependently) nixlGpuXferStatusH *d_xfer_status[2]; for (int i = 0; i < 2; ++i) { d_put_status[i] = deviceAlloc(); - ASSERT_EQ(cudaMalloc(&d_xfer_status[i], sizeof(nixlGpuXferStatusH)), - cudaSuccess); - ASSERT_EQ(cudaMemset(d_xfer_status[i], 0, sizeof(nixlGpuXferStatusH)), - cudaSuccess); + ASSERT_EQ(cudaMalloc(&d_xfer_status[i], sizeof(nixlGpuXferStatusH)), cudaSuccess); + ASSERT_EQ(cudaMemset(d_xfer_status[i], 0, sizeof(nixlGpuXferStatusH)), cudaSuccess); } proxyPutAsyncKernel<<<1, 1>>>(mvhs.src, mvhs.dst, 0, d_put_status[0], d_xfer_status[0]); @@ -1119,9 +1106,8 @@ TEST_F(ProxyDeviceApiTest, ChannelCompletionsAdvanceIndependently) EXPECT_EQ(deviceGet(d_put_status[0]), NIXL_IN_PROG); EXPECT_EQ(deviceGet(d_put_status[1]), NIXL_IN_PROG); ASSERT_TRUE(waitForCondition([adapter]() { - return adapter->pendingCount() == 2 - && adapter->hasPendingForChannel(0) - && adapter->hasPendingForChannel(1); + return adapter->pendingCount() == 2 && adapter->hasPendingForChannel(0) && + adapter->hasPendingForChannel(1); })); uint64_t completed_token = 0; diff --git a/test/gtest/unit/proxy_runtime/proxy_runtime.cpp b/test/gtest/unit/proxy_runtime/proxy_runtime.cpp index ae51ea060d..4490aff505 100644 --- a/test/gtest/unit/proxy_runtime/proxy_runtime.cpp +++ b/test/gtest/unit/proxy_runtime/proxy_runtime.cpp @@ -32,12 +32,12 @@ namespace gtest { namespace proxy_runtime { -class DummyBackendMD : public nixlBackendMD { + class DummyBackendMD : public nixlBackendMD { public: DummyBackendMD() : nixlBackendMD(false) {} -}; + }; -class StubBackend : public nixlDeviceProxyBackendAdapter { + class StubBackend : public nixlDeviceProxyBackendAdapter { public: nixl_status_t init(uint32_t worker_count, uint32_t channel_count) override { @@ -86,9 +86,9 @@ class StubBackend : public nixlDeviceProxyBackendAdapter { mutable std::mutex submit_mutex_; std::vector submissions_; uint64_t next_request_token_ = 0; -}; + }; -class ProxyRuntimeTest : public testing::Test { + class ProxyRuntimeTest : public testing::Test { protected: nixl_status_t initRuntime(uint32_t channel_count, @@ -108,6 +108,22 @@ class ProxyRuntimeTest : public testing::Test { nixlProxyRuntime runtime_; }; +static nixlProxyWorkRing +copyDeviceWorkRing(const nixlProxyChannelView &view) { + nixlProxyWorkRing ring{}; + EXPECT_EQ(cudaMemcpy(&ring, view.work_ring, sizeof(ring), cudaMemcpyDeviceToHost), + cudaSuccess); + return ring; +} + +static nixlProxySubmission * +hostRecordsFromDeviceAlias(nixlProxySubmission *records_host_dev) { + cudaPointerAttributes attrs{}; + EXPECT_EQ(cudaPointerGetAttributes(&attrs, records_host_dev), cudaSuccess); + EXPECT_NE(attrs.hostPointer, nullptr); + return static_cast(attrs.hostPointer); +} + TEST_F(ProxyRuntimeTest, InitCallsBackendInit) { ASSERT_EQ(initRuntime(4, 2), NIXL_SUCCESS); EXPECT_TRUE(backend_->init_called_); @@ -143,10 +159,12 @@ TEST_F(ProxyRuntimeTest, DeviceChannelViewsPopulated) { for (uint32_t i = 0; i < 3; ++i) { EXPECT_EQ(views[i].channel_id, i); EXPECT_NE(views[i].work_ring, nullptr); - EXPECT_NE(views[i].work_ring->producer_idx, nullptr); - EXPECT_NE(views[i].work_ring->consumer_idx, nullptr); + const nixlProxyWorkRing ring = copyDeviceWorkRing(views[i]); + EXPECT_NE(ring.records, nullptr); + EXPECT_NE(ring.producer_idx, nullptr); + EXPECT_NE(ring.consumer_idx, nullptr); EXPECT_NE(views[i].completion_slot, nullptr); - EXPECT_EQ(views[i].work_ring->depth, kDefaultProxyRingDepth); + EXPECT_EQ(ring.depth, kDefaultProxyRingDepth); } } @@ -154,15 +172,16 @@ TEST_F(ProxyRuntimeTest, WorkRingIndicesStartAtZero) { ASSERT_EQ(initRuntime(2, 1), NIXL_SUCCESS); const nixlProxyChannelView *views = runtime_.deviceChannelViews(); for (uint32_t i = 0; i < 2; ++i) { - uint32_t producer = 0; - uint32_t consumer = 0; + const nixlProxyWorkRing ring = copyDeviceWorkRing(views[i]); + uint64_t producer = 0; + uint64_t consumer = 0; ASSERT_EQ(cudaMemcpy(&producer, - views[i].work_ring->producer_idx, + ring.producer_idx, sizeof(producer), cudaMemcpyDeviceToHost), cudaSuccess); ASSERT_EQ(cudaMemcpy(&consumer, - views[i].work_ring->consumer_idx, + ring.consumer_idx, sizeof(consumer), cudaMemcpyDeviceToHost), cudaSuccess); @@ -195,11 +214,13 @@ TEST_F(ProxyRuntimeTest, WorkerCountClampedToChannels) { TEST_F(ProxyRuntimeTest, DeviceContextPopulated) { ASSERT_EQ(initRuntime(3, 1), NIXL_SUCCESS); - auto *ctx = runtime_.deviceContext(); - ASSERT_NE(ctx, nullptr); - EXPECT_EQ(ctx->num_channels, 3u); - EXPECT_EQ(ctx->channels, runtime_.deviceChannelViews()); - EXPECT_NE(ctx->shutdown_word, nullptr); + auto *device_ctx = runtime_.deviceContext(); + ASSERT_NE(device_ctx, nullptr); + nixlProxyDeviceContextData ctx{}; + ASSERT_EQ(cudaMemcpy(&ctx, device_ctx, sizeof(ctx), cudaMemcpyDeviceToHost), cudaSuccess); + EXPECT_EQ(ctx.num_channels, 3u); + EXPECT_NE(ctx.channels, nullptr); + EXPECT_NE(ctx.shutdown_word, nullptr); } TEST_F(ProxyRuntimeTest, DeviceContextNullAfterShutdown) { @@ -376,10 +397,12 @@ TEST_F(ProxyRuntimeTest, WorkerSubmitsPreparedTransportDescriptors) { submission.dst_offset = 8; submission.size = 32; - auto *ring = runtime_.deviceChannelViews()[0].work_ring; + const nixlProxyWorkRing ring = copyDeviceWorkRing(runtime_.deviceChannelViews()[0]); + auto *records = hostRecordsFromDeviceAlias(ring.records); + ASSERT_NE(records, nullptr); submission.op_idx = 0; - ring->records[0] = submission; - __atomic_store_n(&ring->records[0].op_idx, uint64_t{11}, __ATOMIC_RELEASE); + records[0] = submission; + __atomic_store_n(&records[0].op_idx, uint64_t{11}, __ATOMIC_RELEASE); const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(250); while (std::chrono::steady_clock::now() < deadline) { @@ -445,10 +468,12 @@ TEST_F(ProxyRuntimeTest, WorkerSubmitsPreparedAtomicAddDescriptor) { submission.size = sizeof(uint64_t); submission.value = 42; - auto *ring = runtime_.deviceChannelViews()[0].work_ring; + const nixlProxyWorkRing ring = copyDeviceWorkRing(runtime_.deviceChannelViews()[0]); + auto *records = hostRecordsFromDeviceAlias(ring.records); + ASSERT_NE(records, nullptr); submission.op_idx = 0; - ring->records[0] = submission; - __atomic_store_n(&ring->records[0].op_idx, uint64_t{11}, __ATOMIC_RELEASE); + records[0] = submission; + __atomic_store_n(&records[0].op_idx, uint64_t{11}, __ATOMIC_RELEASE); const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(250); while (std::chrono::steady_clock::now() < deadline) { From e23270ff3db5b3a069a603732b46f44a2618d7f9 Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Sun, 31 May 2026 11:51:44 +0300 Subject: [PATCH 16/18] gpu: enqueue: cache consumer index on device Cache the host-pinned consumer index in device memory and refresh from host only when the ring appears full. Signed-off-by: Tomer Davidor --- src/api/gpu/proxy/nixl_device_proxy.cuh | 16 +- src/core/device_proxy/proxy_protocol.h | 2 + src/core/device_proxy/proxy_runtime.cpp | 322 ++++++++--------- src/core/device_proxy/proxy_runtime.h | 457 ++++++++++++------------ 4 files changed, 395 insertions(+), 402 deletions(-) diff --git a/src/api/gpu/proxy/nixl_device_proxy.cuh b/src/api/gpu/proxy/nixl_device_proxy.cuh index 492abb116f..280291908e 100644 --- a/src/api/gpu/proxy/nixl_device_proxy.cuh +++ b/src/api/gpu/proxy/nixl_device_proxy.cuh @@ -80,6 +80,8 @@ static_assert(sizeof(*nixlProxyWorkRing{}.producer_idx) == 8, "producer_idx must be 64-bit to avoid wrap-around false completions"); static_assert(sizeof(*nixlProxyWorkRing{}.consumer_idx) == 8, "consumer_idx must be 64-bit to match producer_idx"); +static_assert(sizeof(*nixlProxyWorkRing{}.consumer_idx_cache) == 8, + "consumer_idx_cache must be 64-bit to match producer_idx"); static_assert(sizeof(nixlProxyCompletionSlot::completed_idx) == 8, "completed_idx must be 64-bit to match producer_idx"); @@ -117,7 +119,8 @@ struct ProxyDeviceContext : nixlProxyDeviceContextData { // // producer_idx lives in device memory and only needs device-scope atomicity. // consumer_idx lives in pinned host memory (accessible from device via - // UVA mapped pointer), so full-ring polling uses system-scope atomics. + // UVA mapped pointer). The device cache keeps the non-full path from + // repeatedly touching host memory. __device__ inline nixl_status_t enqueue(nixlProxySubmission submission, nixlGpuXferStatusH *xfer_status = nullptr) { if (submission.channel_id >= num_channels) { @@ -135,9 +138,14 @@ struct ProxyDeviceContext : nixlProxyDeviceContextData { // Atomically claim a unique slot in the ring. const uint64_t ticket = producer_idx.fetch_add(1, cuda::memory_order_relaxed); - // Spin until the claimed slot has space (consumer has freed it). - while (ticket - cons.load(cuda::memory_order_acquire) >= ring->depth) { - if (shut.load(cuda::memory_order_acquire) + // Fast path: use the device cache. Refresh from host only if the ring + // appears full, since mapped-host loads are much slower than HBM loads. + uint64_t cached_consumer_idx = *ring->consumer_idx_cache; + while (ticket - cached_consumer_idx >= ring->depth) { + cached_consumer_idx = cons.load(cuda::memory_order_acquire); + *ring->consumer_idx_cache = cached_consumer_idx; + + if (shut.load(cuda::memory_order_relaxed) == static_cast(nixl_proxy_control_state_t::SHUTDOWN)) { return NIXL_ERR_BACKEND; } diff --git a/src/core/device_proxy/proxy_protocol.h b/src/core/device_proxy/proxy_protocol.h index 5f3add710d..06a55af0a6 100644 --- a/src/core/device_proxy/proxy_protocol.h +++ b/src/core/device_proxy/proxy_protocol.h @@ -60,6 +60,8 @@ struct nixlProxyWorkRing { uint64_t *producer_idx = nullptr; /** Mapped pinned consumer; host proxy uses __atomic_* on host alias (nixlProxyChannelState). */ uint64_t *consumer_idx = nullptr; + /** Device-resident cached consumer index; GPU refreshes from consumer_idx only when full. */ + uint64_t *consumer_idx_cache = nullptr; /** The depth of the work ring. */ uint32_t depth = 0; }; diff --git a/src/core/device_proxy/proxy_runtime.cpp b/src/core/device_proxy/proxy_runtime.cpp index bcb0d41b7f..ee3d6cf742 100644 --- a/src/core/device_proxy/proxy_runtime.cpp +++ b/src/core/device_proxy/proxy_runtime.cpp @@ -26,7 +26,7 @@ nixl_status_t nixlProxyMemViewRegistry::registerProxyMemView(nixlMemViewH backend_memview, - nixlMemViewH *proxy_memview) { + nixlMemViewH *proxy_memview) { if (proxy_memview == nullptr) { return NIXL_ERR_INVALID_PARAM; } @@ -38,29 +38,26 @@ nixlProxyMemViewRegistry::registerProxyMemView(nixlMemViewH backend_memview, *proxy_memview = reinterpret_cast(entry.proxy_memview_id); ++next_proxy_memview_id_; - NIXL_DEBUG << "nixlProxyMemViewRegistry::register: backend_mvh=" - << backend_memview << " -> proxy_id=" - << (next_proxy_memview_id_ - 1); + NIXL_DEBUG << "nixlProxyMemViewRegistry::register: backend_mvh=" << backend_memview + << " -> proxy_id=" << (next_proxy_memview_id_ - 1); return NIXL_SUCCESS; } nixl_status_t -nixlProxyMemViewRegistry::prepMemView(const nixl_meta_dlist_t &dlist, - nixlMemViewH *proxy_memview) { +nixlProxyMemViewRegistry::prepMemView(const nixl_meta_dlist_t &dlist, nixlMemViewH *proxy_memview) { return prepMemView(nullptr, dlist, proxy_memview); } nixl_status_t -nixlProxyMemViewRegistry::prepMemView( - const nixl_remote_meta_dlist_t &dlist, - nixlMemViewH *proxy_memview) { +nixlProxyMemViewRegistry::prepMemView(const nixl_remote_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview) { return prepMemView(nullptr, dlist, proxy_memview); } nixl_status_t nixlProxyMemViewRegistry::prepMemView(nixlMemViewH backend_memview, - const nixl_meta_dlist_t &dlist, - nixlMemViewH *proxy_memview) { + const nixl_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview) { if (proxy_memview == nullptr) { return NIXL_ERR_INVALID_PARAM; } @@ -82,10 +79,9 @@ nixlProxyMemViewRegistry::prepMemView(nixlMemViewH backend_memview, } nixl_status_t -nixlProxyMemViewRegistry::prepMemView( - nixlMemViewH backend_memview, - const nixl_remote_meta_dlist_t &dlist, - nixlMemViewH *proxy_memview) { +nixlProxyMemViewRegistry::prepMemView(nixlMemViewH backend_memview, + const nixl_remote_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview) { if (proxy_memview == nullptr) { return NIXL_ERR_INVALID_PARAM; } @@ -113,14 +109,13 @@ nixlProxyMemViewRegistry::unregisterProxyMemView(nixlMemViewH proxy_memview) { return NIXL_ERR_INVALID_PARAM; } entry->state = ProxyMemViewRegEntryState::ENTRY_RETIRED; - NIXL_DEBUG << "nixlProxyMemViewRegistry::unregister: proxy_id=" - << entry->proxy_memview_id; + NIXL_DEBUG << "nixlProxyMemViewRegistry::unregister: proxy_id=" << entry->proxy_memview_id; return NIXL_SUCCESS; } bool nixlProxyMemViewRegistry::resolveProxyMemView(nixlMemViewH proxy_memview, - nixlMemViewH &backend_memview) const { + nixlMemViewH &backend_memview) const { const RegistryEntry *entry = getEntryForHandle(proxy_memview); if (entry == nullptr || entry->state == ProxyMemViewRegEntryState::ENTRY_RETIRED) { return false; @@ -131,7 +126,7 @@ nixlProxyMemViewRegistry::resolveProxyMemView(nixlMemViewH proxy_memview, bool nixlProxyMemViewRegistry::resolveProxyMemViewId(uint64_t proxy_memview_id, - nixlMemViewH &backend_memview) const { + nixlMemViewH &backend_memview) const { const RegistryEntry *entry = getEntryForId(proxy_memview_id); if (entry == nullptr || entry->state == ProxyMemViewRegEntryState::ENTRY_RETIRED) { return false; @@ -142,7 +137,7 @@ nixlProxyMemViewRegistry::resolveProxyMemViewId(uint64_t proxy_memview_id, nixl_status_t nixlProxyMemViewRegistry::storeMetadata(nixlMemViewH proxy_memview, - const nixl_meta_dlist_t &dlist) { + const nixl_meta_dlist_t &dlist) { RegistryEntry *entry = getEntryForHandle(proxy_memview); if (entry == nullptr || entry->state == ProxyMemViewRegEntryState::ENTRY_RETIRED) { return NIXL_ERR_NOT_FOUND; @@ -160,7 +155,7 @@ nixlProxyMemViewRegistry::storeMetadata(nixlMemViewH proxy_memview, nixl_status_t nixlProxyMemViewRegistry::storeMetadata(nixlMemViewH proxy_memview, - const nixl_remote_meta_dlist_t &dlist) { + const nixl_remote_meta_dlist_t &dlist) { RegistryEntry *entry = getEntryForHandle(proxy_memview); if (entry == nullptr || entry->state == ProxyMemViewRegEntryState::ENTRY_RETIRED) { return NIXL_ERR_NOT_FOUND; @@ -178,7 +173,7 @@ nixlProxyMemViewRegistry::storeMetadata(nixlMemViewH proxy_memview, nixl_status_t nixlProxyMemViewRegistry::prepareSubmission(const nixlProxySubmission &submission, - nixlBackendProxySubmission &prepared_submission) const { + nixlBackendProxySubmission &prepared_submission) const { bool needs_source = false; size_t transfer_size = 0; switch (submission.opcode) { @@ -197,13 +192,12 @@ nixlProxyMemViewRegistry::prepareSubmission(const nixlProxySubmission &submissio const RemoteMetadata *remote_metadata = nullptr; const ProxyMemViewRegStoredEntry *dst_metadata = nullptr; - nixl_status_t status = getRemoteEntryForSubmission( - submission.dst_proxy_memview_id, - submission.dst_index, - submission.dst_offset, - transfer_size, - remote_metadata, - dst_metadata); + nixl_status_t status = getRemoteEntryForSubmission(submission.dst_proxy_memview_id, + submission.dst_index, + submission.dst_offset, + transfer_size, + remote_metadata, + dst_metadata); if (status != NIXL_SUCCESS) { return status; } @@ -217,32 +211,29 @@ nixlProxyMemViewRegistry::prepareSubmission(const nixlProxySubmission &submissio prepared.value = submission.value; prepared.remote_agent = remote_metadata->remote_agent; prepared.remote.mem_type = remote_metadata->mem_type; - prepared.remote.desc = nixlMetaDesc( - dst_metadata->base_addr + submission.dst_offset, - transfer_size, - dst_metadata->dev_id, - dst_metadata->metadata); + prepared.remote.desc = nixlMetaDesc(dst_metadata->base_addr + submission.dst_offset, + transfer_size, + dst_metadata->dev_id, + dst_metadata->metadata); if (needs_source) { const LocalMetadata *local_metadata = nullptr; const ProxyMemViewRegStoredEntry *src_metadata = nullptr; - status = getLocalEntryForSubmission( - submission.src_proxy_memview_id, - submission.src_index, - submission.src_offset, - transfer_size, - local_metadata, - src_metadata); + status = getLocalEntryForSubmission(submission.src_proxy_memview_id, + submission.src_index, + submission.src_offset, + transfer_size, + local_metadata, + src_metadata); if (status != NIXL_SUCCESS) { return status; } prepared.local.mem_type = local_metadata->mem_type; - prepared.local.desc = nixlMetaDesc( - src_metadata->base_addr + submission.src_offset, - transfer_size, - src_metadata->dev_id, - src_metadata->metadata); + prepared.local.desc = nixlMetaDesc(src_metadata->base_addr + submission.src_offset, + transfer_size, + src_metadata->dev_id, + src_metadata->metadata); } prepared_submission = prepared; @@ -268,9 +259,8 @@ nixlProxyMemViewRegistry::getEntryForHandle(nixlMemViewH proxy_memview) const { nixlProxyMemViewRegistry::RegistryEntry * nixlProxyMemViewRegistry::getEntryForId(uint64_t proxy_memview_id) { - if (proxy_memview_id < 1 - || proxy_memview_id >= next_proxy_memview_id_ - || proxy_memview_id > entries_.size()) { + if (proxy_memview_id < 1 || proxy_memview_id >= next_proxy_memview_id_ || + proxy_memview_id > entries_.size()) { return nullptr; } return &entries_[proxy_memview_id - 1]; @@ -278,26 +268,27 @@ nixlProxyMemViewRegistry::getEntryForId(uint64_t proxy_memview_id) { const nixlProxyMemViewRegistry::RegistryEntry * nixlProxyMemViewRegistry::getEntryForId(uint64_t proxy_memview_id) const { - if (proxy_memview_id < 1 - || proxy_memview_id >= next_proxy_memview_id_ - || proxy_memview_id > entries_.size()) { + if (proxy_memview_id < 1 || proxy_memview_id >= next_proxy_memview_id_ || + proxy_memview_id > entries_.size()) { return nullptr; } return &entries_[proxy_memview_id - 1]; } nixl_status_t -nixlProxyMemViewRegistry::getRemoteEntryForSubmission(uint64_t proxy_memview_id, - size_t index, - size_t offset, - size_t size, - const RemoteMetadata *&metadata, - const ProxyMemViewRegStoredEntry *&entry) const { +nixlProxyMemViewRegistry::getRemoteEntryForSubmission( + uint64_t proxy_memview_id, + size_t index, + size_t offset, + size_t size, + const RemoteMetadata *&metadata, + const ProxyMemViewRegStoredEntry *&entry) const { metadata = nullptr; entry = nullptr; const RegistryEntry *registry_entry = getEntryForId(proxy_memview_id); - if (registry_entry == nullptr || registry_entry->state != ProxyMemViewRegEntryState::ENTRY_READY) { + if (registry_entry == nullptr || + registry_entry->state != ProxyMemViewRegEntryState::ENTRY_READY) { NIXL_DEBUG << "nixlProxyMemViewRegistry::prepareSubmission: dst not ready" << " dst_proxy_id=" << proxy_memview_id; return NIXL_ERR_NOT_FOUND; @@ -324,17 +315,19 @@ nixlProxyMemViewRegistry::getRemoteEntryForSubmission(uint64_t proxy_memview_id, } nixl_status_t -nixlProxyMemViewRegistry::getLocalEntryForSubmission(uint64_t proxy_memview_id, - size_t index, - size_t offset, - size_t size, - const LocalMetadata *&metadata, - const ProxyMemViewRegStoredEntry *&entry) const { +nixlProxyMemViewRegistry::getLocalEntryForSubmission( + uint64_t proxy_memview_id, + size_t index, + size_t offset, + size_t size, + const LocalMetadata *&metadata, + const ProxyMemViewRegStoredEntry *&entry) const { metadata = nullptr; entry = nullptr; const RegistryEntry *registry_entry = getEntryForId(proxy_memview_id); - if (registry_entry == nullptr || registry_entry->state != ProxyMemViewRegEntryState::ENTRY_READY) { + if (registry_entry == nullptr || + registry_entry->state != ProxyMemViewRegEntryState::ENTRY_READY) { NIXL_DEBUG << "nixlProxyMemViewRegistry::prepareSubmission: src not ready" << " src_proxy_id=" << proxy_memview_id; return NIXL_ERR_NOT_FOUND; @@ -361,24 +354,26 @@ nixlProxyMemViewRegistry::getLocalEntryForSubmission(uint64_t proxy_memview_id, } bool -nixlProxyMemViewRegistry::rangeFits(const ProxyMemViewRegStoredEntry &entry, size_t offset, size_t size) { +nixlProxyMemViewRegistry::rangeFits(const ProxyMemViewRegStoredEntry &entry, + size_t offset, + size_t size) { return offset <= entry.len && size <= entry.len - offset; } void -nixlProxyMemViewRegistry::fillLocalMetadata(const nixl_meta_dlist_t &dlist, - LocalMetadata &out) { +nixlProxyMemViewRegistry::fillLocalMetadata(const nixl_meta_dlist_t &dlist, LocalMetadata &out) { out = LocalMetadata{}; out.mem_type = dlist.getType(); out.entries.reserve(dlist.descCount()); for (const auto &desc : dlist) { - out.entries.push_back(ProxyMemViewRegStoredEntry{desc.addr, desc.len, desc.devId, desc.metadataP}); + out.entries.push_back( + ProxyMemViewRegStoredEntry{desc.addr, desc.len, desc.devId, desc.metadataP}); } } void nixlProxyMemViewRegistry::fillRemoteMetadata(const nixl_remote_meta_dlist_t &dlist, - RemoteMetadata &out) { + RemoteMetadata &out) { out = RemoteMetadata{}; out.mem_type = dlist.getType(); out.entries.reserve(dlist.descCount()); @@ -386,23 +381,25 @@ nixlProxyMemViewRegistry::fillRemoteMetadata(const nixl_remote_meta_dlist_t &dli if (out.remote_agent.empty() && desc.remoteAgent != nixl_null_agent) { out.remote_agent = desc.remoteAgent; } - out.entries.push_back(ProxyMemViewRegStoredEntry{desc.addr, desc.len, desc.devId, desc.metadataP}); + out.entries.push_back( + ProxyMemViewRegStoredEntry{desc.addr, desc.len, desc.devId, desc.metadataP}); } } nixl_status_t nixlProxyChannelState::allocate(uint32_t channel_id, uint32_t depth) { - NIXL_INFO << "nixlProxyChannelState::allocate: channel_id=" << channel_id - << " depth=" << depth; + NIXL_INFO << "nixlProxyChannelState::allocate: channel_id=" << channel_id << " depth=" << depth; ring_depth_ = depth; - if (cudaMalloc(reinterpret_cast(&work_ring_dev_), - sizeof(nixlProxyWorkRing)) != cudaSuccess - || cudaMalloc(reinterpret_cast(&producer_idx_dev_), - sizeof(uint64_t)) != cudaSuccess - || cudaMallocHost(&records_host_, sizeof(nixlProxySubmission) * depth) != cudaSuccess - || cudaMallocHost(reinterpret_cast(&consumer_idx_host_), - sizeof(uint64_t)) != cudaSuccess - || cudaMallocHost(&completion_slot_host_, sizeof(nixlProxyCompletionSlot)) != cudaSuccess) { + if (cudaMalloc(reinterpret_cast(&work_ring_dev_), sizeof(nixlProxyWorkRing)) != + cudaSuccess || + cudaMalloc(reinterpret_cast(&producer_idx_dev_), sizeof(uint64_t)) != + cudaSuccess || + cudaMalloc(reinterpret_cast(&consumer_idx_cache_dev_), sizeof(uint64_t)) != + cudaSuccess || + cudaMallocHost(&records_host_, sizeof(nixlProxySubmission) * depth) != cudaSuccess || + cudaMallocHost(reinterpret_cast(&consumer_idx_host_), sizeof(uint64_t)) != + cudaSuccess || + cudaMallocHost(&completion_slot_host_, sizeof(nixlProxyCompletionSlot)) != cudaSuccess) { NIXL_ERROR << "nixlProxyChannelState::allocate: CUDA allocation failed for channel " << channel_id; deallocate(); @@ -433,37 +430,35 @@ nixlProxyChannelState::allocate(uint32_t channel_id, uint32_t depth) { for (uint32_t i = 0; i < depth; ++i) { records_host_[i] = nixlProxySubmission{}; } - if (cudaMemset(producer_idx_dev_, 0, sizeof(*producer_idx_dev_)) != cudaSuccess) { + if (cudaMemset(producer_idx_dev_, 0, sizeof(*producer_idx_dev_)) != cudaSuccess || + cudaMemset(consumer_idx_cache_dev_, 0, sizeof(*consumer_idx_cache_dev_)) != cudaSuccess) { deallocate(); return NIXL_ERR_BACKEND; } __atomic_store_n(consumer_idx_host_, uint64_t{0}, __ATOMIC_RELEASE); completion_slot_host_->next_status = NIXL_IN_PROG; - __atomic_store_n(&completion_slot_host_->completed_idx, - uint64_t{0}, __ATOMIC_RELEASE); + __atomic_store_n(&completion_slot_host_->completed_idx, uint64_t{0}, __ATOMIC_RELEASE); nixlProxyWorkRing work_ring{ records_dev_ptr, producer_idx_dev_, consumer_idx_dev, + consumer_idx_cache_dev_, depth, }; - if (cudaMemcpy(work_ring_dev_, - &work_ring, - sizeof(work_ring), - cudaMemcpyHostToDevice) != cudaSuccess) { + if (cudaMemcpy(work_ring_dev_, &work_ring, sizeof(work_ring), cudaMemcpyHostToDevice) != + cudaSuccess) { deallocate(); return NIXL_ERR_BACKEND; } - device_view = nixlProxyChannelView{ work_ring_dev_, completion_slot_dev_, channel_id }; + device_view = nixlProxyChannelView{work_ring_dev_, completion_slot_dev_, channel_id}; inflight_requests.clear(); NIXL_INFO << "nixlProxyChannelState::allocate: channel " << channel_id << " ready" - << " work_ring(dev)=" << work_ring_dev_ - << " records=" << records_host_ - << " records(dev)=" << records_dev_ptr - << " producer_idx(dev)=" << producer_idx_dev_ + << " work_ring(dev)=" << work_ring_dev_ << " records=" << records_host_ + << " records(dev)=" << records_dev_ptr << " producer_idx(dev)=" << producer_idx_dev_ << " consumer_idx(host)=" << consumer_idx_host_ << " consumer_idx(dev)=" << consumer_idx_dev + << " consumer_idx_cache(dev)=" << consumer_idx_cache_dev_ << " completion_slot(host)=" << completion_slot_host_ << " completion_slot(dev)=" << completion_slot_dev_; return NIXL_SUCCESS; @@ -474,12 +469,16 @@ nixlProxyChannelState::deallocate() noexcept { if (completion_slot_host_) { cudaFreeHost(completion_slot_host_); completion_slot_host_ = nullptr; - completion_slot_dev_ = nullptr; + completion_slot_dev_ = nullptr; } if (producer_idx_dev_) { cudaFree(producer_idx_dev_); producer_idx_dev_ = nullptr; } + if (consumer_idx_cache_dev_) { + cudaFree(consumer_idx_cache_dev_); + consumer_idx_cache_dev_ = nullptr; + } if (consumer_idx_host_) { cudaFreeHost(consumer_idx_host_); consumer_idx_host_ = nullptr; @@ -508,23 +507,25 @@ nixlProxyChannelState & nixlProxyChannelState::operator=(nixlProxyChannelState &&other) noexcept { if (this != &other) { deallocate(); - device_view = other.device_view; + device_view = other.device_view; inflight_requests = std::move(other.inflight_requests); - work_ring_dev_ = other.work_ring_dev_; - records_host_ = other.records_host_; + work_ring_dev_ = other.work_ring_dev_; + records_host_ = other.records_host_; producer_idx_dev_ = other.producer_idx_dev_; - consumer_idx_host_ = other.consumer_idx_host_; - ring_depth_ = other.ring_depth_; - completion_slot_host_ = other.completion_slot_host_; - completion_slot_dev_ = other.completion_slot_dev_; - other.work_ring_dev_ = nullptr; - other.records_host_ = nullptr; - other.producer_idx_dev_ = nullptr; - other.consumer_idx_host_ = nullptr; - other.ring_depth_ = 0; + consumer_idx_host_ = other.consumer_idx_host_; + consumer_idx_cache_dev_ = other.consumer_idx_cache_dev_; + ring_depth_ = other.ring_depth_; + completion_slot_host_ = other.completion_slot_host_; + completion_slot_dev_ = other.completion_slot_dev_; + other.work_ring_dev_ = nullptr; + other.records_host_ = nullptr; + other.producer_idx_dev_ = nullptr; + other.consumer_idx_host_ = nullptr; + other.consumer_idx_cache_dev_ = nullptr; + other.ring_depth_ = 0; other.completion_slot_host_ = nullptr; - other.completion_slot_dev_ = nullptr; - other.device_view = nixlProxyChannelView{}; + other.completion_slot_dev_ = nullptr; + other.device_view = nixlProxyChannelView{}; } return *this; } @@ -539,12 +540,11 @@ nixlProxyRuntime::~nixlProxyRuntime() { nixl_status_t nixlProxyRuntime::init(std::unique_ptr backend, - uint32_t channel_count, - uint32_t worker_count, - uint64_t pthr_delay_us) { + uint32_t channel_count, + uint32_t worker_count, + uint64_t pthr_delay_us) { NIXL_INFO << "ProxyRuntime::init: channel_count=" << channel_count - << " worker_count=" << worker_count - << " pthr_delay_us=" << pthr_delay_us + << " worker_count=" << worker_count << " pthr_delay_us=" << pthr_delay_us << " backend=" << backend.get(); if (backend == nullptr || channel_count == 0 || worker_count == 0) { NIXL_ERROR << "ProxyRuntime::init: invalid params"; @@ -554,8 +554,8 @@ nixlProxyRuntime::init(std::unique_ptr backend, backend_ = std::move(backend); memview_registry_.clear(); - if (cudaMallocHost(reinterpret_cast(&shutdown_word_host_), - sizeof(uint32_t)) != cudaSuccess) { + if (cudaMallocHost(reinterpret_cast(&shutdown_word_host_), sizeof(uint32_t)) != + cudaSuccess) { NIXL_ERROR << "ProxyRuntime::init: failed to allocate shutdown_word"; shutdown_word_host_ = nullptr; backend_.reset(); @@ -582,7 +582,7 @@ nixlProxyRuntime::init(std::unique_ptr backend, NIXL_ERROR << "ProxyRuntime::init: backend init failed: " << rc; cudaFreeHost(shutdown_word_host_); shutdown_word_host_ = nullptr; - shutdown_word_dev_ = nullptr; + shutdown_word_dev_ = nullptr; backend_.reset(); return rc; } @@ -598,7 +598,7 @@ nixlProxyRuntime::init(std::unique_ptr backend, backend_->shutdown(); cudaFreeHost(shutdown_word_host_); shutdown_word_host_ = nullptr; - shutdown_word_dev_ = nullptr; + shutdown_word_dev_ = nullptr; backend_.reset(); return rc; } @@ -611,7 +611,7 @@ nixlProxyRuntime::init(std::unique_ptr backend, backend_->shutdown(); cudaFreeHost(shutdown_word_host_); shutdown_word_host_ = nullptr; - shutdown_word_dev_ = nullptr; + shutdown_word_dev_ = nullptr; backend_.reset(); return NIXL_ERR_BACKEND; } @@ -629,21 +629,17 @@ nixlProxyRuntime::init(std::unique_ptr backend, backend_->shutdown(); cudaFreeHost(shutdown_word_host_); shutdown_word_host_ = nullptr; - shutdown_word_dev_ = nullptr; + shutdown_word_dev_ = nullptr; backend_.reset(); return NIXL_ERR_BACKEND; } nixlProxyDeviceContextData device_context{ - device_channel_views_dev_, - channel_count, - shutdown_word_dev_ - }; + device_channel_views_dev_, channel_count, shutdown_word_dev_}; if (cudaMalloc(reinterpret_cast(&device_context_), - sizeof(nixlProxyDeviceContextData)) != cudaSuccess - || cudaMemcpy(device_context_, - &device_context, - sizeof(device_context), - cudaMemcpyHostToDevice) != cudaSuccess) { + sizeof(nixlProxyDeviceContextData)) != cudaSuccess || + cudaMemcpy( + device_context_, &device_context, sizeof(device_context), cudaMemcpyHostToDevice) != + cudaSuccess) { if (device_context_) { cudaFree(device_context_); device_context_ = nullptr; @@ -655,7 +651,7 @@ nixlProxyRuntime::init(std::unique_ptr backend, backend_->shutdown(); cudaFreeHost(shutdown_word_host_); shutdown_word_host_ = nullptr; - shutdown_word_dev_ = nullptr; + shutdown_word_dev_ = nullptr; backend_.reset(); return NIXL_ERR_BACKEND; } @@ -666,30 +662,27 @@ nixlProxyRuntime::init(std::unique_ptr backend, for (uint32_t w = 0; w < worker_count; ++w) { uint32_t first_ch = (w * channel_count) / worker_count; - uint32_t end_ch = ((w + 1) * channel_count) / worker_count; - uint32_t n_ch = end_ch - first_ch; - - NIXL_INFO << "ProxyRuntime::init: worker " << w - << " assigned channels [" << first_ch << ", " << end_ch << ")"; - workers_.push_back(std::make_unique( - backend_.get(), - &memview_registry_, - shutdown_word_host_, - &channels_[first_ch], - n_ch, - pthr_delay_us)); - } - - NIXL_INFO << "ProxyRuntime::init: complete — " - << channel_count << " channels, " - << worker_count << " workers, " + uint32_t end_ch = ((w + 1) * channel_count) / worker_count; + uint32_t n_ch = end_ch - first_ch; + + NIXL_INFO << "ProxyRuntime::init: worker " << w << " assigned channels [" << first_ch + << ", " << end_ch << ")"; + workers_.push_back(std::make_unique(backend_.get(), + &memview_registry_, + shutdown_word_host_, + &channels_[first_ch], + n_ch, + pthr_delay_us)); + } + + NIXL_INFO << "ProxyRuntime::init: complete — " << channel_count << " channels, " << worker_count + << " workers, " << "device_context(dev)=" << device_context_; return NIXL_SUCCESS; } nixl_status_t -nixlProxyRuntime::loadRemoteConnInfo(const std::string &remote_name, - const nixl_blob_t &conn_info) { +nixlProxyRuntime::loadRemoteConnInfo(const std::string &remote_name, const nixl_blob_t &conn_info) { NIXL_INFO << "ProxyRuntime::loadRemoteConnInfo: remote='" << remote_name << "' conn_info_size=" << conn_info.size(); if (backend_ == nullptr) { @@ -702,20 +695,17 @@ nixlProxyRuntime::loadRemoteConnInfo(const std::string &remote_name, } nixl_status_t -nixlProxyRuntime::registerProxyMemView(nixlMemViewH backend_memview, - nixlMemViewH *proxy_memview) { +nixlProxyRuntime::registerProxyMemView(nixlMemViewH backend_memview, nixlMemViewH *proxy_memview) { return memview_registry_.registerProxyMemView(backend_memview, proxy_memview); } nixl_status_t -nixlProxyRuntime::prepMemView(const nixl_meta_dlist_t &dlist, - nixlMemViewH *proxy_memview) { +nixlProxyRuntime::prepMemView(const nixl_meta_dlist_t &dlist, nixlMemViewH *proxy_memview) { return memview_registry_.prepMemView(dlist, proxy_memview); } nixl_status_t -nixlProxyRuntime::prepMemView(const nixl_remote_meta_dlist_t &dlist, - nixlMemViewH *proxy_memview) { +nixlProxyRuntime::prepMemView(const nixl_remote_meta_dlist_t &dlist, nixlMemViewH *proxy_memview) { return memview_registry_.prepMemView(dlist, proxy_memview); } @@ -727,10 +717,9 @@ nixlProxyRuntime::prepMemView(nixlMemViewH backend_memview, } nixl_status_t -nixlProxyRuntime::prepMemView( - nixlMemViewH backend_memview, - const nixl_remote_meta_dlist_t &dlist, - nixlMemViewH *proxy_memview) { +nixlProxyRuntime::prepMemView(nixlMemViewH backend_memview, + const nixl_remote_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview) { return memview_registry_.prepMemView(backend_memview, dlist, proxy_memview); } @@ -740,33 +729,30 @@ nixlProxyRuntime::unregisterProxyMemView(nixlMemViewH proxy_memview) { } nixl_status_t -nixlProxyRuntime::storeMetadata(nixlMemViewH proxy_memview, - const nixl_meta_dlist_t &dlist) { +nixlProxyRuntime::storeMetadata(nixlMemViewH proxy_memview, const nixl_meta_dlist_t &dlist) { return memview_registry_.storeMetadata(proxy_memview, dlist); } nixl_status_t -nixlProxyRuntime::storeMetadata(nixlMemViewH proxy_memview, - const nixl_remote_meta_dlist_t &dlist) { +nixlProxyRuntime::storeMetadata(nixlMemViewH proxy_memview, const nixl_remote_meta_dlist_t &dlist) { return memview_registry_.storeMetadata(proxy_memview, dlist); } bool nixlProxyRuntime::resolveProxyMemView(nixlMemViewH proxy_memview, - nixlMemViewH &backend_memview) const { + nixlMemViewH &backend_memview) const { return memview_registry_.resolveProxyMemView(proxy_memview, backend_memview); } bool nixlProxyRuntime::resolveProxyMemViewId(uint64_t proxy_memview_id, - nixlMemViewH &backend_memview) const { + nixlMemViewH &backend_memview) const { return memview_registry_.resolveProxyMemViewId(proxy_memview_id, backend_memview); } nixl_status_t nixlProxyRuntime::startWorkers() { - NIXL_INFO << "ProxyRuntime::startWorkers: launching " - << workers_.size() << " worker thread(s)"; + NIXL_INFO << "ProxyRuntime::startWorkers: launching " << workers_.size() << " worker thread(s)"; if (shutdown_word_host_ == nullptr) { NIXL_ERROR << "ProxyRuntime::startWorkers: runtime not initialized"; return NIXL_ERR_NOT_SUPPORTED; @@ -837,7 +823,7 @@ nixlProxyRuntime::shutdown() { if (shutdown_word_host_) { cudaFreeHost(shutdown_word_host_); shutdown_word_host_ = nullptr; - shutdown_word_dev_ = nullptr; + shutdown_word_dev_ = nullptr; } if (device_channel_views_dev_) { cudaFree(device_channel_views_dev_); diff --git a/src/core/device_proxy/proxy_runtime.h b/src/core/device_proxy/proxy_runtime.h index 1c9e1eefd6..c12b1b2ab0 100644 --- a/src/core/device_proxy/proxy_runtime.h +++ b/src/core/device_proxy/proxy_runtime.h @@ -45,22 +45,26 @@ struct alignas(64) nixlProxyChannelState { nixlProxyWorkRing *work_ring_dev_ = nullptr; nixlProxySubmission *records_host_ = nullptr; /** Device-resident producer index; only the GPU updates it. */ - uint64_t *producer_idx_dev_ = nullptr; + uint64_t *producer_idx_dev_ = nullptr; /** Consumer count: host pinned; proxy uses __atomic_* on consumer_idx_host_. */ - uint64_t *consumer_idx_host_ = nullptr; + uint64_t *consumer_idx_host_ = nullptr; + /** Device-resident cache of consumer_idx_host_ used by GPU enqueue backpressure. */ + uint64_t *consumer_idx_cache_dev_ = nullptr; /** Host-side ring depth for the CPU worker; nixlProxyWorkRing itself is device-only. */ - uint32_t ring_depth_ = 0; + uint32_t ring_depth_ = 0; /** Mapped pinned host memory; proxy worker writes directly via host alias. */ - nixlProxyCompletionSlot *completion_slot_host_ = nullptr; + nixlProxyCompletionSlot *completion_slot_host_ = nullptr; /** Device-mapped alias of completion_slot_host_ for nixlProxyChannelView. */ - nixlProxyCompletionSlot *completion_slot_dev_ = nullptr; + nixlProxyCompletionSlot *completion_slot_dev_ = nullptr; nixlProxyChannelState() = default; ~nixlProxyChannelState(); nixlProxyChannelState(nixlProxyChannelState &&) noexcept; - nixlProxyChannelState &operator=(nixlProxyChannelState &&) noexcept; + nixlProxyChannelState & + operator=(nixlProxyChannelState &&) noexcept; nixlProxyChannelState(const nixlProxyChannelState &) = delete; - nixlProxyChannelState &operator=(const nixlProxyChannelState &) = delete; + nixlProxyChannelState & + operator=(const nixlProxyChannelState &) = delete; nixl_status_t allocate(uint32_t channel_id, uint32_t depth); @@ -70,232 +74,225 @@ struct alignas(64) nixlProxyChannelState { }; class nixlProxyMemViewRegistry { - public: - nixl_status_t - registerProxyMemView(nixlMemViewH backend_memview, - nixlMemViewH *proxy_memview); - - nixl_status_t - prepMemView(const nixl_meta_dlist_t &dlist, - nixlMemViewH *proxy_memview); - - nixl_status_t - prepMemView(const nixl_remote_meta_dlist_t &dlist, - nixlMemViewH *proxy_memview); - - nixl_status_t - prepMemView(nixlMemViewH backend_memview, - const nixl_meta_dlist_t &dlist, - nixlMemViewH *proxy_memview); - - nixl_status_t - prepMemView(nixlMemViewH backend_memview, - const nixl_remote_meta_dlist_t &dlist, - nixlMemViewH *proxy_memview); - - nixl_status_t - unregisterProxyMemView(nixlMemViewH proxy_memview); - - nixl_status_t - storeMetadata(nixlMemViewH proxy_memview, - const nixl_meta_dlist_t &dlist); - - nixl_status_t - storeMetadata(nixlMemViewH proxy_memview, - const nixl_remote_meta_dlist_t &dlist); - - bool - resolveProxyMemView(nixlMemViewH proxy_memview, - nixlMemViewH &backend_memview) const; - - bool - resolveProxyMemViewId(uint64_t proxy_memview_id, - nixlMemViewH &backend_memview) const; - - nixl_status_t - prepareSubmission(const nixlProxySubmission &submission, - nixlBackendProxySubmission &prepared_submission) const; - - void - clear() noexcept; - - private: - struct ProxyMemViewRegStoredEntry { - uintptr_t base_addr = 0; - size_t len = 0; - uint64_t dev_id = 0; - nixlBackendMD *metadata = nullptr; - }; - - struct LocalMetadata { - nixl_mem_t mem_type = DRAM_SEG; - std::vector entries; - }; - - struct RemoteMetadata { - nixl_mem_t mem_type = DRAM_SEG; - std::string remote_agent; - std::vector entries; - }; - - enum class ProxyMemViewRegEntryState : uint8_t { - ENTRY_ALLOCATED, - ENTRY_READY, - ENTRY_RETIRED, - }; - - enum class ProxyMemViewRegMetadataKind : uint8_t { - METADATA_KIND_NONE, - METADATA_KIND_LOCAL, - METADATA_KIND_REMOTE, - }; - - struct RegistryEntry { - uint64_t proxy_memview_id = 0; - nixlMemViewH backend_memview = nullptr; - ProxyMemViewRegEntryState state = ProxyMemViewRegEntryState::ENTRY_ALLOCATED; - ProxyMemViewRegMetadataKind metadata_kind = ProxyMemViewRegMetadataKind::METADATA_KIND_NONE; - LocalMetadata local_metadata{}; - RemoteMetadata remote_metadata{}; - }; - - RegistryEntry * - getEntryForHandle(nixlMemViewH proxy_memview); - - const RegistryEntry * - getEntryForHandle(nixlMemViewH proxy_memview) const; - - RegistryEntry * - getEntryForId(uint64_t proxy_memview_id); - - const RegistryEntry * - getEntryForId(uint64_t proxy_memview_id) const; - - nixl_status_t - getRemoteEntryForSubmission(uint64_t proxy_memview_id, - size_t index, - size_t offset, - size_t size, - const RemoteMetadata *&metadata, - const ProxyMemViewRegStoredEntry *&entry) const; - - nixl_status_t - getLocalEntryForSubmission(uint64_t proxy_memview_id, - size_t index, - size_t offset, - size_t size, - const LocalMetadata *&metadata, - const ProxyMemViewRegStoredEntry *&entry) const; - - static bool - rangeFits(const ProxyMemViewRegStoredEntry &entry, size_t offset, size_t size); - - static void - fillLocalMetadata(const nixl_meta_dlist_t &dlist, LocalMetadata &out); - - static void - fillRemoteMetadata(const nixl_remote_meta_dlist_t &dlist, RemoteMetadata &out); - - std::vector entries_; - uint64_t next_proxy_memview_id_ = 1; +public: + nixl_status_t + registerProxyMemView(nixlMemViewH backend_memview, nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(const nixl_meta_dlist_t &dlist, nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(const nixl_remote_meta_dlist_t &dlist, nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(nixlMemViewH backend_memview, + const nixl_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(nixlMemViewH backend_memview, + const nixl_remote_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview); + + nixl_status_t + unregisterProxyMemView(nixlMemViewH proxy_memview); + + nixl_status_t + storeMetadata(nixlMemViewH proxy_memview, const nixl_meta_dlist_t &dlist); + + nixl_status_t + storeMetadata(nixlMemViewH proxy_memview, const nixl_remote_meta_dlist_t &dlist); + + bool + resolveProxyMemView(nixlMemViewH proxy_memview, nixlMemViewH &backend_memview) const; + + bool + resolveProxyMemViewId(uint64_t proxy_memview_id, nixlMemViewH &backend_memview) const; + + nixl_status_t + prepareSubmission(const nixlProxySubmission &submission, + nixlBackendProxySubmission &prepared_submission) const; + + void + clear() noexcept; + +private: + struct ProxyMemViewRegStoredEntry { + uintptr_t base_addr = 0; + size_t len = 0; + uint64_t dev_id = 0; + nixlBackendMD *metadata = nullptr; + }; + + struct LocalMetadata { + nixl_mem_t mem_type = DRAM_SEG; + std::vector entries; + }; + + struct RemoteMetadata { + nixl_mem_t mem_type = DRAM_SEG; + std::string remote_agent; + std::vector entries; + }; + + enum class ProxyMemViewRegEntryState : uint8_t { + ENTRY_ALLOCATED, + ENTRY_READY, + ENTRY_RETIRED, + }; + + enum class ProxyMemViewRegMetadataKind : uint8_t { + METADATA_KIND_NONE, + METADATA_KIND_LOCAL, + METADATA_KIND_REMOTE, + }; + + struct RegistryEntry { + uint64_t proxy_memview_id = 0; + nixlMemViewH backend_memview = nullptr; + ProxyMemViewRegEntryState state = ProxyMemViewRegEntryState::ENTRY_ALLOCATED; + ProxyMemViewRegMetadataKind metadata_kind = ProxyMemViewRegMetadataKind::METADATA_KIND_NONE; + LocalMetadata local_metadata{}; + RemoteMetadata remote_metadata{}; + }; + + RegistryEntry * + getEntryForHandle(nixlMemViewH proxy_memview); + + const RegistryEntry * + getEntryForHandle(nixlMemViewH proxy_memview) const; + + RegistryEntry * + getEntryForId(uint64_t proxy_memview_id); + + const RegistryEntry * + getEntryForId(uint64_t proxy_memview_id) const; + + nixl_status_t + getRemoteEntryForSubmission(uint64_t proxy_memview_id, + size_t index, + size_t offset, + size_t size, + const RemoteMetadata *&metadata, + const ProxyMemViewRegStoredEntry *&entry) const; + + nixl_status_t + getLocalEntryForSubmission(uint64_t proxy_memview_id, + size_t index, + size_t offset, + size_t size, + const LocalMetadata *&metadata, + const ProxyMemViewRegStoredEntry *&entry) const; + + static bool + rangeFits(const ProxyMemViewRegStoredEntry &entry, size_t offset, size_t size); + + static void + fillLocalMetadata(const nixl_meta_dlist_t &dlist, LocalMetadata &out); + + static void + fillRemoteMetadata(const nixl_remote_meta_dlist_t &dlist, RemoteMetadata &out); + + std::vector entries_; + uint64_t next_proxy_memview_id_ = 1; }; class nixlProxyRuntime { - public: - nixlProxyRuntime(); - ~nixlProxyRuntime(); - - nixlProxyRuntime(nixlProxyRuntime &&) = delete; - nixlProxyRuntime(const nixlProxyRuntime &) = delete; - nixlProxyRuntime& operator=(nixlProxyRuntime &&) = delete; - nixlProxyRuntime& operator=(const nixlProxyRuntime &) = delete; - - nixl_status_t - init(std::unique_ptr backend, - uint32_t channel_count, - uint32_t worker_count, - uint64_t pthr_delay_us = 0); - - nixl_status_t - loadRemoteConnInfo(const std::string &remote_name, - const nixl_blob_t &conn_info); - - nixl_status_t - registerProxyMemView(nixlMemViewH backend_memview, - nixlMemViewH *proxy_memview); - - nixl_status_t - prepMemView(const nixl_meta_dlist_t &dlist, - nixlMemViewH *proxy_memview); - - nixl_status_t - prepMemView(const nixl_remote_meta_dlist_t &dlist, - nixlMemViewH *proxy_memview); - - nixl_status_t - prepMemView(nixlMemViewH backend_memview, - const nixl_meta_dlist_t &dlist, - nixlMemViewH *proxy_memview); - - nixl_status_t - prepMemView(nixlMemViewH backend_memview, - const nixl_remote_meta_dlist_t &dlist, - nixlMemViewH *proxy_memview); - - nixl_status_t - unregisterProxyMemView(nixlMemViewH proxy_memview); - - nixl_status_t - storeMetadata(nixlMemViewH proxy_memview, - const nixl_meta_dlist_t &dlist); - - nixl_status_t - storeMetadata(nixlMemViewH proxy_memview, - const nixl_remote_meta_dlist_t &dlist); - - bool - resolveProxyMemView(nixlMemViewH proxy_memview, - nixlMemViewH &backend_memview) const; - - bool - resolveProxyMemViewId(uint64_t proxy_memview_id, - nixlMemViewH &backend_memview) const; - - nixl_status_t - startWorkers(); - - nixl_status_t - shutdown(); - - const nixlProxyMemViewRegistry & - memviewRegistry() const { return memview_registry_; } - - uint32_t - channelCount() const { return static_cast(channels_.size()); } - - const nixlProxyChannelView * - deviceChannelViews() const { - return device_channel_views_.empty() ? nullptr : device_channel_views_.data(); - } - - nixlProxyDeviceContextData * - deviceContext() const { return device_context_; } - - private: - void - joinWorkerThreads() noexcept; - - std::vector channels_; - std::vector device_channel_views_; - nixlProxyChannelView *device_channel_views_dev_ = nullptr; - nixlProxyDeviceContextData *device_context_ = nullptr; - std::vector> workers_; - nixlProxyMemViewRegistry memview_registry_; - std::unique_ptr backend_; - uint32_t *shutdown_word_host_ = nullptr; - uint32_t *shutdown_word_dev_ = nullptr; - uint32_t ring_depth_ = kDefaultProxyRingDepth; - bool workers_started_ = false; +public: + nixlProxyRuntime(); + ~nixlProxyRuntime(); + + nixlProxyRuntime(nixlProxyRuntime &&) = delete; + nixlProxyRuntime(const nixlProxyRuntime &) = delete; + nixlProxyRuntime & + operator=(nixlProxyRuntime &&) = delete; + nixlProxyRuntime & + operator=(const nixlProxyRuntime &) = delete; + + nixl_status_t + init(std::unique_ptr backend, + uint32_t channel_count, + uint32_t worker_count, + uint64_t pthr_delay_us = 0); + + nixl_status_t + loadRemoteConnInfo(const std::string &remote_name, const nixl_blob_t &conn_info); + + nixl_status_t + registerProxyMemView(nixlMemViewH backend_memview, nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(const nixl_meta_dlist_t &dlist, nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(const nixl_remote_meta_dlist_t &dlist, nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(nixlMemViewH backend_memview, + const nixl_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview); + + nixl_status_t + prepMemView(nixlMemViewH backend_memview, + const nixl_remote_meta_dlist_t &dlist, + nixlMemViewH *proxy_memview); + + nixl_status_t + unregisterProxyMemView(nixlMemViewH proxy_memview); + + nixl_status_t + storeMetadata(nixlMemViewH proxy_memview, const nixl_meta_dlist_t &dlist); + + nixl_status_t + storeMetadata(nixlMemViewH proxy_memview, const nixl_remote_meta_dlist_t &dlist); + + bool + resolveProxyMemView(nixlMemViewH proxy_memview, nixlMemViewH &backend_memview) const; + + bool + resolveProxyMemViewId(uint64_t proxy_memview_id, nixlMemViewH &backend_memview) const; + + nixl_status_t + startWorkers(); + + nixl_status_t + shutdown(); + + const nixlProxyMemViewRegistry & + memviewRegistry() const { + return memview_registry_; + } + + uint32_t + channelCount() const { + return static_cast(channels_.size()); + } + + const nixlProxyChannelView * + deviceChannelViews() const { + return device_channel_views_.empty() ? nullptr : device_channel_views_.data(); + } + + nixlProxyDeviceContextData * + deviceContext() const { + return device_context_; + } + +private: + void + joinWorkerThreads() noexcept; + + std::vector channels_; + std::vector device_channel_views_; + nixlProxyChannelView *device_channel_views_dev_ = nullptr; + nixlProxyDeviceContextData *device_context_ = nullptr; + std::vector> workers_; + nixlProxyMemViewRegistry memview_registry_; + std::unique_ptr backend_; + uint32_t *shutdown_word_host_ = nullptr; + uint32_t *shutdown_word_dev_ = nullptr; + uint32_t ring_depth_ = kDefaultProxyRingDepth; + bool workers_started_ = false; }; #endif // NIXL_SRC_CORE_DEVICE_PROXY_PROXY_RUNTIME_H From 5f662b7529fb3712d5e5f9d456857105dc032a5f Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Wed, 3 Jun 2026 21:42:22 +0300 Subject: [PATCH 17/18] gpu: enqueue: use GPU-scope release for D2H FIFO publish Keep the ready-word publish at GPU scope for host-pinned proxy records. The enqueue path writes the payload first and commits op_idx last; the CPU worker acquire-polls op_idx before copying the record. This avoids the enqueue overhead of a system-scope release while documenting the D2H FIFO contract explicitly. Signed-off-by: Tomer Davidor --- src/api/gpu/proxy/nixl_device_proxy.cuh | 41 +++++++++++++++---------- 1 file changed, 24 insertions(+), 17 deletions(-) diff --git a/src/api/gpu/proxy/nixl_device_proxy.cuh b/src/api/gpu/proxy/nixl_device_proxy.cuh index 280291908e..d25c3c1f5a 100644 --- a/src/api/gpu/proxy/nixl_device_proxy.cuh +++ b/src/api/gpu/proxy/nixl_device_proxy.cuh @@ -28,9 +28,10 @@ struct ProxyDeviceContext; // Overlay struct written into nixlGpuXferStatusH::storage by enqueue() // and read back by pollXferStatus(). Must fit within the 64-byte opaque blob. struct ProxyXferStatus { - nixlProxyCompletionSlot *slot; // device pointer to the channel's nixlProxyCompletionSlot - uint64_t op_idx; + nixlProxyCompletionSlot *slot; // device pointer to the channel's nixlProxyCompletionSlot + uint64_t op_idx; }; + static_assert(sizeof(ProxyXferStatus) <= sizeof(nixlGpuXferStatusH), "ProxyXferStatus must fit in nixlGpuXferStatusH::storage"); @@ -43,7 +44,8 @@ extern __device__ __constant__ ProxyDeviceContext *g_nixl_proxy_ctx; __host__ inline cudaError_t nixlProxyPublishContext(nixlProxyDeviceContextData *ctx) { ProxyDeviceContext *device_ctx = reinterpret_cast(ctx); - cudaError_t err = cudaMemcpyToSymbol(g_nixl_proxy_ctx, &device_ctx, sizeof(ProxyDeviceContext *)); + cudaError_t err = + cudaMemcpyToSymbol(g_nixl_proxy_ctx, &device_ctx, sizeof(ProxyDeviceContext *)); if (err != cudaSuccess) { fprintf(stderr, "nixlProxyPublishContext: cudaMemcpyToSymbol failed: code=%d msg=%s\n", @@ -66,12 +68,12 @@ nixlProxyClearContext() { return err; } -__device__ __forceinline__ uint64_t +__device__ __forceinline__ uint64_t proxyMemViewIdFromHandle(nixlMemViewH mvh) { return static_cast(reinterpret_cast(mvh)); } -__device__ __forceinline__ ProxyDeviceContext * +__device__ __forceinline__ ProxyDeviceContext * load_proxy_context() { return g_nixl_proxy_ctx; } @@ -86,7 +88,8 @@ static_assert(sizeof(nixlProxyCompletionSlot::completed_idx) == 8, "completed_idx must be 64-bit to match producer_idx"); template -__device__ inline void nixlProxyExecInit(uint32_t &lane_id) { +__device__ inline void +nixlProxyExecInit(uint32_t &lane_id) { static_assert(level != nixl_gpu_level_t::GRID, "Proxy GPU backend does not support GRID-level operations"); @@ -100,7 +103,8 @@ __device__ inline void nixlProxyExecInit(uint32_t &lane_id) { } template -__device__ inline void nixlProxySync() { +__device__ inline void +nixlProxySync() { static_assert(level != nixl_gpu_level_t::GRID, "Proxy GPU backend does not support GRID-level operations"); @@ -128,10 +132,9 @@ struct ProxyDeviceContext : nixlProxyDeviceContextData { } nixlProxyChannelView &channel_view = channels[submission.channel_id]; - nixlProxyWorkRing *ring = channel_view.work_ring; + nixlProxyWorkRing *ring = channel_view.work_ring; - cuda::atomic_ref producer_idx( - *ring->producer_idx); + cuda::atomic_ref producer_idx(*ring->producer_idx); cuda::atomic_ref cons(*ring->consumer_idx); cuda::atomic_ref shut(*shutdown_word); @@ -145,8 +148,8 @@ struct ProxyDeviceContext : nixlProxyDeviceContextData { cached_consumer_idx = cons.load(cuda::memory_order_acquire); *ring->consumer_idx_cache = cached_consumer_idx; - if (shut.load(cuda::memory_order_relaxed) - == static_cast(nixl_proxy_control_state_t::SHUTDOWN)) { + if (shut.load(cuda::memory_order_relaxed) == + static_cast(nixl_proxy_control_state_t::SHUTDOWN)) { return NIXL_ERR_BACKEND; } } @@ -157,7 +160,13 @@ struct ProxyDeviceContext : nixlProxyDeviceContextData { // Signal this slot is ready for the consumer. The release // guarantees the record write above is visible before the // consumer reads op_idx via an acquire load. op_idx == 0 means empty. - cuda::atomic_ref record_op_idx( + submission.op_idx = 0; + ring->records[slot] = submission; + + // Avoiding system-scope release keeps enqueue from paying + // a global GPU memory drain; the CPU worker acquire-polls op_idx + // before copying the record. + cuda::atomic_ref record_op_idx( ring->records[slot].op_idx); record_op_idx.store(submission_op_idx, cuda::memory_order_release); @@ -179,11 +188,9 @@ struct ProxyDeviceContext : nixlProxyDeviceContextData { // latched the channel __device__ inline nixl_status_t pollXferStatus(const nixlGpuXferStatusH &xfer_status) const { - const ProxyXferStatus *pxs = - reinterpret_cast(xfer_status.storage); + const ProxyXferStatus *pxs = reinterpret_cast(xfer_status.storage); - cuda::atomic_ref comp_idx( - pxs->slot->completed_idx); + cuda::atomic_ref comp_idx(pxs->slot->completed_idx); const uint64_t completed_idx = comp_idx.load(cuda::memory_order_acquire); if (completed_idx > pxs->op_idx) { From 95a5503c803427ac2dd0f4a8f64a9f441159c38d Mon Sep 17 00:00:00 2001 From: Tomer Davidor Date: Mon, 13 Jul 2026 17:30:08 +0300 Subject: [PATCH 18/18] misc: meson.build: Update copyright date Signed-off-by: Tomer Davidor --- src/api/meson.build | 2 +- test/gtest/device_api/meson.build | 2 +- test/gtest/device_api/utils.cuh | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/api/meson.build b/src/api/meson.build index c59466b809..c8a561649d 100644 --- a/src/api/meson.build +++ b/src/api/meson.build @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/test/gtest/device_api/meson.build b/test/gtest/device_api/meson.build index 1be6c91867..6bd70ece95 100644 --- a/test/gtest/device_api/meson.build +++ b/test/gtest/device_api/meson.build @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/test/gtest/device_api/utils.cuh b/test/gtest/device_api/utils.cuh index 3d9ebc36e7..8d26b31d52 100644 --- a/test/gtest/device_api/utils.cuh +++ b/test/gtest/device_api/utils.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License");