m_time_steppers.fpp

!>
!! @file
!! @brief Contains module m_time_steppers

#:include 'macros.fpp'
#:include 'case.fpp'

!> @brief Total-variation-diminishing (TVD) Runge--Kutta time integrators (1st-, 2nd-, and 3rd-order SSP)
module m_time_steppers

    use m_derived_types
    use m_global_parameters
    use m_rhs
    use m_pressure_relaxation
    use m_data_output
    use m_bubbles_EE
    use m_bubbles_EL
    use m_ibm
    use m_hyperelastic
    use m_mpi_proxy
    use m_boundary_common
    use m_helper
    use m_sim_helpers
    use m_fftw
    use m_nvtx
    use m_thermochem, only: num_species
    use m_body_forces
    use m_derived_variables

    implicit none

    type(vector_field), allocatable, dimension(:)    :: q_cons_ts  !< Cell-average conservative variables at each time-stage (TS)
    type(scalar_field), allocatable, dimension(:)    :: q_prim_vf  !< Cell-average primitive variables at the current time-stage
    type(scalar_field), allocatable, dimension(:)    :: rhs_vf     !< Cell-average RHS variables at the current time-stage
    type(integer_field), allocatable, dimension(:,:) :: bc_type    !< Boundary condition identifiers
    !> Cell-average primitive variables at consecutive TIMESTEPS
    type(vector_field), allocatable, dimension(:) :: q_prim_ts1, q_prim_ts2
    real(wp), allocatable, dimension(:,:,:,:,:)   :: rhs_pb
    type(scalar_field)                            :: q_T_sf  !< Cell-average temperature variables at the current time-stage
    real(wp), allocatable, dimension(:,:,:,:,:)   :: rhs_mv
    real(wp), allocatable, dimension(:,:,:)       :: max_dt
    integer, private                              :: num_ts  !< Number of time stages in the time-stepping scheme
    integer                                       :: stor    !< storage index
    real(wp), allocatable, dimension(:,:)         :: rk_coef
    integer, private                              :: num_probe_ts

    $:GPU_DECLARE(create='[q_cons_ts, q_prim_vf, q_T_sf, rhs_vf, q_prim_ts1, q_prim_ts2, rhs_mv, rhs_pb, max_dt, rk_coef, stor, bc_type]')

    !> @cond
#if defined(__NVCOMPILER_GPU_UNIFIED_MEM)
    real(stp), allocatable, dimension(:,:,:,:), pinned, target :: q_cons_ts_pool_host
#elif defined(FRONTIER_UNIFIED)
    real(stp), pointer, contiguous, dimension(:,:,:,:) :: q_cons_ts_pool_host, q_cons_ts_pool_device
    integer(kind=8)                                    :: pool_dims(4), pool_starts(4)
    integer(kind=8)                                    :: pool_size
    type(c_ptr)                                        :: cptr_host, cptr_device
#endif
    !> @endcond

contains

    !> Initialize the time steppers module
    impure subroutine s_initialize_time_steppers_module

#ifdef FRONTIER_UNIFIED
        use hipfort
        use hipfort_hipmalloc
        use hipfort_check
#if defined(MFC_OpenACC)
        use openacc
#endif
#endif
        integer :: i, j  !< Generic loop iterators
        ! Setting number of time-stages for selected time-stepping scheme

        if (time_stepper == 1) then
            num_ts = 1
        else if (any(time_stepper == (/2, 3/))) then
            num_ts = 2
        end if

        if (probe_wrt) then
            num_probe_ts = 2
        end if

        ! Allocating the cell-average conservative variables
        @:ALLOCATE(q_cons_ts(1:num_ts))
        @:PREFER_GPU(q_cons_ts)

        do i = 1, num_ts
            @:ALLOCATE(q_cons_ts(i)%vf(1:sys_size))
            @:PREFER_GPU(q_cons_ts(i)%vf)
        end do

        !> @cond
#if defined(__NVCOMPILER_GPU_UNIFIED_MEM)
        if (num_ts == 2 .and. nv_uvm_out_of_core) then
            ! host allocation for q_cons_ts(2)%vf(j)%sf for all j
            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

        do j = 1, sys_size
            ! q_cons_ts(1) lives on the device
            @:ALLOCATE(q_cons_ts(1)%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(1)%vf(j)%sf)
            if (num_ts == 2) then
                if (nv_uvm_out_of_core) then
                    ! q_cons_ts(2) lives on the host
                    q_cons_ts(2)%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)
                else
                    @:ALLOCATE(q_cons_ts(2)%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(2)%vf(j)%sf)
                end if
            end if
        end do

        do i = 1, num_ts
            @:ACC_SETUP_VFs(q_cons_ts(i))
        end do
#elif defined(FRONTIER_UNIFIED)
        ! Allocate to memory regions using hip calls that we will attach pointers to
        do i = 1, 3
            pool_dims(i) = idwbuff(i)%end - idwbuff(i)%beg + 1
            pool_starts(i) = idwbuff(i)%beg
        end do
        pool_dims(4) = sys_size
        pool_starts(4) = 1
#ifdef MFC_MIXED_PRECISION
        pool_size = 1_8*(idwbuff(1)%end - idwbuff(1)%beg + 1)*(idwbuff(2)%end - idwbuff(2)%beg + 1)*(idwbuff(3)%end &
                         & - idwbuff(3)%beg + 1)*sys_size
        call hipCheck(hipMalloc_(cptr_device, pool_size*2_8))
        call c_f_pointer(cptr_device, q_cons_ts_pool_device, shape=pool_dims)
        q_cons_ts_pool_device(idwbuff(1)%beg:,idwbuff(2)%beg:,idwbuff(3)%beg:,1:) => q_cons_ts_pool_device

        call hipCheck(hipMallocManaged_(cptr_host, pool_size*2_8, hipMemAttachGlobal))
        call c_f_pointer(cptr_host, q_cons_ts_pool_host, shape=pool_dims)
        q_cons_ts_pool_host(idwbuff(1)%beg:,idwbuff(2)%beg:,idwbuff(3)%beg:,1:) => q_cons_ts_pool_host
