From 746f327d3f3f28f9780236cabe0de88ee1eb5afb Mon Sep 17 00:00:00 2001
From: Ed Hartnett <edwardjameshartnett@gmail.com>
Date: Fri, 5 Jul 2019 15:30:21 -0600
Subject: [PATCH] added nc_init_decomp

---
 src/clib/pio.h            |  4 ++++
 src/ncint/ncint_pio.c     | 17 +++++++++++++++++
 tests/ncint/tst_pio_udf.c | 29 ++++++++++++++++++++++-------
 3 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/src/clib/pio.h b/src/clib/pio.h
index 8c74a375e87..94c2f1281e0 100644
--- a/src/clib/pio.h
+++ b/src/clib/pio.h
@@ -1252,6 +1252,10 @@ extern "C" {
 
     int nc_free_iosystem(int iosysid);
 
+    int nc_init_decomp(int iosysid, int pio_type, int ndims, const int *gdimlen,
+                       int maplen, const size_t *compmap, int *ioidp,
+                       int rearranger, const size_t *iostart,
+                       const size_t *iocount);
 #if defined(__cplusplus)
 }
 #endif
diff --git a/src/ncint/ncint_pio.c b/src/ncint/ncint_pio.c
index 32ae5396699..495d593a023 100644
--- a/src/ncint/ncint_pio.c
+++ b/src/ncint/ncint_pio.c
@@ -44,3 +44,20 @@ nc_free_iosystem(int iosysid)
 {
     return PIOc_free_iosystem(iosysid);
 }
+
+/**
+ * Same as PIOc_init_decomp()
+ *
+ * @author Ed Hartnett
+ */
+int
+nc_init_decomp(int iosysid, int pio_type, int ndims, const int *gdimlen,
+               int maplen, const size_t *compmap, int *ioidp,
+               int rearranger, const size_t *iostart,
+               const size_t *iocount)
+{
+    return PIOc_init_decomp(iosysid, pio_type, ndims, gdimlen, maplen,
+                            (const PIO_Offset *)compmap, ioidp, rearranger,
+                            (const PIO_Offset *)iostart,
+                            (const PIO_Offset *)iocount);
+}
diff --git a/tests/ncint/tst_pio_udf.c b/tests/ncint/tst_pio_udf.c
index 8564778a202..9d79ae573ed 100644
--- a/tests/ncint/tst_pio_udf.c
+++ b/tests/ncint/tst_pio_udf.c
@@ -37,16 +37,18 @@ main(int argc, char **argv)
     printf("\n*** Testing netCDF integration layer.\n");
     printf("*** testing simple use of netCDF integration layer format...");
     {
-        int ncid;
+        int ncid, ioid;
         int dimid[NDIM2], varid;
+        int dimlen[NDIM2] = {DIM_LEN_X, DIM_LEN_Y};
         int iosysid;
         NC_Dispatch *disp_in;
+        size_t elements_per_pe;
+        size_t *compdof; /* The decomposition mapping. */
+        int i;
 
-        /* Create an empty file to play with. */
-        /* if (nc_create(FILE_NAME, NC_CLOBBER, &ncid)) ERR; */
-        /* if (nc_close(ncid)) ERR; */
-
+        /* Turn on logging for PIO library. */
         PIOc_set_log_level(3);
+
         /* Initialize the intracomm. */
         if (nc_init_intracomm(MPI_COMM_WORLD, 1, 1, 0, 0, &iosysid)) ERR;
 
@@ -55,9 +57,22 @@ main(int argc, char **argv)
 
         /* Create a file to play with. */
         if (nc_create(FILE_NAME, NC_UDF0, &ncid)) ERR;
-        if (nc_def_dim(ncid, DIM_NAME_X, DIM_LEN_X, &dimid[0])) ERR;
-        if (nc_def_dim(ncid, DIM_NAME_Y, DIM_LEN_Y, &dimid[1])) ERR;
+        if (nc_def_dim(ncid, DIM_NAME_X, dimlen[0], &dimid[0])) ERR;
+        if (nc_def_dim(ncid, DIM_NAME_Y, dimlen[1], &dimid[1])) ERR;
         if (nc_def_var(ncid, VAR_NAME, NC_INT, NDIM2, dimid, &varid)) ERR;
+
+        /* Create a decomposition for distributed arrays. */
+        elements_per_pe = DIM_LEN_X * DIM_LEN_Y / ntasks;
+        if (!(compdof = malloc(elements_per_pe * sizeof(size_t))))
+            ERR;
+        for (i = 0; i < elements_per_pe; i++)
+            compdof[i] = my_rank * elements_per_pe + i;
+
+        /* Create the PIO decomposition for this test. */
+        if (nc_init_decomp(iosysid, PIO_INT, NDIM2, dimlen, elements_per_pe,
+                           compdof, &ioid, 0, NULL, NULL)) ERR;
+        free(compdof);
+
         if (nc_close(ncid)) ERR;
 
         /* Check that our user-defined format has been added. */