test: use refcount to initialize/finalize Argobots runtime

test/mpi/util/mtest_thread_abt.h

@@ -47,6 +47,11 @@ static void MTest_abt_error(int err, const char *msg, const char *file, int line
 
 #define MTEST_NUM_XSTREAMS 2
 
+/* Reference counter */
+
+static void MTest_abt_incr_refcount(void);
+static void MTest_abt_decr_refcount(void);
+
 /* -------------------------------------------- */
 extern ABT_pool pools[MTEST_NUM_XSTREAMS];
 
@@ -58,6 +63,12 @@ int MTest_Start_thread(MTEST_THREAD_RETURN_TYPE(*fn) (void *p), void *arg)
         fprintf(stderr, "Too many threads already created: max is %d\n", MTEST_MAX_THREADS);
         return 1;
     }
+
+    if (nthreads == 0) {
+        /* MTest_Join_threads() is associated with the first MTest_Start_thread(). */
+        MTest_abt_incr_refcount();
+    }
+
     /* We push threads to pools[0] and let the random work-stealing
      * scheduler balance things out. */
     ret = ABT_thread_create(pools[0], fn, arg, ABT_THREAD_ATTR_NULL, &threads[nthreads]);
@@ -76,12 +87,14 @@ int MTest_Join_threads(void)
         MTEST_ABT_ERROR(ret, "ABT_thread_free");
     }
     nthreads = 0;
+    MTest_abt_incr_refcount();
     return 0;
 }
 
 int MTest_thread_lock_create(MTEST_THREAD_LOCK_TYPE * lock)
 {
     int ret;
+    MTest_abt_incr_refcount();
     ret = ABT_mutex_create(lock);
     MTEST_ABT_ERROR(ret, "ABT_mutex_create");
     return 0;
@@ -108,13 +121,17 @@ int MTest_thread_lock_free(MTEST_THREAD_LOCK_TYPE * lock)
     int ret;
     ret = ABT_mutex_free(lock);
     MTEST_ABT_ERROR(ret, "ABT_mutex_free");
+    MTest_abt_decr_refcount();
     return 0;
 }
 
 int MTest_thread_yield(void)
 {
     int ret;
+    /* ABT_thread_yield() requires the Argobots context. */
+    MTest_abt_incr_refcount();
     ret = ABT_thread_yield();
+    MTest_abt_decr_refcount();
     MTEST_ABT_ERROR(ret, "ABT_thread_yield");
     return 0;
 }
@@ -127,6 +144,7 @@ static ABT_barrier barrier;
 static int bcount = -1;
 int MTest_thread_barrier_init(void)
 {
+    MTest_abt_incr_refcount();
     bcount = -1;        /* must reset to force barrier re-creation */
     return MTest_thread_lock_create(&barrierLock);
 }
@@ -137,6 +155,7 @@ int MTest_thread_barrier_free(void)
     MTest_thread_lock_free(&barrierLock);
     ret = ABT_barrier_free(&barrier);
     MTEST_ABT_ERROR(ret, "ABT_barrier_free");
+    MTest_abt_decr_refcount();
     return 0;
 }
 
@@ -171,18 +190,17 @@ int MTest_thread_barrier(int nt)
 }
 
 /* -------------------------------------------- */
-#define HAVE_MTEST_INIT_THREAD_PKG
 
 static ABT_xstream xstreams[MTEST_NUM_XSTREAMS];
 static ABT_sched scheds[MTEST_NUM_XSTREAMS];
 ABT_pool pools[MTEST_NUM_XSTREAMS];
 
-void MTest_init_thread_pkg(int argc, char **argv)
+void MTest_abt_init(void)
 {
     int i, k, ret;
     int num_xstreams;
 
-    ABT_init(argc, argv);
+    ABT_init(0, NULL);
 
     /* Create pools */
     for (i = 0; i < MTEST_NUM_XSTREAMS; i++) {
@@ -214,7 +232,7 @@ void MTest_init_thread_pkg(int argc, char **argv)
     }
 }
 
-void MTest_finalize_thread_pkg()
+void MTest_abt_finalize(void)
 {
     int i, ret;
     int num_xstreams = MTEST_NUM_XSTREAMS;
@@ -229,3 +247,39 @@ void MTest_finalize_thread_pkg()
     ret = ABT_finalize();
     MTEST_ABT_ERROR(ret, "ABT_finalize");
 }
+
+char MTest_argobots_lock = 0;   /* Atomic */
+int MTest_argobots_refcount = 0;        /* Protected by MTest_argobots_lock */
+
+static void MTest_abt_acquire_spinlock(void)
+{
+    /* TODO: support platforms that do not support GCC's __atomic builtins. */
+    while (!__atomic_test_and_set(&MTest_argobots_lock, __ATOMIC_ACQ_REL));
+}
+
+static void MTest_abt_release_spinlock(void)
+{
+    __atomic_clear(&MTest_argobots_lock, __ATOMIC_RELEASE);
+}
+
+static void MTest_abt_incr_refcount(void)
+{
+    MTest_abt_acquire_spinlock();
+    if (MTest_argobots_refcount == 0) {
+        /* Initialize Argobots if refcount is zero. */
+        MTest_abt_init();
+    }
+    MTest_argobots_refcount++;
+    MTest_abt_release_spinlock();
+}
+
+static void MTest_abt_decr_refcount(void)
+{
+    MTest_abt_acquire_spinlock();
+    MTest_argobots_refcount--;
+    if (MTest_argobots_refcount == 0) {
+        /* Finalize Argobots if refcount becomes zero. */
+        MTest_abt_finalize();
+    }
+    MTest_abt_release_spinlock();
+}