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

@@ -610,9 +610,11 @@ 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, and host-located VMM reports a
+     * zero base address; keep the caller's extent in both cases. */
+    if (is_vmm && ((base_address == 0) ||
+                   ((base_address + alloc_length) < (address + length)))) {
         goto out_ctx_pop;
     }
 

test/gtest/uct/cuda/test_cuda_ipc_md.cc

@@ -807,4 +807,33 @@ UCS_TEST_P(test_cuda_copy_md, vmm_locality_domain_mem_type) {
 #endif
 }
 
+UCS_TEST_P(test_cuda_copy_md, vmm_multi_handle_unaligned_range,
+           "CUDA_COPY_REG_WHOLE_ALLOC=on") {
+    cuda_vmm_mem_buffer buffer(1, UCS_MEMORY_TYPE_CUDA, 2);
+    /* Starts in the middle of the first handle and ends in the second one,
+     * while being shorter than a single handle */
+    size_t chunk_size = buffer.size() / 2;
+    void *address     = UCS_PTR_BYTE_OFFSET(buffer.ptr(), chunk_size / 2);
+    size_t length     = (chunk_size / 2) + 1;
+    CUdeviceptr alloc_base;
+    size_t alloc_length;
+    uct_md_mem_attr_t mem_attr;
+
+    /* If the query fails, the MD keeps the requested extent as well */
+    if ((cuMemGetAddressRange(&alloc_base, &alloc_length,
+                              (CUdeviceptr)address) == CUDA_SUCCESS) &&
+        ((alloc_base + alloc_length) >= ((CUdeviceptr)address + length))) {
+        UCS_TEST_SKIP_R("allocation range covers the requested range");
+    }
+
+    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));
+
+    /* The containing allocation is too short, so the requested extent must be
+     * kept */
+    EXPECT_EQ(address, mem_attr.base_address);
+    EXPECT_EQ(length, mem_attr.alloc_length);
+}
+
 _UCT_MD_INSTANTIATE_TEST_CASE(test_cuda_copy_md, cuda_cpy);

test/gtest/uct/cuda/test_switch_cuda_device.cc

@@ -647,13 +647,8 @@ UCS_TEST_P(test_mem_alloc_device, no_current_context_vmm_mem_registrable,
 }
 
 #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 +658,30 @@ 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);
+    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(), buffer.ptr(), buffer.size(),
+                                   &mem_attr));
+    EXPECT_EQ(buffer.ptr(), mem_attr.base_address);
+    EXPECT_EQ(buffer.size(), mem_attr.alloc_length);
+}
 #endif
 
 UCS_TEST_P(test_mem_alloc_device, legacy_cuda_no_memtype_copy_inter_node)