#else
        ! Doing hipMalloc then mapping should be most performant
        call hipCheck(hipMalloc(q_cons_ts_pool_device, dims8=pool_dims, lbounds8=pool_starts))
        ! Without this map CCE will still create a device copy, because it's silly like that
#if defined(MFC_OpenACC)
        call acc_map_data(q_cons_ts_pool_device, c_loc(q_cons_ts_pool_device), c_sizeof(q_cons_ts_pool_device))
#endif
        ! CCE see it can access this and will leave it on the host. It will stay on the host so long as HSA_XNACK=1 NOTE: WE CANNOT
        ! DO ATOMICS INTO THIS MEMORY. We have to change a property to use atomics here Otherwise leaving this as fine-grained will
        ! actually help performance since it can't be cached in GPU L2
        if (num_ts == 2) then
            call hipCheck(hipMallocManaged(q_cons_ts_pool_host, dims8=pool_dims, lbounds8=pool_starts, flags=hipMemAttachGlobal))
#if defined(MFC_OpenMP)
            call hipCheck(hipMemAdvise(c_loc(q_cons_ts_pool_host), c_sizeof(q_cons_ts_pool_host), &
                          & hipMemAdviseSetPreferredLocation, -1))
#endif
        end if
#endif

        do j = 1, sys_size
            ! q_cons_ts(1) lives on the device
            q_cons_ts(1)%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_device(:,:,:,j)
            if (num_ts == 2) then
                ! q_cons_ts(2) lives on the host
                q_cons_ts(2)%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
        end do

        do i = 1, num_ts
            @:ACC_SETUP_VFs(q_cons_ts(i))
            do j = 1, sys_size
                $:GPU_UPDATE(device='[q_cons_ts(i)%vf(j)]')
            end do
        end do
#else
        !> @endcond
        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))
            end do
            @:ACC_SETUP_VFs(q_cons_ts(i))
        end do
        !> @cond
