UCT/ZE/COPY: Initialize memtype cache when opening MD

src/uct/ze/copy/ze_copy_md.c

@@ -341,6 +341,10 @@ uct_ze_copy_md_open(uct_component_h component, const char *md_name,
     md->super.component = &uct_ze_copy_component;
 
     *md_p = (uct_md_h)md;
+
+    /* Create the memtype cache, so that its UCM handler sees allocations made
+     * through this MD. */
+    ucs_memtype_cache_global_create();
     return UCS_OK;
 }
 

test/gtest/uct/ze/test_ze_copy_md.cc

@@ -9,6 +9,8 @@ extern "C" {
 #include <uct/api/uct.h>
 #include <uct/ze/base/ze_base.h>
 #include <uct/ze/copy/ze_copy_md.h>
+#include <ucs/config/global_opts.h>
+#include <ucs/memory/memtype_cache.h>
 }
 
 
@@ -57,6 +59,22 @@ class test_ze_copy_md : public ucs::test {
 
         return (status == UCS_OK) ? md : NULL;
     }
+
+    /* Allocate ZE_DEVICE memory through the MD; skip the test if the MD
+     * cannot allocate it on this system. */
+    static void alloc_device_or_skip(uct_md_h md, size_t *length_p,
+                                     void **addr_p, uct_mem_h *memh_p)
+    {
+        ucs_status_t status = md->ops->mem_alloc(md, length_p, addr_p,
+                                                 UCS_MEMORY_TYPE_ZE_DEVICE,
+                                                 UCS_SYS_DEVICE_ID_UNKNOWN, 0,
+                                                 "test_ze_copy", memh_p);
+        if (status == UCS_ERR_UNSUPPORTED) {
+            uct_md_close(md);
+            UCS_TEST_SKIP_R("ZE_DEVICE alloc unsupported on this system");
+        }
+        ASSERT_UCS_OK(status);
+    }
 };
 
 
@@ -122,14 +140,7 @@ UCS_TEST_F(test_ze_copy_md, mem_alloc_free_device) {
     void *addr     = NULL;
     uct_mem_h memh = NULL;
 
-    ucs_status_t status = md->ops->mem_alloc(md, &length, &addr,
-                                             UCS_MEMORY_TYPE_ZE_DEVICE,
-                                             UCS_SYS_DEVICE_ID_UNKNOWN, 0,
-                                             "test_ze_copy", &memh);
-    if (status == UCS_ERR_UNSUPPORTED) {
-        UCS_TEST_SKIP_R("ZE_DEVICE alloc unsupported on this system");
-    }
-    ASSERT_UCS_OK(status);
+    alloc_device_or_skip(md, &length, &addr, &memh);
     ASSERT_TRUE(addr != NULL);
     ASSERT_GE(length, 4096u);
 
@@ -147,17 +158,10 @@ UCS_TEST_F(test_ze_copy_md, detect_memory_type_device) {
         UCS_TEST_SKIP_R("Could not open ze_cpy MD on this system");
     }
 
-    size_t length       = 4096;
-    void *addr          = NULL;
-    uct_mem_h memh      = NULL;
-    ucs_status_t status = md->ops->mem_alloc(md, &length, &addr,
-                                             UCS_MEMORY_TYPE_ZE_DEVICE,
-                                             UCS_SYS_DEVICE_ID_UNKNOWN, 0,
-                                             "test_ze_copy", &memh);
-    if (status == UCS_ERR_UNSUPPORTED) {
-        UCS_TEST_SKIP_R("ZE_DEVICE alloc unsupported on this system");
-    }
-    ASSERT_UCS_OK(status);
+    size_t length  = 4096;
+    void *addr     = NULL;
+    uct_mem_h memh = NULL;
+    alloc_device_or_skip(md, &length, &addr, &memh);
 
     ucs_memory_type_t mem_type = UCS_MEMORY_TYPE_HOST;
     EXPECT_UCS_OK(md->ops->detect_memory_type(md, addr, length, &mem_type));
@@ -168,6 +172,54 @@ UCS_TEST_F(test_ze_copy_md, detect_memory_type_device) {
 }
 
 
+/*
+ * The first allocation through this MD must be tracked in the memtype cache.
+ */
+UCS_TEST_F(test_ze_copy_md, memtype_cache_created_at_md_open) {
+    if (ucs_global_opts.enable_memtype_cache == UCS_NO) {
+        UCS_TEST_SKIP_R("memtype cache is disabled");
+    }
+
+    /* Creation fails permanently once it has failed in this process, so make
+     * sure the cache can be created at all; otherwise a NULL instance after
+     * uct_md_open() would prove nothing. */
+    ucs_memtype_cache_global_create();
+    if (ucs_memtype_cache_global_instance == NULL) {
+        UCS_TEST_SKIP_R("memtype cache cannot be created");
+    }
+
+    /* Drop the cache, so that the MD open has to create it again. */
+    ucs_memtype_cache_cleanup();
+    ucs_memtype_cache_global_init();
+    ASSERT_TRUE(ucs_memtype_cache_global_instance == NULL);
+
+    uct_md_h md = open_first_md();
+    if (md == NULL) {
+        UCS_TEST_SKIP_R("Could not open ze_cpy MD on this system");
+    }
+
+    EXPECT_TRUE(ucs_memtype_cache_global_instance != NULL);
+
+    size_t length  = 4096;
+    void *addr     = NULL;
+    uct_mem_h memh = NULL;
+    alloc_device_or_skip(md, &length, &addr, &memh);
+
+    ucs_memory_info_t mem_info;
+    ucs_status_t status = ucs_memtype_cache_lookup(addr, length, &mem_info);
+    EXPECT_UCS_OK(status);
+    if (status == UCS_OK) {
+        EXPECT_EQ(UCS_MEMORY_TYPE_ZE_DEVICE, mem_info.type);
+    }
+
+    EXPECT_UCS_OK(md->ops->mem_free(md, memh));
+    EXPECT_EQ(UCS_ERR_NO_ELEM,
+              ucs_memtype_cache_lookup(addr, length, &mem_info));
+
+    uct_md_close(md);
+}
+
+
 /*
  * mem_query reports base address and allocation length back.
  */
@@ -177,17 +229,10 @@ UCS_TEST_F(test_ze_copy_md, mem_query_base_and_length) {
         UCS_TEST_SKIP_R("Could not open ze_cpy MD on this system");
     }
 
-    size_t length       = 8192;
-    void *addr          = NULL;
-    uct_mem_h memh      = NULL;
-    ucs_status_t status = md->ops->mem_alloc(md, &length, &addr,
-                                             UCS_MEMORY_TYPE_ZE_DEVICE,
-                                             UCS_SYS_DEVICE_ID_UNKNOWN, 0,
-                                             "test_ze_copy", &memh);
-    if (status == UCS_ERR_UNSUPPORTED) {
-        UCS_TEST_SKIP_R("ZE_DEVICE alloc unsupported on this system");
-    }
-    ASSERT_UCS_OK(status);
+    size_t length  = 8192;
+    void *addr     = NULL;
+    uct_mem_h memh = NULL;
+    alloc_device_or_skip(md, &length, &addr, &memh);
 
     uct_md_mem_attr_v2_t mem_attr = {};
     mem_attr.field_mask           = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE |