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.
pmodels · Oct 2, 2024 · c2df86d · c2df86d
1 parent d6a78a9
commit c2df86d
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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