#endif
        !> @endcond

        ! Allocating the cell-average primitive ts variables
        if (probe_wrt) then
            @:ALLOCATE(q_prim_ts1(1:num_probe_ts))

            do i = 1, num_probe_ts
                @:ALLOCATE(q_prim_ts1(i)%vf(1:sys_size))
            end do

            do i = 1, num_probe_ts
                do j = 1, sys_size
                    @:ALLOCATE(q_prim_ts1(i)%vf(j)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, &
                               & idwbuff(3)%beg:idwbuff(3)%end))
                end do
                @:ACC_SETUP_VFs(q_prim_ts1(i))
            end do

            @:ALLOCATE(q_prim_ts2(1:num_probe_ts))

            do i = 1, num_probe_ts
                @:ALLOCATE(q_prim_ts2(i)%vf(1:sys_size))
            end do

            do i = 1, num_probe_ts
                do j = 1, sys_size
                    @:ALLOCATE(q_prim_ts2(i)%vf(j)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, &
                               & idwbuff(3)%beg:idwbuff(3)%end))
                end do
                @:ACC_SETUP_VFs(q_prim_ts2(i))
            end do
        end if

        ! Allocating the cell-average primitive variables
        @:ALLOCATE(q_prim_vf(1:sys_size))

        if (.not. igr) then
            do i = 1, eqn_idx%adv%end
                @:ALLOCATE(q_prim_vf(i)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, &
                           & idwbuff(3)%beg:idwbuff(3)%end))
                @:ACC_SETUP_SFs(q_prim_vf(i))
            end do

            if (bubbles_euler) then
                do i = eqn_idx%bub%beg, eqn_idx%bub%end
                    @:ALLOCATE(q_prim_vf(i)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, &
                               & idwbuff(3)%beg:idwbuff(3)%end))
                    @:ACC_SETUP_SFs(q_prim_vf(i))
                end do
                if (adv_n) then
                    @:ALLOCATE(q_prim_vf(eqn_idx%n)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, &
                               & idwbuff(3)%beg:idwbuff(3)%end))
                    @:ACC_SETUP_SFs(q_prim_vf(eqn_idx%n))
                end if
            end if

            if (mhd) then
                do i = eqn_idx%B%beg, eqn_idx%B%end
                    @:ALLOCATE(q_prim_vf(i)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, &
                               & idwbuff(3)%beg:idwbuff(3)%end))
                    @:ACC_SETUP_SFs(q_prim_vf(i))
                end do
            end if

            if (elasticity) then
                do i = eqn_idx%stress%beg, eqn_idx%stress%end
                    @:ALLOCATE(q_prim_vf(i)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, &
                               & idwbuff(3)%beg:idwbuff(3)%end))
                    @:ACC_SETUP_SFs(q_prim_vf(i))
                end do
            end if

            if (hyperelasticity) then
                do i = eqn_idx%xi%beg, eqn_idx%xi%end + 1
                    @:ALLOCATE(q_prim_vf(i)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, &
                               & idwbuff(3)%beg:idwbuff(3)%end))
                    @:ACC_SETUP_SFs(q_prim_vf(i))
                end do
            end if

            if (cont_damage) then
                @:ALLOCATE(q_prim_vf(eqn_idx%damage)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, &
                           & idwbuff(3)%beg:idwbuff(3)%end))
                @:ACC_SETUP_SFs(q_prim_vf(eqn_idx%damage))
            end if

            if (hyper_cleaning) then
                @:ALLOCATE(q_prim_vf(eqn_idx%psi)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, &
                           & idwbuff(3)%beg:idwbuff(3)%end))
                @:ACC_SETUP_SFs(q_prim_vf(eqn_idx%psi))
            end if

            if (model_eqns == 3) then
                do i = eqn_idx%int_en%beg, eqn_idx%int_en%end
                    @:ALLOCATE(q_prim_vf(i)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, &
                               & idwbuff(3)%beg:idwbuff(3)%end))
                    @:ACC_SETUP_SFs(q_prim_vf(i))
                end do
            end if

            if (surface_tension) then
                @:ALLOCATE(q_prim_vf(eqn_idx%c)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, &
                           & idwbuff(3)%beg:idwbuff(3)%end))
                @:ACC_SETUP_SFs(q_prim_vf(eqn_idx%c))
            end if

            if (chemistry) then
                do i = eqn_idx%species%beg, eqn_idx%species%end
                    @:ALLOCATE(q_prim_vf(i)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, &
                               & idwbuff(3)%beg:idwbuff(3)%end))
                    @:ACC_SETUP_SFs(q_prim_vf(i))
                end do

                @:ALLOCATE(q_T_sf%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, idwbuff(3)%beg:idwbuff(3)%end))
                @:ACC_SETUP_SFs(q_T_sf)
            end if
        end if

        @:ALLOCATE(pb_ts(1:2))
        ! Initialize bubble variables pb and mv at all quadrature nodes for all R0 bins
        if (qbmm .and. (.not. polytropic)) then
            @:ALLOCATE(pb_ts(1)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, idwbuff(3)%beg:idwbuff(3)%end, &
                       & 1:nnode, 1:nb))
            @:ACC_SETUP_SFs(pb_ts(1))

            @:ALLOCATE(pb_ts(2)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, idwbuff(3)%beg:idwbuff(3)%end, &
                       & 1:nnode, 1:nb))
            @:ACC_SETUP_SFs(pb_ts(2))

            @:ALLOCATE(rhs_pb(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, idwbuff(3)%beg:idwbuff(3)%end, &
                       & 1:nnode, 1:nb))
        else if (qbmm .and. polytropic) then
            @:ALLOCATE(pb_ts(1)%sf(idwbuff(1)%beg:idwbuff(1)%beg + 1, idwbuff(2)%beg:idwbuff(2)%beg + 1, &
                       & idwbuff(3)%beg:idwbuff(3)%beg + 1, 1:nnode, 1:nb))
            @:ACC_SETUP_SFs(pb_ts(1))

            @:ALLOCATE(pb_ts(2)%sf(idwbuff(1)%beg:idwbuff(1)%beg + 1, idwbuff(2)%beg:idwbuff(2)%beg + 1, &
                       & idwbuff(3)%beg:idwbuff(3)%beg + 1, 1:nnode, 1:nb))
            @:ACC_SETUP_SFs(pb_ts(2))

            @:ALLOCATE(rhs_pb(idwbuff(1)%beg:idwbuff(1)%beg + 1, idwbuff(2)%beg:idwbuff(2)%beg + 1, &
                       & idwbuff(3)%beg:idwbuff(3)%beg + 1, 1:nnode, 1:nb))
        else
            @:ALLOCATE(pb_ts(1)%sf(0,0,0,0,0))
            @:ACC_SETUP_SFs(pb_ts(1))

            @:ALLOCATE(pb_ts(2)%sf(0,0,0,0,0))
            @:ACC_SETUP_SFs(pb_ts(2))

            @:ALLOCATE(rhs_pb(0,0,0,0,0))
        end if

        @:ALLOCATE(mv_ts(1:2))

        if (qbmm .and. (.not. polytropic)) then
            @:ALLOCATE(mv_ts(1)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, idwbuff(3)%beg:idwbuff(3)%end, &
                       & 1:nnode, 1:nb))
            @:ACC_SETUP_SFs(mv_ts(1))

            @:ALLOCATE(mv_ts(2)%sf(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, idwbuff(3)%beg:idwbuff(3)%end, &
                       & 1:nnode, 1:nb))
            @:ACC_SETUP_SFs(mv_ts(2))

            @:ALLOCATE(rhs_mv(idwbuff(1)%beg:idwbuff(1)%end, idwbuff(2)%beg:idwbuff(2)%end, idwbuff(3)%beg:idwbuff(3)%end, &
                       & 1:nnode, 1:nb))
        else if (qbmm .and. polytropic) then
            @:ALLOCATE(mv_ts(1)%sf(idwbuff(1)%beg:idwbuff(1)%beg + 1, idwbuff(2)%beg:idwbuff(2)%beg + 1, &
                       & idwbuff(3)%beg:idwbuff(3)%beg + 1, 1:nnode, 1:nb))
            @:ACC_SETUP_SFs(mv_ts(1))

            @:ALLOCATE(mv_ts(2)%sf(idwbuff(1)%beg:idwbuff(1)%beg + 1, idwbuff(2)%beg:idwbuff(2)%beg + 1, &
                       & idwbuff(3)%beg:idwbuff(3)%beg + 1, 1:nnode, 1:nb))
            @:ACC_SETUP_SFs(mv_ts(2))

            @:ALLOCATE(rhs_mv(idwbuff(1)%beg:idwbuff(1)%beg + 1, idwbuff(2)%beg:idwbuff(2)%beg + 1, &
                       & idwbuff(3)%beg:idwbuff(3)%beg + 1, 1:nnode, 1:nb))
        else
            @:ALLOCATE(mv_ts(1)%sf(0,0,0,0,0))
            @:ACC_SETUP_SFs(mv_ts(1))

            @:ALLOCATE(mv_ts(2)%sf(0,0,0,0,0))
            @:ACC_SETUP_SFs(mv_ts(2))

            @:ALLOCATE(rhs_mv(0,0,0,0,0))
        end if

        ! Allocating the cell-average RHS variables
        @:ALLOCATE(rhs_vf(1:sys_size))
        @:PREFER_GPU(rhs_vf)

        if (igr) then
            do i = 1, sys_size
                @:ALLOCATE(rhs_vf(i)%sf(-1:m+1,-1:n+1,-1:p+1))
                @:ACC_SETUP_SFs(rhs_vf(i))
                @:PREFER_GPU(rhs_vf(i)%sf)
            end do
        else
            do i = 1, sys_size
                @:ALLOCATE(rhs_vf(i)%sf(0:m, 0:n, 0:p))
                @:ACC_SETUP_SFs(rhs_vf(i))
            end do
        end if

        ! Opening and writing the header of the run-time information file
        if (proc_rank == 0 .and. run_time_info) then
            call s_open_run_time_information_file()
        end if

        if (cfl_dt) then
            @:ALLOCATE(max_dt(0:m, 0:n, 0:p))
        end if

        ! Allocating arrays to store the bc types
        @:ALLOCATE(bc_type(1:num_dims,1:2))

        @:ALLOCATE(bc_type(1,1)%sf(0:0,0:n,0:p))
        @:ALLOCATE(bc_type(1,2)%sf(0:0,0:n,0:p))
        #:if not MFC_CASE_OPTIMIZATION or num_dims > 1
            if (n > 0) then
                @:ALLOCATE(bc_type(2,1)%sf(-buff_size:m+buff_size,0:0,0:p))
                @:ALLOCATE(bc_type(2,2)%sf(-buff_size:m+buff_size,0:0,0:p))
                #:if not MFC_CASE_OPTIMIZATION or num_dims > 2
                    if (p > 0) then
                        @:ALLOCATE(bc_type(3,1)%sf(-buff_size:m+buff_size,-buff_size:n+buff_size,0:0))
                        @:ALLOCATE(bc_type(3,2)%sf(-buff_size:m+buff_size,-buff_size:n+buff_size,0:0))
                    end if
                #:endif
            end if
        #:endif

        do i = 1, num_dims
            do j = 1, 2
                @:ACC_SETUP_SFs(bc_type(i,j))
            end do
        end do

        if (any(time_stepper == (/1, 2, 3/))) then
            ! temporary array index for TVD RK
            if (time_stepper == 1) then
                stor = 1
            else
                stor = 2
            end if

            ! TVD RK coefficients
            @:ALLOCATE(rk_coef(time_stepper, 4))
            if (time_stepper == 1) then
                rk_coef(1,:) = (/1._wp, 0._wp, 1._wp, 1._wp/)
            else if (time_stepper == 2) then
                rk_coef(1,:) = (/1._wp, 0._wp, 1._wp, 1._wp/)
                rk_coef(2,:) = (/1._wp, 1._wp, 1._wp, 2._wp/)
            else if (time_stepper == 3) then
                rk_coef(1,:) = (/1._wp, 0._wp, 1._wp, 1._wp/)
                rk_coef(2,:) = (/1._wp, 3._wp, 1._wp, 4._wp/)
                rk_coef(3,:) = (/2._wp, 1._wp, 2._wp, 3._wp/)
            end if
            $:GPU_UPDATE(device='[rk_coef, stor]')
        end if

    end subroutine s_initialize_time_steppers_module

    !> Advance the solution one full step using a TVD Runge-Kutta time integrator
    impure subroutine s_tvd_rk(t_step, time_avg, nstage)

