Allow keeping q_cons_ts(2) on CPU using pinned allocations

Author: ntselepidis

src/simulation/m_time_steppers.fpp

@@ -75,8 +75,14 @@ module m_time_steppers
     integer, private :: num_ts !<
     !! Number of time stages in the time-stepping scheme
 
+    integer, private :: out_of_core
+
     $:GPU_DECLARE(create='[q_cons_ts,q_prim_vf,q_T_sf,rhs_vf,q_prim_ts,rhs_mv,rhs_pb,max_dt]')
 
+#ifdef __NVCOMPILER_GPU_UNIFIED_MEM
+    real(wp), allocatable, dimension(:, :, :, :), pinned, target :: q_cons_ts_pool_host
+#endif
+
 contains
 
     !> The computation of parameters, the allocation of memory,
@@ -86,6 +92,21 @@ contains
 
         integer :: i, j !< Generic loop iterators
 
+        character(len=10) :: out_of_core_str
+        out_of_core = 0
+
+#ifdef __NVCOMPILER_GPU_UNIFIED_MEM
+        call get_environment_variable("MFC_OUT_OF_CORE", out_of_core_str)
+
+        if (trim(out_of_core_str) == "0") then
+            out_of_core = 0
+        elseif (trim(out_of_core_str) == "1") then
+            out_of_core = 1
+        else ! default
+            out_of_core = 0
+        endif
+#endif
+
         ! Setting number of time-stages for selected time-stepping scheme
         if (time_stepper == 1) then
             num_ts = 1
@@ -102,12 +123,33 @@ contains
             @:PREFER_GPU(q_cons_ts(i)%vf)
         end do
 
+#ifdef __NVCOMPILER_GPU_UNIFIED_MEM
+        if ( out_of_core == 1 ) then
+           allocate(q_cons_ts_pool_host(idwbuff(1)%beg:idwbuff(1)%end, &
+                                        idwbuff(2)%beg:idwbuff(2)%end, &
+                                        idwbuff(3)%beg:idwbuff(3)%end, &
+                                        1:sys_size))
+        end if
+#endif
+
         do i = 1, num_ts
             do j = 1, sys_size
-                @:ALLOCATE(q_cons_ts(i)%vf(j)%sf(idwbuff(1)%beg:idwbuff(1)%end, &
-                    idwbuff(2)%beg:idwbuff(2)%end, &
-                    idwbuff(3)%beg:idwbuff(3)%end))
-                @:PREFER_GPU(q_cons_ts(i)%vf(j)%sf)
+#ifdef __NVCOMPILER_GPU_UNIFIED_MEM
+                if ( i <= (num_ts - out_of_core) ) then
+                    !print*, "q_cons_ts", i, j, "on GPU"
+#endif
+                    @:ALLOCATE(q_cons_ts(i)%vf(j)%sf(idwbuff(1)%beg:idwbuff(1)%end, &
+                        idwbuff(2)%beg:idwbuff(2)%end, &
+                        idwbuff(3)%beg:idwbuff(3)%end))
+                    @:PREFER_GPU(q_cons_ts(i)%vf(j)%sf)
+#ifdef __NVCOMPILER_GPU_UNIFIED_MEM
+                else
+                    !print*, "q_cons_ts", i, j, "on CPU"
+                    q_cons_ts(i)%vf(j)%sf(idwbuff(1)%beg:idwbuff(1)%end, &
+                        idwbuff(2)%beg:idwbuff(2)%end, &
+                        idwbuff(3)%beg:idwbuff(3)%end) => q_cons_ts_pool_host(:,:,:,j)
+                end if
+#endif
             end do
             @:ACC_SETUP_VFs(q_cons_ts(i))
         end do
@@ -1205,14 +1247,30 @@ contains
         ! Deallocating the cell-average conservative variables
         do i = 1, num_ts
             do j = 1, sys_size
-                @:DEALLOCATE(q_cons_ts(i)%vf(j)%sf)
+#ifdef __NVCOMPILER_GPU_UNIFIED_MEM
+                if ( i <= (num_ts - out_of_core) ) then
+                    !print*, "q_cons_ts", i, j, "dealloc"
+#endif
+                    @:DEALLOCATE(q_cons_ts(i)%vf(j)%sf)
+#ifdef __NVCOMPILER_GPU_UNIFIED_MEM
+                else
+                    !print*, "q_cons_ts", i, j, "nullify"
+                    nullify(q_cons_ts(i)%vf(j)%sf)
+                end if
+#endif
             end do
 
             @:DEALLOCATE(q_cons_ts(i)%vf)
         end do
 
         @:DEALLOCATE(q_cons_ts)
 
+#ifdef __NVCOMPILER_GPU_UNIFIED_MEM
+        if ( out_of_core == 1 ) then
+            deallocate(q_cons_ts_pool_host)
+        end if
+#endif
+
         ! Deallocating the cell-average primitive ts variables
         if (probe_wrt) then
             do i = 0, 3

toolchain/templates/santis.mako

@@ -38,10 +38,9 @@ export FI_MR_CACHE_MONITOR=disabled
 export MPICH_NO_BUFFER_ALIAS_CHECK=1
 
 # CUSTOM env vars to MFC
-export NVIDIA_ALLOC_MODE=0                    # do nothing
+export MFC_OUT_OF_CORE=1                      # out of core
 export NVIDIA_MANUAL_GPU_HINTS=1              # prefloc GPU on some
 export NVIDIA_IGR_TEMPS_ON_GPU=3              # jac, jac_rhs, and jac_old on GPU
-export NVIDIA_VARS_ON_GPU=7                   # q_cons_ts(1)%vf%sf for j=1-7 on GPU
 
 # NSYS
 export NSYS=1                                 # enable nsys profiling