From 405759344308b3d57dedd32326e5786f73319128 Mon Sep 17 00:00:00 2001
From: Hui Zhou <hzhou321@anl.gov>
Date: Tue, 2 Jul 2024 16:41:21 -0500
Subject: [PATCH] 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.
---
 src/binding/c/misc_api.txt     | 19 +++++++++++++++++++
 src/binding/custom_mapping.txt |  7 +++++++
 src/include/mpi.h.in           |  8 ++++++++
 src/mpl/include/mpl_gpu.h      |  3 ++-
 4 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/src/binding/c/misc_api.txt b/src/binding/c/misc_api.txt
index ee4dd6d1bd9..ff22c13877d 100644
--- 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
index d3efc9eb545..57e3b9bef55 100644
--- a/src/binding/custom_mapping.txt
+++ b/src/binding/custom_mapping.txt
@@ -7,6 +7,7 @@ LIS_KIND_MAP:
     STREAM: handle
     IOVEC: None
     ASYNC_THING: None
+    POINTER_TYPE: integer
 
 SMALL_F90_KIND_MAP:
     GPU_TYPE: INTEGER
@@ -14,6 +15,7 @@ SMALL_F90_KIND_MAP:
     STREAM: INTEGER
     IOVEC: None
     ASYNC_THING: INTEGER(KIND=MPI_ADDRESS_KIND)
+    POINTER_TYPE: INTEGER
 
 BIG_F90_KIND_MAP:
     GPU_TYPE: INTEGER
@@ -21,6 +23,7 @@ BIG_F90_KIND_MAP:
     STREAM: INTEGER
     IOVEC: None
     ASYNC_THING: INTEGER(KIND=MPI_ADDRESS_KIND)
+    POINTER_TYPE: INTEGER
 
 SMALL_F08_KIND_MAP:
     GPU_TYPE: INTEGER
@@ -28,6 +31,7 @@ SMALL_F08_KIND_MAP:
     STREAM: TYPE(MPIX_Stream)
     IOVEC: None
     ASYNC_THING: INTEGER(KIND=MPI_ADDRESS_KIND)
+    POINTER_TYPE: INTEGER
 
 BIG_F08_KIND_MAP:
     GPU_TYPE: INTEGER
@@ -35,6 +39,7 @@ BIG_F08_KIND_MAP:
     STREAM: TYPE(MPIX_Stream)
     IOVEC: None
     ASYNC_THING: INTEGER(KIND=MPI_ADDRESS_KIND)
+    POINTER_TYPE: INTEGER
 
 SMALL_C_KIND_MAP:
     GPU_TYPE: int
@@ -42,6 +47,7 @@ SMALL_C_KIND_MAP:
     STREAM: MPIX_Stream
     IOVEC: MPIX_Iov
     ASYNC_THING: MPIX_Async_thing
+    POINTER_TYPE: int
 
 BIG_C_KIND_MAP:
     GPU_TYPE: int
@@ -49,3 +55,4 @@ BIG_C_KIND_MAP:
     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
index 1180c3f17a1..df972147e8d 100644
--- 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
index b81cf2ddf0c..c84892dc752 100644
--- 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