checkpoint.f90

Go to the documentation of this file.

 
!
module simulation_checkpoint
  use num_types, only: rp, sp, dp
  use case, only: case_t
  use json_file_module, only: json_file
  use json_utils, only: json_get, json_get_or_default
  use scalar_scheme, only: scalar_scheme_t
  use time_state, only: time_state_t
  use chkp_output, only: chkp_output_t
  use field, only: field_t
  use field_list, only: field_list_t
  use logger, only: neko_log, log_size, neko_log_debug
  use mpi_f08, only: mpi_wtime, mpi_barrier
  use comm, only: neko_comm, pe_rank
  use utils, only: neko_error
  use math, only: copy, rzero
  use profiler, only: profiler_start_region, profiler_end_region
  use neko_config, only: neko_bcknd_device
  use device, only: device_memcpy, device_to_host, host_to_device
  use registry, only: neko_registry
  implicit none
  private
 
  type, public :: simulation_checkpoint_t
     private
 
     ! ----------------------------------------------------------------------- !
     ! User parameters
 
     logical :: enabled = .false.
     character(len=256) :: algorithm = "linear"
     character(len=256) :: filename = "forward_checkpoint"
     character(len=256) :: path = "checkpoints/"
     character(len=8) :: fmt = "chkp"
     integer :: n_saves_memory = 10
     logical :: keep_checkpoints = .false.
 
     ! Internal parameters
     integer :: n_saves_disc = 0
     integer :: n_timesteps = 0
     integer :: first_valid_timestep = 2
     integer :: loaded_checkpoint = -1
 
     ! Field pointers
     type(field_list_t) :: state_list
     type(host_array), dimension(:,:), allocatable :: state_storage
 
     ! Structures to hold the checkpoint data
     type(chkp_output_t) :: chkp_output
 
   contains
     generic, public :: init => init_from_json, init_from_components
     procedure, public, pass(this) :: init_from_json => &
          checkpoint_init_from_json
     procedure, public, pass(this) :: init_from_components => &
          checkpoint_init_from_components
     procedure, public, pass(this) :: free => checkpoint_free
     procedure, public, pass(this) :: reset => checkpoint_reset
     procedure, public, pass(this) :: save => checkpoint_save
     procedure, public, pass(this) :: restore => checkpoint_restore
 
     procedure, pass(this) :: save_data => checkpoint_save_data
     procedure, pass(this) :: load_data => checkpoint_load_data
  end type simulation_checkpoint_t
 
  type :: host_array
     real(kind=rp), allocatable :: data(:)
     integer :: size = 0
   contains
     procedure, pass(this) :: init => host_array_init
     procedure, pass(this) :: free => host_array_free
     procedure, pass(this) :: is_allocated => host_array_is_allocated
  end type host_array
 
  ! ========================================================================== !
  ! Module procedures for our algorithm implementations.
 
  interface
 
     module subroutine checkpoint_save_linear(this, neko_case)
       class(simulation_checkpoint_t), intent(inout) :: this
       class(case_t), intent(inout) :: neko_case
     end subroutine checkpoint_save_linear
 
     module subroutine checkpoint_restore_linear(this, neko_case, tstep)
       class(simulation_checkpoint_t), intent(inout) :: this
       class(case_t), target, intent(inout) :: neko_case
       integer, intent(in) :: tstep
     end subroutine checkpoint_restore_linear
  end interface
 
