opal/cuda: avoid direct access to cumem host numa memory

Signed-off-by: Akshay Venkatesh <[email protected]>

Author: Akshay-Venkatesh

config/opal_check_cuda.m4

@@ -1,5 +1,6 @@
 dnl -*- autoconf -*-
 dnl
+dnl Copyright (c) 2024      NVIDIA Corporation.  All rights reserved.
 dnl Copyright (c) 2004-2010 The Trustees of Indiana University and Indiana
 dnl                         University Research and Technology
 dnl                         Corporation.  All rights reserved.
@@ -118,6 +119,12 @@ AS_IF([test "$opal_check_cuda_happy" = "yes"],
         [#include <$opal_cuda_incdir/cuda.h>])],
     [])
 
+# If we have CUDA support, check to see if we have support for cuMemCreate memory on host NUMA.
+AS_IF([test "$opal_check_cuda_happy"="yes"],
+    [AC_CHECK_DECL([CU_MEM_LOCATION_TYPE_HOST_NUMA], [CUDA_HOST_NUMA_SUPPORT=1], [CUDA_HOST_NUMA_SUPPORT=0],
+        [#include <$opal_cuda_incdir/cuda.h>])],
+    [])
+
 # If we have CUDA support, check to see if we have support for SYNC_MEMOPS
 # which was first introduced in CUDA 6.0.
 AS_IF([test "$opal_check_cuda_happy" = "yes"],
@@ -160,6 +167,10 @@ AM_CONDITIONAL([OPAL_cuda_support], [test "x$CUDA_SUPPORT" = "x1"])
 AC_DEFINE_UNQUOTED([OPAL_CUDA_SUPPORT],$CUDA_SUPPORT,
                    [Whether we want cuda device pointer support])
 
+AM_CONDITIONAL([OPAL_cuda_host_numa_support], [test "x$CUDA_HOST_NUMA_SUPPORT" = "x1"])
+AC_DEFINE_UNQUOTED([OPAL_CUDA_HOST_NUMA_SUPPORT],$CUDA_HOST_NUMA_SUPPORT,
+                   [Whether we have CU_MEM_LOCATION_TYPE_HOST_NUMA support available])
+
 AM_CONDITIONAL([OPAL_cuda_sync_memops], [test "x$CUDA_SYNC_MEMOPS" = "x1"])
 AC_DEFINE_UNQUOTED([OPAL_CUDA_SYNC_MEMOPS],$CUDA_SYNC_MEMOPS,
                    [Whether we have CUDA CU_POINTER_ATTRIBUTE_SYNC_MEMOPS support available])

opal/mca/accelerator/cuda/accelerator_cuda.c

@@ -1,12 +1,12 @@
 /*
+ * Copyright (c) 2024      NVIDIA Corporation.  All rights reserved.
  * Copyright (c) 2014-2015 Intel, Inc.  All rights reserved.
  * Copyright (c) 2014      Research Organization for Information Science
  *                         and Technology (RIST). All rights reserved.
  * Copyright (c) 2014      Mellanox Technologies, Inc.
  *                         All rights reserved.
  * Copyright (c)           Amazon.com, Inc. or its affiliates.
  *                         All Rights reserved.
- * Copyright (c) 2024      NVIDIA Corporation.  All rights reserved.
  * Copyright (c) 2024      The University of Tennessee and The University
  *                         of Tennessee Research Foundation.  All rights
  *                         reserved.
@@ -154,6 +154,46 @@ static int accelerator_cuda_get_device_id(CUcontext mem_ctx) {
     return dev_id;
 }
 
+static CUmemorytype accelerator_cuda_check_host_numa(CUdeviceptr dbuf)
+{
+#if OPAL_CUDA_HOST_NUMA_SUPPORT
+    CUmemAllocationProp prop;
+    CUmemLocation location;
+    CUresult result;
+    unsigned long long flags;
+    CUmemGenericAllocationHandle alloc_handle;
+    /* Check if memory is allocated using VMM API and see if host memory needs
+     * to be treated as pinned device memory */
+    result = cuMemRetainAllocationHandle(&alloc_handle, (void*)dbuf);
+    if (result != CUDA_SUCCESS) {
+        return CU_MEMORYTYPE_HOST;
+    }
+
+    result = cuMemGetAllocationPropertiesFromHandle(&prop, alloc_handle);
+    if (result != CUDA_SUCCESS) {
+        return CU_MEMORYTYPE_HOST;
+    }
+
+    if ((CU_MEM_LOCATION_TYPE_HOST == prop.location.type) ||
+        (CU_MEM_LOCATION_TYPE_HOST_NUMA == prop.location.type) ||
+        (CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT == prop.location.type)) {
+        /* If host has read-write access, then range is accessible by CPU */
+        result = cuMemGetAccess(&flags, &location, dbuf);
+        if ((CUDA_SUCCESS == result) &&
+            ((CU_MEM_LOCATION_TYPE_HOST == location.type) ||
+             (CU_MEM_LOCATION_TYPE_HOST_NUMA == location.type) ||
+             (CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT == location.type)) &&
+            (CU_MEM_ACCESS_FLAGS_PROT_READWRITE == flags)) {
+            return CU_MEMORYTYPE_HOST;
+        } else {
+            return CU_MEMORYTYPE_DEVICE;
+        }
+    }
+#else
+    return CU_MEMORYTYPE_HOST;
+#endif
+}
+
 static int accelerator_cuda_check_addr(const void *addr, int *dev_id, uint64_t *flags)
 {
     CUresult result;
@@ -197,8 +237,11 @@ static int accelerator_cuda_check_addr(const void *addr, int *dev_id, uint64_t *
             return OPAL_ERROR;
         }
     } else if (CU_MEMORYTYPE_HOST == mem_type) {
-        /* Host memory, nothing to do here */
-        return 0;
+        mem_type = accelerator_cuda_check_host_numa(dbuf);
+        if (CU_MEMORYTYPE_HOST == mem_type) {
+            /* Host memory, nothing to do here */
+            return 0;
+        }
     } else if (0 == mem_type) {
         /* This can happen when CUDA is initialized but dbuf is not valid CUDA pointer */
         return 0;
@@ -218,8 +261,11 @@ static int accelerator_cuda_check_addr(const void *addr, int *dev_id, uint64_t *
             return OPAL_ERROR;
         }
     } else if (CU_MEMORYTYPE_HOST == mem_type) {
-        /* Host memory, nothing to do here */
-        return 0;
+        mem_type = accelerator_cuda_check_host_numa(dbuf);
+        if (CU_MEMORYTYPE_HOST == mem_type) {
+            /* Host memory, nothing to do here */
+            return 0;
+        }
     } else {
         result = cuPointerGetAttribute(&mem_ctx, CU_POINTER_ATTRIBUTE_CONTEXT, dbuf);
         /* query the device from the context */