#ifdef _CRAYFTN
        ! DIR$ OPTIMIZE (-haggress)
#endif
        integer, intent(in)     :: t_step
        real(wp), intent(inout) :: time_avg
        integer, intent(in)     :: nstage
        integer                 :: i, j, k, l, q, s  !< Generic loop iterator
        real(wp)                :: start, finish

        call cpu_time(start)
        call nvtxStartRange("TIMESTEP")

        ! Adaptive dt: initial stage
        if (adap_dt) call s_adaptive_dt_bubble(1)

        do s = 1, nstage
            call s_compute_rhs(q_cons_ts(1)%vf, q_T_sf, q_prim_vf, bc_type, rhs_vf, pb_ts(1)%sf, rhs_pb, mv_ts(1)%sf, rhs_mv, &
                               & t_step, time_avg, s)

            if (s == 1) then
                if (run_time_info) then
                    if (igr) then
                        call s_write_run_time_information(q_cons_ts(1)%vf, t_step)
                    end if
                    if (.not. igr) then
                        call s_write_run_time_information(q_prim_vf, t_step)
                    end if
                end if

                if (probe_wrt) then
                    call s_time_step_cycling(t_step)
                    call s_compute_derived_variables(t_step, q_cons_ts(1)%vf, q_prim_ts1, q_prim_ts2)
                end if

                if (cfl_dt) then
                    if (mytime >= t_stop) return
                else
                    if (t_step == t_step_stop) return
                end if
            end if

            if (bubbles_lagrange .and. .not. adap_dt) call s_update_lagrange_tdv_rk(stage=s)
            $:GPU_PARALLEL_LOOP(collapse=4)
            do i = 1, sys_size
                do l = 0, p
                    do k = 0, n
                        do j = 0, m
                            if (s == 1 .and. nstage > 1) then
                                q_cons_ts(stor)%vf(i)%sf(j, k, l) = q_cons_ts(1)%vf(i)%sf(j, k, l)
                            end if
                            if (igr) then
                                q_cons_ts(1)%vf(i)%sf(j, k, l) = (rk_coef(s, 1)*q_cons_ts(1)%vf(i)%sf(j, k, l) + rk_coef(s, &
                                          & 2)*q_cons_ts(stor)%vf(i)%sf(j, k, l) + rk_coef(s, 3)*rhs_vf(i)%sf(j, k, &
                                          & l))/rk_coef(s, 4)
                            else
                                q_cons_ts(1)%vf(i)%sf(j, k, l) = (rk_coef(s, 1)*q_cons_ts(1)%vf(i)%sf(j, k, l) + rk_coef(s, &
                                          & 2)*q_cons_ts(stor)%vf(i)%sf(j, k, l) + rk_coef(s, 3)*dt*rhs_vf(i)%sf(j, k, &
                                          & l))/rk_coef(s, 4)
                            end if
                        end do
                    end do
                end do
            end do
            $:END_GPU_PARALLEL_LOOP()
            ! Evolve pb and mv for non-polytropic qbmm
            if (qbmm .and. (.not. polytropic)) then
                $:GPU_PARALLEL_LOOP(collapse=5)
                do i = 1, nb
                    do l = 0, p
                        do k = 0, n
                            do j = 0, m
                                do q = 1, nnode
                                    if (s == 1 .and. nstage > 1) then
                                        pb_ts(stor)%sf(j, k, l, q, i) = pb_ts(1)%sf(j, k, l, q, i)
                                        mv_ts(stor)%sf(j, k, l, q, i) = mv_ts(1)%sf(j, k, l, q, i)
                                    end if
                                    pb_ts(1)%sf(j, k, l, q, i) = (rk_coef(s, 1)*pb_ts(1)%sf(j, k, l, q, i) + rk_coef(s, &
                                          & 2)*pb_ts(stor)%sf(j, k, l, q, i) + rk_coef(s, 3)*dt*rhs_pb(j, k, l, q, i))/rk_coef(s, 4)
                                    mv_ts(1)%sf(j, k, l, q, i) = (rk_coef(s, 1)*mv_ts(1)%sf(j, k, l, q, i) + rk_coef(s, &
                                          & 2)*mv_ts(stor)%sf(j, k, l, q, i) + rk_coef(s, 3)*dt*rhs_mv(j, k, l, q, i))/rk_coef(s, 4)
                                end do
                            end do
                        end do
                    end do
                end do
                $:END_GPU_PARALLEL_LOOP()
            end if

            if (bodyForces) call s_apply_bodyforces(q_cons_ts(1)%vf, q_prim_vf, rhs_vf, rk_coef(s, 3)*dt/rk_coef(s, 4))

            if (grid_geometry == 3) call s_apply_fourier_filter(q_cons_ts(1)%vf)

            if (model_eqns == 3 .and. (.not. relax)) then
                call s_pressure_relaxation_procedure(q_cons_ts(1)%vf)
            end if

            if (adv_n) call s_comp_alpha_from_n(q_cons_ts(1)%vf)

            if (ib) then
                ! check if any IBMS are moving, and if so, update the markers, ghost points, levelsets, and levelset norms
                if (moving_immersed_boundary_flag) then
                    call s_propagate_immersed_boundaries(s)
                end if

                ! update the ghost fluid properties point values based on IB state
                if (qbmm .and. .not. polytropic) then
                    call s_ibm_correct_state(q_cons_ts(1)%vf, q_prim_vf, pb_ts(1)%sf, mv_ts(1)%sf)
                else
                    call s_ibm_correct_state(q_cons_ts(1)%vf, q_prim_vf)
                end if
            end if
        end do

        if (ib) then
            if (moving_immersed_boundary_flag) then
                call s_wrap_periodic_ibs()  ! wraps the positions of IBs to the local proc
                call s_handoff_ib_ownership()  ! recomputes which ranks own which IBs and communicate to neighbors
            else if (ib_state_wrt) then
                call s_compute_ib_forces(q_prim_vf, fluid_pp)
            end if
        end if

        ! Adaptive dt: final stage
        if (adap_dt) call s_adaptive_dt_bubble(3)

        call nvtxEndRange
        call cpu_time(finish)

        wall_time = abs(finish - start)

        if (t_step - t_step_start >= 2) then
            wall_time_avg = (wall_time + (t_step - t_step_start - 2)*wall_time_avg)/(t_step - t_step_start - 1)
        else
            wall_time_avg = 0._wp
        end if

    end subroutine s_tvd_rk

    !> Bubble source part in Strang operator splitting scheme
    impure subroutine s_adaptive_dt_bubble(stage)

        integer, intent(in) :: stage

        call s_convert_conservative_to_primitive_variables(q_cons_ts(1)%vf, q_T_sf, q_prim_vf, idwint)

        if (bubbles_euler) then
            call s_compute_bubble_EE_source(q_cons_ts(1)%vf, q_prim_vf, rhs_vf, divu)
            call s_comp_alpha_from_n(q_cons_ts(1)%vf)
        else if (bubbles_lagrange) then
            call s_populate_variables_buffers(bc_type, q_prim_vf, pb_ts(1)%sf, mv_ts(1)%sf, q_T_sf)
            call s_compute_bubble_EL_dynamics(q_prim_vf, stage)
            call s_transfer_data_to_tmp()
            call s_smear_voidfraction()
            if (stage == 3) then
                if (lag_params%write_bubbles_stats) call s_calculate_lag_bubble_stats()
                if (lag_params%write_bubbles) then
                    $:GPU_UPDATE(host='[gas_p, gas_mv, intfc_rad, intfc_vel]')
                    call s_write_lag_particles(mytime)
                end if
                call s_write_void_evol(mytime)
            end if
        end if

    end subroutine s_adaptive_dt_bubble

    !> Compute the global time step size from CFL stability constraints across all cells
    impure subroutine s_compute_dt()

        real(wp) :: rho  !< Cell-avg. density

        #:if not MFC_CASE_OPTIMIZATION and USING_AMD
            real(wp), dimension(3) :: vel    !< Cell-avg. velocity
            real(wp), dimension(3) :: alpha  !< Cell-avg. volume fraction
        #:else
            real(wp), dimension(num_vels)   :: vel    !< Cell-avg. velocity
            real(wp), dimension(num_fluids) :: alpha  !< Cell-avg. volume fraction
        #:endif
        real(wp)               :: vel_sum  !< Cell-avg. velocity sum
        real(wp)               :: pres     !< Cell-avg. pressure
        real(wp)               :: gamma    !< Cell-avg. sp. heat ratio
        real(wp)               :: pi_inf   !< Cell-avg. liquid stiffness function
        real(wp)               :: qv       !< Cell-avg. fluid reference energy
        real(wp)               :: c        !< Cell-avg. sound speed
        real(wp)               :: H        !< Cell-avg. enthalpy
        real(wp), dimension(2) :: Re       !< Cell-avg. Reynolds numbers
        real(wp)               :: dt_local
        integer                :: j, k, l  !< Generic loop iterators

        if (.not. igr) then
            call s_convert_conservative_to_primitive_variables(q_cons_ts(1)%vf, q_T_sf, q_prim_vf, idwint)
        end if

        $:GPU_PARALLEL_LOOP(collapse=3, private='[vel, alpha, Re, rho, vel_sum, pres, gamma, pi_inf, c, H, qv]')
        do l = 0, p
            do k = 0, n
                do j = 0, m
                    if (igr) then
                        call s_compute_enthalpy(q_cons_ts(1)%vf, pres, rho, gamma, pi_inf, Re, H, alpha, vel, vel_sum, qv, j, k, l)
                    else
                        call s_compute_enthalpy(q_prim_vf, pres, rho, gamma, pi_inf, Re, H, alpha, vel, vel_sum, qv, j, k, l)
                    end if

                    ! Compute mixture sound speed
                    call s_compute_speed_of_sound(pres, rho, gamma, pi_inf, H, alpha, vel_sum, 0._wp, c, qv)

                    call s_compute_dt_from_cfl(vel, c, max_dt, rho, Re, j, k, l)
                end do
            end do
        end do
        $:END_GPU_PARALLEL_LOOP()

        #:call GPU_PARALLEL(copyout='[dt_local]', copyin='[max_dt]')
            dt_local = minval(max_dt)
        #:endcall GPU_PARALLEL

        if (num_procs == 1) then
            dt = dt_local
        else
            call s_mpi_allreduce_min(dt_local, dt)
        end if

        $:GPU_UPDATE(device='[dt]')

    end subroutine s_compute_dt

    !> Apply the body forces source term at each Runge-Kutta stage
    subroutine s_apply_bodyforces(q_cons_vf, q_prim_vf_in, rhs_vf_in, ldt)

        type(scalar_field), dimension(1:sys_size), intent(inout) :: q_cons_vf
        type(scalar_field), dimension(1:sys_size), intent(in)    :: q_prim_vf_in
        type(scalar_field), dimension(1:sys_size), intent(inout) :: rhs_vf_in
        real(wp), intent(in)                                     :: ldt  !< local dt
        integer                                                  :: i, j, k, l

        call nvtxStartRange("RHS-BODYFORCES")
        call s_compute_body_forces_rhs(q_prim_vf_in, q_cons_vf, rhs_vf_in)

        $:GPU_PARALLEL_LOOP(collapse=4)
        do i = eqn_idx%mom%beg, eqn_idx%E
            do l = 0, p
                do k = 0, n
                    do j = 0, m
                        q_cons_vf(i)%sf(j, k, l) = q_cons_vf(i)%sf(j, k, l) + ldt*rhs_vf_in(i)%sf(j, k, l)
                    end do
                end do
            end do
        end do
        $:END_GPU_PARALLEL_LOOP()

        call nvtxEndRange

    end subroutine s_apply_bodyforces

    !> Update immersed boundary positions and velocities at the current Runge-Kutta stage
    subroutine s_propagate_immersed_boundaries(s)

        integer, intent(in) :: s
        integer             :: i
        integer             :: gbl_id  ! used for analytic ib patch motion

        call nvtxStartRange("PROPAGATE-IMMERSED-BOUNDARIES")

        if (moving_immersed_boundary_flag) call s_compute_ib_forces(q_prim_vf, fluid_pp)

        do i = 1, num_ibs
            if (s == 1) then
                patch_ib(i)%step_vel = patch_ib(i)%vel
                patch_ib(i)%step_angular_vel = patch_ib(i)%angular_vel
                patch_ib(i)%step_angles = patch_ib(i)%angles
                patch_ib(i)%step_x_centroid = patch_ib(i)%x_centroid
                patch_ib(i)%step_y_centroid = patch_ib(i)%y_centroid
                patch_ib(i)%step_z_centroid = patch_ib(i)%z_centroid
            end if

            ! Compute forces BEFORE the RK velocity blend so the device copy of patch_ib%vel matches the host (pre-blend) when
            ! velocity-dependent collision damping forces are evaluated on the GPU.
            if (patch_ib(i)%moving_ibm > 0) then
                patch_ib(i)%vel = (rk_coef(s, 1)*patch_ib(i)%step_vel + rk_coef(s, 2)*patch_ib(i)%vel)/rk_coef(s, 4)
                patch_ib(i)%angular_vel = (rk_coef(s, 1)*patch_ib(i)%step_angular_vel + rk_coef(s, &
                         & 2)*patch_ib(i)%angular_vel)/rk_coef(s, 4)

                if (patch_ib(i)%moving_ibm == 1) then
                    ! plug in analytic velocities for 1-way coupling, if it exists
                    @:mib_analytical()
                else if (patch_ib(i)%moving_ibm == 2) then  ! if we are using two-way coupling, apply force and torque
                    ! update the velocity from the force value
                    patch_ib(i)%vel = patch_ib(i)%vel + rk_coef(s, 3)*dt*(patch_ib(i)%force/patch_ib(i)%mass)/rk_coef(s, 4)

                    ! update the angular velocity with the torque value
                    patch_ib(i)%angular_vel = (patch_ib(i)%angular_vel*patch_ib(i)%moment) + (rk_coef(s, &
                             & 3)*dt*patch_ib(i)%torque/rk_coef(s, 4))  ! add the torque to the angular momentum
                    call s_compute_moment_of_inertia(i, patch_ib(i)%angular_vel)
                    ! update the moment of inertia to be based on the direction of the angular momentum
                    patch_ib(i)%angular_vel = patch_ib(i)%angular_vel/patch_ib(i)%moment
                end if

                ! Update the angle of the IB
                patch_ib(i)%angles = (rk_coef(s, 1)*patch_ib(i)%step_angles + rk_coef(s, 2)*patch_ib(i)%angles + rk_coef(s, &
                         & 3)*patch_ib(i)%angular_vel*dt)/rk_coef(s, 4)

                ! Update the position of the IB
                patch_ib(i)%x_centroid = (rk_coef(s, 1)*patch_ib(i)%step_x_centroid + rk_coef(s, &
                         & 2)*patch_ib(i)%x_centroid + rk_coef(s, 3)*patch_ib(i)%vel(1)*dt)/rk_coef(s, 4)
                patch_ib(i)%y_centroid = (rk_coef(s, 1)*patch_ib(i)%step_y_centroid + rk_coef(s, &
                         & 2)*patch_ib(i)%y_centroid + rk_coef(s, 3)*patch_ib(i)%vel(2)*dt)/rk_coef(s, 4)
                patch_ib(i)%z_centroid = (rk_coef(s, 1)*patch_ib(i)%step_z_centroid + rk_coef(s, &
                         & 2)*patch_ib(i)%z_centroid + rk_coef(s, 3)*patch_ib(i)%vel(3)*dt)/rk_coef(s, 4)
            end if
        end do

        $:GPU_UPDATE(device='[patch_ib]')
        call s_update_mib(num_ibs)

        call nvtxEndRange

    end subroutine s_propagate_immersed_boundaries

    !> Save the temporary q_prim_vf vector into q_prim_ts for use in p_main
    subroutine s_time_step_cycling(t_step)

        integer, intent(in) :: t_step
        integer             :: i, j, k, l  !< Generic loop iterator

        if (t_step == t_step_start) then
            $:GPU_PARALLEL_LOOP(collapse=4)
            do i = 1, sys_size
                do l = 0, p
                    do k = 0, n
                        do j = 0, m
                            q_prim_ts2(2)%vf(i)%sf(j, k, l) = q_prim_vf(i)%sf(j, k, l)
                        end do
                    end do
                end do
            end do
            $:END_GPU_PARALLEL_LOOP()
        else if (t_step == t_step_start + 1) then
            $:GPU_PARALLEL_LOOP(collapse=4)
            do i = 1, sys_size
                do l = 0, p
                    do k = 0, n
                        do j = 0, m
                            q_prim_ts2(1)%vf(i)%sf(j, k, l) = q_prim_vf(i)%sf(j, k, l)
                        end do
                    end do
                end do
            end do
            $:END_GPU_PARALLEL_LOOP()
        else if (t_step == t_step_start + 2) then
            $:GPU_PARALLEL_LOOP(collapse=4)
            do i = 1, sys_size
                do l = 0, p
                    do k = 0, n
                        do j = 0, m
                            q_prim_ts1(2)%vf(i)%sf(j, k, l) = q_prim_vf(i)%sf(j, k, l)
                        end do
                    end do
                end do
            end do
            $:END_GPU_PARALLEL_LOOP()
        else if (t_step == t_step_start + 3) then
            $:GPU_PARALLEL_LOOP(collapse=4)
            do i = 1, sys_size
                do l = 0, p
                    do k = 0, n
                        do j = 0, m
                            q_prim_ts1(1)%vf(i)%sf(j, k, l) = q_prim_vf(i)%sf(j, k, l)
                        end do
                    end do
                end do
            end do
            $:END_GPU_PARALLEL_LOOP()
        else  ! All other timesteps
            $:GPU_PARALLEL_LOOP(collapse=4)
            do i = 1, sys_size
                do l = 0, p
                    do k = 0, n
                        do j = 0, m
                            q_prim_ts2(2)%vf(i)%sf(j, k, l) = q_prim_ts2(1)%vf(i)%sf(j, k, l)
                            q_prim_ts2(1)%vf(i)%sf(j, k, l) = q_prim_ts1(2)%vf(i)%sf(j, k, l)
                            q_prim_ts1(2)%vf(i)%sf(j, k, l) = q_prim_ts1(1)%vf(i)%sf(j, k, l)
                            q_prim_ts1(1)%vf(i)%sf(j, k, l) = q_prim_vf(i)%sf(j, k, l)
                        end do
                    end do
                end do
            end do
            $:END_GPU_PARALLEL_LOOP()
        end if

    end subroutine s_time_step_cycling

    !> Module deallocation and/or disassociation procedures
    impure subroutine s_finalize_time_steppers_module

