cudaMallocAsync not supported by UCX, may cause failure in OpenMPI+Kokkos+cuda applications

[lfmeadow]

See UCX issue openucx/ucx#7194
cudaMallocAsync was added to Kokkos/Cuda sometime in the last few months and results in a application crash in lammps+kokkos+cuda using OpenMPI+UCX.
The UCX issue comments mention 3 PRs to UCX that should resolve the problem.
In the meantime the following patch to Kokkos will fix it:

diff --git a/core/src/Cuda/Kokkos_CudaSpace.cpp b/core/src/Cuda/Kokkos_CudaSpace.cpp
index 39da8de6b..13c35eb87 100644
--- a/core/src/Cuda/Kokkos_CudaSpace.cpp
+++ b/core/src/Cuda/Kokkos_CudaSpace.cpp
@@ -229,7 +229,7 @@ void *CudaSpace::impl_allocate(
 
 #ifndef CUDART_VERSION
 #error CUDART_VERSION undefined!
-#elif (CUDART_VERSION >= 11020)
+#elif 0 // (CUDART_VERSION >= 11020) UCX issue
   cudaError_t error_code;
   if (arg_alloc_size >= memory_threshold_g) {
     error_code = cudaMallocAsync(&ptr, arg_alloc_size, 0);
@@ -358,7 +358,7 @@ void CudaSpace::impl_deallocate(
   try {
 #ifndef CUDART_VERSION
 #error CUDART_VERSION undefined!
-#elif (CUDART_VERSION >= 11020)
+#elif 0 // (CUDART_VERSION >= 11020) UCX issue
     if (arg_alloc_size >= memory_threshold_g) {
       CUDA_SAFE_CALL(cudaDeviceSynchronize());
       CUDA_SAFE_CALL(cudaFreeAsync(arg_alloc_ptr, 0));

[maxpkatz]

#4026 is the PR that does this (@matt-stack). We should make the use of cudaMallocAsync opt-in with a CMake flag at least until UCX fully supports that allocator.

[BenWibking]

This affects our code as well. It looks like UCX 1.14rc1 may fix this issue: https://github.com/openucx/ucx/blob/d83ef403646473d76be42cbef03fb38652169f78/NEWS#L49.

Do you have a reproducer for the issue that can be tested with the latest UCX?

[lfmeadow]

Geez I found that bug ages ago I forgot all about it. So cray seems to have the same problem. There’s also an easy kokkos patch. I can test it when Perlmutter is back. Thanks!

…

Sent from my iPhone On Jan 25, 2023, at 9:12 AM, Ben Wibking ***@***.***> wrote:  This affects my code as well. It looks like UCX 1.14rc1 may fix this issue: https://github.com/openucx/ucx/blob/d83ef403646473d76be42cbef03fb38652169f78/NEWS#L49. Do you have a reproducer for the issue that can be tested with the latest UCX? — Reply to this email directly, view it on GitHub<#4228 (comment)>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/ABH5YK7HRQ6ZXPBWLSENKSTWUFNFVANCNFSM5B66AZFQ>. You are receiving this because you authored the thread.Message ID: ***@***.***>

[masterleinad]

Fixed in #4233.