From 91a94201d2b275705c7f9ab9250b9cfbcf03aa9a Mon Sep 17 00:00:00 2001
From: Joseph Schuchart <schuchart@hlrs.de>
Date: Mon, 28 Sep 2020 15:12:24 +0200
Subject: [PATCH] PML UCX: add SPC instrumentation for message size
 sent/received

Signed-off-by: Joseph Schuchart <schuchart@icl.utk.edu>
---
 ompi/mca/pml/ucx/pml_ucx.c         | 22 ++++++++++++++++++++++
 ompi/mca/pml/ucx/pml_ucx_request.c |  4 ++++
 2 files changed, 26 insertions(+)

diff --git a/ompi/mca/pml/ucx/pml_ucx.c b/ompi/mca/pml/ucx/pml_ucx.c
index 7e5b566ff9..fc2c43e5fe 100644
--- a/ompi/mca/pml/ucx/pml_ucx.c
+++ b/ompi/mca/pml/ucx/pml_ucx.c
@@ -20,6 +20,7 @@
 #include "opal/mca/pmix/pmix-internal.h"
 #include "ompi/attribute/attribute.h"
 #include "ompi/message/message.h"
+#include "ompi/runtime/ompi_spc.h"
 #include "ompi/mca/pml/base/pml_base_bsend.h"
 #include "opal/mca/common/ucx/common_ucx.h"
 #if OPAL_CUDA_SUPPORT
@@ -627,6 +628,13 @@ int mca_pml_ucx_recv(void *buf, size_t count, ompi_datatype_t *datatype, int src
         status = ucp_request_test(req, &info);
         if (status != UCS_INPROGRESS) {
             mca_pml_ucx_set_recv_status_safe(mpi_status, status, &info);
+
+#if SPC_ENABLE == 1
+            size_t dt_size;
+            ompi_datatype_type_size(datatype, &dt_size);
+            SPC_USER_OR_MPI(tag, dt_size*count,
+                            OMPI_SPC_BYTES_RECEIVED_USER, OMPI_SPC_BYTES_RECEIVED_MPI);
+#endif
             return OMPI_SUCCESS;
         }
     }
@@ -827,6 +835,13 @@ int mca_pml_ucx_isend(const void *buf, size_t count, ompi_datatype_t *datatype,
                                                    mca_pml_ucx_send_completion);
 #endif
 
+#if SPC_ENABLE == 1
+    size_t dt_size;
+    ompi_datatype_type_size(datatype, &dt_size);
+    SPC_USER_OR_MPI(tag, dt_size*count,
+                    OMPI_SPC_BYTES_SENT_USER, OMPI_SPC_BYTES_SENT_MPI);
+#endif
+
     if (req == NULL) {
         PML_UCX_VERBOSE(8, "returning completed request");
         *request = &ompi_pml_ucx.completed_send_req;
@@ -918,6 +933,13 @@ int mca_pml_ucx_send(const void *buf, size_t count, ompi_datatype_t *datatype, i
         return OMPI_ERROR;
     }
 
+#if SPC_ENABLE == 1
+    size_t dt_size;
+    ompi_datatype_type_size(datatype, &dt_size);
+    SPC_USER_OR_MPI(tag, dt_size*count,
+                    OMPI_SPC_BYTES_SENT_USER, OMPI_SPC_BYTES_SENT_MPI);
+#endif
+
 #if HAVE_DECL_UCP_TAG_SEND_NBR
     if (OPAL_LIKELY((MCA_PML_BASE_SEND_BUFFERED != mode) &&
                     (MCA_PML_BASE_SEND_SYNCHRONOUS != mode))) {
diff --git a/ompi/mca/pml/ucx/pml_ucx_request.c b/ompi/mca/pml/ucx/pml_ucx_request.c
index b072afc66e..744d580358 100644
--- a/ompi/mca/pml/ucx/pml_ucx_request.c
+++ b/ompi/mca/pml/ucx/pml_ucx_request.c
@@ -13,6 +13,7 @@
 #include "pml_ucx_request.h"
 #include "ompi/mca/pml/base/pml_base_bsend.h"
 #include "ompi/message/message.h"
+#include "ompi/runtime/ompi_spc.h"
 #include <inttypes.h>
 
 
@@ -71,6 +72,9 @@ mca_pml_ucx_recv_completion_internal(void *request, ucs_status_t status,
                     (void*)req, ucs_status_string(status), info->sender_tag,
                     info->length);
 
+    SPC_USER_OR_MPI(PML_UCX_TAG_GET_MPI_TAG(info->sender_tag), info->length,
+                    OMPI_SPC_BYTES_RECEIVED_USER, OMPI_SPC_BYTES_RECEIVED_MPI);
+
     mca_pml_ucx_set_recv_status(&req->req_status, status, info);
     PML_UCX_ASSERT( !(REQUEST_COMPLETE(req)));
     ompi_request_complete(req, true);