UCT/GDR_COPY: Fix review commnets

openucx · Nov 16, 2018 · bbd0058 · bbd0058
1 parent 3b7f72e
commit bbd0058
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/src/uct/cuda/gdr_copy/gdr_copy_md.c b/src/uct/cuda/gdr_copy/gdr_copy_md.c
@@ -9,6 +9,7 @@
 #include <limits.h>
 #include <ucs/debug/log.h>
 #include <ucs/sys/sys.h>
+#include <ucs/sys/math.h>
 #include <ucs/debug/memtrack.h>
 #include <ucs/type/class.h>
 #include <ucm/api/ucm.h>
@@ -178,7 +179,7 @@ static ucs_status_t uct_gdr_copy_mem_reg(uct_md_h uct_md, void *address, size_t
                                          unsigned flags, uct_mem_h *memh_p)
 {
     uct_gdr_copy_mem_t *mem_hndl = NULL;
-    uintptr_t start, end;
+    void *start, *end;
     ucs_status_t status;
 
     mem_hndl = ucs_malloc(sizeof(uct_gdr_copy_mem_t), "gdr_copy handle");
@@ -187,10 +188,11 @@ static ucs_status_t uct_gdr_copy_mem_reg(uct_md_h uct_md, void *address, size_t
         return UCS_ERR_NO_MEMORY;
     }
 
-    start = ucs_align_down_pow2((uintptr_t)address, GPU_PAGE_SIZE);
-    end   = ucs_align_up_pow2((uintptr_t)address + length, GPU_PAGE_SIZE);
+    start = ucs_align_down_pow2_ptr(address, GPU_PAGE_SIZE);
+    end   = ucs_align_up_pow2_ptr(address + length, GPU_PAGE_SIZE);
+    ucs_assert_always(start <= end);
 
-    status = uct_gdr_copy_mem_reg_internal(uct_md, (void *)start, (end - start), 0, mem_hndl);
+    status = uct_gdr_copy_mem_reg_internal(uct_md, start, end - start, 0, mem_hndl);
     if (status != UCS_OK) {
         ucs_free(mem_hndl);
         return status;