UCP/RMA: Report error if doing RMA for non-HOST memory

openucx · Nov 11, 2019 · c37ffa4 · c37ffa4
1 parent c9527ad
commit c37ffa4
Showing 1 changed file with 32 additions and 0 deletions.
diff --git a/src/ucp/rma/rma_send.c b/src/ucp/rma/rma_send.c
@@ -118,11 +118,33 @@ static void ucp_rma_request_zcopy_completion(uct_completion_t *self,
     }
 }
 
+static UCS_F_ALWAYS_INLINE ucs_status_t
+ucp_rma_check_mem_type(ucp_ep_h ep, const void *buffer, size_t length)
+{
+    ucs_memory_type_t mem_type = ucp_memory_type_detect(ep->worker->context,
+                                                        buffer, length);
+
+    /* TODO: remove when support for non-HOST memory types will be added */
+    if (ucs_unlikely(mem_type != UCS_MEMORY_TYPE_HOST)) {
+        ucs_error("UCP doesn't support RMA for \"%s\" memory type",
+                  ucs_memory_type_names[mem_type]);
+        return UCS_ERR_INVALID_PARAM;
+    }
+
+    return UCS_OK;
+}
+
 static UCS_F_ALWAYS_INLINE ucs_status_t
 ucp_rma_request_init(ucp_request_t *req, ucp_ep_h ep, const void *buffer,
                      size_t length, uint64_t remote_addr, ucp_rkey_h rkey,
                      uct_pending_callback_t cb, size_t zcopy_thresh, int flags)
 {
+    ucs_status_t status = ucp_rma_check_mem_type(ep, buffer, length);
+
+    if (ucs_unlikely(status != UCS_OK)) {
+        return status;
+    }
+
     req->flags                = flags; /* Implicit release */
     req->send.ep              = ep;
     req->send.buffer          = (void*)buffer;
@@ -215,6 +237,11 @@ ucs_status_t ucp_put_nbi(ucp_ep_h ep, const void *buffer, size_t length,
 
     /* Fast path for a single short message */
     if (ucs_likely((ssize_t)length <= (int)rkey->cache.max_put_short)) {
+        status = ucp_rma_check_mem_type(ep, buffer, length);
+        if (ucs_unlikely(status != UCS_OK)) {
+            return status;
+        }
+
         status = UCS_PROFILE_CALL(uct_ep_put_short, ep->uct_eps[rkey->cache.rma_lane],
                                   buffer, length, remote_addr, rkey->cache.rma_rkey);
         if (ucs_likely(status != UCS_ERR_NO_RESOURCE)) {
@@ -253,6 +280,11 @@ ucs_status_ptr_t ucp_put_nb(ucp_ep_h ep, const void *buffer, size_t length,
 
     /* Fast path for a single short message */
     if (ucs_likely((ssize_t)length <= (int)rkey->cache.max_put_short)) {
+        status = ucp_rma_check_mem_type(ep, buffer, length);
+        if (ucs_unlikely(status != UCS_OK)) {
+            return UCS_STATUS_PTR(status);
+        }
+
         status = UCS_PROFILE_CALL(uct_ep_put_short, ep->uct_eps[rkey->cache.rma_lane],
                                   buffer, length, remote_addr, rkey->cache.rma_rkey);
         if (ucs_likely(status != UCS_ERR_NO_RESOURCE)) {