contains
 
  ! ========================================================================== !
  ! Initialization and deallocation
 
  subroutine checkpoint_init_from_json(this, neko_case, params)
    class(simulation_checkpoint_t), intent(inout) :: this
    class(case_t), target, intent(inout) :: neko_case
    type(json_file), target, intent(inout) :: params
    integer :: n_saves_memory
    character(len=:), allocatable :: path, filename, algorithm, fmt
    character(len=256), dimension(:), allocatable :: extra_field_names
    type(field_list_t) :: extra_fields
    type(field_t), pointer :: fi
    integer :: i
    logical :: enabled, keep_checkpoints
 
    call json_get_or_default(params, "enabled", enabled, .false.)
    if (.not. enabled) return
 
    call json_get_or_default(params, "algorithm", algorithm, "linear")
    call json_get_or_default(params, "n_memory", n_saves_memory, 10)
    call json_get_or_default(params, "path", path, "checkpoints/")
    call json_get_or_default(params, "filename", filename, "checkpoint")
    call json_get_or_default(params, "format", fmt, "chkp")
    call json_get_or_default(params, "keep_checkpoints", keep_checkpoints, &
         .false.)
 
    if ("extra_fields" .in. params) then
       allocate(extra_field_names(0))
       call json_get(params, "extra_fields", extra_field_names)
       call extra_fields%init(size(extra_field_names))
       do i = 1, size(extra_field_names)
          fi => neko_registry%get_field(extra_field_names(i))
          call extra_fields%assign(i, fi)
       end do
       ! Create a field list for the extra fields
       call this%init_from_components(neko_case, algorithm, n_saves_memory, &
            path, filename, fmt, keep_checkpoints, extra_fields)
    else
       ! Create a field list without the extra fields
       call this%init_from_components(neko_case, algorithm, n_saves_memory, &
            path, filename, fmt, keep_checkpoints)
    end if
 
  end subroutine checkpoint_init_from_json
 
  subroutine checkpoint_init_from_components(this, neko_case, algorithm, &
       n_saves_memory, path, filename, fmt, keep_checkpoints, extra_fields)
    class(simulation_checkpoint_t), intent(inout), target :: this
    class(case_t), target, intent(inout) :: neko_case
    character(len=*), optional, intent(in) :: algorithm
    integer, optional, intent(in) :: n_saves_memory
    character(len=*), optional, intent(in) :: path
    character(len=*), optional, intent(in) :: filename
    character(len=*), optional, intent(in) :: fmt
    logical, optional, intent(in) :: keep_checkpoints
    type(field_list_t), optional, intent(inout) :: extra_fields
    type(field_t), pointer :: si
    character(len=LOG_SIZE) :: msg
    integer :: i, n_states
    logical :: exists
 
    call this%free()
 
    ! Assign parameters from arguments or defaults
    this%enabled = .true.
    if (present(algorithm)) this%algorithm = algorithm
    if (present(n_saves_memory)) this%n_saves_memory = n_saves_memory
    if (present(path)) this%path = trim(path)
    if (present(filename)) this%filename = trim(filename)
    if (present(fmt)) this%fmt = trim(fmt)
    if (present(keep_checkpoints)) this%keep_checkpoints = keep_checkpoints
 
    inquire(file = trim(this%path), exist = exists)
    if (.not. exists) then
       call mpi_barrier(neko_comm)
       if (pe_rank .eq. 0) then
          call execute_command_line("mkdir -p '" // trim(this%path) // "'")
       end if
       call mpi_barrier(neko_comm)
    end if
 
    ! Initialize the Neko checkpoint output
    call this%chkp_output%init(neko_case%chkp, this%filename, &
         fmt = this%fmt, path = this%path, overwrite = .true.)
 
    n_states = 4
    if (allocated(neko_case%scalars)) then
       n_states = n_states + size(neko_case%scalars%scalar_fields)
    end if
    if (present(extra_fields)) then
       n_states = n_states + extra_fields%size()
    end if
 
    call this%state_list%init(n_states)
 
    ! Assign fluid pointers
    call this%state_list%assign(1, neko_case%fluid%p)
    call this%state_list%assign(2, neko_case%fluid%u)
    call this%state_list%assign(3, neko_case%fluid%v)
    call this%state_list%assign(4, neko_case%fluid%w)
    n_states = 4
 
    ! Assign scalar pointers
    if (allocated(neko_case%scalars)) then
       do i = 1, size(neko_case%scalars%scalar_fields)
          si => neko_case%scalars%scalar_fields(i)%scalar%s
          call this%state_list%assign(n_states + i, si)
       end do
       n_states = n_states + size(neko_case%scalars%scalar_fields)
    end if
 
    ! Assign any extra fields specified by the user
    if (present(extra_fields)) then
       do i = 1, extra_fields%size()
          si => extra_fields%get_by_index(i)
          call this%state_list%assign(n_states + i, si)
       end do
       n_states = n_states + extra_fields%size()
    end if
 
    ! Allocate the storage for the RAM checkpoints
    allocate(this%state_storage(this%n_saves_memory, this%state_list%size()))
 
    ! Write a status message with the parameters set
    call neko_log%section("Checkpointing")
 
    write(msg, '(A, A)') "Algorithm:                    ", trim(this%algorithm)
    call neko_log%message(trim(msg))
    write(msg, '(A,I0)') "Number of checkpoints in RAM: ", this%n_saves_memory
    call neko_log%message(trim(msg))
    write(msg, '(A, A)') "Checkpoint file path:         ", trim(this%path)
    call neko_log%message(trim(msg))
    write(msg, '(A, A)') "Checkpoint file name:         ", trim(this%filename)
    call neko_log%message(trim(msg))
    write(msg, '(A, A)') "Checkpoint file format:       ", trim(this%fmt)
    call neko_log%message(trim(msg))
 
    if (.not. this%keep_checkpoints) then
       call neko_log%message("Checkpoint files will be deleted.")
    else
       call neko_log%message("Checkpoint files will be kept.")
    end if
 
    call neko_log%message("Fields in checkpoint:", neko_log_debug)
    do i = 1, this%state_list%size()
       si => this%state_list%get(i)
       call neko_log%message("  - " // trim(si%name), neko_log_debug)
    end do
 
    call neko_log%end_section()
 
  end subroutine checkpoint_init_from_components
 
  subroutine checkpoint_free(this)
    class(simulation_checkpoint_t), intent(inout) :: this
    integer :: i, j
    character(len=1024) :: file_name
    logical :: exists
    integer :: stat, unit
 
    ! Free the RAM Checkpoints
    if (allocated(this%state_storage)) then
       do i = 1, this%n_saves_memory
          do j = 1, this%state_list%size()
             call this%state_storage(i, j)%free()
          end do
       end do
    end if
 
    call this%state_list%free()
    if (allocated(this%state_storage)) deallocate(this%state_storage)
 
    ! Delete the checkpoint file list
    if (.not. this%keep_checkpoints .and. pe_rank .eq. 0) then
       do i = this%n_timesteps, 1, -1
          call this%chkp_output%set_counter(i)
          file_name = this%chkp_output%file_%get_fname()
          inquire(file = trim(file_name), exist = exists)
          if (exists) then
             open(newunit = unit, file = trim(file_name), iostat = stat, &
                  status='old')
             if (stat .eq. 0) close(unit, status = 'delete')
          end if
       end do
    end if
    call mpi_barrier(neko_comm)
 
    ! Reset to default values
    this%enabled = .false.
    this%filename = "checkpoint"
    this%fmt = "chkp"
    this%algorithm = "linear"
    this%n_saves_memory = 10
    this%keep_checkpoints = .false.
 
    this%n_saves_disc = 0
    this%n_timesteps = 0
    this%first_valid_timestep = 2
    this%loaded_checkpoint = -1
 
  end subroutine checkpoint_free
 
  ! ========================================================================== !
  ! Saving and Restoring
 
  subroutine checkpoint_save(this, neko_case)
    class(simulation_checkpoint_t), intent(inout) :: this
    class(case_t), intent(inout) :: neko_case
 
    if (.not. this%enabled) return
 
    call profiler_start_region("Checkpoint save")
 
    ! Update the number of recorded timesteps
    this%n_timesteps = this%n_timesteps + 1
 
    select case (this%algorithm)
    case ("linear")
       call checkpoint_save_linear(this, neko_case)
    case default
       call neko_error("Unknown checkpoint algorithm: " // this%algorithm)
    end select
 
    call profiler_end_region("Checkpoint save")
  end subroutine checkpoint_save
 
  subroutine checkpoint_restore(this, neko_case, tstep)
    class(simulation_checkpoint_t), intent(inout) :: this
    class(case_t), target, intent(inout) :: neko_case
    integer, intent(in) :: tstep
    character(len=256) :: msg
 
    if (.not. this%enabled) return
 
    call profiler_start_region("Checkpoint restore")
 
    if (tstep .lt. 1 .or. tstep .gt. this%n_timesteps) then
       write(msg, '(A,I0,A,I0,A)') "Requested timestep ", tstep, &
            " is out of range [1, ", this%n_timesteps, "]"
       call neko_error(trim(msg))
    end if
 
    select case (this%algorithm)
    case ("linear")
       call checkpoint_restore_linear(this, neko_case, tstep)
    case default
       call neko_error("Unknown checkpoint algorithm: " // this%algorithm)
    end select
 
    call profiler_end_region("Checkpoint restore")
  end subroutine checkpoint_restore
 
  subroutine checkpoint_save_data(this, index)
    class(simulation_checkpoint_t), intent(inout) :: this
    integer, intent(in) :: index
    type(field_t), pointer :: si
    integer :: i
    character(len=1024) :: msg
 
    if (index .lt. 1 .or. index .gt. this%n_saves_memory) then
       write(msg, '(A,I0,A,I0,A)') "Checkpoint save index ", index, &
            " is out of range [1, ", this%n_saves_memory, "]"
       call neko_error(trim(msg))
    end if
 
    ! Allocate the RAM checkpoint if not already allocated
    do i = 1, this%state_list%size()
       if (.not. this%state_storage(index, i)%is_allocated()) then
          si => this%state_list%get(i)
          call this%state_storage(index, i)%init(si%size())
       end if
    end do
 
    ! Save the current iterates to memory
    if (neko_bcknd_device .eq. 0) then
       do i = 1, this%state_list%size()
          si => this%state_list%get(i)
          call copy(this%state_storage(index, i)%data, si%x, si%size())
       end do
    else
       do i = 1, this%state_list%size()
          si => this%state_list%get(i)
          call device_memcpy(this%state_storage(index, i)%data, si%x_d, &
               si%size(), device_to_host, this%state_list%size() .eq. i)
       end do
    end if
  end subroutine checkpoint_save_data
 
  subroutine checkpoint_load_data(this, index)
    class(simulation_checkpoint_t), intent(inout) :: this
    integer, intent(in) :: index
    type(field_t), pointer :: si
    character(len=1024) :: msg
    integer :: i
 
    if (index .lt. 1 .or. index .gt. this%n_saves_memory) then
       write(msg, '(A,I0,A,I0,A)') "Checkpoint save index ", index, &
            " is out of range [1, ", this%n_saves_memory, "]"
       call neko_error(trim(msg))
    end if
 
    ! Save the current iterates to memory
    if (neko_bcknd_device .eq. 0) then
       do i = 1, this%state_list%size()
          si => this%state_list%get(i)
          call copy(si%x, this%state_storage(index, i)%data, si%size())
       end do
    else
       do i = 1, this%state_list%size()
          si => this%state_list%get(i)
          call device_memcpy(this%state_storage(index, i)%data, si%x_d, &
               si%size(), host_to_device, this%state_list%size() .eq. i)
       end do
 
    end if
  end subroutine checkpoint_load_data
 
  ! ========================================================================== !
  ! Meta handling
 
  subroutine checkpoint_reset(this)
    class(simulation_checkpoint_t), intent(inout) :: this
    integer :: i, j
 
    if (.not. this%enabled) return
 
    ! Reset our checkpoints
    this%loaded_checkpoint = -1
    this%n_saves_disc = 0
    this%n_timesteps = 0
 
    do i = 1, size(this%state_storage, 1)
       do j = 1, size(this%state_storage, 2)
          call rzero(this%state_storage(i, j)%data, &
               this%state_storage(i, j)%size)
       end do
    end do
 
  end subroutine checkpoint_reset
 
  ! -------------------------------------------------------------------------- !
  ! Host array routines
 
  subroutine host_array_init(this, size)
    class(host_array), intent(inout) :: this
    integer, intent(in) :: size
 
    call this%free()
    this%size = size
    allocate(this%data(size))
    call rzero(this%data, this%size)
 
  end subroutine host_array_init
 
  subroutine host_array_free(this)
    class(host_array), intent(inout) :: this
 
    this%size = 0
    if (allocated(this%data)) deallocate(this%data)
 
  end subroutine host_array_free
 
  pure function host_array_is_allocated(this) result(is_alloc)
    class(host_array), intent(in) :: this
    logical :: is_alloc
 
    is_alloc = allocated(this%data)
 
  end function host_array_is_allocated
 
end module simulation_checkpoint