#ifdef FRONTIER_UNIFIED
        use hipfort
        use hipfort_hipmalloc
        use hipfort_check
#endif
        integer :: i, j  !< Generic loop iterators
        ! Deallocating the cell-average conservative variables
#if defined(__NVCOMPILER_GPU_UNIFIED_MEM)
        do j = 1, sys_size
            @:DEALLOCATE(q_cons_ts(1)%vf(j)%sf)
            if (num_ts == 2) then
                if (nv_uvm_out_of_core) then
                    nullify (q_cons_ts(2)%vf(j)%sf)
                else
                    @:DEALLOCATE(q_cons_ts(2)%vf(j)%sf)
                end if
            end if
        end do
        if (num_ts == 2 .and. nv_uvm_out_of_core) then
            deallocate (q_cons_ts_pool_host)
        end if
#elif defined(FRONTIER_UNIFIED)
        do i = 1, num_ts
            do j = 1, sys_size
                nullify (q_cons_ts(i)%vf(j)%sf)
            end do
        end do
#ifdef MFC_MIXED_PRECISION
        call hipCheck(hipHostFree_(c_loc(q_cons_ts_pool_host)))
        nullify (q_cons_ts_pool_host)
        call hipCheck(hipFree_(c_loc(q_cons_ts_pool_device)))
        nullify (q_cons_ts_pool_device)
