hmem: Support dev reg with FI_HMEM_CUDA

This is a wrapper around GDRcopy.

Signed-off-by: Ian Ziemba <[email protected]>

include/ofi_hmem.h

@@ -168,13 +168,23 @@ int cuda_hmem_cleanup(void);
 bool cuda_is_addr_valid(const void *addr, uint64_t *device, uint64_t *flags);
 int cuda_host_register(void *ptr, size_t size);
 int cuda_host_unregister(void *ptr);
-int cuda_dev_register(struct fi_mr_attr *mr_attr, uint64_t *handle);
+int cuda_dev_register(const void *addr, size_t size, uint64_t *handle);
 int cuda_dev_unregister(uint64_t handle);
+int cuda_dev_reg_copy_to_hmem(uint64_t handle, void *dest, const void *src,
+			      size_t size);
+int cuda_dev_reg_copy_from_hmem(uint64_t handle, void *dest, const void *src,
+				size_t size);
 int cuda_get_handle(void *dev_buf, size_t size, void **handle);
 int cuda_open_handle(void **handle, size_t size, uint64_t device,
 		     void **ipc_ptr);
 int cuda_close_handle(void *ipc_ptr);
 int cuda_get_base_addr(const void *ptr, size_t len, void **base, size_t *size);
+int cuda_dev_register(const void *addr, size_t size, uint64_t *handle);
+int cuda_dev_unregister(uint64_t handle);
+int cuda_dev_reg_copy_to_hmem(uint64_t handle, void *dest, const void *src,
+			      size_t size);
+int cuda_dev_reg_copy_from_hmem(uint64_t handle, void *dest, const void *src,
+				size_t size);
 
 bool cuda_is_ipc_enabled(void);
 int cuda_get_ipc_handle_size(size_t *size);

src/hmem.c

@@ -161,10 +161,10 @@ struct ofi_hmem_ops hmem_ops[] = {
 		.get_base_addr = cuda_get_base_addr,
 		.is_ipc_enabled = cuda_is_ipc_enabled,
 		.get_ipc_handle_size = cuda_get_ipc_handle_size,
-		.dev_register = ofi_hmem_no_dev_register,
-		.dev_unregister = ofi_hmem_no_dev_unregister,
-		.dev_reg_copy_to_hmem = ofi_hmem_no_dev_reg_copy_to_hmem,
-		.dev_reg_copy_from_hmem = ofi_hmem_no_dev_reg_copy_from_hmem,
+		.dev_register = cuda_dev_register,
+		.dev_unregister = cuda_dev_unregister,
+		.dev_reg_copy_to_hmem = cuda_dev_reg_copy_to_hmem,
+		.dev_reg_copy_from_hmem = cuda_dev_reg_copy_from_hmem,
 	},
 	[FI_HMEM_ROCR] = {
 		.initialized = false,

src/hmem_cuda.c

@@ -312,20 +312,35 @@ int cuda_copy_from_dev(uint64_t device, void *dst, const void *src, size_t size)
 	return -FI_EIO;
 }
 
-int cuda_dev_register(struct fi_mr_attr *mr_attr, uint64_t *handle)
+int cuda_dev_register(const void *addr, size_t size, uint64_t *handle)
 {
-	if (cuda_attr.use_gdrcopy)
-		return cuda_gdrcopy_dev_register(mr_attr, handle);
+	struct fi_mr_attr mr_attr = {};
+	struct iovec iov = {};
 
-	*handle = mr_attr->device.cuda;
-	return FI_SUCCESS;
+	iov.iov_base = (void *) addr;
+	iov.iov_len = size;
+	mr_attr.mr_iov = &iov;
+	mr_attr.iov_count = 1;
+
+	return cuda_gdrcopy_dev_register(&mr_attr, handle);
 }
 
 int cuda_dev_unregister(uint64_t handle)
 {
-	if (cuda_attr.use_gdrcopy)
-		return cuda_gdrcopy_dev_unregister(handle);
+	return cuda_gdrcopy_dev_unregister(handle);
+}
 
+int cuda_dev_reg_copy_to_hmem(uint64_t handle, void *dest, const void *src,
+			      size_t size)
+{
+	cuda_gdrcopy_to_dev(handle, dest, src, size);
+	return FI_SUCCESS;
+}
+
+int cuda_dev_reg_copy_from_hmem(uint64_t handle, void *dest, const void *src,
+				size_t size)
+{
+	cuda_gdrcopy_from_dev(handle, dest, src, size);
 	return FI_SUCCESS;
 }
 
@@ -811,14 +826,26 @@ int cuda_host_unregister(void *ptr)
 	return -FI_ENOSYS;
 }
 
-int cuda_dev_register(struct fi_mr_attr *mr_attr, uint64_t *handle)
+int cuda_dev_register(const void *addr, size_t size, uint64_t *handle)
 {
-	return FI_SUCCESS;
+	return -FI_ENOSYS;
 }
 
 int cuda_dev_unregister(uint64_t handle)
 {
-	return FI_SUCCESS;
+	return -FI_ENOSYS;
+}
+
+int cuda_dev_reg_copy_to_hmem(uint64_t handle, void *dest, const void *src,
+			      size_t size)
+{
+	return -FI_ENOSYS;
+}
+
+int cuda_dev_reg_copy_from_hmem(uint64_t handle, void *dest, const void *src,
+				size_t size)
+{
+	return -FI_ENOSYS;
 }
 
 int cuda_get_handle(void *dev_buf, size_t size, void **handle)