opal/accelerator: Add get_buffer_id API

Signed-off-by: William Zhang <wilzhang@amazon.com>

Author: wckzhang

opal/mca/accelerator/accelerator.h

@@ -79,7 +79,6 @@
 
 BEGIN_C_DECLS
 
-
 #define MCA_ACCELERATOR_NO_DEVICE_ID -1
 /**
  * Accelerator flags
@@ -103,6 +102,8 @@ typedef enum {
     MCA_ACCELERATOR_TRANSFER_DTOD,
 } opal_accelerator_transfer_type_t;
 
+typedef uint64_t opal_accelerator_buffer_id_t;
+
 struct opal_accelerator_stream_t {
     opal_object_t super;
     /* Stream object */
@@ -359,6 +360,20 @@ typedef int (*opal_accelerator_base_module_get_device_fn_t)(
 typedef int (*opal_accelerator_base_module_device_can_access_peer_fn_t)(
     int *access, int dev1, int dev2);
 
+/**
+ * Retrieves current device id for a device associated with the local process.
+ * If MCA_ACCELERATOR_NO_DEVICE_ID is provided, there is no device/process pairing.
+ *
+ * @param[IN] dev_id         ID of the device or MCA_ACCELERATOR_NO_DEVICE_ID
+ * @param[IN] addr           Buffer pointer to check
+ * @param[OUT] buf_id        ID of the given buffer
+ *
+ *
+ * @return                   OPAL_SUCCESS or error status on failure
+ */
+typedef int (*opal_accelerator_base_module_get_buffer_id_fn_t)(
+    int dev_id, const void *addr, opal_accelerator_buffer_id_t *buf_id);
+
 /*
  * the standard public API data structure
  */
@@ -384,6 +399,8 @@ typedef struct {
 
     opal_accelerator_base_module_get_device_fn_t get_device;
     opal_accelerator_base_module_device_can_access_peer_fn_t device_can_access_peer;
+
+    opal_accelerator_base_module_get_buffer_id_fn_t get_buffer_id;
 } opal_accelerator_base_module_t;
 
 /**

opal/mca/accelerator/cuda/accelerator_cuda.c

@@ -47,6 +47,8 @@ static int accelerator_cuda_host_unregister(int dev_id, void *ptr);
 static int accelerator_cuda_get_device(int *dev_id);
 static int accelerator_cuda_device_can_access_peer( int *access, int dev1, int dev2);
 
+static int accelerator_cuda_get_buffer_id(int dev_id, const void *addr, opal_accelerator_buffer_id_t *buf_id);
+
 opal_accelerator_base_module_t opal_accelerator_cuda_module =
 {
     accelerator_cuda_check_addr,
@@ -68,7 +70,9 @@ opal_accelerator_base_module_t opal_accelerator_cuda_module =
     accelerator_cuda_host_unregister,
 
     accelerator_cuda_get_device,
-    accelerator_cuda_device_can_access_peer
+    accelerator_cuda_device_can_access_peer,
+
+    accelerator_cuda_get_buffer_id
 };
 
 static int accelerator_cuda_check_addr(const void *addr, int *dev_id, uint64_t *flags)
@@ -538,3 +542,30 @@ static int accelerator_cuda_device_can_access_peer(int *access, int dev1, int de
     }
     return 0;
 }
+
+/*
+ * Get the buffer ID from the memory.
+ * This is needed to ensure the cached registration is not stale.  If
+ * we fail to get buffer ID, print an error and set buffer ID to 0.
+ * Also set SYNC_MEMOPS on any GPU registration to ensure that
+ * synchronous copies complete before the buffer is accessed.
+ */
+static int accelerator_cuda_get_buffer_id(int dev_id, const void *addr, opal_accelerator_buffer_id_t *buf_id)
+{
+    CUresult result;
+    int enable = 1;
+    result = opal_accelerator_cuda_func.cuPointerGetAttribute((unsigned long long *)buf_id, CU_POINTER_ATTRIBUTE_BUFFER_ID, (CUdeviceptr) addr);
+    if (OPAL_UNLIKELY(result != CUDA_SUCCESS)) {
+        opal_show_help("help-accelerator-cuda.txt", "bufferID failed", true, OPAL_PROC_MY_HOSTNAME,
+                       result);
+        return result;
+    }
+    result = opal_accelerator_cuda_func.cuPointerSetAttribute(&enable, CU_POINTER_ATTRIBUTE_SYNC_MEMOPS,
+                                       (CUdeviceptr) addr);
+    if (OPAL_UNLIKELY(CUDA_SUCCESS != result)) {
+        opal_show_help("help-accelerator-cuda.txt", "cuPointerSetAttribute failed", true,
+                       OPAL_PROC_MY_HOSTNAME, result, addr);
+        return result;
+    }
+    return OPAL_SUCCESS;
+}

opal/mca/accelerator/null/accelerator_null_component.c

@@ -61,6 +61,8 @@ static int accelerator_null_host_unregister(int dev_id, void *ptr);
 static int accelerator_null_get_device(int *dev_id);
 static int accelerator_null_device_can_access_peer(int *access, int dev1, int dev2);
 
+static int accelerator_null_get_buffer_id(int dev_id, const void *addr, opal_accelerator_buffer_id_t *buf_id);
+
 /*
  * Instantiate the public struct with all of our public information
  * and pointers to our public functions in it
@@ -120,7 +122,9 @@ opal_accelerator_base_module_t opal_accelerator_null_module =
     accelerator_null_host_unregister,
 
     accelerator_null_get_device,
-    accelerator_null_device_can_access_peer
+    accelerator_null_device_can_access_peer,
+
+    accelerator_null_get_buffer_id
 };
 
 static int accelerator_null_open(void)
@@ -235,3 +239,8 @@ static int accelerator_null_device_can_access_peer( int *access, int dev1, int d
 {
     return OPAL_ERR_NOT_IMPLEMENTED;
 }
+
+static int accelerator_null_get_buffer_id(int dev_id, const void *addr, opal_accelerator_buffer_id_t *buf_id)
+{
+    return OPAL_ERR_NOT_IMPLEMENTED;
+}

opal/mca/accelerator/rocm/accelerator_rocm_module.c

@@ -37,6 +37,7 @@ static int mca_accelerator_rocm_host_unregister(int dev_id, void *ptr);
 static int mca_accelerator_rocm_get_device(int *dev_id);
 static int mca_accelerator_rocm_device_can_access_peer( int *access, int dev1, int dev2);
 
+static int mca_accelerator_rocm_get_buffer_id(int dev_id, const void *addr, opal_accelerator_buffer_id_t *buf_id);
 
 opal_accelerator_base_module_t opal_accelerator_rocm_module =
 {
@@ -59,7 +60,9 @@ opal_accelerator_base_module_t opal_accelerator_rocm_module =
     mca_accelerator_rocm_host_unregister,
 
     mca_accelerator_rocm_get_device,
-    mca_accelerator_rocm_device_can_access_peer
+    mca_accelerator_rocm_device_can_access_peer,
+
+    mca_accelerator_rocm_get_buffer_id
 };
 
 
@@ -483,3 +486,9 @@ static int mca_accelerator_rocm_device_can_access_peer(int *access, int dev1, in
 
     return OPAL_SUCCESS;
 }
+
+static int accelerator_rocm_get_buffer_id(int dev_id, const void *addr, opal_accelerator_buffer_id_t *buf_id)
+{
+    *buf_id = 0;
+    return OPAL_SUCCESS;
+}