From e559864e35651aa8e2469b3c95c67e9676160570 Mon Sep 17 00:00:00 2001
From: Jean-Luc Fattebert <fattebertj@ornl.gov>
Date: Thu, 11 Nov 2021 12:13:03 -0500
Subject: [PATCH] Add get_sparsity function to distributed2d

---
 src/C-interface/bml_introspection.c           |  5 +++
 .../bml_introspection_distributed2d.c         | 28 +++++++++++++
 .../bml_introspection_distributed2d.h         |  4 ++
 tests/C-tests/CMakeLists.txt                  |  1 +
 tests/C-tests/get_sparsity_typed.c            | 41 ++++++++++++-------
 5 files changed, 65 insertions(+), 14 deletions(-)

diff --git a/src/C-interface/bml_introspection.c b/src/C-interface/bml_introspection.c
index 2246f59f3..a4d6ba037 100644
--- a/src/C-interface/bml_introspection.c
+++ b/src/C-interface/bml_introspection.c
@@ -339,6 +339,11 @@ bml_get_sparsity(
         case csr:
             return bml_get_sparsity_csr(A, threshold);
             break;
+#ifdef DO_MPI
+        case distributed2d:
+            return bml_get_sparsity_distributed2d(A, threshold);
+            break;
+#endif
         default:
             LOG_ERROR("unknown matrix type in bml_get_sparsity\n");
             break;
diff --git a/src/C-interface/distributed2d/bml_introspection_distributed2d.c b/src/C-interface/distributed2d/bml_introspection_distributed2d.c
index f6171b558..df0d49e4f 100644
--- a/src/C-interface/distributed2d/bml_introspection_distributed2d.c
+++ b/src/C-interface/distributed2d/bml_introspection_distributed2d.c
@@ -3,6 +3,7 @@
 #include "../bml_logger.h"
 #include "../bml_logger.h"
 #include "../bml_types.h"
+#include "../bml_parallel.h"
 #include "bml_introspection_distributed2d.h"
 #include "bml_types_distributed2d.h"
 
@@ -96,3 +97,30 @@ bml_get_local_matrix_distributed2d(
         return NULL;
     }
 }
+
+/** Return the sparsity of a matrix.
+ *
+ *  Note that the the sparsity of a matrix is defined
+ *  as NumberOfZeroes/N*N where N is the matrix dimension.
+ *  The density of matrix A will be defined as 1-sparsity(A)
+ *
+ * \ingroup introspection_group_C
+ *
+ * \param A The bml matrix.
+ * \param threshold The threshold used to compute the sparsity.
+ * \return The sparsity of A.
+ */
+double
+bml_get_sparsity_distributed2d(
+    bml_matrix_distributed2d_t * A,
+    double threshold)
+{
+    double sp = bml_get_sparsity(A->matrix, threshold);
+    int nloc = A->N / A->nprows;
+    sp *= nloc * nloc;
+
+    bml_sumRealReduce(&sp);
+    sp /= (A->N * A->N);
+
+    return sp;
+}
diff --git a/src/C-interface/distributed2d/bml_introspection_distributed2d.h b/src/C-interface/distributed2d/bml_introspection_distributed2d.h
index 2a12e9ff5..115e26afb 100644
--- a/src/C-interface/distributed2d/bml_introspection_distributed2d.h
+++ b/src/C-interface/distributed2d/bml_introspection_distributed2d.h
@@ -17,4 +17,8 @@ int bml_get_M_distributed2d(
 
 bml_matrix_t *bml_get_local_matrix_distributed2d(
     bml_matrix_distributed2d_t * A);
+
+double bml_get_sparsity_distributed2d(
+    bml_matrix_distributed2d_t * A,
+    double threshold);
 #endif
diff --git a/tests/C-tests/CMakeLists.txt b/tests/C-tests/CMakeLists.txt
index e47b04dda..be18903c2 100644
--- a/tests/C-tests/CMakeLists.txt
+++ b/tests/C-tests/CMakeLists.txt
@@ -128,6 +128,7 @@ set(testlist-mpi
     convert
     copy
     diagonalize
+    get_sparsity
     import_export
     introspection
     io_matrix
diff --git a/tests/C-tests/get_sparsity_typed.c b/tests/C-tests/get_sparsity_typed.c
index 15c3b997d..ac9825305 100644
--- a/tests/C-tests/get_sparsity_typed.c
+++ b/tests/C-tests/get_sparsity_typed.c
@@ -23,36 +23,49 @@ int TYPED_FUNC(
     double threshold = 0.1;
     REAL_T *A_dense;
 
-    A = bml_random_matrix(matrix_type, matrix_precision, N, M, sequential);
+    bml_distribution_mode_t distrib_mode = sequential;
+#ifdef DO_MPI
+    if (bml_getNRanks() > 1)
+    {
+        LOG_INFO("Use distributed matrix\n");
+        distrib_mode = distributed;
+    }
+#endif
+
+    A = bml_random_matrix(matrix_type, matrix_precision, N, M, distrib_mode);
 
     sparsity = bml_get_sparsity(A, threshold);
 
     // Form bml to dense_array
     A_dense = bml_export_to_dense(A, dense_row_major);
 
-    for (int i = 0; i < N; i++)
+    if (bml_getMyRank() == 0)
     {
-        for (int j = 0; j < N; j++)
+        for (int i = 0; i < N; i++)
         {
-            if (ABS(A_dense[ROWMAJOR(i, j, N, N)]) > threshold)
+            for (int j = 0; j < N; j++)
             {
-                nnzs++;
+                if (ABS(A_dense[ROWMAJOR(i, j, N, N)]) > threshold)
+                {
+                    nnzs++;
+                }
             }
         }
-    }
 
-    sparsity_ref = (1.0 - (double) nnzs / ((double) (N * N)));
+        sparsity_ref = (1.0 - (double) nnzs / ((double) (N * N)));
 
-    if (fabs(sparsity - sparsity_ref) > 1e-12)
-    {
-        LOG_ERROR("bml_get_sparsity is corrupted\n");
-        return -1;
-    }
+        if (fabs(sparsity - sparsity_ref) > 1e-12)
+        {
+            LOG_ERROR("bml_get_sparsity is corrupted\n");
+            return -1;
+        }
 
-    printf("Sparsity = %f\n", sparsity);
+        printf("Sparsity = %f\n", sparsity);
+    }
 
     bml_deallocate(&A);
-    bml_free_memory(A_dense);
+    if (bml_getMyRank() == 0)
+        bml_free_memory(A_dense);
 
     LOG_INFO("bml_get_sparsity passed\n");