mpix: add MPIX_Query_pointer_attr

It is not always easy to check what kind of memory type user has as seen by MPI. For example, a memory from a Python library. Add MPIX_Query_pointer_attr for the debugging purpose.
pmodels · Jul 9, 2024 · 4057593 · 4057593
1 parent 842f4cf
commit 4057593
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 1 deletion.
diff --git a/src/binding/c/misc_api.txt b/src/binding/c/misc_api.txt
@@ -208,3 +208,22 @@ MPIX_Query_hip_support:
 
     return is_supported;
 }
+
+MPIX_Query_pointer_attr:
+    ptr: BUFFER, [Pointer address to be queried]
+    type: POINTER_TYPE, direction=out, [Memory type (enum MPIX_Pointer_type)]
+    device: INDEX, direction=out, [Device id (if the pointer address is associated with a device, -1 otherwise)]
+    .desc: return buffer type (enum MPIX_Pointer_type) and device id (if it is a device buffer or device registered buffer)
+    .impl: direct
+    .skip: global_cs
+{
+    int mpi_errno = MPI_SUCCESS;
+
+    MPL_pointer_attr_t attr;
+    mpi_errno = MPL_gpu_query_pointer_attr(ptr, &attr);
+
+    *type = attr.type;
+    *device = attr.device;
+
+    return mpi_errno;
+}
diff --git a/src/binding/custom_mapping.txt b/src/binding/custom_mapping.txt
@@ -7,45 +7,52 @@ LIS_KIND_MAP:
     STREAM: handle
     IOVEC: None
     ASYNC_THING: None
+    POINTER_TYPE: integer
 
 SMALL_F90_KIND_MAP:
     GPU_TYPE: INTEGER
     GREQUEST_CLASS: INTEGER
     STREAM: INTEGER
     IOVEC: None
     ASYNC_THING: INTEGER(KIND=MPI_ADDRESS_KIND)
+    POINTER_TYPE: INTEGER
 
 BIG_F90_KIND_MAP:
     GPU_TYPE: INTEGER
     GREQUEST_CLASS: INTEGER
     STREAM: INTEGER
     IOVEC: None
     ASYNC_THING: INTEGER(KIND=MPI_ADDRESS_KIND)
+    POINTER_TYPE: INTEGER
 
 SMALL_F08_KIND_MAP:
     GPU_TYPE: INTEGER
     GREQUEST_CLASS: INTEGER
     STREAM: TYPE(MPIX_Stream)
     IOVEC: None
     ASYNC_THING: INTEGER(KIND=MPI_ADDRESS_KIND)
+    POINTER_TYPE: INTEGER
 
 BIG_F08_KIND_MAP:
     GPU_TYPE: INTEGER
     GREQUEST_CLASS: INTEGER
     STREAM: TYPE(MPIX_Stream)
     IOVEC: None
     ASYNC_THING: INTEGER(KIND=MPI_ADDRESS_KIND)
+    POINTER_TYPE: INTEGER
 
 SMALL_C_KIND_MAP:
     GPU_TYPE: int
     GREQUEST_CLASS: MPIX_Grequest_class
     STREAM: MPIX_Stream
     IOVEC: MPIX_Iov
     ASYNC_THING: MPIX_Async_thing
+    POINTER_TYPE: int
 
 BIG_C_KIND_MAP:
     GPU_TYPE: int
     GREQUEST_CLASS: MPIX_Grequest_class
     STREAM: MPIX_Stream
     IOVEC: MPIX_Iov
     ASYNC_THING: MPIX_Async_thing
+    POINTER_TYPE: int
diff --git a/src/include/mpi.h.in b/src/include/mpi.h.in
@@ -763,6 +763,14 @@ enum MPIR_Combiner_enum {
 #define MPIX_GPU_SUPPORT_ZE    (1)
 #define MPIX_GPU_SUPPORT_HIP   (2)
 
+/* MPIX_Query_pointer_attr */
+enum MPIX_Pointer_type {
+    MPIX_POINTER_UNREGISTERED_HOST = 0,
+    MPIX_POINTER_REGISTERED_HOST,
+    MPIX_POINTER_DEV,
+    MPIX_POINTER_MANAGED
+};
+
 /* feature advertisement */
 #define MPIIMPL_ADVERTISES_FEATURES 1
 #define MPIIMPL_HAVE_MPI_INFO 1

diff --git a/src/mpl/include/mpl_gpu.h b/src/mpl/include/mpl_gpu.h
@@ -22,6 +22,7 @@
 #include "mpl_gpu_fallback.h"
 #endif
 
+/* NOTE: must agree with MPIX_Pointer_type in mpi.h */
 typedef enum {
     MPL_GPU_POINTER_UNREGISTERED_HOST = 0,
     MPL_GPU_POINTER_REGISTERED_HOST,
@@ -61,7 +62,7 @@ typedef enum {
     MPL_GPU_COPY_H2D,
     MPL_GPU_COPY_D2D_INCOMING,  /* copy from remote to local */
     MPL_GPU_COPY_D2D_OUTGOING,  /* copy from local to remote */
-    MPL_GPU_COPY_DIRECTION_NONE,  /* copy in any direction and to/from any buffer type */
+    MPL_GPU_COPY_DIRECTION_NONE,        /* copy in any direction and to/from any buffer type */
 } MPL_gpu_copy_direction_t;
 
 #define MPL_GPU_COPY_DIRECTION_TYPES 4