UCM/MEM_ATTR: Make memory attribute opaque and add unit tests

Author: SeyedMir

src/ucm/Makefile.am

@@ -15,13 +15,15 @@ libucm_la_CPPFLAGS = $(BASE_CPPFLAGS) -DUCM_MALLOC_PREFIX=ucm_dl
 libucm_la_CFLAGS   = $(BASE_CFLAGS) $(CFLAGS_NO_DEPRECATED)
 
 nobase_dist_libucm_la_HEADERS = \
-	api/ucm.h
+	api/ucm.h \
+	mem_attr/mem_attr.h
 
 noinst_HEADERS = \
 	event/event.h \
 	malloc/malloc_hook.h \
 	malloc/allocator.h \
 	mmap/mmap.h \
+	mem_attr/mem_attr_int.h \
 	util/khash_safe.h \
 	util/replace.h \
 	util/log.h \
@@ -37,6 +39,7 @@ libucm_la_SOURCES = \
 	event/event.c \
 	malloc/malloc_hook.c \
 	mmap/install.c \
+	mem_attr/mem_attr.c \
 	util/replace.c \
 	util/log.c \
 	util/reloc.c \

src/ucm/cuda/cudamem.c

@@ -16,11 +16,11 @@
 #include <ucm/util/reloc.h>
 #include <ucm/util/replace.h>
 #include <ucm/util/sys.h>
+#include <ucm/mem_attr/mem_attr_int.h>
+#include <ucs/debug/memtrack.h>
 #include <ucs/debug/assert.h>
 #include <ucs/sys/compiler.h>
 #include <ucs/sys/preprocessor.h>
-#include <ucs/memory/memory_type.h>
-#include <ucs/type/status.h>
 
 #include <sys/mman.h>
 #include <pthread.h>
@@ -457,12 +457,35 @@ static void ucm_cudamem_get_existing_alloc(ucm_event_handler_t *handler)
     }
 }
 
