Python Tests: w/o Managed Memory (#340)

* Python Tests: w/o Managed Memory * pyAMReX: dtoh copy functions * Limit D2H Copies to GPU runs
ECP-WarpX · Apr 13, 2023 · 30f93ec · 30f93ec
1 parent 2bf9a6d
commit 30f93ec
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 6 deletions.
diff --git a/cmake/dependencies/pyAMReX.cmake b/cmake/dependencies/pyAMReX.cmake
@@ -76,7 +76,7 @@ option(ImpactX_pyamrex_internal "Download & build pyAMReX" ON)
 set(ImpactX_pyamrex_repo "https://github.com/AMReX-Codes/pyamrex.git"
     CACHE STRING
     "Repository URI to pull and build pyamrex from if(ImpactX_pyamrex_internal)")
-set(ImpactX_pyamrex_branch "23.04"
+set(ImpactX_pyamrex_branch "bc24ec54bbb982ce3c388b1caf5b9a0a573ce0a8"
     CACHE STRING
     "Repository branch for ImpactX_pyamrex_repo if(ImpactX_pyamrex_internal)")
 

diff --git a/tests/python/conftest.py b/tests/python/conftest.py
@@ -26,8 +26,8 @@ def amrex_init():
             "amrex.signal_handling=0",
             # abort GPU runs if out-of-memory instead of swapping to host RAM
             "amrex.abort_on_out_of_gpu_memory=1",
-            # allow implicit host-device memory transfers
-            "amrex.the_arena_is_managed=1",
+            # do not rely on implicit host-device memory transfers
+            "amrex.the_arena_is_managed=0",
         ]
     )
     yield

diff --git a/tests/python/test_charge_deposition.py b/tests/python/test_charge_deposition.py
@@ -56,14 +56,28 @@ def test_charge_deposition(save_png=True):
 
     half_z = sim.n_cell[2] // 2  # order: x,y,z
 
+    # for GPU runs, copy data from device to host for plotting
+    if impactx.Config.have_gpu:
+        rho_host = amrex.MultiFab(
+            rho.box_array(),
+            rho.dm(),
+            rho.n_comp(),
+            rho.n_grow_vect(),
+            amrex.MFInfo().set_arena(amrex.The_Pinned_Arena()),
+        )
+        amrex.dtoh_memcpy(rho_host, rho)
+    else:
+        rho_host = rho
+
+    # plot data slices
     f = plt.figure()
     ax = f.gca()
-    ng = rho.nGrowVect
-    for mfi in rho:
+    ng = rho_host.nGrowVect
+    for mfi in rho_host:
         bx = mfi.validbox()
         rbx = amrex.RealBox(bx, dr, gm.ProbLo())
 
-        arr = rho.array(mfi)  # TODO: explicit device-to-host memory copy
+        arr = rho_host.array(mfi)
         arr_np = np.array(arr, copy=False)  # indices: comp, z, y, x
 
         # shift box to zero-based local mfi index space