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,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; keep the caller's extent. */
+    if (is_vmm && ((base_address + alloc_length) < (address + length))) {
         goto out_ctx_pop;
     }
 

test/gtest/uct/cuda/test_cuda_ipc_md.cc

@@ -807,4 +807,23 @@ 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;
+    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_LE(mem_attr.base_address, address);
+    EXPECT_GE(UCS_PTR_BYTE_OFFSET(mem_attr.base_address, mem_attr.alloc_length),
+              UCS_PTR_BYTE_OFFSET(address, length));
+}
+
 _UCT_MD_INSTANTIATE_TEST_CASE(test_cuda_copy_md, cuda_cpy);