+typedef struct ucm_cudamem_attr {
+    ucm_mem_attr_t super;
+    unsigned long long buf_id;
+} ucm_cudamem_attr_t;
+
+static int ucm_cudamem_attr_cmp(const ucm_mem_attr_h mem_attr1,
+                                const ucm_mem_attr_h mem_attr2)
+{
+    ucm_cudamem_attr_t *cuda_mem_attr1, *cuda_mem_attr2;
+    cuda_mem_attr1 = ucs_derived_of(mem_attr1, ucm_cudamem_attr_t);
+    cuda_mem_attr2 = ucs_derived_of(mem_attr2, ucm_cudamem_attr_t);
+    if (cuda_mem_attr1->buf_id == cuda_mem_attr2->buf_id) return 0;
+    return 1;
+}
+
+static void ucm_cudamem_attr_destroy(ucm_mem_attr_h mem_attr)
+{
+    ucm_cudamem_attr_t *cuda_mem_attr;
+    cuda_mem_attr = ucs_derived_of(mem_attr,ucm_cudamem_attr_t);
+    ucs_free(cuda_mem_attr);
+}
+
 static ucs_status_t ucm_cudamem_attr_get(const void *address, size_t length,
-                                         ucs_memory_attr_t *mem_attr)
+                                         ucm_mem_attr_h *mem_attr_p)
 {
 #define UCM_CUDA_MEM_QUERY_NUM_ATTRS 3
     CUmemorytype cuda_mem_type = (CUmemorytype)0;
     uint32_t is_managed        = 0;
+    unsigned long long buf_id  = 0;
     CUpointer_attribute attr_type[UCM_CUDA_MEM_QUERY_NUM_ATTRS];
     void *attr_data[UCM_CUDA_MEM_QUERY_NUM_ATTRS];
     CUresult cu_err;
@@ -474,23 +497,31 @@ static ucs_status_t ucm_cudamem_attr_get(const void *address, size_t length,
     attr_type[1] = CU_POINTER_ATTRIBUTE_IS_MANAGED;
     attr_data[1] = &is_managed;
     attr_type[2] = CU_POINTER_ATTRIBUTE_BUFFER_ID;
-    attr_data[2] = &mem_attr->cuda.buf_id;
+    attr_data[2] = &buf_id;
 
     cu_err = cuPointerGetAttributes(ucs_static_array_size(attr_data),
                                     attr_type, attr_data,
                                     (CUdeviceptr)address);
     if (cu_err == CUDA_SUCCESS) {
+        ucm_cudamem_attr_t *mem_attr;
+        mem_attr = ucs_malloc(sizeof(*mem_attr),"cudamem_attr");
+        if (mem_attr == NULL) return UCS_ERR_NO_MEMORY;
         switch (cuda_mem_type) {
             case CU_MEMORYTYPE_DEVICE:
-                mem_attr->mem_type = is_managed ? UCS_MEMORY_TYPE_CUDA_MANAGED
-                                                : UCS_MEMORY_TYPE_CUDA;
+                mem_attr->super.mem_type = is_managed
+                                           ? UCS_MEMORY_TYPE_CUDA_MANAGED
+                                           : UCS_MEMORY_TYPE_CUDA;
                 break;
             case CU_MEMORYTYPE_HOST:
-                mem_attr->mem_type = UCS_MEMORY_TYPE_HOST;
+                mem_attr->super.mem_type = UCS_MEMORY_TYPE_HOST;
                 break;
             default:
                 return UCS_ERR_INVALID_ADDR;
         }
+        mem_attr->buf_id        = buf_id;
+        mem_attr->super.cmp     = &ucm_cudamem_attr_cmp;
+        mem_attr->super.destroy = &ucm_cudamem_attr_destroy;
+        *mem_attr_p = &mem_attr->super;
         return UCS_OK;
     }
 

src/ucm/event/event.h

@@ -9,6 +9,7 @@
 
 #include <ucm/api/ucm.h>
 #include <ucm/util/log.h>
+#include <ucm/mem_attr/mem_attr.h>
 #include <ucs/datastruct/list.h>
 #include <ucs/type/status.h>
 
@@ -34,7 +35,7 @@ typedef struct ucm_event_installer {
     ucs_status_t          (*install)(int events);
     void                  (*get_existing_alloc)(ucm_event_handler_t *handler);
     ucs_status_t          (*get_mem_attr)(const void *address, size_t length,
-                                          ucs_memory_attr_t *mem_attr);
+                                          ucm_mem_attr_h *mem_attr_p);
     ucs_list_link_t       list;
 } ucm_event_installer_t;
 

src/ucm/mem_attr/mem_attr.c

@@ -0,0 +1,72 @@
+/**
+ * Copyright (C) NVIDIA Corporation. 2020.  ALL RIGHTS RESERVED.
+ * See file LICENSE for terms.
+ */
+
+#ifdef HAVE_CONFIG_H
+#  include "config.h"
+#endif
+
+#include "mem_attr.h"
+#include "mem_attr_int.h"
+
+#include <ucs/memory/memory_type.h>
+#include <ucm/event/event.h>
+
+
+static int ucm_mem_attr_cmp_type(ucm_mem_attr_h mem_attr1,
+                                 ucm_mem_attr_h mem_attr2)
+{
+    if (mem_attr1->mem_type == mem_attr2->mem_type) return 0;
+    return 1;
+}
+
+static void ucm_mem_attr_destroy_host(ucm_mem_attr_h mem_attr)
+{
+    /* Nothing to be done for host memory attributes */
+}
+
+/* All host memory will have the same attributes (only a type).
+ * So, they will all point to this static struct */
+static ucm_mem_attr_t mem_attr_host = {
+    .mem_type = UCS_MEMORY_TYPE_HOST,
+    .cmp      = &ucm_mem_attr_cmp_type,
+    .destroy  = &ucm_mem_attr_destroy_host
+};
+
+ucs_status_t ucm_mem_attr_get(const void *address, size_t length,
+                              ucm_mem_attr_h *mem_attr_p)
+{
+    ucm_event_installer_t *event_installer;
+    ucs_status_t status;
+    int failure = 0;
+
+    ucs_list_for_each(event_installer, &ucm_event_installer_list, list) {
+        if (event_installer->get_mem_attr == NULL) continue;
+        status = event_installer->get_mem_attr(address, length, mem_attr_p);
+        if (status == UCS_OK) return UCS_OK;
+        if (status != UCS_ERR_INVALID_ADDR) failure = 1;
+    }
+
+    if (failure) return UCS_ERR_NO_RESOURCE;
+
+    /* none of the installers recognized the address. So, it must be HOST */
+    *mem_attr_p = &mem_attr_host;
+    return UCS_OK;
+}
+
+ucs_memory_type_t ucm_mem_attr_get_type(ucm_mem_attr_h mem_attr)
+{
+    return mem_attr->mem_type;
+}
+
+int ucm_mem_attr_cmp(ucm_mem_attr_h mem_attr1, ucm_mem_attr_h mem_attr2)
+{
+    if (ucm_mem_attr_cmp_type(mem_attr1, mem_attr2)) return 1;
+    return mem_attr1->cmp(mem_attr1, mem_attr2);
+}
+
+void ucm_mem_attr_destroy(ucm_mem_attr_h mem_attr)
+{
+    mem_attr->destroy(mem_attr);
+}

src/ucm/mem_attr/mem_attr.h

@@ -0,0 +1,21 @@
+/**
+ * Copyright (C) NVIDIA Corporation. 2020.  ALL RIGHTS RESERVED.
+ * See file LICENSE for terms.
+ */
+
+#ifndef UCM_MEM_ATTR_H_
+#define UCM_MEM_ATTR_H_
+
+#include <ucs/memory/memory_type.h>
+#include <ucs/type/status.h>
+#include <stddef.h>
+
+typedef struct ucm_mem_attr *ucm_mem_attr_h;
+
+ucs_status_t ucm_mem_attr_get(const void *address, size_t length,
+                              ucm_mem_attr_h *mem_attr_p);
+ucs_memory_type_t ucm_mem_attr_get_type(ucm_mem_attr_h mem_attr);
+int ucm_mem_attr_cmp(ucm_mem_attr_h mem_attr1, ucm_mem_attr_h mem_attr2);
+void ucm_mem_attr_destroy(ucm_mem_attr_h mem_attr);
+
+#endif

src/ucm/mem_attr/mem_attr_int.h

@@ -0,0 +1,23 @@
+/**
+ * Copyright (C) NVIDIA Corporation. 2020.  ALL RIGHTS RESERVED.
+ * See file LICENSE for terms.
+ */
+
+#ifndef UCM_MEM_ATTR_INT_H_
+#define UCM_MEM_ATTR_INT_H_
+
+
+#include "mem_attr.h"
+
+#include <ucs/memory/memory_type.h>
+#include <ucs/type/status.h>
+
+
+typedef struct ucm_mem_attr {
+    ucs_memory_type_t mem_type;
+    int (*cmp)(ucm_mem_attr_h mem_attr1, ucm_mem_attr_h mem_attr2);
+    void (*destroy)(ucm_mem_attr_h mem_attr);
+} ucm_mem_attr_t;
+
+
+#endif

src/ucm/rocm/rocmmem.c

@@ -13,6 +13,8 @@
 #include <ucm/util/log.h>
 #include <ucm/util/reloc.h>
 #include <ucm/util/replace.h>
+#include <ucm/mem_attr/mem_attr_int.h>
+#include <ucs/debug/memtrack.h>
 #include <ucs/debug/assert.h>
 #include <ucs/sys/compiler.h>
 #include <ucs/sys/preprocessor.h>
@@ -185,38 +187,62 @@ static void ucm_rocmmem_get_existing_alloc(ucm_event_handler_t *handler)
 {
 }
 
-static ucs_status_t ucm_rocmmem_attr_get(const void *address, size_t length,
-                                         ucs_memory_attr_t *mem_attr)
+typedef struct ucm_rocmmem_attr {
+    ucm_mem_attr_t super;
+    /* TODO what rocm memory attributes do we need? */
+} ucm_rocmmem_attr_t;
+
+static int ucm_rocmmem_attr_cmp(const ucm_mem_attr_h mem_attr1,
+                                const ucm_mem_attr_h mem_attr2)
 {
-    if (address == NULL) return UCS_ERR_INVALID_ADDR;
+    /* TODO how should we compare two rocm memory attributes? */
+    return 0;
+}
+
+static void ucm_rocmmem_attr_destroy(ucm_mem_attr_h mem_attr)
+{
+    ucm_rocmmem_attr_t *rocm_mem_attr;
+    rocm_mem_attr = ucs_derived_of(mem_attr,ucm_rocmmem_attr_t);
+    ucs_free(rocm_mem_attr);
+}
 
+static ucs_status_t ucm_rocmmem_attr_get(const void *address, size_t length,
+                                         ucm_mem_attr_h *mem_attr_p)
+{
     hsa_status_t status;
     hsa_amd_pointer_info_t info = {
         .size = sizeof(hsa_amd_pointer_info_t),
     };
 
-    status = hsa_amd_pointer_info((void*)addr, &info, NULL, NULL, NULL);
-    if (status == HSA_STATUS_SUCCESS) {
-        if (info.type == HSA_EXT_POINTER_TYPE_UNKNOWN) return UCS_ERR_INVALID_ADDR;
+    if (address == NULL) return UCS_ERR_INVALID_ADDR;
 
+    status = hsa_amd_pointer_info((void*)address, &info, NULL, NULL, NULL);
+    if (status == HSA_STATUS_SUCCESS) {
         hsa_device_type_t dev_type;
+        if (info.type == HSA_EXT_POINTER_TYPE_UNKNOWN) return UCS_ERR_INVALID_ADDR;
         status = hsa_agent_get_info(info.agentOwner, HSA_AGENT_INFO_DEVICE, &dev_type);
         if (status == HSA_STATUS_SUCCESS) {
+            ucm_rocmmem_attr_t *mem_attr;
+            mem_attr = ucs_malloc(sizeof(*mem_attr),"rocmmem_attr");
+            if (mem_attr == NULL) return UCS_ERR_NO_MEMORY;
             switch (dev_type) {
                 case HSA_DEVICE_TYPE_GPU:
-                    mem_attr->mem_type = UCS_MEMORY_TYPE_ROCM;
+                    mem_attr->super.mem_type = UCS_MEMORY_TYPE_ROCM;
                     break;
                 case HSA_DEVICE_TYPE_CPU:
-                    mem_attr->mem_type = UCS_MEMORY_TYPE_HOST;
+                    mem_attr->super.mem_type = UCS_MEMORY_TYPE_HOST;
                     break;
                 default:
                     return UCS_ERR_INVALID_ADDR;
             }
+            mem_attr->super.cmp     = &ucm_rocmmem_attr_cmp;
+            mem_attr->super.destroy = &ucm_rocmmem_attr_destroy;
+            *mem_attr_p = &mem_attr->super;
             return UCS_OK;
         }
     }
 
-    ucs_error("failed to get AMD pointer info");
+    ucm_error("failed to get AMD pointer info");
     return UCS_ERR_NO_RESOURCE;
 }
 

src/ucm/util/sys.c

@@ -16,7 +16,6 @@
 
 #include <ucm/api/ucm.h>
 #include <ucm/util/log.h>
-#include <ucm/event/event.h>
 #include <ucm/mmap/mmap.h>
 #include <ucs/sys/math.h>
 #include <linux/mman.h>
@@ -350,26 +349,3 @@ char *ucm_concat_path(char *buffer, size_t max, const char *dir, const char *fil
 
     return buffer;
 }
-
-ucs_status_t ucm_mem_attr_get(const void *address, size_t length,
-                              ucs_memory_attr_t *mem_attr)
-{
-    ucm_event_installer_t *event_installer;
-    ucs_status_t status;
-    int failure = 0;
-
-    ucs_list_for_each(event_installer, &ucm_event_installer_list, list) {
-        if (NULL == event_installer->get_mem_attr) continue;
-        status = event_installer->get_mem_attr(address, length, mem_attr);
-        if (status == UCS_OK) return UCS_OK;
-        if (status != UCS_ERR_INVALID_ADDR) failure = 1;
-    }
-
-    if (failure) {
-        mem_attr->mem_type = UCS_MEMORY_TYPE_UNKNOWN;
-        return UCS_ERR_NO_RESOURCE;
-    }
-
-    mem_attr->mem_type = UCS_MEMORY_TYPE_HOST;
-    return UCS_OK;
-}

src/ucm/util/sys.h

@@ -9,8 +9,6 @@
 #define UCM_UTIL_SYS_H_
 
 #include <stddef.h>
-#include <ucs/memory/memory_type.h>
-#include <ucs/type/status.h>
 
 
 /*
@@ -90,17 +88,4 @@ void ucm_prevent_dl_unload();
 char *ucm_concat_path(char *buffer, size_t max, const char *dir, const char *file);
 
 
-/*
- * Get the memory attributes of a given address and length.
- *
- * @param [in]  address   Buffer address.
- * @param [in]  length    Buffer length.
- * @param [out] mem_attr  Memory attributes as defined by @ref ucs_memory_attr_t.
- *
- * @return Error code as defined by @ref ucs_status_t.
- */
-ucs_status_t ucm_mem_attr_get(const void *address, size_t length,
-                              ucs_memory_attr_t *mem_attr);
-
-
 #endif