From a9d274ce420c9a25113d8f0c463df35739ad3bcc Mon Sep 17 00:00:00 2001 From: romerojosh Date: Tue, 25 Aug 2026 18:31:29 +0000 Subject: [PATCH 1/2] Add automatic workspace management Signed-off-by: romerojosh --- docs/api/f_api.rst | 14 +- docs/basic_usage.rst | 36 +- include/cudecomp.h | 45 +- include/internal/common.h | 17 + src/cudecomp.cc | 871 ++++++++++++++--------- src/cudecomp_m.cuf | 66 +- tests/ctest/CMakeLists.txt | 36 + tests/ctest/api_tests.cc | 6 - tests/ctest/automatic_workspace_tests.cc | 240 +++++++ tests/ctest/fortran_halo_case.inc | 24 +- tests/ctest/fortran_transpose_case.inc | 8 +- tests/ctest/halo_tests.cc | 62 +- tests/ctest/transpose_tests.cc | 61 +- 13 files changed, 1063 insertions(+), 423 deletions(-) create mode 100644 tests/ctest/automatic_workspace_tests.cc diff --git a/docs/api/f_api.rst b/docs/api/f_api.rst index abbe449..088ab5d 100644 --- a/docs/api/f_api.rst +++ b/docs/api/f_api.rst @@ -471,7 +471,7 @@ _____________________ :p cudecompGridDesc grid_desc [in]: A cuDecomp grid descriptor. :p T input(*) [in]: Device array containing input X-axis aligned pencil data. :p T output(*) [out]: Device array to write output Y-axis aligned pencil data. If :code:`input` and :code:`output` are the same, operation is performed in-place - :p T work(*) [in]: Device array to use for transpose workspace. + :p T work(*) [in]: Caller-owned device array to use for transpose workspace, or :code:`CUDECOMP_WORKSPACE_AUTO` to use handle-owned workspace management. :p cudecompDataType dtype [in]: The :code:`cudecompDataType` to use for the operation. :p integer input_halo_extents(3) [in,optional]: An array of three integers to define halo region extents of the input data, in global order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the i-th global domain axis. Symmetric halos are assumed (e.g. a value of one in halo_extents means there are 2 halo elements, one element on each side). If not provided, input data is assumed to have no halos. :p integer output_halo_extents(3) [in,optional]: Similar to :code:`input_halo_extents` but for the output data. If not provided, output data is assumed to have no halos. @@ -497,7 +497,7 @@ _____________________ :p cudecompGridDesc grid_desc [in]: A cuDecomp grid descriptor. :p T input(*) [in]: Device array containing input Y-axis aligned pencil data. :p T output(*) [out]: Device array to write output Z-axis aligned pencil data. If :code:`input` and :code:`output` are the same, operation is performed in-place - :p T work(*) [in]: Device array to use for transpose workspace. + :p T work(*) [in]: Caller-owned device array to use for transpose workspace, or :code:`CUDECOMP_WORKSPACE_AUTO` to use handle-owned workspace management. :p cudecompDataType dtype [in]: The :code:`cudecompDataType` to use for the operation. :p integer input_halo_extents(3) [in,optional]: An array of three integers to define halo region extents of the input data, in global order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the i-th global domain axis. Symmetric halos are assumed (e.g. a value of one in halo_extents means there are 2 halo elements, one element on each side). If not provided, input data is assumed to have no halos. :p integer output_halo_extents(3) [in,optional]: Similar to :code:`intput_halo_extents` but for the output data. If not provided, output data is assumed to have no halos. @@ -523,7 +523,7 @@ _____________________ :p cudecompGridDesc grid_desc [in]: A cuDecomp grid descriptor. :p T input(*) [in]: Device array containing input Z-axis aligned pencil data. :p T output(*) [out]: Device array to write output Y-axis aligned pencil data. If :code:`input` and :code:`output` are the same, operation is performed in-place - :p T work(*) [in]: Device array to use for transpose workspace. + :p T work(*) [in]: Caller-owned device array to use for transpose workspace, or :code:`CUDECOMP_WORKSPACE_AUTO` to use handle-owned workspace management. :p cudecompDataType dtype [in]: The :code:`cudecompDataType` to use for the operation. :p integer input_halo_extents(3) [in,optional]: An array of three integers to define halo region extents of the input data, in global order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the i-th global domain axis. Symmetric halos are assumed (e.g. a value of one in halo_extents means there are 2 halo elements, one element on each side). If not provided, input data is assumed to have no halos. :p integer output_halo_extents(3) [in,optional]: Similar to :code:`intput_halo_extents` but for the output data. If not provided, output data is assumed to have no halos. @@ -550,7 +550,7 @@ _____________________ :p cudecompGridDesc grid_desc [in]: A cuDecomp grid descriptor. :p T input(*) [in]: Device array containing input Y-axis aligned pencil data. :p T output(*) [out]: Device array to write output X-axis aligned pencil data. If :code:`input` and :code:`output` are the same, operation is performed in-place - :p T work(*) [in]: Device array to use for transpose workspace. + :p T work(*) [in]: Caller-owned device array to use for transpose workspace, or :code:`CUDECOMP_WORKSPACE_AUTO` to use handle-owned workspace management. :p cudecompDataType dtype [in]: The :code:`cudecompDataType` to use for the operation. :p integer input_halo_extents(3) [in,optional]: An array of three integers to define halo region extents of the input data, in global order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the i-th global domain axis. Symmetric halos are assumed (e.g. a value of one in halo_extents means there are 2 halo elements, one element on each side). If not provided, input data is assumed to have no halos. :p integer output_halo_extents(3) [in,optional]: Similar to :code:`intput_halo_extents` but for the output data. If not provided, output data is assumed to have no halos. @@ -578,7 +578,7 @@ ____________________ :p cudecompHandle handle [in]: The initialized cuDecomp library handle :p cudecompGridDesc grid_desc [in]: A cuDecomp grid descriptor. :p T input(*) [in,out]: Device array containing input X-axis aligned pencil data. On successful completion, this buffer will contain the input X-axis aligned pencil data with the specified halo regions updated. - :p T work(*) [in]: Device array to use for halo workspace. + :p T work(*) [in]: Caller-owned device array to use for halo workspace, or :code:`CUDECOMP_WORKSPACE_AUTO` to use handle-owned workspace management. :p cudecompDataType dtype [in]: The :code:`cudecompDataType` to use for the operation. :p integer halo_extents(3) [in]: An array of three integers to define halo region extents of the input data, in global order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the i-th global domain axis. Symmetric halos are assumed (e.g. a value of one in halo_extents means there are 2 halo elements, one element on each side). :p logical halo_periods(3) [in]: An array of three boolean values to define halo periodicity of the input data, in global order. If the i-th entry in this array is true, the domain is treated periodically along the i-th global domain axis. @@ -603,7 +603,7 @@ ____________________ :p cudecompHandle handle [in]: The initialized cuDecomp library handle :p cudecompGridDesc grid_desc [in]: A cuDecomp grid descriptor. :p T input(*) [in,out]: Device array containing input Y-axis aligned pencil data. On successful completion, this buffer will contain the input X-axis aligned pencil data with the specified halo regions updated. - :p T work(*) [in]: Device array to use for halo workspace. + :p T work(*) [in]: Caller-owned device array to use for halo workspace, or :code:`CUDECOMP_WORKSPACE_AUTO` to use handle-owned workspace management. :p cudecompDataType dtype [in]: The :code:`cudecompDataType` to use for the operation. :p integer halo_extents(3) [in]: An array of three integers to define halo region extents of the input data, in global order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the i-th global domain axis. Symmetric halos are assumed (e.g. a value of one in halo_extents means there are 2 halo elements, one element on each side). :p logical halo_periods(3) [in]: An array of three boolean values to define halo periodicity of the input data, in global order. If the i-th entry in this array is true, the domain is treated periodically along the i-th global domain axis. @@ -628,7 +628,7 @@ ____________________ :p cudecompHandle handle [in]: The initialized cuDecomp library handle :p cudecompGridDesc grid_desc [in]: A cuDecomp grid descriptor. :p T input(*) [in,out]: Device array containing input Z-axis aligned pencil data. On successful completion, this buffer will contain the input X-axis aligned pencil data with the specified halo regions updated. - :p T work(*) [in]: Device array to use for halo workspace. + :p T work(*) [in]: Caller-owned device array to use for halo workspace, or :code:`CUDECOMP_WORKSPACE_AUTO` to use handle-owned workspace management. :p cudecompDataType dtype [in]: The :code:`cudecompDataType` to use for the operation. :p integer halo_extents(3) [in]: An array of three integers to define halo region extents of the input data, in global order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the i-th global domain axis. Symmetric halos are assumed (e.g. a value of one in halo_extents means there are 2 halo elements, one element on each side). :p logical halo_periods(3) [in]: An array of three boolean values to define halo periodicity of the input data, in global order. If the i-th entry in this array is true, the domain is treated periodically along the i-th global domain axis. diff --git a/docs/basic_usage.rst b/docs/basic_usage.rst index b92262a..2631e13 100644 --- a/docs/basic_usage.rst +++ b/docs/basic_usage.rst @@ -522,7 +522,9 @@ Besides device memory to store pencil data, cuDecomp also requires workspace buf is used to facilitate local packing/unpacking and transposition operations (which are currently performed out-of-place). As a result, this workspace buffer will be approximately 2x the size of the largest pencil assigned to this process. For halo communication, the workspace is used to facilitate local packing of non-contiguous -halo elements. We can query the required workspace sizes, in number of elements, using the +halo elements. + +Applications can allocate caller-owned workspaces. Query the required workspace sizes, in number of elements, using the :ref:`cudecompGetTransposeWorkspaceSize-ref` and :ref:`cudecompGetHaloWorkspaceSize-ref` functions. .. tabs:: @@ -582,6 +584,35 @@ is required for NVSHMEM operations (see NVSHMEM documentation for more details). istat = cudecompMalloc(handle, grid_desc, halo_work_d, halo_work_num_elements) call CHECK_CUDECOMP_EXIT(istat) +Alternatively, applications may let cuDecomp manage workspace memory by skipping the size queries and allocations +above and passing :code:`CUDECOMP_WORKSPACE_AUTO` directly to each operation. cuDecomp then allocates, grows, and reuses +handle-owned workspace memory and preserves ordering with the stream passed to each operation. All participating ranks +must consistently choose automatic or explicit workspace management for a given operation. Automatic workspace +management cannot be used while the caller's stream is being captured by a CUDA Graph; provide an explicit workspace +in that case. + +.. tabs:: + + .. code-tab:: c++ + + CHECK_CUDECOMP_EXIT(cudecompTransposeXToY(handle, grid_desc, data_d, data_d, + CUDECOMP_WORKSPACE_AUTO, CUDECOMP_DOUBLE, + pinfo_x.halo_extents, nullptr, nullptr, nullptr, 0)); + + CHECK_CUDECOMP_EXIT(cudecompUpdateHalosX(handle, grid_desc, data_d, CUDECOMP_WORKSPACE_AUTO, + CUDECOMP_DOUBLE, pinfo_x.halo_extents, halo_periods, + 0, nullptr, 0)); + + .. code-tab:: fortran + + istat = cudecompTransposeXToY(handle, grid_desc, data_d, data_d, CUDECOMP_WORKSPACE_AUTO, & + CUDECOMP_DOUBLE, pinfo_x%halo_extents, [0,0,0]) + call CHECK_CUDECOMP_EXIT(istat) + + istat = cudecompUpdateHalosX(handle, grid_desc, data_d, CUDECOMP_WORKSPACE_AUTO, CUDECOMP_DOUBLE, & + pinfo_x%halo_extents, halo_periods, 1) + call CHECK_CUDECOMP_EXIT(istat) + Transposing the data -------------------- @@ -681,7 +712,8 @@ them unspecified in Fortran. Cleaning up and finalizing the library -------------------------------------- Finally, we can clean up resources. Note the usage of :ref:`cudecompFree-ref` to deallocate the workspace arrays -allocated with :ref:`cudecompMalloc-ref`. +allocated with :ref:`cudecompMalloc-ref`. Automatically managed workspaces are owned by the handle and are released by +:code:`cudecompFinalize`; applications must not pass them to :ref:`cudecompFree-ref`. .. tabs:: diff --git a/include/cudecomp.h b/include/cudecomp.h index 2774996..790e7c4 100644 --- a/include/cudecomp.h +++ b/include/cudecomp.h @@ -38,6 +38,9 @@ #define CUDECOMP_PENCIL_INFO_MAGIC INT32_C(0x50494e46) /** @endcond */ +/** Workspace argument sentinel that selects handle-owned workspace management. */ +#define CUDECOMP_WORKSPACE_AUTO ((void*)0) + #ifdef __cplusplus extern "C" { #endif @@ -525,7 +528,11 @@ cudecompResult_t cudecompGetShiftedRank(cudecompHandle_t handle, cudecompGridDes * @param[in] input A pointer to the memory buffer to read input X-axis aligned pencil data * @param[out] output A pointer to the memory buffer to write output Y-axis aligned pencil data. If input and output are * the same, operation is performed in-place - * @param[in] work A pointer to the transpose workspace memory + * @param[in] work A pointer to transpose workspace memory. For handle-owned workspace management, a NULL pointer (or + * CUDECOMP_WORKSPACE_AUTO) can be provided. In this case, cuDecomp allocates and grows the workspace as needed and + * serializes the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same + * workspace-management choice on all participating ranks. Handle-owned workspace management is not supported while + * stream is being captured. * @param[in] dtype The cuDecomp datatype to use for the transpose operation * @param[in] input_halo_extents An array of three integers to define halo region extents of the input data, in global * order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along @@ -554,7 +561,11 @@ cudecompResult_t cudecompTransposeXToY(cudecompHandle_t handle, cudecompGridDesc * @param[in] input A pointer to the memory buffer to read input Y-axis aligned pencil data * @param[out] output A pointer to the memory buffer to write output Z-axis aligned pencil data. If input and output are * the same, operation is performed in-place - * @param[in] work A pointer to the transpose workspace memory + * @param[in] work A pointer to transpose workspace memory. For handle-owned workspace management, a NULL pointer (or + * CUDECOMP_WORKSPACE_AUTO) can be provided. In this case, cuDecomp allocates and grows the workspace as needed and + * serializes the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same + * workspace-management choice on all participating ranks. Handle-owned workspace management is not supported while + * stream is being captured. * @param[in] dtype The cuDecomp datatype to use for the transpose operation * @param[in] input_halo_extents An array of three integers to define halo region extents of the input data, in global * order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along @@ -583,7 +594,11 @@ cudecompResult_t cudecompTransposeYToZ(cudecompHandle_t handle, cudecompGridDesc * @param[in] input A pointer to the memory buffer to read input Z-axis aligned pencil data * @param[out] output A pointer to the memory buffer to write output Y-axis aligned pencil data. If input and output are * the same, operation is performed in-place - * @param[in] work A pointer to the transpose workspace memory + * @param[in] work A pointer to transpose workspace memory. For handle-owned workspace management, a NULL pointer (or + * CUDECOMP_WORKSPACE_AUTO) can be provided. In this case, cuDecomp allocates and grows the workspace as needed and + * serializes the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same + * workspace-management choice on all participating ranks. Handle-owned workspace management is not supported while + * stream is being captured. * @param[in] dtype The cuDecomp datatype to use for the transpose operation * @param[in] input_halo_extents An array of three integers to define halo region extents of the input data, in global * order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along @@ -612,7 +627,11 @@ cudecompResult_t cudecompTransposeZToY(cudecompHandle_t handle, cudecompGridDesc * @param[in] input A pointer to the memory buffer to read input Y-axis aligned pencil data * @param[out] output A pointer to the memory buffer to write output X-axis aligned pencil data. If input and output are * the same, operation is performed in-place - * @param[in] work A pointer to the transpose workspace memory + * @param[in] work A pointer to transpose workspace memory. For handle-owned workspace management, a NULL pointer (or + * CUDECOMP_WORKSPACE_AUTO) can be provided. In this case, cuDecomp allocates and grows the workspace as needed and + * serializes the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same + * workspace-management choice on all participating ranks. Handle-owned workspace management is not supported while + * stream is being captured. * @param[in] dtype The cuDecomp datatype to use for the transpose operation * @param[in] input_halo_extents An array of three integers to define halo region extents of the input data, in global * order. The i-th entry in this array should contain the number of halo elements (per direction) expected in the along @@ -641,7 +660,11 @@ cudecompResult_t cudecompTransposeYToX(cudecompHandle_t handle, cudecompGridDesc * @param[in] grid_desc A cuDecomp grid descriptor * @param[in,out] input A pointer to the memory buffer to read input X-axis aligned pencil data. On successful * completion, this buffer will contain the input X-axis aligned pencil data with the specified halo regions updated. - * @param[in] work A pointer to the halo workspace memory + * @param[in] work A pointer to halo workspace memory. For handle-owned workspace management, a NULL pointer (or + * CUDECOMP_WORKSPACE_AUTO) can be provided. In this case, cuDecomp allocates and grows the workspace as needed and + * serializes the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same + * workspace-management choice on all participating ranks. Handle-owned workspace management is not supported while + * stream is being captured. * @param[in] dtype The cuDecomp datatype to use for the halo operation * @param[in] halo_extents An array of three integers to define halo region extents of the input data, in global order. * The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the @@ -668,7 +691,11 @@ cudecompResult_t cudecompUpdateHalosX(cudecompHandle_t handle, cudecompGridDesc_ * @param[in] grid_desc A cuDecomp grid descriptor * @param[in,out] input A pointer to the memory buffer to read input Y-axis aligned pencil data. On successful * completion, this buffer will contain the input Y-axis aligned pencil data with the specified halo regions updated. - * @param[in] work A pointer to the halo workspace memory + * @param[in] work A pointer to halo workspace memory. For handle-owned workspace management, a NULL pointer (or + * CUDECOMP_WORKSPACE_AUTO) can be provided. In this case, cuDecomp allocates and grows the workspace as needed and + * serializes the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same + * workspace-management choice on all participating ranks. Handle-owned workspace management is not supported while + * stream is being captured. * @param[in] dtype The cuDecomp datatype to use for the halo operation * @param[in] halo_extents An array of three integers to define halo region extents of the input data, in global order. * The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the @@ -695,7 +722,11 @@ cudecompResult_t cudecompUpdateHalosY(cudecompHandle_t handle, cudecompGridDesc_ * @param[in] grid_desc A cuDecomp grid descriptor * @param[in,out] input A pointer to the memory buffer to read input Z-axis aligned pencil data. On successful * completion, this buffer will contain the input Z-axis aligned pencil data with the specified halo regions updated. - * @param[in] work A pointer to the halo workspace memory + * @param[in] work A pointer to halo workspace memory. For handle-owned workspace management, a NULL pointer (or + * CUDECOMP_WORKSPACE_AUTO) can be provided. In this case, cuDecomp allocates and grows the workspace as needed and + * serializes the operation on a handle-owned stream while preserving ordering with stream. Callers must make the same + * workspace-management choice on all participating ranks. Handle-owned workspace management is not supported while + * stream is being captured. * @param[in] dtype The cuDecomp datatype to use for the halo operation * @param[in] halo_extents An array of three integers to define halo region extents of the input data, in global order. * The i-th entry in this array should contain the number of halo elements (per direction) expected in the along the diff --git a/include/internal/common.h b/include/internal/common.h index 18c9bad..547953f 100644 --- a/include/internal/common.h +++ b/include/internal/common.h @@ -26,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -121,6 +122,22 @@ struct cudecompHandle { std::vector streams; // internal streams for concurrent scheduling + // Automatic workspaces are split by allocation domain but share one execution stream so operations submitted with + // a null workspace cannot race with each other. + struct ManagedWorkspace { + void* ptr = nullptr; + size_t size = 0; +#ifdef ENABLE_NVSHMEM + cudecomp::nvshmemRuntime nvshmem_runtime; +#endif + }; + ManagedWorkspace ordinary_workspace; + ManagedWorkspace nvshmem_workspace; + std::unique_ptr workspace_stream; + std::unique_ptr workspace_ingress_event; + std::unique_ptr workspace_egress_event; + std::mutex workspace_mutex; + #if CUTENSOR_MAJOR >= 2 cutensorHandle_t cutensor_handle = nullptr; // cuTENSOR handle; cutensorPlanPreference_t cutensor_plan_pref = nullptr; // cuTENSOR plan preference; diff --git a/src/cudecomp.cc b/src/cudecomp.cc index 5471ccd..dd9bd95 100644 --- a/src/cudecomp.cc +++ b/src/cudecomp.cc @@ -856,6 +856,191 @@ struct cuMemAllocationGuard { }; #endif +static size_t getDataTypeSizeBytes(cudecompDataType_t dtype) { + switch (dtype) { + case CUDECOMP_FLOAT: return 4; + case CUDECOMP_DOUBLE: + case CUDECOMP_FLOAT_COMPLEX: return 8; + case CUDECOMP_DOUBLE_COMPLEX: return 16; + default: THROW_INVALID_USAGE("unknown data type"); + } +} + +static size_t getWorkspaceSizeBytes(int64_t num_elements, cudecompDataType_t dtype) { + if (num_elements < 0) { THROW_INTERNAL_ERROR("workspace element count cannot be negative"); } + return static_cast(num_elements) * getDataTypeSizeBytes(dtype); +} + +static size_t getWorkspaceAllocationSize(cudecompHandle_t handle, size_t requested_size, bool use_nvshmem) { +#ifdef ENABLE_NVSHMEM + if (use_nvshmem) { + CHECK_MPI(MPI_Allreduce(MPI_IN_PLACE, &requested_size, 1, mpiSizeTDatatype(), MPI_MAX, handle->mpi_comm)); + } +#else + (void)handle; + if (use_nvshmem) { THROW_NOT_SUPPORTED("build does not support NVSHMEM communication backends."); } +#endif + return requested_size; +} + +static void deregisterNcclBuffer(cudecompHandle_t handle, void* buffer) { +#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 19, 0) + auto entry = handle->nccl_ubr_handles.find(buffer); + if (entry == handle->nccl_ubr_handles.end()) return; + for (const auto& registration : entry->second) { + CHECK_NCCL(ncclCommDeregister(*registration.first, registration.second)); + } + handle->nccl_ubr_handles.erase(entry); +#else + (void)handle; + (void)buffer; +#endif +} + +static void freeOrdinaryWorkspace(cudecompHandle_t handle, void* buffer) { + if (!buffer) return; + deregisterNcclBuffer(handle, buffer); + if (handle->cuda_cumem_enable) { +#if CUDART_VERSION >= 11030 + CUmemGenericAllocationHandle cumem_handle; + CHECK_CUDA_DRV(cuMemRetainAllocationHandle(&cumem_handle, buffer)); + CHECK_CUDA_DRV(cuMemRelease(cumem_handle)); + size_t size = 0; + CHECK_CUDA_DRV(cuMemGetAddressRange(nullptr, &size, reinterpret_cast(buffer))); + CHECK_CUDA_DRV(cuMemUnmap(reinterpret_cast(buffer), size)); + CHECK_CUDA_DRV(cuMemRelease(cumem_handle)); + CHECK_CUDA_DRV(cuMemAddressFree(reinterpret_cast(buffer), size)); +#endif + } else { + CHECK_CUDA(cudaFree(buffer)); + } +} + +#ifdef ENABLE_NVSHMEM +static void freeNvshmemWorkspace(const nvshmemRuntime& runtime, void* buffer) { + if (!buffer) return; + if (!runtime || !runtime->initialized) { THROW_INVALID_USAGE("NVSHMEM runtime is unavailable"); } + nvshmem_free(buffer); + auto entry = runtime->nvshmem_allocations.find(buffer); + if (entry != runtime->nvshmem_allocations.end()) { + runtime->nvshmem_allocation_size -= entry->second; + runtime->nvshmem_allocations.erase(entry); + } +} +#endif + +static void releaseManagedWorkspaces(cudecompHandle_t handle) { + if (handle->workspace_stream) { CHECK_CUDA(cudaStreamSynchronize(*handle->workspace_stream)); } + + freeOrdinaryWorkspace(handle, handle->ordinary_workspace.ptr); + handle->ordinary_workspace.ptr = nullptr; + handle->ordinary_workspace.size = 0; + +#ifdef ENABLE_NVSHMEM + freeNvshmemWorkspace(handle->nvshmem_workspace.nvshmem_runtime, handle->nvshmem_workspace.ptr); + handle->nvshmem_workspace.nvshmem_runtime.reset(); +#endif + handle->nvshmem_workspace.ptr = nullptr; + handle->nvshmem_workspace.size = 0; + + handle->workspace_ingress_event.reset(); + handle->workspace_egress_event.reset(); + handle->workspace_stream.reset(); +} + +// NVSHMEM allocation sizes must be normalized with getWorkspaceAllocationSize before calling this function. +static void allocateWorkspace(cudecompHandle_t handle, cudecompGridDesc_t grid_desc, void** buffer, + size_t buffer_size_bytes, bool use_nvshmem); + +static void registerWorkspaceWithNccl(cudecompHandle_t handle, cudecompGridDesc_t grid_desc, void* buffer, + size_t size) { +#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 19, 0) + if (!handle->nccl_enable_ubr || !buffer) return; + + auto register_comm = [&](const ncclComm& comm) { + if (!comm) return; + auto& registrations = handle->nccl_ubr_handles[buffer]; + auto existing = std::find_if(registrations.begin(), registrations.end(), + [&](const auto& registration) { return registration.first.get() == comm.get(); }); + if (existing != registrations.end()) return; + void* registration_handle; + CHECK_NCCL(ncclCommRegister(*comm, buffer, size, ®istration_handle)); + registrations.emplace_back(comm, registration_handle); + }; + + register_comm(grid_desc->nccl_comm); + register_comm(grid_desc->nccl_local_comm); +#else + (void)handle; + (void)grid_desc; + (void)buffer; + (void)size; +#endif +} + +static void* prepareManagedWorkspace(cudecompHandle_t handle, cudecompGridDesc_t grid_desc, size_t required_size, + bool use_nvshmem) { + if (required_size == 0) return nullptr; + required_size = getWorkspaceAllocationSize(handle, required_size, use_nvshmem); + + auto& workspace = use_nvshmem ? handle->nvshmem_workspace : handle->ordinary_workspace; + if (workspace.size < required_size) { + CHECK_CUDA(cudaStreamSynchronize(*handle->workspace_stream)); + if (workspace.ptr) { +#ifdef ENABLE_NVSHMEM + if (use_nvshmem) { + freeNvshmemWorkspace(workspace.nvshmem_runtime, workspace.ptr); + workspace.nvshmem_runtime.reset(); + } else +#endif + { + freeOrdinaryWorkspace(handle, workspace.ptr); + } + workspace.ptr = nullptr; + workspace.size = 0; + } + + allocateWorkspace(handle, grid_desc, &workspace.ptr, required_size, use_nvshmem); + workspace.size = required_size; +#ifdef ENABLE_NVSHMEM + if (use_nvshmem) { workspace.nvshmem_runtime = grid_desc->nvshmem_runtime; } +#endif + } else if (!use_nvshmem) { + registerWorkspaceWithNccl(handle, grid_desc, workspace.ptr, workspace.size); + } + + return workspace.ptr; +} + +template +static void runWithWorkspace(cudecompHandle_t handle, cudecompGridDesc_t grid_desc, void* work, size_t required_size, + bool use_nvshmem, cudaStream_t caller_stream, Function&& function) { + if (work) { + function(work, caller_stream); + return; + } + + cudaStreamCaptureStatus capture_status; + CHECK_CUDA(cudaStreamIsCapturing(caller_stream, &capture_status)); + if (capture_status != cudaStreamCaptureStatusNone) { + THROW_NOT_SUPPORTED( + "automatic workspace management is not supported during CUDA stream capture; use an explicit workspace"); + } + + std::lock_guard lock(handle->workspace_mutex); + if (!handle->workspace_stream) { handle->workspace_stream = std::make_unique(); } + if (!handle->workspace_ingress_event) { handle->workspace_ingress_event = std::make_unique(); } + if (!handle->workspace_egress_event) { handle->workspace_egress_event = std::make_unique(); } + void* managed_work = prepareManagedWorkspace(handle, grid_desc, required_size, use_nvshmem); + auto managed_stream = handle->workspace_stream->get(); + + CHECK_CUDA(cudaEventRecord(*handle->workspace_ingress_event, caller_stream)); + CHECK_CUDA(cudaStreamWaitEvent(managed_stream, *handle->workspace_ingress_event, 0)); + function(managed_work, managed_stream); + CHECK_CUDA(cudaEventRecord(*handle->workspace_egress_event, managed_stream)); + CHECK_CUDA(cudaStreamWaitEvent(caller_stream, *handle->workspace_egress_event, 0)); +} + } // namespace } // namespace cudecomp @@ -880,6 +1065,10 @@ void warnIfNvshmemBufferUsedWithMpi(cudecompHandle_t handle, const void* send_bu #endif cudecompHandle::~cudecompHandle() noexcept { + try { + cudecomp::releaseManagedWorkspaces(this); + } catch (...) {} + #if NCCL_VERSION_CODE >= NCCL_VERSION(2, 19, 0) for (auto& entry : nccl_ubr_handles) { for (const auto& ubr_handle : entry.second) { @@ -1012,14 +1201,11 @@ cudecompResult_t cudecompFinalize(cudecompHandle_t handle) { try { checkHandle(handle); -#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 19, 0) - for (auto& entry : handle->nccl_ubr_handles) { - for (const auto& ubr_handle : entry.second) { - CHECK_NCCL(ncclCommDeregister(*ubr_handle.first, ubr_handle.second)); - } + releaseManagedWorkspaces(handle); + + while (!handle->nccl_ubr_handles.empty()) { + deregisterNcclBuffer(handle, handle->nccl_ubr_handles.begin()->first); } - handle->nccl_ubr_handles.clear(); -#endif delete handle; } @@ -1458,150 +1644,141 @@ cudecompResult_t cudecompGetHaloWorkspaceSize(cudecompHandle_t handle, cudecompG return CUDECOMP_RESULT_SUCCESS; } -cudecompResult_t cudecompMalloc(cudecompHandle_t handle, cudecompGridDesc_t grid_desc, void** buffer, - size_t buffer_size_bytes) { - using namespace cudecomp; - try { - checkHandle(handle); - checkGridDesc(handle, grid_desc); - if (!buffer) { THROW_INVALID_USAGE("buffer argument cannot be null"); } - if (buffer_size_bytes == 0) { THROW_INVALID_USAGE("buffer size cannot be zero"); } - - if (transposeBackendRequiresNvshmem(grid_desc->config.transpose_comm_backend) || - haloBackendRequiresNvshmem(grid_desc->config.halo_comm_backend)) { +namespace cudecomp { +namespace { +static void allocateWorkspace(cudecompHandle_t handle, cudecompGridDesc_t grid_desc, void** buffer, + size_t buffer_size_bytes, bool use_nvshmem) { + if (use_nvshmem) { #ifdef ENABLE_NVSHMEM - // NVSHMEM requires allocations to be the same size for all ranks. Find maximum. - CHECK_MPI(MPI_Allreduce(MPI_IN_PLACE, &buffer_size_bytes, 1, mpiSizeTDatatype(), MPI_MAX, handle->mpi_comm)); + auto nvshmem_runtime = grid_desc->nvshmem_runtime; + if (!nvshmem_runtime || !nvshmem_runtime->initialized) { THROW_INVALID_USAGE("NVSHMEM runtime is unavailable"); } - auto nvshmem_runtime = grid_desc->nvshmem_runtime; - if (!nvshmem_runtime || !nvshmem_runtime->initialized) { THROW_INVALID_USAGE("NVSHMEM runtime is unavailable"); } - - size_t nvshmem_free_size = 0; - if (nvshmem_runtime->nvshmem_symmetric_size > nvshmem_runtime->nvshmem_allocation_size) { - nvshmem_free_size = nvshmem_runtime->nvshmem_symmetric_size - nvshmem_runtime->nvshmem_allocation_size; - } - if (!nvshmem_runtime->nvshmem_vmm && handle->rank == 0 && buffer_size_bytes > nvshmem_free_size) { - fprintf(stderr, - "CUDECOMP:WARN: Attempting an NVSHMEM allocation of %zu bytes but *approximately* " - "%zu free bytes of %zu total bytes of symmetric heap space available. If the allocation fails, " - "set NVSHMEM_SYMMETRIC_SIZE >= %zu and try again.\n", - buffer_size_bytes, nvshmem_free_size, nvshmem_runtime->nvshmem_symmetric_size, - nvshmem_runtime->nvshmem_symmetric_size + (buffer_size_bytes - nvshmem_free_size)); - } + size_t nvshmem_free_size = 0; + if (nvshmem_runtime->nvshmem_symmetric_size > nvshmem_runtime->nvshmem_allocation_size) { + nvshmem_free_size = nvshmem_runtime->nvshmem_symmetric_size - nvshmem_runtime->nvshmem_allocation_size; + } + if (!nvshmem_runtime->nvshmem_vmm && handle->rank == 0 && buffer_size_bytes > nvshmem_free_size) { + fprintf(stderr, + "CUDECOMP:WARN: Attempting an NVSHMEM allocation of %zu bytes but *approximately* " + "%zu free bytes of %zu total bytes of symmetric heap space available. If the allocation fails, " + "set NVSHMEM_SYMMETRIC_SIZE >= %zu and try again.\n", + buffer_size_bytes, nvshmem_free_size, nvshmem_runtime->nvshmem_symmetric_size, + nvshmem_runtime->nvshmem_symmetric_size + (buffer_size_bytes - nvshmem_free_size)); + } - *buffer = nvshmem_malloc(buffer_size_bytes); - if (buffer_size_bytes != 0 && *buffer == nullptr) { THROW_NVSHMEM_ERROR("nvshmem_malloc failed"); } - // Record NVSHMEM allocation details - nvshmem_runtime->nvshmem_allocations[*buffer] = buffer_size_bytes; - nvshmem_runtime->nvshmem_allocation_size += buffer_size_bytes; + *buffer = nvshmem_malloc(buffer_size_bytes); + if (buffer_size_bytes != 0 && *buffer == nullptr) { THROW_NVSHMEM_ERROR("nvshmem_malloc failed"); } + // Record NVSHMEM allocation details + nvshmem_runtime->nvshmem_allocations[*buffer] = buffer_size_bytes; + nvshmem_runtime->nvshmem_allocation_size += buffer_size_bytes; #else - THROW_NOT_SUPPORTED("build does not support NVSHMEM communication backends."); + THROW_NOT_SUPPORTED("build does not support NVSHMEM communication backends."); #endif - } else { - if (handle->cuda_cumem_enable) { + } else { + if (handle->cuda_cumem_enable) { #if CUDART_VERSION >= 11030 - int dev; - CUdevice cu_dev; - CHECK_CUDA(cudaGetDevice(&dev)); - CHECK_CUDA_DRV(cuDeviceGet(&cu_dev, dev)); + int dev; + CUdevice cu_dev; + CHECK_CUDA(cudaGetDevice(&dev)); + CHECK_CUDA_DRV(cuDeviceGet(&cu_dev, dev)); - int requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; + int requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; #if CUDART_VERSION >= 12030 - int driverVersion; - CHECK_CUDA(cudaDriverGetVersion(&driverVersion)); - if (driverVersion >= 12030) { - int fabric_supported = 0; - CHECK_CUDA_DRV( - cuDeviceGetAttribute(&fabric_supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, cu_dev)); - if (fabric_supported) requestedHandleTypes |= CU_MEM_HANDLE_TYPE_FABRIC; - } + int driverVersion; + CHECK_CUDA(cudaDriverGetVersion(&driverVersion)); + if (driverVersion >= 12030) { + int fabric_supported = 0; + CHECK_CUDA_DRV( + cuDeviceGetAttribute(&fabric_supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, cu_dev)); + if (fabric_supported) requestedHandleTypes |= CU_MEM_HANDLE_TYPE_FABRIC; + } #endif - CUmemAllocationProp prop = {}; - prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; - prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + CUmemAllocationProp prop = {}; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.requestedHandleTypes = static_cast(requestedHandleTypes); + prop.location.id = cu_dev; + + // Check for RDMA support + int flag; + CHECK_CUDA_DRV(cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED, cu_dev)); + if (flag) prop.allocFlags.gpuDirectRDMACapable = 1; + + // Keep the caller-requested size so any retry can realign from the original value. + size_t original_buffer_size_bytes = buffer_size_bytes; + size_t granularity; + CHECK_CUDA_DRV(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); + buffer_size_bytes = (original_buffer_size_bytes + granularity - 1) / granularity * granularity; + + // Allocate memory + cuMemAllocationGuard cumem_guard; + cumem_guard.size = buffer_size_bytes; + CUresult err = cuFnTable.pfn_cuMemCreate(&cumem_guard.handle, buffer_size_bytes, &prop, 0); +#if CUDART_VERSION >= 12030 + if ((requestedHandleTypes & CU_MEM_HANDLE_TYPE_FABRIC) && + (err == CUDA_ERROR_NOT_PERMITTED || err == CUDA_ERROR_NOT_SUPPORTED)) { + // Fabric handles are useful when the platform supports them, but regular NCCL user buffer registration only + // requires POSIX FD export support. If Fabric creation is unavailable at runtime, keep VMM enabled and fall + // back to a POSIX-FD-only allocation. + requestedHandleTypes &= ~CU_MEM_HANDLE_TYPE_FABRIC; prop.requestedHandleTypes = static_cast(requestedHandleTypes); - prop.location.id = cu_dev; - - // Check for RDMA support - int flag; - CHECK_CUDA_DRV( - cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED, cu_dev)); - if (flag) prop.allocFlags.gpuDirectRDMACapable = 1; - - // Keep the caller-requested size so any retry can realign from the original value. - size_t original_buffer_size_bytes = buffer_size_bytes; - size_t granularity; CHECK_CUDA_DRV(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); buffer_size_bytes = (original_buffer_size_bytes + granularity - 1) / granularity * granularity; - - // Allocate memory - cuMemAllocationGuard cumem_guard; cumem_guard.size = buffer_size_bytes; - CUresult err = cuFnTable.pfn_cuMemCreate(&cumem_guard.handle, buffer_size_bytes, &prop, 0); -#if CUDART_VERSION >= 12030 - if ((requestedHandleTypes & CU_MEM_HANDLE_TYPE_FABRIC) && - (err == CUDA_ERROR_NOT_PERMITTED || err == CUDA_ERROR_NOT_SUPPORTED)) { - // Fabric handles are useful when the platform supports them, but regular NCCL user buffer registration only - // requires POSIX FD export support. If Fabric creation is unavailable at runtime, keep VMM enabled and fall - // back to a POSIX-FD-only allocation. - requestedHandleTypes &= ~CU_MEM_HANDLE_TYPE_FABRIC; - prop.requestedHandleTypes = static_cast(requestedHandleTypes); - CHECK_CUDA_DRV(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); - buffer_size_bytes = (original_buffer_size_bytes + granularity - 1) / granularity * granularity; - cumem_guard.size = buffer_size_bytes; - err = cuFnTable.pfn_cuMemCreate(&cumem_guard.handle, buffer_size_bytes, &prop, 0); - } -#endif - if (CUDA_SUCCESS != err) { - const char* error_str; - cuFnTable.pfn_cuGetErrorString(err, &error_str); - throw cudecomp::CudaError(__FILE__, __LINE__, error_str); - } - cumem_guard.handle_created = true; - CHECK_CUDA_DRV(cuMemAddressReserve(&cumem_guard.ptr, buffer_size_bytes, granularity, 0, 0)); - cumem_guard.address_reserved = true; - CHECK_CUDA_DRV(cuMemMap(cumem_guard.ptr, buffer_size_bytes, 0, cumem_guard.handle, 0)); - cumem_guard.mapped = true; - - // Set read/write access - CUmemAccessDesc accessDesc = {}; - accessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - accessDesc.location.id = cu_dev; - accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; - CHECK_CUDA_DRV(cuMemSetAccess(cumem_guard.ptr, buffer_size_bytes, &accessDesc, 1)); - - *buffer = reinterpret_cast(cumem_guard.ptr); - cumem_guard.release(); -#endif - } else { - CHECK_CUDA(cudaMalloc(buffer, buffer_size_bytes)); + err = cuFnTable.pfn_cuMemCreate(&cumem_guard.handle, buffer_size_bytes, &prop, 0); } -#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 19, 0) - if (transposeBackendRequiresNccl(grid_desc->config.transpose_comm_backend) || - haloBackendRequiresNccl(grid_desc->config.halo_comm_backend)) { - - if (handle->nccl_enable_ubr) { - try { - void* nccl_ubr_handle; - if (grid_desc->nccl_comm) { - CHECK_NCCL(ncclCommRegister(*grid_desc->nccl_comm, *buffer, buffer_size_bytes, &nccl_ubr_handle)); - handle->nccl_ubr_handles[*buffer].push_back(std::make_pair(grid_desc->nccl_comm, nccl_ubr_handle)); - } - if (grid_desc->nccl_local_comm) { - CHECK_NCCL(ncclCommRegister(*grid_desc->nccl_local_comm, *buffer, buffer_size_bytes, &nccl_ubr_handle)); - handle->nccl_ubr_handles[*buffer].push_back(std::make_pair(grid_desc->nccl_local_comm, nccl_ubr_handle)); - } - } catch (...) { - cudecompFree(handle, grid_desc, *buffer); - *buffer = nullptr; - throw; - } - } +#endif + if (CUDA_SUCCESS != err) { + const char* error_str; + cuFnTable.pfn_cuGetErrorString(err, &error_str); + throw cudecomp::CudaError(__FILE__, __LINE__, error_str); } + cumem_guard.handle_created = true; + CHECK_CUDA_DRV(cuMemAddressReserve(&cumem_guard.ptr, buffer_size_bytes, granularity, 0, 0)); + cumem_guard.address_reserved = true; + CHECK_CUDA_DRV(cuMemMap(cumem_guard.ptr, buffer_size_bytes, 0, cumem_guard.handle, 0)); + cumem_guard.mapped = true; + + // Set read/write access + CUmemAccessDesc accessDesc = {}; + accessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + accessDesc.location.id = cu_dev; + accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + CHECK_CUDA_DRV(cuMemSetAccess(cumem_guard.ptr, buffer_size_bytes, &accessDesc, 1)); + + *buffer = reinterpret_cast(cumem_guard.ptr); + cumem_guard.release(); #endif + } else { + CHECK_CUDA(cudaMalloc(buffer, buffer_size_bytes)); + } + try { + registerWorkspaceWithNccl(handle, grid_desc, *buffer, buffer_size_bytes); + } catch (...) { + freeOrdinaryWorkspace(handle, *buffer); + *buffer = nullptr; + throw; } } +} +} // namespace +} // namespace cudecomp + +cudecompResult_t cudecompMalloc(cudecompHandle_t handle, cudecompGridDesc_t grid_desc, void** buffer, + size_t buffer_size_bytes) { + using namespace cudecomp; + try { + checkHandle(handle); + checkGridDesc(handle, grid_desc); + if (!buffer) { THROW_INVALID_USAGE("buffer argument cannot be null"); } + if (buffer_size_bytes == 0) { THROW_INVALID_USAGE("buffer size cannot be zero"); } + + bool use_nvshmem = transposeBackendRequiresNvshmem(grid_desc->config.transpose_comm_backend) || + haloBackendRequiresNvshmem(grid_desc->config.halo_comm_backend); + buffer_size_bytes = getWorkspaceAllocationSize(handle, buffer_size_bytes, use_nvshmem); + allocateWorkspace(handle, grid_desc, buffer, buffer_size_bytes, use_nvshmem); + } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; } @@ -1612,54 +1789,16 @@ cudecompResult_t cudecompFree(cudecompHandle_t handle, cudecompGridDesc_t grid_d checkHandle(handle); checkGridDesc(handle, grid_desc); -#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 19, 0) - if (handle->nccl_ubr_handles.count(buffer) != 0) { - for (const auto& entry : handle->nccl_ubr_handles[buffer]) { - CHECK_NCCL(ncclCommDeregister(*entry.first, entry.second)); - } - handle->nccl_ubr_handles.erase(buffer); - } -#endif - if (transposeBackendRequiresNvshmem(grid_desc->config.transpose_comm_backend) || haloBackendRequiresNvshmem(grid_desc->config.halo_comm_backend)) { #ifdef ENABLE_NVSHMEM - if (buffer) { - auto nvshmem_runtime = grid_desc->nvshmem_runtime; - if (!nvshmem_runtime || !nvshmem_runtime->initialized) { - THROW_INVALID_USAGE("NVSHMEM runtime is unavailable"); - } - - nvshmem_free(buffer); - - // Record NVSHMEM deallocation details - auto entry = nvshmem_runtime->nvshmem_allocations.find(buffer); - if (entry != nvshmem_runtime->nvshmem_allocations.end()) { - nvshmem_runtime->nvshmem_allocation_size -= entry->second; - nvshmem_runtime->nvshmem_allocations.erase(entry); - } - } + freeNvshmemWorkspace(grid_desc->nvshmem_runtime, buffer); #else THROW_NOT_SUPPORTED("build does not support NVSHMEM communication backends."); #endif } else { - if (handle->cuda_cumem_enable) { -#if CUDART_VERSION >= 11030 - if (buffer) { - CUmemGenericAllocationHandle cumem_handle; - CHECK_CUDA_DRV(cuMemRetainAllocationHandle(&cumem_handle, buffer)); - CHECK_CUDA_DRV(cuMemRelease(cumem_handle)); - size_t size = 0; - CHECK_CUDA_DRV(cuMemGetAddressRange(NULL, &size, (CUdeviceptr)buffer)); - CHECK_CUDA_DRV(cuMemUnmap((CUdeviceptr)buffer, size)); - CHECK_CUDA_DRV(cuMemRelease(cumem_handle)); - CHECK_CUDA_DRV(cuMemAddressFree((CUdeviceptr)buffer, size)); - } -#endif - } else { - if (buffer) { CHECK_CUDA(cudaFree(buffer)); } - } + freeOrdinaryWorkspace(handle, buffer); } } CUDECOMP_CATCH_C_API_ERRORS() @@ -1696,12 +1835,7 @@ cudecompResult_t cudecompGetDataTypeSize(cudecompDataType_t dtype, int64_t* dtyp try { checkDataType(dtype); if (!dtype_size) { THROW_INVALID_USAGE("dtype_size cannot be null."); } - switch (dtype) { - case CUDECOMP_FLOAT: *dtype_size = 4; break; - case CUDECOMP_DOUBLE: - case CUDECOMP_FLOAT_COMPLEX: *dtype_size = 8; break; - case CUDECOMP_DOUBLE_COMPLEX: *dtype_size = 16; break; - } + *dtype_size = static_cast(getDataTypeSizeBytes(dtype)); } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; @@ -1765,31 +1899,40 @@ cudecompResult_t cudecompTransposeXToY(cudecompHandle_t handle, cudecompGridDesc checkDataType(dtype); if (!input) { THROW_INVALID_USAGE("input argument cannot be null"); } if (!output) { THROW_INVALID_USAGE("output argument cannot be null"); } - if (!work) { THROW_INVALID_USAGE("work argument cannot be null"); } - switch (dtype) { - case CUDECOMP_FLOAT: - cudecompTransposeXToY(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), - reinterpret_cast(work), input_halo_extents, output_halo_extents, input_padding, - output_padding, stream); - break; - case CUDECOMP_DOUBLE: - cudecompTransposeXToY(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), - reinterpret_cast(work), input_halo_extents, output_halo_extents, input_padding, - output_padding, stream); - break; - case CUDECOMP_FLOAT_COMPLEX: - cudecompTransposeXToY(handle, grid_desc, reinterpret_cast*>(input), - reinterpret_cast*>(output), - reinterpret_cast*>(work), input_halo_extents, output_halo_extents, - input_padding, output_padding, stream); - break; - case CUDECOMP_DOUBLE_COMPLEX: - cudecompTransposeXToY(handle, grid_desc, reinterpret_cast*>(input), - reinterpret_cast*>(output), - reinterpret_cast*>(work), input_halo_extents, - output_halo_extents, input_padding, output_padding, stream); - break; + size_t workspace_size = 0; + if (!work) { + int64_t workspace_elements; + CHECK_CUDECOMP(cudecompGetTransposeWorkspaceSize(handle, grid_desc, &workspace_elements)); + workspace_size = getWorkspaceSizeBytes(workspace_elements, dtype); } + auto execute = [&](void* actual_work, cudaStream_t actual_stream) { + switch (dtype) { + case CUDECOMP_FLOAT: + cudecompTransposeXToY(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(actual_work), input_halo_extents, output_halo_extents, + input_padding, output_padding, actual_stream); + break; + case CUDECOMP_DOUBLE: + cudecompTransposeXToY(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(actual_work), input_halo_extents, output_halo_extents, + input_padding, output_padding, actual_stream); + break; + case CUDECOMP_FLOAT_COMPLEX: + cudecompTransposeXToY(handle, grid_desc, reinterpret_cast*>(input), + reinterpret_cast*>(output), + reinterpret_cast*>(actual_work), input_halo_extents, + output_halo_extents, input_padding, output_padding, actual_stream); + break; + case CUDECOMP_DOUBLE_COMPLEX: + cudecompTransposeXToY(handle, grid_desc, reinterpret_cast*>(input), + reinterpret_cast*>(output), + reinterpret_cast*>(actual_work), input_halo_extents, + output_halo_extents, input_padding, output_padding, actual_stream); + break; + } + }; + bool use_nvshmem = transposeBackendRequiresNvshmem(grid_desc->config.transpose_comm_backend); + runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, stream, execute); } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; @@ -1806,31 +1949,40 @@ cudecompResult_t cudecompTransposeYToZ(cudecompHandle_t handle, cudecompGridDesc checkDataType(dtype); if (!input) { THROW_INVALID_USAGE("input argument cannot be null"); } if (!output) { THROW_INVALID_USAGE("output argument cannot be null"); } - if (!work) { THROW_INVALID_USAGE("work argument cannot be null"); } - switch (dtype) { - case CUDECOMP_FLOAT: - cudecompTransposeYToZ(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), - reinterpret_cast(work), input_halo_extents, output_halo_extents, input_padding, - output_padding, stream); - break; - case CUDECOMP_DOUBLE: - cudecompTransposeYToZ(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), - reinterpret_cast(work), input_halo_extents, output_halo_extents, input_padding, - output_padding, stream); - break; - case CUDECOMP_FLOAT_COMPLEX: - cudecompTransposeYToZ(handle, grid_desc, reinterpret_cast*>(input), - reinterpret_cast*>(output), - reinterpret_cast*>(work), input_halo_extents, output_halo_extents, - input_padding, output_padding, stream); - break; - case CUDECOMP_DOUBLE_COMPLEX: - cudecompTransposeYToZ(handle, grid_desc, reinterpret_cast*>(input), - reinterpret_cast*>(output), - reinterpret_cast*>(work), input_halo_extents, - output_halo_extents, input_padding, output_padding, stream); - break; + size_t workspace_size = 0; + if (!work) { + int64_t workspace_elements; + CHECK_CUDECOMP(cudecompGetTransposeWorkspaceSize(handle, grid_desc, &workspace_elements)); + workspace_size = getWorkspaceSizeBytes(workspace_elements, dtype); } + auto execute = [&](void* actual_work, cudaStream_t actual_stream) { + switch (dtype) { + case CUDECOMP_FLOAT: + cudecompTransposeYToZ(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(actual_work), input_halo_extents, output_halo_extents, + input_padding, output_padding, actual_stream); + break; + case CUDECOMP_DOUBLE: + cudecompTransposeYToZ(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(actual_work), input_halo_extents, output_halo_extents, + input_padding, output_padding, actual_stream); + break; + case CUDECOMP_FLOAT_COMPLEX: + cudecompTransposeYToZ(handle, grid_desc, reinterpret_cast*>(input), + reinterpret_cast*>(output), + reinterpret_cast*>(actual_work), input_halo_extents, + output_halo_extents, input_padding, output_padding, actual_stream); + break; + case CUDECOMP_DOUBLE_COMPLEX: + cudecompTransposeYToZ(handle, grid_desc, reinterpret_cast*>(input), + reinterpret_cast*>(output), + reinterpret_cast*>(actual_work), input_halo_extents, + output_halo_extents, input_padding, output_padding, actual_stream); + break; + } + }; + bool use_nvshmem = transposeBackendRequiresNvshmem(grid_desc->config.transpose_comm_backend); + runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, stream, execute); } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; @@ -1847,31 +1999,40 @@ cudecompResult_t cudecompTransposeZToY(cudecompHandle_t handle, cudecompGridDesc checkDataType(dtype); if (!input) { THROW_INVALID_USAGE("input argument cannot be null"); } if (!output) { THROW_INVALID_USAGE("output argument cannot be null"); } - if (!work) { THROW_INVALID_USAGE("work argument cannot be null"); } - switch (dtype) { - case CUDECOMP_FLOAT: - cudecompTransposeZToY(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), - reinterpret_cast(work), input_halo_extents, output_halo_extents, input_padding, - output_padding, stream); - break; - case CUDECOMP_DOUBLE: - cudecompTransposeZToY(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), - reinterpret_cast(work), input_halo_extents, output_halo_extents, input_padding, - output_padding, stream); - break; - case CUDECOMP_FLOAT_COMPLEX: - cudecompTransposeZToY(handle, grid_desc, reinterpret_cast*>(input), - reinterpret_cast*>(output), - reinterpret_cast*>(work), input_halo_extents, output_halo_extents, - input_padding, output_padding, stream); - break; - case CUDECOMP_DOUBLE_COMPLEX: - cudecompTransposeZToY(handle, grid_desc, reinterpret_cast*>(input), - reinterpret_cast*>(output), - reinterpret_cast*>(work), input_halo_extents, - output_halo_extents, input_padding, output_padding, stream); - break; + size_t workspace_size = 0; + if (!work) { + int64_t workspace_elements; + CHECK_CUDECOMP(cudecompGetTransposeWorkspaceSize(handle, grid_desc, &workspace_elements)); + workspace_size = getWorkspaceSizeBytes(workspace_elements, dtype); } + auto execute = [&](void* actual_work, cudaStream_t actual_stream) { + switch (dtype) { + case CUDECOMP_FLOAT: + cudecompTransposeZToY(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(actual_work), input_halo_extents, output_halo_extents, + input_padding, output_padding, actual_stream); + break; + case CUDECOMP_DOUBLE: + cudecompTransposeZToY(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(actual_work), input_halo_extents, output_halo_extents, + input_padding, output_padding, actual_stream); + break; + case CUDECOMP_FLOAT_COMPLEX: + cudecompTransposeZToY(handle, grid_desc, reinterpret_cast*>(input), + reinterpret_cast*>(output), + reinterpret_cast*>(actual_work), input_halo_extents, + output_halo_extents, input_padding, output_padding, actual_stream); + break; + case CUDECOMP_DOUBLE_COMPLEX: + cudecompTransposeZToY(handle, grid_desc, reinterpret_cast*>(input), + reinterpret_cast*>(output), + reinterpret_cast*>(actual_work), input_halo_extents, + output_halo_extents, input_padding, output_padding, actual_stream); + break; + } + }; + bool use_nvshmem = transposeBackendRequiresNvshmem(grid_desc->config.transpose_comm_backend); + runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, stream, execute); } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; @@ -1888,31 +2049,40 @@ cudecompResult_t cudecompTransposeYToX(cudecompHandle_t handle, cudecompGridDesc checkDataType(dtype); if (!input) { THROW_INVALID_USAGE("input argument cannot be null"); } if (!output) { THROW_INVALID_USAGE("output argument cannot be null"); } - if (!work) { THROW_INVALID_USAGE("work argument cannot be null"); } - switch (dtype) { - case CUDECOMP_FLOAT: - cudecompTransposeYToX(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), - reinterpret_cast(work), input_halo_extents, output_halo_extents, input_padding, - output_padding, stream); - break; - case CUDECOMP_DOUBLE: - cudecompTransposeYToX(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), - reinterpret_cast(work), input_halo_extents, output_halo_extents, input_padding, - output_padding, stream); - break; - case CUDECOMP_FLOAT_COMPLEX: - cudecompTransposeYToX(handle, grid_desc, reinterpret_cast*>(input), - reinterpret_cast*>(output), - reinterpret_cast*>(work), input_halo_extents, output_halo_extents, - input_padding, output_padding, stream); - break; - case CUDECOMP_DOUBLE_COMPLEX: - cudecompTransposeYToX(handle, grid_desc, reinterpret_cast*>(input), - reinterpret_cast*>(output), - reinterpret_cast*>(work), input_halo_extents, - output_halo_extents, input_padding, output_padding, stream); - break; + size_t workspace_size = 0; + if (!work) { + int64_t workspace_elements; + CHECK_CUDECOMP(cudecompGetTransposeWorkspaceSize(handle, grid_desc, &workspace_elements)); + workspace_size = getWorkspaceSizeBytes(workspace_elements, dtype); } + auto execute = [&](void* actual_work, cudaStream_t actual_stream) { + switch (dtype) { + case CUDECOMP_FLOAT: + cudecompTransposeYToX(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(actual_work), input_halo_extents, output_halo_extents, + input_padding, output_padding, actual_stream); + break; + case CUDECOMP_DOUBLE: + cudecompTransposeYToX(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(output), + reinterpret_cast(actual_work), input_halo_extents, output_halo_extents, + input_padding, output_padding, actual_stream); + break; + case CUDECOMP_FLOAT_COMPLEX: + cudecompTransposeYToX(handle, grid_desc, reinterpret_cast*>(input), + reinterpret_cast*>(output), + reinterpret_cast*>(actual_work), input_halo_extents, + output_halo_extents, input_padding, output_padding, actual_stream); + break; + case CUDECOMP_DOUBLE_COMPLEX: + cudecompTransposeYToX(handle, grid_desc, reinterpret_cast*>(input), + reinterpret_cast*>(output), + reinterpret_cast*>(actual_work), input_halo_extents, + output_halo_extents, input_padding, output_padding, actual_stream); + break; + } + }; + bool use_nvshmem = transposeBackendRequiresNvshmem(grid_desc->config.transpose_comm_backend); + runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, stream, execute); } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; @@ -1932,29 +2102,38 @@ cudecompResult_t cudecompUpdateHalosX(cudecompHandle_t handle, cudecompGridDesc_ return CUDECOMP_RESULT_SUCCESS; } if (!input) { THROW_INVALID_USAGE("input argument cannot be null"); } - if (!work) { THROW_INVALID_USAGE("work argument cannot be null"); } if (dim < 0 || dim > 2) { THROW_INVALID_USAGE("dim argument out of range"); } - - switch (dtype) { - case CUDECOMP_FLOAT: - cudecompUpdateHalosX(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(work), - halo_extents, halo_periods, dim, padding, stream); - break; - case CUDECOMP_DOUBLE: - cudecompUpdateHalosX(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(work), - halo_extents, halo_periods, dim, padding, stream); - break; - case CUDECOMP_FLOAT_COMPLEX: - cudecompUpdateHalosX(handle, grid_desc, reinterpret_cast*>(input), - reinterpret_cast*>(work), halo_extents, halo_periods, dim, padding, - stream); - break; - case CUDECOMP_DOUBLE_COMPLEX: - cudecompUpdateHalosX(handle, grid_desc, reinterpret_cast*>(input), - reinterpret_cast*>(work), halo_extents, halo_periods, dim, - padding, stream); - break; + size_t workspace_size = 0; + if (!work) { + int64_t workspace_elements; + CHECK_CUDECOMP(cudecompGetHaloWorkspaceSize(handle, grid_desc, 0, halo_extents, &workspace_elements)); + workspace_size = getWorkspaceSizeBytes(workspace_elements, dtype); } + auto execute = [&](void* actual_work, cudaStream_t actual_stream) { + switch (dtype) { + case CUDECOMP_FLOAT: + cudecompUpdateHalosX(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(actual_work), + halo_extents, halo_periods, dim, padding, actual_stream); + break; + case CUDECOMP_DOUBLE: + cudecompUpdateHalosX(handle, grid_desc, reinterpret_cast(input), + reinterpret_cast(actual_work), halo_extents, halo_periods, dim, padding, + actual_stream); + break; + case CUDECOMP_FLOAT_COMPLEX: + cudecompUpdateHalosX(handle, grid_desc, reinterpret_cast*>(input), + reinterpret_cast*>(actual_work), halo_extents, halo_periods, dim, + padding, actual_stream); + break; + case CUDECOMP_DOUBLE_COMPLEX: + cudecompUpdateHalosX(handle, grid_desc, reinterpret_cast*>(input), + reinterpret_cast*>(actual_work), halo_extents, halo_periods, + dim, padding, actual_stream); + break; + } + }; + bool use_nvshmem = haloBackendRequiresNvshmem(grid_desc->config.halo_comm_backend); + runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, stream, execute); } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; @@ -1974,29 +2153,38 @@ cudecompResult_t cudecompUpdateHalosY(cudecompHandle_t handle, cudecompGridDesc_ return CUDECOMP_RESULT_SUCCESS; } if (!input) { THROW_INVALID_USAGE("input argument cannot be null"); } - if (!work) { THROW_INVALID_USAGE("work argument cannot be null"); } if (dim < 0 || dim > 2) { THROW_INVALID_USAGE("dim argument out of range"); } - - switch (dtype) { - case CUDECOMP_FLOAT: - cudecompUpdateHalosY(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(work), - halo_extents, halo_periods, dim, padding, stream); - break; - case CUDECOMP_DOUBLE: - cudecompUpdateHalosY(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(work), - halo_extents, halo_periods, dim, padding, stream); - break; - case CUDECOMP_FLOAT_COMPLEX: - cudecompUpdateHalosY(handle, grid_desc, reinterpret_cast*>(input), - reinterpret_cast*>(work), halo_extents, halo_periods, dim, padding, - stream); - break; - case CUDECOMP_DOUBLE_COMPLEX: - cudecompUpdateHalosY(handle, grid_desc, reinterpret_cast*>(input), - reinterpret_cast*>(work), halo_extents, halo_periods, dim, - padding, stream); - break; + size_t workspace_size = 0; + if (!work) { + int64_t workspace_elements; + CHECK_CUDECOMP(cudecompGetHaloWorkspaceSize(handle, grid_desc, 1, halo_extents, &workspace_elements)); + workspace_size = getWorkspaceSizeBytes(workspace_elements, dtype); } + auto execute = [&](void* actual_work, cudaStream_t actual_stream) { + switch (dtype) { + case CUDECOMP_FLOAT: + cudecompUpdateHalosY(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(actual_work), + halo_extents, halo_periods, dim, padding, actual_stream); + break; + case CUDECOMP_DOUBLE: + cudecompUpdateHalosY(handle, grid_desc, reinterpret_cast(input), + reinterpret_cast(actual_work), halo_extents, halo_periods, dim, padding, + actual_stream); + break; + case CUDECOMP_FLOAT_COMPLEX: + cudecompUpdateHalosY(handle, grid_desc, reinterpret_cast*>(input), + reinterpret_cast*>(actual_work), halo_extents, halo_periods, dim, + padding, actual_stream); + break; + case CUDECOMP_DOUBLE_COMPLEX: + cudecompUpdateHalosY(handle, grid_desc, reinterpret_cast*>(input), + reinterpret_cast*>(actual_work), halo_extents, halo_periods, + dim, padding, actual_stream); + break; + } + }; + bool use_nvshmem = haloBackendRequiresNvshmem(grid_desc->config.halo_comm_backend); + runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, stream, execute); } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; @@ -2016,29 +2204,38 @@ cudecompResult_t cudecompUpdateHalosZ(cudecompHandle_t handle, cudecompGridDesc_ return CUDECOMP_RESULT_SUCCESS; } if (!input) { THROW_INVALID_USAGE("input argument cannot be null"); } - if (!work) { THROW_INVALID_USAGE("work argument cannot be null"); } if (dim < 0 || dim > 2) { THROW_INVALID_USAGE("dim argument out of range"); } - - switch (dtype) { - case CUDECOMP_FLOAT: - cudecompUpdateHalosZ(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(work), - halo_extents, halo_periods, dim, padding, stream); - break; - case CUDECOMP_DOUBLE: - cudecompUpdateHalosZ(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(work), - halo_extents, halo_periods, dim, padding, stream); - break; - case CUDECOMP_FLOAT_COMPLEX: - cudecompUpdateHalosZ(handle, grid_desc, reinterpret_cast*>(input), - reinterpret_cast*>(work), halo_extents, halo_periods, dim, padding, - stream); - break; - case CUDECOMP_DOUBLE_COMPLEX: - cudecompUpdateHalosZ(handle, grid_desc, reinterpret_cast*>(input), - reinterpret_cast*>(work), halo_extents, halo_periods, dim, - padding, stream); - break; + size_t workspace_size = 0; + if (!work) { + int64_t workspace_elements; + CHECK_CUDECOMP(cudecompGetHaloWorkspaceSize(handle, grid_desc, 2, halo_extents, &workspace_elements)); + workspace_size = getWorkspaceSizeBytes(workspace_elements, dtype); } + auto execute = [&](void* actual_work, cudaStream_t actual_stream) { + switch (dtype) { + case CUDECOMP_FLOAT: + cudecompUpdateHalosZ(handle, grid_desc, reinterpret_cast(input), reinterpret_cast(actual_work), + halo_extents, halo_periods, dim, padding, actual_stream); + break; + case CUDECOMP_DOUBLE: + cudecompUpdateHalosZ(handle, grid_desc, reinterpret_cast(input), + reinterpret_cast(actual_work), halo_extents, halo_periods, dim, padding, + actual_stream); + break; + case CUDECOMP_FLOAT_COMPLEX: + cudecompUpdateHalosZ(handle, grid_desc, reinterpret_cast*>(input), + reinterpret_cast*>(actual_work), halo_extents, halo_periods, dim, + padding, actual_stream); + break; + case CUDECOMP_DOUBLE_COMPLEX: + cudecompUpdateHalosZ(handle, grid_desc, reinterpret_cast*>(input), + reinterpret_cast*>(actual_work), halo_extents, halo_periods, + dim, padding, actual_stream); + break; + } + }; + bool use_nvshmem = haloBackendRequiresNvshmem(grid_desc->config.halo_comm_backend); + runWithWorkspace(handle, grid_desc, work, workspace_size, use_nvshmem, stream, execute); } CUDECOMP_CATCH_C_API_ERRORS() return CUDECOMP_RESULT_SUCCESS; diff --git a/src/cudecomp_m.cuf b/src/cudecomp_m.cuf index cf7ca63..1ee1a6e 100644 --- a/src/cudecomp_m.cuf +++ b/src/cudecomp_m.cuf @@ -162,6 +162,9 @@ module cudecomp enumerator :: CUDECOMP_RESULT_NVML_ERROR = 9 end enum + ! Sentinel used to select handle-owned workspace management. Its pointer association must not be changed. + real(c_float), pointer, public, protected :: CUDECOMP_WORKSPACE_AUTO(:) => null() + ! types ! Opaque handle to cuDecomp handle @@ -408,8 +411,9 @@ module cudecomp import type(cudecompHandle), value :: handle type(cudecompGridDesc), value :: grid_desc - !dir$ ignore_tkr input, output, work - real(c_float), device :: input(*), output(*), work(*) + !dir$ ignore_tkr input, output + real(c_float), device :: input(*), output(*) + type(c_devptr), value :: work integer(c_int), value :: dtype integer(c_int32_t) :: input_halo_extents(3), output_halo_extents(3) integer(c_int32_t) :: input_padding(3), output_padding(3) @@ -426,8 +430,9 @@ module cudecomp import type(cudecompHandle), value :: handle type(cudecompGridDesc), value :: grid_desc - !dir$ ignore_tkr input, output, work - real(c_float), device :: input(*), output(*), work(*) + !dir$ ignore_tkr input, output + real(c_float), device :: input(*), output(*) + type(c_devptr), value :: work integer(c_int), value :: dtype integer(c_int32_t) :: input_halo_extents(3), output_halo_extents(3) integer(c_int32_t) :: input_padding(3), output_padding(3) @@ -444,8 +449,9 @@ module cudecomp import type(cudecompHandle), value :: handle type(cudecompGridDesc), value :: grid_desc - !dir$ ignore_tkr input, output, work - real(c_float), device :: input(*), output(*), work(*) + !dir$ ignore_tkr input, output + real(c_float), device :: input(*), output(*) + type(c_devptr), value :: work integer(c_int), value :: dtype integer(c_int32_t) :: input_halo_extents(3), output_halo_extents(3) integer(c_int32_t) :: input_padding(3), output_padding(3) @@ -462,8 +468,9 @@ module cudecomp import type(cudecompHandle), value :: handle type(cudecompGridDesc), value :: grid_desc - !dir$ ignore_tkr input, output, work - real(c_float), device :: input(*), output(*), work(*) + !dir$ ignore_tkr input, output + real(c_float), device :: input(*), output(*) + type(c_devptr), value :: work integer(c_int), value :: dtype integer(c_int32_t) :: input_halo_extents(3), output_halo_extents(3) integer(c_int32_t) :: input_padding(3), output_padding(3) @@ -480,8 +487,9 @@ module cudecomp import type(cudecompHandle), value :: handle type(cudecompGridDesc), value :: grid_desc - !dir$ ignore_tkr input, work - real(c_float), device :: input(*), work(*) + !dir$ ignore_tkr input + real(c_float), device :: input(*) + type(c_devptr), value :: work integer(c_int), value :: dtype integer(c_int32_t) :: halo_extents(3) logical(c_bool) :: halo_periods(3) @@ -499,8 +507,9 @@ module cudecomp import type(cudecompHandle), value :: handle type(cudecompGridDesc), value :: grid_desc - !dir$ ignore_tkr input, work - real(c_float), device :: input(*), work(*) + !dir$ ignore_tkr input + real(c_float), device :: input(*) + type(c_devptr), value :: work integer(c_int), value :: dtype integer(c_int32_t) :: halo_extents(3) logical(c_bool) :: halo_periods(3) @@ -518,8 +527,9 @@ module cudecomp import type(cudecompHandle), value :: handle type(cudecompGridDesc), value :: grid_desc - !dir$ ignore_tkr input, work - real(c_float), device :: input(*), work(*) + !dir$ ignore_tkr input + real(c_float), device :: input(*) + type(c_devptr), value :: work integer(c_int), value :: dtype integer(c_int32_t) :: halo_extents(3) logical(c_bool) :: halo_periods(3) @@ -842,6 +852,7 @@ contains integer :: output_halo_extents_(3) integer :: input_padding_(3) integer :: output_padding_(3) + type(c_devptr) :: work_ stream_ = 0 input_halo_extents_(:) = [0, 0, 0] @@ -853,8 +864,9 @@ contains if (present(output_halo_extents)) output_halo_extents_ = output_halo_extents if (present(input_padding)) input_padding_ = input_padding if (present(output_padding)) output_padding_ = output_padding + work_ = c_devloc(work) res = cudecompTransposeXToY_C(handle, grid_desc, & - input, output, work, dtype, input_halo_extents_, output_halo_extents_, & + input, output, work_, dtype, input_halo_extents_, output_halo_extents_, & input_padding_, output_padding_, stream_) end function cudecompTransposeXToY @@ -879,6 +891,7 @@ contains integer :: output_halo_extents_(3) integer :: input_padding_(3) integer :: output_padding_(3) + type(c_devptr) :: work_ stream_ = 0 input_halo_extents_(:) = [0, 0, 0] @@ -890,8 +903,9 @@ contains if (present(output_halo_extents)) output_halo_extents_ = output_halo_extents if (present(input_padding)) input_padding_ = input_padding if (present(output_padding)) output_padding_ = output_padding + work_ = c_devloc(work) res = cudecompTransposeYToZ_C(handle, grid_desc, & - input, output, work, dtype, input_halo_extents_, output_halo_extents_, & + input, output, work_, dtype, input_halo_extents_, output_halo_extents_, & input_padding_, output_padding_, stream_) end function cudecompTransposeYToZ @@ -916,6 +930,7 @@ contains integer :: output_halo_extents_(3) integer :: input_padding_(3) integer :: output_padding_(3) + type(c_devptr) :: work_ stream_ = 0 input_halo_extents_(:) = [0, 0, 0] @@ -927,8 +942,9 @@ contains if (present(output_halo_extents)) output_halo_extents_ = output_halo_extents if (present(input_padding)) input_padding_ = input_padding if (present(output_padding)) output_padding_ = output_padding + work_ = c_devloc(work) res = cudecompTransposeZToY_C(handle, grid_desc, & - input, output, work, dtype, input_halo_extents_, output_halo_extents_, & + input, output, work_, dtype, input_halo_extents_, output_halo_extents_, & input_padding_, output_padding_, stream_) end function cudecompTransposeZToY @@ -953,6 +969,7 @@ contains integer :: output_halo_extents_(3) integer :: input_padding_(3) integer :: output_padding_(3) + type(c_devptr) :: work_ stream_ = 0 input_halo_extents_(:) = [0, 0, 0] @@ -964,8 +981,9 @@ contains if (present(output_halo_extents)) output_halo_extents_ = output_halo_extents if (present(input_padding)) input_padding_ = input_padding if (present(output_padding)) output_padding_ = output_padding + work_ = c_devloc(work) res = cudecompTransposeYToX_C(handle, grid_desc, & - input, output, work, dtype, input_halo_extents_, output_halo_extents_, & + input, output, work_, dtype, input_halo_extents_, output_halo_extents_, & input_padding_, output_padding_, stream_) end function cudecompTransposeYToX @@ -989,6 +1007,7 @@ contains integer(cuda_stream_kind) :: stream_ logical(c_bool) :: halo_periods_c(3) integer :: padding_(3) + type(c_devptr) :: work_ halo_periods_c(:) = halo_periods @@ -996,8 +1015,9 @@ contains padding_ = [0, 0, 0] if (present(stream)) stream_ = stream if (present(padding)) padding_ = padding + work_ = c_devloc(work) res = cudecompUpdateHalosX_C(handle, grid_desc, & - input, work, dtype, halo_extents, halo_periods_c, & + input, work_, dtype, halo_extents, halo_periods_c, & dim - 1, padding_, stream_) end function cudecompUpdateHalosX @@ -1020,6 +1040,7 @@ contains integer(cuda_stream_kind) :: stream_ logical(c_bool) :: halo_periods_c(3) integer :: padding_(3) + type(c_devptr) :: work_ halo_periods_c(:) = halo_periods @@ -1027,8 +1048,9 @@ contains padding_ = [0, 0, 0] if (present(stream)) stream_ = stream if (present(padding)) padding_ = padding + work_ = c_devloc(work) res = cudecompUpdateHalosY_C(handle, grid_desc, & - input, work, dtype, halo_extents, halo_periods_c, & + input, work_, dtype, halo_extents, halo_periods_c, & dim - 1, padding_, stream_) end function cudecompUpdateHalosY @@ -1051,6 +1073,7 @@ contains integer(cuda_stream_kind) :: stream_ logical(c_bool) :: halo_periods_c(3) integer :: padding_(3) + type(c_devptr) :: work_ halo_periods_c(:) = halo_periods @@ -1058,8 +1081,9 @@ contains padding_ = [0, 0, 0] if (present(stream)) stream_ = stream if (present(padding)) padding_ = padding + work_ = c_devloc(work) res = cudecompUpdateHalosZ_C(handle, grid_desc, & - input, work, dtype, halo_extents, halo_periods_c, & + input, work_, dtype, halo_extents, halo_periods_c, & dim - 1, padding_, stream_) end function cudecompUpdateHalosZ diff --git a/tests/ctest/CMakeLists.txt b/tests/ctest/CMakeLists.txt index d1b1b50..9603e58 100644 --- a/tests/ctest/CMakeLists.txt +++ b/tests/ctest/CMakeLists.txt @@ -115,6 +115,42 @@ if (CUDECOMP_ENABLE_NVSHMEM) ) endif() +add_executable(cudecomp_test_automatic_workspace) +target_sources(cudecomp_test_automatic_workspace + PRIVATE + automatic_workspace_tests.cc + mpi_test_main.cc +) +target_link_libraries(cudecomp_test_automatic_workspace + PRIVATE + cudecomp + cudecomp_test_support +) +set_target_properties(cudecomp_test_automatic_workspace PROPERTIES + RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/tests/ctest +) + +set(CUDECOMP_TEST_AUTOMATIC_WORKSPACE_MPI_RANKS 4) +add_test( + NAME cudecomp_automatic_workspace + COMMAND ${MPIEXEC_EXECUTABLE} ${CUDECOMP_TEST_MPIEXEC_PREFLAGS_LIST} + ${MPIEXEC_NUMPROC_FLAG} ${CUDECOMP_TEST_AUTOMATIC_WORKSPACE_MPI_RANKS} + $ +) +set_tests_properties(cudecomp_automatic_workspace PROPERTIES + LABELS "workspace;mpi;nccl;nccl_ubr" + TIMEOUT ${CUDECOMP_TEST_TIMEOUT} + ENVIRONMENT "CUDECOMP_ENABLE_NCCL_UBR=1" +) +if (CUDECOMP_ENABLE_NVSHMEM) + set_property(TEST cudecomp_automatic_workspace APPEND PROPERTY + LABELS "nvshmem" + ) + set_property(TEST cudecomp_automatic_workspace APPEND PROPERTY + ENVIRONMENT "NVSHMEM_DISABLE_NCCL=1" + ) +endif() + add_executable(cudecomp_test_transpose) target_sources(cudecomp_test_transpose PRIVATE diff --git a/tests/ctest/api_tests.cc b/tests/ctest/api_tests.cc index e1778f1..0e9a15b 100644 --- a/tests/ctest/api_tests.cc +++ b/tests/ctest/api_tests.cc @@ -1482,9 +1482,6 @@ TEST_F(ApiTransposeTest, RejectsInvalidArguments) { EXPECT_EQ(CUDECOMP_RESULT_INVALID_USAGE, cudecompTransposeXToY(handle_, grid_desc, valid_pointer, nullptr, valid_pointer, CUDECOMP_FLOAT, nullptr, nullptr, nullptr, nullptr, 0)); - EXPECT_EQ(CUDECOMP_RESULT_INVALID_USAGE, - cudecompTransposeXToY(handle_, grid_desc, valid_pointer, valid_pointer, nullptr, CUDECOMP_FLOAT, nullptr, - nullptr, nullptr, nullptr, 0)); EXPECT_EQ(CUDECOMP_RESULT_INVALID_USAGE, cudecompTransposeXToY(handle_, grid_desc, valid_pointer, valid_pointer, valid_pointer, static_cast(999), nullptr, nullptr, nullptr, nullptr, 0)); @@ -1521,9 +1518,6 @@ TEST_F(ApiHaloTest, RejectsInvalidArguments) { EXPECT_EQ(CUDECOMP_RESULT_INVALID_USAGE, cudecompUpdateHalosX(handle_, grid_desc, nullptr, valid_pointer, CUDECOMP_FLOAT, kHaloExtents.data(), kHaloPeriods.data(), 0, nullptr, 0)); - EXPECT_EQ(CUDECOMP_RESULT_INVALID_USAGE, - cudecompUpdateHalosX(handle_, grid_desc, valid_pointer, nullptr, CUDECOMP_FLOAT, kHaloExtents.data(), - kHaloPeriods.data(), 0, nullptr, 0)); EXPECT_EQ(CUDECOMP_RESULT_INVALID_USAGE, cudecompUpdateHalosX(handle_, grid_desc, valid_pointer, valid_pointer, CUDECOMP_FLOAT, kHaloExtents.data(), kHaloPeriods.data(), 3, nullptr, 0)); diff --git a/tests/ctest/automatic_workspace_tests.cc b/tests/ctest/automatic_workspace_tests.cc new file mode 100644 index 0000000..f281bae --- /dev/null +++ b/tests/ctest/automatic_workspace_tests.cc @@ -0,0 +1,240 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include + +#include +#include + +#include "cudecomp.h" +#include "internal/common.h" + +#include "gpu_test_utils.h" +#include "mpi_test_utils.h" +#include "test_utils.h" + +namespace { + +constexpr std::array kGdims{9, 10, 11}; +constexpr std::array kPdims{2, 2}; +constexpr std::array kHaloExtents{0, 1, 0}; +constexpr std::array kHaloPeriods{true, true, true}; + +cudecompGridDescConfig_t makeConfig(cudecompTransposeCommBackend_t transpose_backend, + cudecompHaloCommBackend_t halo_backend) { + cudecompGridDescConfig_t config; + EXPECT_EQ(CUDECOMP_RESULT_SUCCESS, cudecompGridDescConfigSetDefaults(&config)); + std::copy(kGdims.begin(), kGdims.end(), config.gdims); + std::copy(kPdims.begin(), kPdims.end(), config.pdims); + config.transpose_comm_backend = transpose_backend; + config.halo_comm_backend = halo_backend; + return config; +} + +} // namespace + +TEST(AutomaticWorkspaceTest, GrowsAcrossHaloAndDtypeRequirements) { + const auto world_comm = cudecomp_test::MpiTestComm::world(); + if (world_comm.size() != 4) { GTEST_SKIP() << "automatic workspace growth test requires exactly four ranks"; } + + const auto setup_decision = cudecomp_test::initializeGpuForTest(world_comm); + ASSERT_FALSE(setup_decision.fail) << setup_decision.reason; + if (setup_decision.skip) { GTEST_SKIP() << setup_decision.reason; } + + cudecompHandle_t handle = nullptr; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompInit(&handle, world_comm.mpiComm())); + cudecomp_test::cudecompHandleGuard handle_guard(handle); + + auto config = makeConfig(CUDECOMP_TRANSPOSE_COMM_MPI_P2P, CUDECOMP_HALO_COMM_MPI); + cudecompGridDesc_t grid_desc = nullptr; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGridDescCreate(handle, &grid_desc, &config, nullptr)); + cudecomp_test::gridDescGuard grid_desc_guard(handle, grid_desc); + + cudecompPencilInfo_t halo_pinfo; + CHECK_CUDECOMP_GLOBAL(world_comm, + cudecompGetPencilInfo(handle, grid_desc, &halo_pinfo, 0, kHaloExtents.data(), nullptr)); + float* halo_data = nullptr; + CHECK_CUDA_GLOBAL(world_comm, cudaMalloc(&halo_data, halo_pinfo.size * sizeof(*halo_data))); + cudecomp_test::cudaBufferGuard halo_data_guard(halo_data); + CHECK_CUDA_GLOBAL(world_comm, cudaMemset(halo_data, 0, halo_pinfo.size * sizeof(*halo_data))); + + int64_t halo_workspace_elements = 0; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGetHaloWorkspaceSize(handle, grid_desc, 0, kHaloExtents.data(), + &halo_workspace_elements)); + CHECK_CUDECOMP_GLOBAL(world_comm, + cudecompUpdateHalosX(handle, grid_desc, halo_data, CUDECOMP_WORKSPACE_AUTO, CUDECOMP_FLOAT, + kHaloExtents.data(), kHaloPeriods.data(), 1, nullptr, nullptr)); + CHECK_CUDA_GLOBAL(world_comm, cudaDeviceSynchronize()); + + ASSERT_NE(handle->ordinary_workspace.ptr, nullptr); + const size_t initial_size = handle->ordinary_workspace.size; + EXPECT_EQ(initial_size, static_cast(halo_workspace_elements) * sizeof(float)); + + cudecompPencilInfo_t x_pinfo; + cudecompPencilInfo_t y_pinfo; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGetPencilInfo(handle, grid_desc, &x_pinfo, 0, nullptr, nullptr)); + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGetPencilInfo(handle, grid_desc, &y_pinfo, 1, nullptr, nullptr)); + const int64_t transpose_data_elements = std::max(x_pinfo.size, y_pinfo.size); + std::complex* transpose_input = nullptr; + std::complex* transpose_output = nullptr; + CHECK_CUDA_GLOBAL(world_comm, cudaMalloc(&transpose_input, transpose_data_elements * sizeof(*transpose_input))); + cudecomp_test::cudaBufferGuard transpose_input_guard(transpose_input); + CHECK_CUDA_GLOBAL(world_comm, cudaMalloc(&transpose_output, transpose_data_elements * sizeof(*transpose_output))); + cudecomp_test::cudaBufferGuard transpose_output_guard(transpose_output); + CHECK_CUDA_GLOBAL(world_comm, cudaMemset(transpose_input, 0, transpose_data_elements * sizeof(*transpose_input))); + + int64_t transpose_workspace_elements = 0; + CHECK_CUDECOMP_GLOBAL(world_comm, + cudecompGetTransposeWorkspaceSize(handle, grid_desc, &transpose_workspace_elements)); + const size_t required_transpose_size = static_cast(transpose_workspace_elements) * sizeof(*transpose_input); + ASSERT_GT(required_transpose_size, initial_size); + + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompTransposeXToY(handle, grid_desc, transpose_input, transpose_output, + CUDECOMP_WORKSPACE_AUTO, CUDECOMP_DOUBLE_COMPLEX, nullptr, + nullptr, nullptr, nullptr, nullptr)); + CHECK_CUDA_GLOBAL(world_comm, cudaDeviceSynchronize()); + + ASSERT_NE(handle->ordinary_workspace.ptr, nullptr); + EXPECT_EQ(handle->ordinary_workspace.size, required_transpose_size); + void* grown_workspace = handle->ordinary_workspace.ptr; + + CHECK_CUDECOMP_GLOBAL(world_comm, + cudecompUpdateHalosX(handle, grid_desc, halo_data, CUDECOMP_WORKSPACE_AUTO, CUDECOMP_FLOAT, + kHaloExtents.data(), kHaloPeriods.data(), 1, nullptr, nullptr)); + CHECK_CUDA_GLOBAL(world_comm, cudaDeviceSynchronize()); + EXPECT_EQ(handle->ordinary_workspace.ptr, grown_workspace); + EXPECT_EQ(handle->ordinary_workspace.size, required_transpose_size); +} + +TEST(AutomaticWorkspaceTest, RegistersCacheReusedByLaterDescriptor) { +#if NCCL_VERSION_CODE < NCCL_VERSION(2, 19, 0) + GTEST_SKIP() << "NCCL user buffer registration requires NCCL 2.19 or newer"; +#else + const auto world_comm = cudecomp_test::MpiTestComm::world(); + if (world_comm.size() != 4) { GTEST_SKIP() << "NCCL workspace registration test requires exactly four ranks"; } + + const auto setup_decision = cudecomp_test::initializeGpuForTest(world_comm, true); + ASSERT_FALSE(setup_decision.fail) << setup_decision.reason; + if (setup_decision.skip) { GTEST_SKIP() << setup_decision.reason; } + + cudecompHandle_t handle = nullptr; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompInit(&handle, world_comm.mpiComm())); + cudecomp_test::cudecompHandleGuard handle_guard(handle); + ASSERT_TRUE(handle->nccl_enable_ubr); + + auto mpi_config = makeConfig(CUDECOMP_TRANSPOSE_COMM_MPI_P2P, CUDECOMP_HALO_COMM_MPI); + cudecompGridDesc_t mpi_grid_desc = nullptr; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGridDescCreate(handle, &mpi_grid_desc, &mpi_config, nullptr)); + cudecomp_test::gridDescGuard mpi_grid_desc_guard(handle, mpi_grid_desc); + + cudecompPencilInfo_t x_pinfo; + cudecompPencilInfo_t y_pinfo; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGetPencilInfo(handle, mpi_grid_desc, &x_pinfo, 0, nullptr, nullptr)); + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGetPencilInfo(handle, mpi_grid_desc, &y_pinfo, 1, nullptr, nullptr)); + const int64_t data_elements = std::max(x_pinfo.size, y_pinfo.size); + float* input = nullptr; + float* output = nullptr; + CHECK_CUDA_GLOBAL(world_comm, cudaMalloc(&input, data_elements * sizeof(*input))); + cudecomp_test::cudaBufferGuard input_guard(input); + CHECK_CUDA_GLOBAL(world_comm, cudaMalloc(&output, data_elements * sizeof(*output))); + cudecomp_test::cudaBufferGuard output_guard(output); + CHECK_CUDA_GLOBAL(world_comm, cudaMemset(input, 0, data_elements * sizeof(*input))); + + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompTransposeXToY(handle, mpi_grid_desc, input, output, CUDECOMP_WORKSPACE_AUTO, + CUDECOMP_FLOAT, nullptr, nullptr, nullptr, nullptr, nullptr)); + CHECK_CUDA_GLOBAL(world_comm, cudaDeviceSynchronize()); + void* cached_workspace = handle->ordinary_workspace.ptr; + ASSERT_NE(cached_workspace, nullptr); + EXPECT_EQ(handle->nccl_ubr_handles.count(cached_workspace), 0); + + auto nccl_config = makeConfig(CUDECOMP_TRANSPOSE_COMM_NCCL, CUDECOMP_HALO_COMM_MPI); + cudecompGridDesc_t nccl_grid_desc = nullptr; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGridDescCreate(handle, &nccl_grid_desc, &nccl_config, nullptr)); + cudecomp_test::gridDescGuard nccl_grid_desc_guard(handle, nccl_grid_desc); + + CHECK_CUDECOMP_GLOBAL(world_comm, + cudecompTransposeXToY(handle, nccl_grid_desc, input, output, CUDECOMP_WORKSPACE_AUTO, + CUDECOMP_FLOAT, nullptr, nullptr, nullptr, nullptr, nullptr)); + CHECK_CUDA_GLOBAL(world_comm, cudaDeviceSynchronize()); + EXPECT_EQ(handle->ordinary_workspace.ptr, cached_workspace); + + auto entry = handle->nccl_ubr_handles.find(cached_workspace); + ASSERT_NE(entry, handle->nccl_ubr_handles.end()); + auto registered_with = [&](const cudecomp::ncclComm& comm) { + return !comm || std::any_of(entry->second.begin(), entry->second.end(), + [&](const auto& registration) { return registration.first.get() == comm.get(); }); + }; + EXPECT_TRUE(registered_with(nccl_grid_desc->nccl_comm)); + EXPECT_TRUE(registered_with(nccl_grid_desc->nccl_local_comm)); +#endif +} + +#ifdef ENABLE_NVSHMEM +TEST(AutomaticWorkspaceTest, KeepsAllocationDomainsSeparate) { + const auto world_comm = cudecomp_test::MpiTestComm::world(); + if (world_comm.size() != 4) { GTEST_SKIP() << "NVSHMEM allocation-domain test requires exactly four ranks"; } + + const auto setup_decision = cudecomp_test::initializeGpuForTest(world_comm); + ASSERT_FALSE(setup_decision.fail) << setup_decision.reason; + if (setup_decision.skip) { GTEST_SKIP() << setup_decision.reason; } + + cudecompHandle_t handle = nullptr; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompInit(&handle, world_comm.mpiComm())); + cudecomp_test::cudecompHandleGuard handle_guard(handle); + + auto config = makeConfig(CUDECOMP_TRANSPOSE_COMM_MPI_P2P, CUDECOMP_HALO_COMM_NVSHMEM); + cudecompGridDesc_t grid_desc = nullptr; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGridDescCreate(handle, &grid_desc, &config, nullptr)); + cudecomp_test::gridDescGuard grid_desc_guard(handle, grid_desc); + + cudecompPencilInfo_t x_pinfo; + cudecompPencilInfo_t y_pinfo; + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGetPencilInfo(handle, grid_desc, &x_pinfo, 0, nullptr, nullptr)); + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompGetPencilInfo(handle, grid_desc, &y_pinfo, 1, nullptr, nullptr)); + const int64_t transpose_data_elements = std::max(x_pinfo.size, y_pinfo.size); + float* x_data = nullptr; + float* y_data = nullptr; + CHECK_CUDA_GLOBAL(world_comm, cudaMalloc(&x_data, transpose_data_elements * sizeof(*x_data))); + cudecomp_test::cudaBufferGuard x_data_guard(x_data); + CHECK_CUDA_GLOBAL(world_comm, cudaMalloc(&y_data, transpose_data_elements * sizeof(*y_data))); + cudecomp_test::cudaBufferGuard y_data_guard(y_data); + CHECK_CUDA_GLOBAL(world_comm, cudaMemset(x_data, 0, transpose_data_elements * sizeof(*x_data))); + + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompTransposeXToY(handle, grid_desc, x_data, y_data, CUDECOMP_WORKSPACE_AUTO, + CUDECOMP_FLOAT, nullptr, nullptr, nullptr, nullptr, nullptr)); + CHECK_CUDA_GLOBAL(world_comm, cudaDeviceSynchronize()); + void* ordinary_workspace = handle->ordinary_workspace.ptr; + ASSERT_NE(ordinary_workspace, nullptr); + EXPECT_EQ(handle->nvshmem_workspace.ptr, nullptr); + + cudecompPencilInfo_t halo_pinfo; + CHECK_CUDECOMP_GLOBAL(world_comm, + cudecompGetPencilInfo(handle, grid_desc, &halo_pinfo, 0, kHaloExtents.data(), nullptr)); + float* halo_data = nullptr; + CHECK_CUDA_GLOBAL(world_comm, cudaMalloc(&halo_data, halo_pinfo.size * sizeof(*halo_data))); + cudecomp_test::cudaBufferGuard halo_data_guard(halo_data); + CHECK_CUDA_GLOBAL(world_comm, cudaMemset(halo_data, 0, halo_pinfo.size * sizeof(*halo_data))); + + CHECK_CUDECOMP_GLOBAL(world_comm, + cudecompUpdateHalosX(handle, grid_desc, halo_data, CUDECOMP_WORKSPACE_AUTO, CUDECOMP_FLOAT, + kHaloExtents.data(), kHaloPeriods.data(), 1, nullptr, nullptr)); + CHECK_CUDA_GLOBAL(world_comm, cudaDeviceSynchronize()); + void* nvshmem_workspace = handle->nvshmem_workspace.ptr; + ASSERT_NE(nvshmem_workspace, nullptr); + EXPECT_EQ(handle->ordinary_workspace.ptr, ordinary_workspace); + EXPECT_NE(nvshmem_workspace, ordinary_workspace); + + CHECK_CUDECOMP_GLOBAL(world_comm, cudecompTransposeYToX(handle, grid_desc, y_data, x_data, CUDECOMP_WORKSPACE_AUTO, + CUDECOMP_FLOAT, nullptr, nullptr, nullptr, nullptr, nullptr)); + CHECK_CUDA_GLOBAL(world_comm, cudaDeviceSynchronize()); + EXPECT_EQ(handle->ordinary_workspace.ptr, ordinary_workspace); + EXPECT_EQ(handle->nvshmem_workspace.ptr, nvshmem_workspace); +} +#endif diff --git a/tests/ctest/fortran_halo_case.inc b/tests/ctest/fortran_halo_case.inc index 9583b56..6cab334 100644 --- a/tests/ctest/fortran_halo_case.inc +++ b/tests/ctest/fortran_halo_case.inc @@ -98,27 +98,31 @@ subroutine RUN_HALO_AXIS(scenario, axis, halo_extents, padding, halo_periods, me select case (axis) case (1) if (use_optional_padding) then - res = cudecompUpdateHalosX(handle, grid_desc, data_d, work_d, DTYPE, pinfo%halo_extents, halo_periods, & - dim, pinfo%padding) + res = cudecompUpdateHalosX(handle, grid_desc, data_d, CUDECOMP_WORKSPACE_AUTO, DTYPE, pinfo%halo_extents, & + halo_periods, dim, pinfo%padding) else - res = cudecompUpdateHalosX(handle, grid_desc, data_d, work_d, DTYPE, pinfo%halo_extents, halo_periods, dim) + res = cudecompUpdateHalosX(handle, grid_desc, data_d, CUDECOMP_WORKSPACE_AUTO, DTYPE, pinfo%halo_extents, & + halo_periods, dim) endif case (2) if (use_optional_padding) then - res = cudecompUpdateHalosY(handle, grid_desc, data_d, work_d, DTYPE, pinfo%halo_extents, halo_periods, & - dim, pinfo%padding) + res = cudecompUpdateHalosY(handle, grid_desc, data_d, CUDECOMP_WORKSPACE_AUTO, DTYPE, pinfo%halo_extents, & + halo_periods, dim, pinfo%padding) else - res = cudecompUpdateHalosY(handle, grid_desc, data_d, work_d, DTYPE, pinfo%halo_extents, halo_periods, dim) + res = cudecompUpdateHalosY(handle, grid_desc, data_d, CUDECOMP_WORKSPACE_AUTO, DTYPE, pinfo%halo_extents, & + halo_periods, dim) endif case (3) if (use_optional_padding) then - res = cudecompUpdateHalosZ(handle, grid_desc, data_d, work_d, DTYPE, pinfo%halo_extents, halo_periods, & - dim, pinfo%padding) + res = cudecompUpdateHalosZ(handle, grid_desc, data_d, CUDECOMP_WORKSPACE_AUTO, DTYPE, pinfo%halo_extents, & + halo_periods, dim, pinfo%padding) else - res = cudecompUpdateHalosZ(handle, grid_desc, data_d, work_d, DTYPE, pinfo%halo_extents, halo_periods, dim) + res = cudecompUpdateHalosZ(handle, grid_desc, data_d, CUDECOMP_WORKSPACE_AUTO, DTYPE, pinfo%halo_extents, & + halo_periods, dim) endif end select - call expect_success(res, DTYPE_NAME//" "//trim(scenario)//" axis "//axis_name(axis)//" cudecompUpdateHalos") + call expect_success(res, DTYPE_NAME//" "//trim(scenario)//" axis "//axis_name(axis)// & + " cudecompUpdateHalos automatic workspace") enddo allocate(actual(pinfo%size)) diff --git a/tests/ctest/fortran_transpose_case.inc b/tests/ctest/fortran_transpose_case.inc index 1fae212..1173032 100644 --- a/tests/ctest/fortran_transpose_case.inc +++ b/tests/ctest/fortran_transpose_case.inc @@ -115,12 +115,12 @@ subroutine RUN_TRANSPOSE_SCENARIO(scenario, halo_extents_x, halo_extents_y, halo endif work_d = 0 if (use_optional_args) then - res = cudecompTransposeXToY(handle, grid_desc, input_d, output_d, work_d, DTYPE, pinfo_x%halo_extents, & - pinfo_y%halo_extents, pinfo_x%padding, pinfo_y%padding) + res = cudecompTransposeXToY(handle, grid_desc, input_d, output_d, CUDECOMP_WORKSPACE_AUTO, DTYPE, & + pinfo_x%halo_extents, pinfo_y%halo_extents, pinfo_x%padding, pinfo_y%padding) else - res = cudecompTransposeXToY(handle, grid_desc, input_d, output_d, work_d, DTYPE) + res = cudecompTransposeXToY(handle, grid_desc, input_d, output_d, CUDECOMP_WORKSPACE_AUTO, DTYPE) endif - call expect_success(res, DTYPE_NAME//" "//trim(scenario)//" cudecompTransposeXToY") + call expect_success(res, DTYPE_NAME//" "//trim(scenario)//" cudecompTransposeXToY automatic workspace") if (res == CUDECOMP_RESULT_SUCCESS) then allocate(actual(pinfo_y%size)) actual = output_d(1:pinfo_y%size) diff --git a/tests/ctest/halo_tests.cc b/tests/ctest/halo_tests.cc index e645709..ddace8b 100644 --- a/tests/ctest/halo_tests.cc +++ b/tests/ctest/halo_tests.cc @@ -16,6 +16,7 @@ #include #include "cudecomp.h" +#include "internal/common.h" #include "backend_test_context.h" #include "backend_utils.h" @@ -48,6 +49,7 @@ struct HaloCase { std::array halo_periods; std::array padding; cudecompRankOrder_t rank_order; + bool automatic_workspace; }; const char* axisName(int axis) { @@ -85,7 +87,8 @@ std::string paramName(const testing::TestParamInfo& info) { const auto& test_case = info.param; return sanitizeParamName(test_case.scenario) + "_Axis" + axisName(test_case.axis) + "_" + sanitizeParamName(test_case.backend.name) + "_" + dtypeName(test_case.dtype) + "_P" + - std::to_string(test_case.pdims[0]) + "x" + std::to_string(test_case.pdims[1]); + std::to_string(test_case.pdims[0]) + "x" + std::to_string(test_case.pdims[1]) + + (test_case.automatic_workspace ? "_AutomaticWorkspace" : ""); } HaloCase makeCase(cudecomp_test::HaloBackend backend, const char* scenario, int axis, @@ -96,8 +99,13 @@ HaloCase makeCase(cudecomp_test::HaloBackend backend, const char* scenario, int std::array halo_extents = kBaselineHaloExtents, std::array halo_periods = kPeriodicHalos, std::array padding = kZeroExtents, cudecompRankOrder_t rank_order = CUDECOMP_RANK_ORDER_DEFAULT) { - return {backend, scenario, axis, gdims, pdims, dtype, - axis_contiguous, mem_order, halo_extents, halo_periods, padding, rank_order}; + return {backend, scenario, axis, gdims, pdims, dtype, axis_contiguous, + mem_order, halo_extents, halo_periods, padding, rank_order, false}; +} + +HaloCase withAutomaticWorkspace(HaloCase test_case) { + test_case.automatic_workspace = true; + return test_case; } void appendBaselineCases(std::vector& cases, const cudecomp_test::HaloBackend& backend) { @@ -122,6 +130,8 @@ void appendBaselineCases(std::vector& cases, const cudecomp_test::Halo } } } + + cases.push_back(withAutomaticWorkspace(makeCase(backend, "AutomaticWorkspaceReuse", 0))); } void appendCoverageCases(std::vector& cases, const cudecomp_test::HaloBackend& backend) { @@ -356,26 +366,42 @@ template void runHaloCase(const HaloCase& test_case) { cudecomp_test::cudaBufferGuard data_buffer(data_d); CHECK_CUDA_GLOBAL(active_comm, data_alloc_result); - void* work_d = nullptr; - const cudecompResult_t work_alloc_result = - cudecompMalloc(handle, grid_desc, &work_d, workspace_num_elements * dtype_size); + void* work_d = CUDECOMP_WORKSPACE_AUTO; + cudecompResult_t work_alloc_result = CUDECOMP_RESULT_SUCCESS; + if (!test_case.automatic_workspace) { + work_alloc_result = cudecompMalloc(handle, grid_desc, &work_d, workspace_num_elements * dtype_size); + } cudecomp_test::cudecompBufferGuard work_buffer(handle, grid_desc, work_d); CHECK_CUDECOMP_GLOBAL(active_comm, work_alloc_result); - CHECK_CUDA_GLOBAL(active_comm, cudaMemset(data_d, 0, pinfo.size * sizeof(*data_d))); - CHECK_CUDA_GLOBAL(active_comm, - cudaMemcpy(data_d, initial.data(), initial.size() * sizeof(*data_d), cudaMemcpyHostToDevice)); - CHECK_CUDA_GLOBAL(active_comm, cudaMemset(work_d, 0, workspace_num_elements * dtype_size)); + auto run_and_verify = [&](void* work) { + CHECK_CUDA_GLOBAL(active_comm, cudaMemset(data_d, 0, pinfo.size * sizeof(*data_d))); + CHECK_CUDA_GLOBAL(active_comm, + cudaMemcpy(data_d, initial.data(), initial.size() * sizeof(*data_d), cudaMemcpyHostToDevice)); + if (work) { CHECK_CUDA_GLOBAL(active_comm, cudaMemset(work, 0, workspace_num_elements * dtype_size)); } - for (int dim = 0; dim < 3; ++dim) { - CHECK_CUDECOMP_GLOBAL(active_comm, runHalo(handle, grid_desc, test_case.axis, data_d, work_d, test_case.dtype, - pinfo, test_case.halo_periods, dim)); - } + for (int dim = 0; dim < 3; ++dim) { + CHECK_CUDECOMP_GLOBAL(active_comm, runHalo(handle, grid_desc, test_case.axis, data_d, work, test_case.dtype, + pinfo, test_case.halo_periods, dim)); + } - std::vector actual(expected.size(), unsetValue()); - CHECK_CUDA_GLOBAL(active_comm, - cudaMemcpy(actual.data(), data_d, actual.size() * sizeof(*data_d), cudaMemcpyDeviceToHost)); - EXPECT_TRUE(pencilMatches(expected, actual, pinfo)); + std::vector actual(expected.size(), unsetValue()); + CHECK_CUDA_GLOBAL(active_comm, + cudaMemcpy(actual.data(), data_d, actual.size() * sizeof(*data_d), cudaMemcpyDeviceToHost)); + EXPECT_TRUE(pencilMatches(expected, actual, pinfo)); + }; + + run_and_verify(work_d); + if (test_case.automatic_workspace) { + void* cached_workspace = cudecomp::haloBackendRequiresNvshmem(test_case.backend.backend) + ? handle->nvshmem_workspace.ptr + : handle->ordinary_workspace.ptr; + ASSERT_NE(cached_workspace, nullptr); + run_and_verify(nullptr); + EXPECT_EQ(cached_workspace, cudecomp::haloBackendRequiresNvshmem(test_case.backend.backend) + ? handle->nvshmem_workspace.ptr + : handle->ordinary_workspace.ptr); + } } TEST_P(HaloCorrectnessTest, UpdateHalos) { diff --git a/tests/ctest/transpose_tests.cc b/tests/ctest/transpose_tests.cc index 385ac21..4433992 100644 --- a/tests/ctest/transpose_tests.cc +++ b/tests/ctest/transpose_tests.cc @@ -58,6 +58,7 @@ struct TransposeCase { std::array output_padding; cudecompRankOrder_t rank_order; std::vector synthetic_host_groups; + bool automatic_workspace; }; const char* operationName(TransposeOperation operation) { @@ -117,7 +118,8 @@ std::string paramName(const testing::TestParamInfo& info) { return sanitizeParamName(test_case.scenario) + "_" + operationName(test_case.operation) + "_" + sanitizeParamName(test_case.backend.name) + "_" + dtypeName(test_case.dtype) + "_P" + std::to_string(test_case.pdims[0]) + "x" + std::to_string(test_case.pdims[1]) + "_" + - (test_case.out_of_place ? "OutOfPlace" : "InPlace"); + (test_case.out_of_place ? "OutOfPlace" : "InPlace") + + (test_case.automatic_workspace ? "_AutomaticWorkspace" : ""); } TransposeCase makeCase(cudecomp_test::TransposeBackend backend, const char* scenario, TransposeOperation operation, @@ -144,7 +146,13 @@ TransposeCase makeCase(cudecomp_test::TransposeBackend backend, const char* scen input_padding, output_padding, rank_order, - {}}; + {}, + false}; +} + +TransposeCase withAutomaticWorkspace(TransposeCase test_case) { + test_case.automatic_workspace = true; + return test_case; } TransposeCase withSyntheticHostGroups(TransposeCase test_case, std::vector synthetic_host_groups) { @@ -185,6 +193,9 @@ void appendBaselineCases(std::vector& cases, const cudecomp_test: } } } + + cases.push_back(withAutomaticWorkspace(makeCase(backend, "AutomaticWorkspaceReuse", TransposeOperation::XToY, + kBaselineGdims, {2, 2}, CUDECOMP_FLOAT, true))); } void appendNcclNativeAlltoAllCases(std::vector& cases, const cudecomp_test::TransposeBackend& backend) { @@ -380,20 +391,21 @@ testing::AssertionResult pencilMatches(const std::vector& expected, const std template cudecompResult_t runTranspose(cudecompHandle_t handle, cudecompGridDesc_t grid_desc, TransposeOperation operation, T* input, T* output, void* work, cudecompDataType_t dtype, - const cudecompPencilInfo_t& input_info, const cudecompPencilInfo_t& output_info) { + const cudecompPencilInfo_t& input_info, const cudecompPencilInfo_t& output_info, + cudaStream_t stream = nullptr) { switch (operation) { case TransposeOperation::XToY: return cudecompTransposeXToY(handle, grid_desc, input, output, work, dtype, input_info.halo_extents, - output_info.halo_extents, input_info.padding, output_info.padding, 0); + output_info.halo_extents, input_info.padding, output_info.padding, stream); case TransposeOperation::YToX: return cudecompTransposeYToX(handle, grid_desc, input, output, work, dtype, input_info.halo_extents, - output_info.halo_extents, input_info.padding, output_info.padding, 0); + output_info.halo_extents, input_info.padding, output_info.padding, stream); case TransposeOperation::YToZ: return cudecompTransposeYToZ(handle, grid_desc, input, output, work, dtype, input_info.halo_extents, - output_info.halo_extents, input_info.padding, output_info.padding, 0); + output_info.halo_extents, input_info.padding, output_info.padding, stream); case TransposeOperation::ZToY: return cudecompTransposeZToY(handle, grid_desc, input, output, work, dtype, input_info.halo_extents, - output_info.halo_extents, input_info.padding, output_info.padding, 0); + output_info.halo_extents, input_info.padding, output_info.padding, stream); } return CUDECOMP_RESULT_INVALID_USAGE; } @@ -410,7 +422,7 @@ void runAndVerifyTranspose(const cudecomp_test::MpiTestComm& active_comm, cudeco } CHECK_CUDA_GLOBAL(active_comm, cudaMemcpy(input_d, input_ref.data(), input_ref.size() * sizeof(*input_d), cudaMemcpyHostToDevice)); - CHECK_CUDA_GLOBAL(active_comm, cudaMemset(work_d, 0, workspace_num_elements * dtype_size)); + if (work_d) { CHECK_CUDA_GLOBAL(active_comm, cudaMemset(work_d, 0, workspace_num_elements * dtype_size)); } CHECK_CUDECOMP_GLOBAL(active_comm, runTranspose(handle, grid_desc, test_case.operation, input_d, output_d, work_d, test_case.dtype, input_info, output_info)); @@ -654,9 +666,11 @@ void runTransposeCase(const TransposeCase& test_case, bool check_cuda_graph_repl output_d = allocated_output_d; } - void* work_d = nullptr; - const cudecompResult_t work_alloc_result = - cudecompMalloc(handle, grid_desc, &work_d, workspace_num_elements * dtype_size); + void* work_d = CUDECOMP_WORKSPACE_AUTO; + cudecompResult_t work_alloc_result = CUDECOMP_RESULT_SUCCESS; + if (!test_case.automatic_workspace) { + work_alloc_result = cudecompMalloc(handle, grid_desc, &work_d, workspace_num_elements * dtype_size); + } cudecomp_test::cudecompBufferGuard work_buffer(handle, grid_desc, work_d); CHECK_CUDECOMP_GLOBAL(active_comm, work_alloc_result); if (check_nccl_user_buffer_registration) { ASSERT_TRUE(ncclUserBufferRegistrationIsActive(handle, work_d)); } @@ -664,6 +678,31 @@ void runTransposeCase(const TransposeCase& test_case, bool check_cuda_graph_repl runAndVerifyTranspose(active_comm, handle, grid_desc, test_case, input_d, output_d, work_d, data_num_elements, workspace_num_elements, dtype_size, input_ref, output_ref, input_info, output_info); + if (test_case.automatic_workspace) { + void* cached_workspace = cudecomp::transposeBackendRequiresNvshmem(test_case.backend.backend) + ? handle->nvshmem_workspace.ptr + : handle->ordinary_workspace.ptr; + ASSERT_NE(cached_workspace, nullptr); + + runAndVerifyTranspose(active_comm, handle, grid_desc, test_case, input_d, output_d, CUDECOMP_WORKSPACE_AUTO, + data_num_elements, workspace_num_elements, dtype_size, input_ref, output_ref, input_info, + output_info); + EXPECT_EQ(cached_workspace, cudecomp::transposeBackendRequiresNvshmem(test_case.backend.backend) + ? handle->nvshmem_workspace.ptr + : handle->ordinary_workspace.ptr); + + cudaStream_t capture_stream = nullptr; + CHECK_CUDA_GLOBAL(active_comm, cudaStreamCreateWithFlags(&capture_stream, cudaStreamNonBlocking)); + CHECK_CUDA_GLOBAL(active_comm, cudaStreamBeginCapture(capture_stream, cudaStreamCaptureModeThreadLocal)); + EXPECT_EQ(CUDECOMP_RESULT_NOT_SUPPORTED, + runTranspose(handle, grid_desc, test_case.operation, input_d, output_d, CUDECOMP_WORKSPACE_AUTO, + test_case.dtype, input_info, output_info, capture_stream)); + cudaGraph_t graph = nullptr; + CHECK_CUDA_GLOBAL(active_comm, cudaStreamEndCapture(capture_stream, &graph)); + if (graph) { CHECK_CUDA_GLOBAL(active_comm, cudaGraphDestroy(graph)); } + CHECK_CUDA_GLOBAL(active_comm, cudaStreamDestroy(capture_stream)); + } + if (check_nccl_user_buffer_registration) { const cudecompResult_t work_free_result = cudecompFree(handle, grid_desc, work_d); if (work_free_result == CUDECOMP_RESULT_SUCCESS) { work_buffer.release(); } From e441d81ec157b05c77ca3baabe9b4637bbb74540 Mon Sep 17 00:00:00 2001 From: romerojosh Date: Tue, 25 Aug 2026 21:43:21 +0000 Subject: [PATCH 2/2] Fix cuTENSOR dependencies for CTest targets Signed-off-by: romerojosh --- tests/ctest/CMakeLists.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/ctest/CMakeLists.txt b/tests/ctest/CMakeLists.txt index 9603e58..150fa84 100644 --- a/tests/ctest/CMakeLists.txt +++ b/tests/ctest/CMakeLists.txt @@ -125,6 +125,7 @@ target_link_libraries(cudecomp_test_automatic_workspace PRIVATE cudecomp cudecomp_test_support + NVHPC::CUTENSOR ) set_target_properties(cudecomp_test_automatic_workspace PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/tests/ctest @@ -208,6 +209,7 @@ target_link_libraries(cudecomp_test_halo PRIVATE cudecomp cudecomp_test_support + NVHPC::CUTENSOR ) set_target_properties(cudecomp_test_halo PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/tests/ctest