From c2df86df4878269933bba40f1c808e6259de7fa2 Mon Sep 17 00:00:00 2001
From: Hui Zhou <hzhou321@anl.gov>
Date: Tue, 1 Oct 2024 18:17:18 -0500
Subject: [PATCH] test/bench: add support for device memory

Add device memory support using mtest_common utilities. This will add
the dependency to utitlity libraries, which the makefile already
imports.

However, this will remove the simpliicity of building single
source with mpicc or mydef_run. If one doesn't need test device memory,
one can simply comment off "$include macros/mtest.def" to restore the
simplicity.
---
 test/mpi/bench/macros/bench_frame.def | 22 ++++++++++++++++++----
 test/mpi/bench/macros/mtest.def       | 14 ++++++++++++++
 test/mpi/bench/p2p_bw.def             |  1 +
 test/mpi/bench/p2p_latency.def        |  1 +
 4 files changed, 34 insertions(+), 4 deletions(-)
 create mode 100644 test/mpi/bench/macros/mtest.def

diff --git a/test/mpi/bench/macros/bench_frame.def b/test/mpi/bench/macros/bench_frame.def
index 7c943fe7714..c34f00f4607 100644
--- a/test/mpi/bench/macros/bench_frame.def
+++ b/test/mpi/bench/macros/bench_frame.def
@@ -9,10 +9,16 @@
 subcode: bench_frame
     $include stdio
     $include stdlib
-    $include mpi
+    $(if:HAS_MTEST)
+        $include mpitest.h
+    $(else)
+        $include mpi
 
     $function main
-        MPI_Init(NULL, NULL);
+        $(if:HAS_MTEST)
+            MTest_Init(NULL, NULL);
+        $(else)
+            MPI_Init(NULL, NULL);
 
         $my grank, gsize: int
         MPI_Comm_rank(MPI_COMM_WORLD, &grank);
@@ -23,7 +29,12 @@ subcode: bench_frame
                 return 1
 
         MPI_Comm comm = MPI_COMM_WORLD;
-        char *buf = malloc(MAX_BUFSIZE)
+
+        $my void *buf
+        $(if:HAS_MTEST)
+            $call mtest_malloc, MAX_BUFSIZE
+        $(else)
+            buf = malloc(MAX_BUFSIZE)
         $if !buf
             printf("! Failed to allocate buffer (size=%d)\n", MAX_BUFSIZE)
             return 1
@@ -35,7 +46,10 @@ subcode: bench_frame
         $if grank == 0
             printf("\n")
 
-        MPI_Finalize();
+        $(if:HAS_MTEST)
+            MTest_Finalize(0);
+        $(else)
+            MPI_Finalize();
 
 macros:
     use_double: 1
diff --git a/test/mpi/bench/macros/mtest.def b/test/mpi/bench/macros/mtest.def
new file mode 100644
index 00000000000..185a35bf77e
--- /dev/null
+++ b/test/mpi/bench/macros/mtest.def
@@ -0,0 +1,14 @@
+macros:
+    HAS_MTEST: 1
+
+subcode: mtest_malloc(size)
+    MTestArgList *head = MTestArgListCreate(argc, argv)
+    int send_rank = 0, recv_rank = 1;
+    $(for:a in send,recv)
+        $if grank == $(a)_rank
+            $my mtest_mem_type_e $(a)_memtype, int $(a)_device
+            $(a)_memtype = MTestArgListGetMemType(head, "$(a)mem")
+            $(a)_device = MTestArgListGetInt_with_default(head, "$(a)dev", 0)
+            MTestMalloc($(size), $(a)_memtype, NULL, &buf, $(a)_device)
+            MTestPrintfMsg(1, "Allocating buffer: memtype=%s, device=%d, size=%d\n", MTest_memtype_name($(a)_memtype), $(a)_device, $(size))
+    MTestArgListDestroy(head)
diff --git a/test/mpi/bench/p2p_bw.def b/test/mpi/bench/p2p_bw.def
index b0c55e7fc94..ebd5e94ab55 100644
--- a/test/mpi/bench/p2p_bw.def
+++ b/test/mpi/bench/p2p_bw.def
@@ -1,5 +1,6 @@
 include: macros/bench_frame.def
 include: macros/bench_p2p.def
+include: macros/mtest.def
 
 subcode: _autoload
     $define WINDOW_SIZE 64
diff --git a/test/mpi/bench/p2p_latency.def b/test/mpi/bench/p2p_latency.def
index 3fa94d547eb..4d18cce0efd 100644
--- a/test/mpi/bench/p2p_latency.def
+++ b/test/mpi/bench/p2p_latency.def
@@ -1,5 +1,6 @@
 include: macros/bench_frame.def
 include: macros/bench_p2p.def
+include: macros/mtest.def
 
 page: p2p_latency, bench_frame
     MULTIPLICITY: 2