UCT/CUDA_COPY/MD: Don't widen VMM range crossing a handle boundary

AUTHORS

@@ -129,6 +129,7 @@ Swen Boehm <boehms@ornl.gov>
 Thomas Vegas <tvegas@nvidia.com>
 Tomer Gilad <tgilad@nvidia.com>
 Tony Curtis <tonycurtis32@gmail.com>
+Tristan Rice <rice@fn.lc>
 Tushar Gohad <tushar.gohad@intel.com>
 Tzafrir Cohen <tzafrirc@nvidia.com>
 Valentin Petrov <valentinp@mellanox.com>

src/uct/cuda/cuda_copy/cuda_copy_md.c

@@ -601,6 +601,11 @@ static void uct_cuda_copy_md_sync_memops_get_address_range(
     ucs_trace("query address 0x%llx: 0x%llx..0x%llx length %zu", address,
               base_address, base_address + alloc_length, alloc_length);
 
+    /* Host-located VMM reports a zero base address */
+    if (base_address == 0) {
+        goto out_ctx_pop;
+    }
+
     if (md->config.alloc_whole_reg == UCS_CONFIG_AUTO) {
         total_bytes = uct_cuda_copy_md_get_total_device_mem(cuda_device);
         if (alloc_length > (total_bytes * md->config.max_reg_ratio)) {
@@ -610,9 +615,9 @@ static void uct_cuda_copy_md_sync_memops_get_address_range(
         ucs_assert(md->config.alloc_whole_reg == UCS_CONFIG_ON);
     }
 
-    /* For multi-handle VMM, the physical allocation may be smaller than the
-     * mapped virtual range; preserve the caller's requested extent. */
-    if (is_vmm && (alloc_length < length)) {
+    /* For multi-handle VMM, the physical allocation containing the address
+     * may not cover the whole requested range */
+    if (is_vmm && ((base_address + alloc_length) < (address + length))) {
         goto out_ctx_pop;
     }
 

test/gtest/uct/cuda/test_cuda_ipc_md.cc

@@ -740,71 +740,3 @@ UCS_TEST_F(test_cuda_ipc_cache_lru, stale_destroy_while_in_use) {
     EXPECT_EQ(0UL, m_cache->total_size);
     EXPECT_TRUE(ucs_list_is_empty(&m_cache->lru_list));
 }
-
-
-class test_cuda_copy_md : public test_md {
-};
-
-UCS_TEST_P(test_cuda_copy_md, vmm_multi_handle_range) {
-    /* A single byte per chunk is rounded up to the allocation granularity, so
-     * the buffer is backed by 3 distinct physical handles */
-    cuda_vmm_mem_buffer buffer(1, UCS_MEMORY_TYPE_CUDA, 3);
-    uct_md_mem_attr_t mem_attr;
-
-    mem_attr.field_mask = UCT_MD_MEM_ATTR_FIELD_MEM_TYPE |
-                          UCT_MD_MEM_ATTR_FIELD_ALLOC_LENGTH;
-
-    /* Query the full multi-handle range; the guard preserves the caller's
-     * requested extent instead of shrinking to the single chunk that
-     * contains the base pointer. */
-    ASSERT_UCS_OK(uct_md_mem_query(md(), buffer.ptr(), buffer.size(),
-                                   &mem_attr));
-
-    EXPECT_EQ(UCS_MEMORY_TYPE_CUDA, mem_attr.mem_type);
-    EXPECT_GE(mem_attr.alloc_length, buffer.size());
-}
-
-UCS_TEST_P(test_cuda_copy_md, vmm_locality_domain_mem_type) {
-#if HAVE_DECL_CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN
-    ASSERT_EQ(cudaSuccess, cudaSetDevice(0));
-
-    auto check_domain = [this](unsigned char domain_id,
-                               ucs_sys_device_t &sys_dev) {
-        cuda_vmm_mem_buffer buffer;
-        CUresult ret = buffer.alloc(64, 0,
-                                    CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN,
-                                    1, domain_id);
-        if (ret != CUDA_SUCCESS) {
-            return false;
-        }
-
-        uct_md_mem_attr_t mem_attr;
-        mem_attr.field_mask = UCT_MD_MEM_ATTR_FIELD_MEM_TYPE |
-                              UCT_MD_MEM_ATTR_FIELD_SYS_DEV;
-
-        ASSERT_UCS_OK(uct_md_mem_query(md(), buffer.ptr(), buffer.size(),
-                                       &mem_attr));
-        EXPECT_EQ(UCS_MEMORY_TYPE_CUDA, mem_attr.mem_type);
-        EXPECT_NE(UCS_SYS_DEVICE_ID_UNKNOWN, mem_attr.sys_dev);
-        sys_dev = mem_attr.sys_dev;
-        return true;
-    };
-
-    ucs_sys_device_t sys_dev0, sys_dev1;
-    if (!check_domain(0, sys_dev0)) {
-        UCS_TEST_SKIP_R("failed to allocate locality-domain 0 VMM memory");
-    }
-
-    if (!check_domain(1, sys_dev1)) {
-        ADD_FAILURE() << "locality domain 1 allocation failed even though "
-                         "domain 0 succeeded";
-        return;
-    }
-
-    EXPECT_EQ(sys_dev0, sys_dev1);
-#else
-    UCS_TEST_SKIP_R("built without locality-domain support");
-#endif
-}
-
-_UCT_MD_INSTANTIATE_TEST_CASE(test_cuda_copy_md, cuda_cpy);

test/gtest/uct/cuda/test_switch_cuda_device.cc

@@ -646,14 +646,83 @@ UCS_TEST_P(test_mem_alloc_device, no_current_context_vmm_mem_registrable,
     query_registrable_no_current_context(buffer.ptr(), size);
 }
 
+UCS_TEST_P(test_mem_alloc_device, vmm_multi_handle_range,
+           "CUDA_COPY_REG_WHOLE_ALLOC=on")
+{
+    cuda_vmm_mem_buffer buffer(4 * UCS_MBYTE, UCS_MEMORY_TYPE_CUDA, 3);
+
+    /* The handle containing the address covers only the first chunk, so the
+     * requested extent must be kept */
+    check_mem_query_exact_range(buffer.ptr(), buffer.size());
+}
+
+UCS_TEST_P(test_mem_alloc_device, vmm_locality_domain_mem_type)
+{
+#if HAVE_DECL_CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN
+    ASSERT_EQ(cudaSuccess, cudaSetDevice(0));
+
+    auto check_domain = [this](unsigned char domain_id,
+                               ucs_sys_device_t &sys_dev) {
+        cuda_vmm_mem_buffer buffer;
+        CUresult ret = buffer.alloc(64, 0,
+                                    CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN,
+                                    1, domain_id);
+        if (ret != CUDA_SUCCESS) {
+            return false;
+        }
+
+        uct_md_mem_attr_t mem_attr;
+        mem_attr.field_mask = UCT_MD_MEM_ATTR_FIELD_MEM_TYPE |
+                              UCT_MD_MEM_ATTR_FIELD_SYS_DEV;
+
+        ASSERT_UCS_OK(uct_md_mem_query(md(), buffer.ptr(), buffer.size(),
+                                       &mem_attr));
+        EXPECT_EQ(UCS_MEMORY_TYPE_CUDA, mem_attr.mem_type);
+        EXPECT_NE(UCS_SYS_DEVICE_ID_UNKNOWN, mem_attr.sys_dev);
+        sys_dev = mem_attr.sys_dev;
+        return true;
+    };
+
+    ucs_sys_device_t sys_dev0, sys_dev1;
+    if (!check_domain(0, sys_dev0)) {
+        UCS_TEST_SKIP_R("failed to allocate locality-domain 0 VMM memory");
+    }
+
+    if (!check_domain(1, sys_dev1)) {
+        ADD_FAILURE() << "locality domain 1 allocation failed even though "
+                         "domain 0 succeeded";
+        return;
+    }
+
+    EXPECT_EQ(sys_dev0, sys_dev1);
+#else
+    UCS_TEST_SKIP_R("built without locality-domain support");
+#endif
+}
+
+UCS_TEST_P(test_mem_alloc_device, vmm_multi_handle_unaligned_range,
+           "CUDA_COPY_REG_WHOLE_ALLOC=on")
+{
+    cuda_vmm_mem_buffer buffer(4 * UCS_MBYTE, UCS_MEMORY_TYPE_CUDA, 2);
+    size_t chunk_size = buffer.size() / 2;
+    void *address     = UCS_PTR_BYTE_OFFSET(buffer.ptr(), chunk_size / 2);
+    uct_md_mem_attr_t mem_attr;
+
+    /* Contained in the first handle: extended to the whole handle */
+    mem_attr.field_mask = UCT_MD_MEM_ATTR_FIELD_BASE_ADDRESS |
+                          UCT_MD_MEM_ATTR_FIELD_ALLOC_LENGTH;
+    ASSERT_UCS_OK(uct_md_mem_query(md(), address, chunk_size / 2, &mem_attr));
+    EXPECT_EQ(buffer.ptr(), mem_attr.base_address);
+    EXPECT_EQ(chunk_size, mem_attr.alloc_length);
+
+    /* Ends in the second handle, while being shorter than a single handle:
+     * the requested extent must be kept */
+    check_mem_query_exact_range(address, (chunk_size / 2) + 1);
+}
+
 #if CUDA_VERSION >= 12020
-/* Host-located VMM is not dmabuf-exportable but is registrable by IB.
- *
- * TODO: REG_WHOLE_ALLOC=off is needed because of an existing issue with
- * base_address being set to 0, causing cache pollution with this type of memory
- */
-UCS_TEST_P(test_mem_alloc_device, host_vmm_mem_registrable,
-           "CUDA_COPY_REG_WHOLE_ALLOC=off")
+/* Host-located VMM is not dmabuf-exportable but is registrable by IB */
+UCS_TEST_P(test_mem_alloc_device, host_vmm_mem_registrable)
 {
     constexpr size_t size = 4 * UCS_MBYTE;
     cuda_host_vmm_mem_buffer buffer(size, UCS_MEMORY_TYPE_CUDA);
@@ -663,14 +732,24 @@ UCS_TEST_P(test_mem_alloc_device, host_vmm_mem_registrable,
 
 /* Host-located VMM is typed as CUDA memory, but cannot be allocated as
  * GPUDirect RDMA capable */
-UCS_TEST_P(test_mem_alloc_device, host_vmm_mem_not_gdr_capable,
-           "CUDA_COPY_REG_WHOLE_ALLOC=off")
+UCS_TEST_P(test_mem_alloc_device, host_vmm_mem_not_gdr_capable)
 {
     constexpr size_t size = 4 * UCS_MBYTE;
     cuda_host_vmm_mem_buffer buffer(size, UCS_MEMORY_TYPE_CUDA);
 
     query_registrable_gdr_capable(buffer.ptr(), size, 1, 0);
 }
+
+/* cuMemGetAddressRange() reports a zero base address for host-located VMM, so
+ * the registration range must not be extended to it */
+UCS_TEST_P(test_mem_alloc_device, host_vmm_mem_range,
+           "CUDA_COPY_REG_WHOLE_ALLOC=on")
+{
+    constexpr size_t size = 4 * UCS_MBYTE;
+    cuda_host_vmm_mem_buffer buffer(size, UCS_MEMORY_TYPE_CUDA);
+
+    check_mem_query_exact_range(buffer.ptr(), buffer.size());
+}
 #endif
 
 UCS_TEST_P(test_mem_alloc_device, legacy_cuda_no_memtype_copy_inter_node)

test/gtest/uct/test_md.cc

@@ -274,6 +274,18 @@ void test_md::check_memory(void *address, void *expect, size_t size,
     EXPECT_TRUE(mem_buffer::compare(expect, address, size, mem_type));
 }
 
+void test_md::check_mem_query_exact_range(void *address, size_t length)
+{
+    uct_md_mem_attr_t mem_attr;
+
+    mem_attr.field_mask = UCT_MD_MEM_ATTR_FIELD_BASE_ADDRESS |
+                          UCT_MD_MEM_ATTR_FIELD_ALLOC_LENGTH;
+    ASSERT_UCS_OK(uct_md_mem_query(md(), address, length, &mem_attr));
+
+    EXPECT_EQ(address, mem_attr.base_address);
+    EXPECT_EQ(length, mem_attr.alloc_length);
+}
+
 void test_md::free_memory(void *address, ucs_memory_type_t mem_type)
 {
     mem_buffer::release(address, mem_type);

test/gtest/uct/test_md.h

@@ -46,6 +46,7 @@ class test_md : public testing::TestWithParam<test_md_param>,
                       ucs_memory_type_t mem_type);
     void check_memory(void *address, void *expect, size_t size,
                       ucs_memory_type_t mem_type);
+    void check_mem_query_exact_range(void *address, size_t length);
     void free_memory(void *address, ucs_memory_type_t mem_type);
     static bool is_device_detected(ucs_memory_type_t mem_type);
     static void* alloc_thread(void *arg);