#else
        call hipCheck(hipHostFree(q_cons_ts_pool_host))
        call hipCheck(hipFree(q_cons_ts_pool_device))
#endif
#else
        do i = 1, num_ts
            do j = 1, sys_size
                @:DEALLOCATE(q_cons_ts(i)%vf(j)%sf)
            end do
        end do
#endif
        do i = 1, num_ts
            @:DEALLOCATE(q_cons_ts(i)%vf)
        end do

        @:DEALLOCATE(q_cons_ts)

        ! Deallocating the cell-average primitive ts variables
        if (probe_wrt) then
            do i = 1, num_probe_ts
                do j = 1, sys_size
                    @:DEALLOCATE(q_prim_ts1(i)%vf(j)%sf,q_prim_ts2(i)%vf(j)%sf )
                end do
                @:DEALLOCATE(q_prim_ts1(i)%vf, q_prim_ts2(i)%vf)
            end do
            @:DEALLOCATE(q_prim_ts1, q_prim_ts2)
        end if

        if (.not. igr) then
            ! Deallocating the cell-average primitive variables
            do i = 1, eqn_idx%adv%end
                @:DEALLOCATE(q_prim_vf(i)%sf)
            end do

            if (mhd) then
                do i = eqn_idx%B%beg, eqn_idx%B%end
                    @:DEALLOCATE(q_prim_vf(i)%sf)
                end do
            end if

            if (elasticity) then
                do i = eqn_idx%stress%beg, eqn_idx%stress%end
                    @:DEALLOCATE(q_prim_vf(i)%sf)
                end do
            end if

            if (hyperelasticity) then
                do i = eqn_idx%xi%beg, eqn_idx%xi%end + 1
                    @:DEALLOCATE(q_prim_vf(i)%sf)
                end do
            end if

            if (cont_damage) then
                @:DEALLOCATE(q_prim_vf(eqn_idx%damage)%sf)
            end if

            if (hyper_cleaning) then
                @:DEALLOCATE(q_prim_vf(eqn_idx%psi)%sf)
            end if

            if (bubbles_euler) then
                do i = eqn_idx%bub%beg, eqn_idx%bub%end
                    @:DEALLOCATE(q_prim_vf(i)%sf)
                end do
            end if

            if (model_eqns == 3) then
                do i = eqn_idx%int_en%beg, eqn_idx%int_en%end
                    @:DEALLOCATE(q_prim_vf(i)%sf)
                end do
            end if
        end if

        @:DEALLOCATE(q_prim_vf)

        ! Deallocating the cell-average RHS variables
        do i = 1, sys_size
            @:DEALLOCATE(rhs_vf(i)%sf)
        end do

        @:DEALLOCATE(rhs_vf)

        ! Writing the footer of and closing the run-time information file
        if (proc_rank == 0 .and. run_time_info) then
            call s_close_run_time_information_file()
        end if

        if (chemistry) then
            @:DEALLOCATE(q_T_sf%sf)
        end if
        @:DEALLOCATE(pb_ts(1)%sf)
        @:DEALLOCATE(pb_ts(2)%sf)
        @:DEALLOCATE(rhs_pb)
        @:DEALLOCATE(pb_ts)
        @:DEALLOCATE(mv_ts(1)%sf)
        @:DEALLOCATE(mv_ts(2)%sf)
        @:DEALLOCATE(rhs_mv)
        @:DEALLOCATE(mv_ts)
        if (cfl_dt) then
            @:DEALLOCATE(max_dt)
        end if
        do i = 1, num_dims
            @:DEALLOCATE(bc_type(i,1)%sf)
            @:DEALLOCATE(bc_type(i,2)%sf)
        end do
        @:DEALLOCATE(bc_type)
        if (any(time_stepper == (/1, 2, 3/))) then
            @:DEALLOCATE(rk_coef)
        end if

    end subroutine s_finalize_time_steppers_module

end module m_time_steppers