checkpoint.f90

Go to the documentation of this file.

 
!
module state_recover_checkpoint
  use num_types, only: rp
  use case, only: case_t
  use json_file_module, only: json_file
  use json_utils, only: json_get, json_get_or_default
  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 utils, only: neko_error, mkdir
  use math, only: copy, rzero
  use host_array, only: host_array_t
  use profiler, only: profiler_start_region, profiler_end_region
  use state_recover, only: state_recover_t
  use comm, only: pe_rank, neko_comm
  use neko_config, only: neko_bcknd_device
  use device, only: device_memcpy, device_to_host, host_to_device
  use registry, only: neko_registry
  use mpi_f08, only: mpi_barrier
  implicit none
  private
 
  ! Backend Implementations
  integer, parameter :: CHECKPOINT_LINEAR = 1
 
  type, public, extends(state_recover_t) :: state_recover_checkpoint_t
     private
 
     ! ----------------------------------------------------------------------- !
     ! User parameters
 
     character(len=256) :: algorithm = ""
     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 :: algorithm_id = 0
     integer :: n_saves_disc = 0
     integer :: first_valid_timestep = 2
     integer :: loaded_checkpoint = -1
 
     ! Field pointers
     type(field_list_t) :: state_list
     type(host_array_t), dimension(:,:), allocatable :: state_storage
 
     ! Structures to hold the checkpoint data
     type(chkp_output_t) :: chkp_output
 
   contains
     procedure, public, pass(this) :: init => 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 state_recover_checkpoint_t
 
  ! ========================================================================== !
  ! Module procedures for our algorithm implementations.
 
  interface
 
     module subroutine checkpoint_save_linear(this)
       class(state_recover_checkpoint_t), intent(inout) :: this
     end subroutine checkpoint_save_linear
 
     module subroutine checkpoint_restore_linear(this, tstep)
       class(state_recover_checkpoint_t), intent(inout) :: this
       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(state_recover_checkpoint_t), intent(inout) :: this
    class(case_t), target, intent(inout) :: neko_case
    type(json_file), 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 :: keep_checkpoints
 
    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
       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(state_recover_checkpoint_t), intent(inout), target :: this
    class(case_t), target, intent(inout) :: neko_case
    character(len=*), 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, j, n_states
    logical :: exists
 
    call this%free()
    this%neko_case => neko_case
 
    ! Assign parameters from arguments or defaults
    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
 
    ! Assign the checkpointing algorithm
    select case (trim(algorithm))
    case ("linear", "LINEAR", "Linear")
       this%algorithm = "linear"
       this%algorithm_id = checkpoint_linear
    case default
       call neko_error("Only the linear checkpoint strategy is supported.")
    end select
 
    inquire(file = trim(this%path), exist = exists)
    if (.not. exists) then
       call mpi_barrier(neko_comm)
       if (pe_rank .eq. 0) then
          call mkdir(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()))
    do i = 1, this%n_saves_memory
       do j = 1, this%state_list%size()
          si => this%state_list%get(j)
          call this%state_storage(i, j)%init(si%size())
       end do
    end do
 
    ! 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(state_recover_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%get_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%filename = "checkpoint"
    this%fmt = "chkp"
    this%algorithm = "linear"
    this%n_saves_memory = 10
    this%keep_checkpoints = .false.
 
    this%n_saves_disc = 0
    call this%set_n_timesteps(0)
    this%first_valid_timestep = 2
    this%loaded_checkpoint = -1
    nullify(this%neko_case)
 
  end subroutine checkpoint_free
 
  ! ========================================================================== !
  ! Saving and Restoring
 
  subroutine checkpoint_save(this)
    class(state_recover_checkpoint_t), intent(inout) :: this
 
    call profiler_start_region("Checkpoint save")
 
    ! Update the number of recorded timesteps
    call this%set_n_timesteps(this%get_n_timesteps() + 1)
 
    select case (this%algorithm_id)
    case (checkpoint_linear)
       call checkpoint_save_linear(this)
    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, tstep)
    class(state_recover_checkpoint_t), intent(inout) :: this
    integer, intent(in) :: tstep
    character(len=256) :: msg
 
    call profiler_start_region("Checkpoint restore")
 
    if (tstep .lt. 1 .or. tstep .gt. this%get_n_timesteps()) then
       write(msg, '(A,I0,A,I0,A)') "Requested timestep ", tstep, &
            " is out of range [1, ", this%get_n_timesteps(), "]"
       call neko_error(trim(msg))
    end if
 
    select case (this%algorithm_id)
    case (checkpoint_linear)
       call checkpoint_restore_linear(this, 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(state_recover_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)%x, 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)%x, si%x_d, &
               si%size(), device_to_host, this%state_list%size() .eq. i)
       end do
    end if
 
    nullify(si)
  end subroutine checkpoint_save_data
 
  subroutine checkpoint_load_data(this, index)
    class(state_recover_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)%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)%x, si%x_d, &
               si%size(), host_to_device, this%state_list%size() .eq. i)
       end do
    end if
 
    nullify(si)
  end subroutine checkpoint_load_data
 
  ! ========================================================================== !
  ! Meta handling
 
  subroutine checkpoint_reset(this)
    class(state_recover_checkpoint_t), intent(inout) :: this
    integer :: i, j
 
    ! Reset our checkpoints
    this%loaded_checkpoint = -1
    this%n_saves_disc = 0
    call this%set_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)%x, this%state_storage(i, j)%size())
       end do
    end do
 
  end subroutine checkpoint_reset
 
end module state_recover_checkpoint