scratch_registry.f90

Go to the documentation of this file.

! Copyright (c) 2025-2026, The Neko Authors
! All rights reserved.
!
! Redistribution and use in source and binary forms, with or without
! modification, are permitted provided that the following conditions
! are met:
!
!   * Redistributions of source code must retain the above copyright
!     notice, this list of conditions and the following disclaimer.
!
!   * Redistributions in binary form must reproduce the above
!     copyright notice, this list of conditions and the following
!     disclaimer in the documentation and/or other materials provided
!     with the distribution.
!
!   * Neither the name of the authors nor the names of its
!     contributors may be used to endorse or promote products derived
!     from this software without specific prior written permission.
!
! THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
! "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
! LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
! FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
! COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
! INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
! BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
! LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
! CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
! LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
! ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
! POSSIBILITY OF SUCH DAMAGE.
!
module scratch_registry
  use registry_entry, only : registry_entry_t
  use host_array, only : host_array_t
  use device_array, only : device_array_t
  use field, only : field_t
  use vector, only : vector_t
  use matrix, only : matrix_t
 
  use math, only : rzero
  use device_math, only : device_rzero
  use field_math, only : field_rzero
  use vector_math, only : vector_rzero
  use matrix_math, only : matrix_rzero
 
  use dofmap, only : dofmap_t
  use utils, only : neko_error
  implicit none
  private
 
  type, public :: scratch_registry_t
     type(registry_entry_t), private, allocatable :: entries(:)
     logical, private, allocatable :: inuse(:)
     integer, private :: n_entries = 0
     integer, private :: n_inuse = 0
     integer, private :: expansion_size = 10
     type(dofmap_t), pointer :: dof => null()
   contains
     procedure, private, pass(this) :: expand
     procedure, pass(this) :: init => scratch_registry_init
     procedure, pass(this) :: free => scratch_registry_free
     procedure, pass(this) :: set_dofmap => scratch_registry_set_dofmap
     procedure, pass(this) :: get_n_entries
     procedure, pass(this) :: get_n_inuse
     procedure, pass(this) :: get_expansion_size
     procedure, pass(this) :: get_size
     procedure, pass(this) :: get_inuse
 
     procedure, pass(this) :: request_host_array
     procedure, pass(this) :: relinquish_host_array_single
     procedure, pass(this) :: relinquish_host_array_multiple
     generic :: relinquish_host_array => relinquish_host_array_single, &
          relinquish_host_array_multiple
 
     procedure, pass(this) :: request_device_array
     procedure, pass(this) :: relinquish_device_array_single
     procedure, pass(this) :: relinquish_device_array_multiple
     generic :: relinquish_device_array => relinquish_device_array_single, &
          relinquish_device_array_multiple
 
     procedure, pass(this) :: request_vector
     procedure, pass(this) :: relinquish_vector_single
     procedure, pass(this) :: relinquish_vector_multiple
     generic :: relinquish_vector => relinquish_vector_single, &
          relinquish_vector_multiple
 
     procedure, pass(this) :: request_matrix
     procedure, pass(this) :: relinquish_matrix_single
     procedure, pass(this) :: relinquish_matrix_multiple
     generic :: relinquish_matrix => relinquish_matrix_single, &
          relinquish_matrix_multiple
 
     procedure, pass(this) :: request_field
     procedure, pass(this) :: relinquish_field_single
     procedure, pass(this) :: relinquish_field_multiple
     generic :: relinquish_field => relinquish_field_single, &
          relinquish_field_multiple
 
     generic :: request => request_host_array, request_device_array, &
          request_vector, request_matrix, request_field
     procedure, pass(this) :: relinquish_single
     procedure, pass(this) :: relinquish_multiple
     generic :: relinquish => relinquish_single, relinquish_multiple
  end type scratch_registry_t
 
  type(scratch_registry_t), public, target :: neko_scratch_registry
 
contains
 
  subroutine scratch_registry_init(this, size, expansion_size, dof)
    class(scratch_registry_t), intent(inout) :: this
    integer, optional, intent(in) :: size
    integer, optional, intent(in) :: expansion_size
    type(dofmap_t), target, intent(in), optional :: dof
    integer :: s
 
    call this%free()
 
    s = 10
    if (present(size)) s = size
    if (present(dof)) this%dof => dof
 
    allocate(this%entries(s))
    allocate(this%inuse(s))
    this%inuse(:) = .false.
 
    this%expansion_size = 10
    if (present(expansion_size)) this%expansion_size = expansion_size
 
  end subroutine scratch_registry_init
 
  subroutine scratch_registry_free(this)
    class(scratch_registry_t), intent(inout):: this
    integer :: i
 
    if (allocated(this%inuse)) then
       if(any(this%inuse)) then
          call neko_error("scratch_registry::free: " // &
               "Cannot free scratch registry with in-use entries.")
       end if
       deallocate(this%inuse)
    end if
 
    if (allocated(this%entries)) then
       do i = 1, this%n_entries
          call this%entries(i)%free()
       end do
 
       deallocate(this%entries)
    end if
 
    if (associated(this%dof)) nullify(this%dof)
 
    ! Reset to default values
    this%n_entries = 0
    this%n_inuse = 0
    this%expansion_size = 10
 
  end subroutine scratch_registry_free
 
  subroutine scratch_registry_set_dofmap(this, dof)
    class(scratch_registry_t), intent(inout) :: this
    type(dofmap_t), target, intent(in) :: dof
 
    if (associated(this%dof, dof)) then
       return
    else if (associated(this%dof)) then
       call neko_error("scratch_registry::set_dofmap: "&
            // "Dofmap is already assigned to scratch registry.")
    end if
 
    this%dof => dof
  end subroutine scratch_registry_set_dofmap
 
  pure function get_n_entries(this) result(n)
    class(scratch_registry_t), intent(in) :: this
    integer :: n
 
    n = this%n_entries
  end function get_n_entries
 
  pure function get_n_inuse(this) result(n)
    class(scratch_registry_t), intent(in) :: this
    integer :: n, i
 
    n = count(this%inuse)
  end function get_n_inuse
 
  pure function get_size(this) result(n)
    class(scratch_registry_t), intent(in) :: this
    integer :: n
 
    if (allocated(this%entries)) then
       n = size(this%entries)
    else
       n = 0
    end if
  end function get_size
 
  pure function get_expansion_size(this) result(n)
    class(scratch_registry_t), intent(in) :: this
    integer :: n
 
    n = this%expansion_size
  end function get_expansion_size
 
  pure logical function get_inuse(this, index)
    class(scratch_registry_t), target, intent(in) :: this
    integer, intent(in) :: index
 
    get_inuse = this%inuse(index)
  end function get_inuse
 
  subroutine expand(this)
    class(scratch_registry_t), intent(inout) :: this
    type(registry_entry_t), allocatable :: temp(:)
    logical, allocatable :: temp2(:)
    integer :: i, n
 
    n = this%get_size()
 
    if (n .gt. 0) then
       call move_alloc(this%entries, temp)
       call move_alloc(this%inuse, temp2)
    end if
 
    allocate(this%entries(n + this%expansion_size))
    allocate(this%inuse(n + this%expansion_size), source = .false.)
 
    if (n .gt. 0) then
       do i = 1, n
          call this%entries(i)%move_from(temp(i))
          this%inuse(i) = temp2(i)
          call temp(i)%free()
       end do
    end if
 
    if (allocated(temp)) deallocate(temp)
    if (allocated(temp2)) deallocate(temp2)
 
  end subroutine expand
 
  subroutine request_host_array(this, v, index, n, clear)
    class(scratch_registry_t), target, intent(inout) :: this
    type(host_array_t), pointer, intent(inout) :: v
    integer, intent(inout) :: index
    integer, intent(in) :: n
    logical, intent(in) :: clear
 
    associate(entries => this%entries, n_entries => this%n_entries, &
         n_inuse => this%n_inuse)
 
      do index = 1, this%get_size()
         if (.not. this%inuse(index)) then
 
            if (.not. entries(index)%is_allocated()) then
               call entries(index)%init_host_array(n)
               n_entries = n_entries + 1
            else if (trim(entries(index)%get_type()) .ne. 'host_array') then
               cycle
            end if
 
            v => entries(index)%get_host_array()
            if (v%size() .ne. n) then
               nullify(v)
               cycle
            end if
 
            if (clear) call rzero(v%x, v%size())
            this%inuse(index) = .true.
            this%n_inuse = this%n_inuse + 1
            return
         end if
      end do
 
      ! all existing host_arrays in use, we need to expand to add a new one
      index = n_entries + 1
      call this%expand()
      n_entries = n_entries + 1
      n_inuse = n_inuse + 1
      this%inuse(n_entries) = .true.
      call this%entries(n_entries)%init_host_array(n)
      v => this%entries(n_entries)%get_host_array()
 
    end associate
  end subroutine request_host_array
 
  subroutine request_device_array(this, v, index, n, clear)
    class(scratch_registry_t), target, intent(inout) :: this
    type(device_array_t), pointer, intent(inout) :: v
    integer, intent(inout) :: index
    integer, intent(in) :: n
    logical, intent(in) :: clear
 
    associate(entries => this%entries, n_entries => this%n_entries, &
         n_inuse => this%n_inuse)
 
      do index = 1, this%get_size()
         if (.not. this%inuse(index)) then
 
            if (.not. entries(index)%is_allocated()) then
               call entries(index)%init_device_array(n)
               n_entries = n_entries + 1
            else if (trim(entries(index)%get_type()) .ne. 'device_array') then
               cycle
            end if
 
            v => entries(index)%get_device_array()
            if (v%size() .ne. n) then
               nullify(v)
               cycle
            end if
 
            if (clear) call device_rzero(v%x_d, v%size())
            this%inuse(index) = .true.
            this%n_inuse = this%n_inuse + 1
            return
         end if
      end do
 
      ! all existing device_arrays in use, we need to expand to add a new one
      index = n_entries + 1
      call this%expand()
      n_entries = n_entries + 1
      n_inuse = n_inuse + 1
      this%inuse(n_entries) = .true.
      call this%entries(n_entries)%init_device_array(n)
      v => this%entries(n_entries)%get_device_array()
 
    end associate
  end subroutine request_device_array
 
  subroutine request_vector(this, v, index, n, clear)
    class(scratch_registry_t), target, intent(inout) :: this
    type(vector_t), pointer, intent(inout) :: v
    integer, intent(inout) :: index
    integer, intent(in) :: n
    logical, intent(in) :: clear
 
    associate(entries => this%entries, n_entries => this%n_entries, &
         n_inuse => this%n_inuse)
 
      do index = 1, this%get_size()
         if (.not. this%inuse(index)) then
 
            if (.not. entries(index)%is_allocated()) then
               call entries(index)%init_vector(n)
               n_entries = n_entries + 1
            else if (trim(entries(index)%get_type()) .ne. 'vector') then
               cycle
            end if
 
            v => entries(index)%get_vector()
            if (v%size() .ne. n) then
               nullify(v)
               cycle
            end if
 
            if (clear) call vector_rzero(v)
            this%inuse(index) = .true.
            this%n_inuse = this%n_inuse + 1
            return
         end if
      end do
 
      ! all existing vectors in use, we need to expand to add a new one
      index = n_entries + 1
      call this%expand()
      n_entries = n_entries + 1
      n_inuse = n_inuse + 1
      this%inuse(n_entries) = .true.
      call this%entries(n_entries)%init_vector(n)
      v => this%entries(n_entries)%get_vector()
 
    end associate
  end subroutine request_vector
 
  subroutine request_matrix(this, m, index, nrows, ncols, clear)
    class(scratch_registry_t), target, intent(inout) :: this
    type(matrix_t), pointer, intent(inout) :: m
    integer, intent(inout) :: index
    integer, intent(in) :: nrows, ncols
    logical, intent(in) :: clear
 
    associate(entries => this%entries, n_entries => this%n_entries, &
         n_inuse => this%n_inuse)
 
      do index = 1, this%get_size()
         if (.not. this%inuse(index)) then
 
            if (.not. entries(index)%is_allocated()) then
               call entries(index)%init_matrix(nrows, ncols)
               n_entries = n_entries + 1
            else if (trim(entries(index)%get_type()) .ne. 'matrix') then
               cycle
            end if
 
            m => entries(index)%get_matrix()
            if (m%get_nrows() .ne. nrows .or. &
                 m%get_ncols() .ne. ncols) then
               nullify(m)
               cycle
            end if
 
            if (clear) call matrix_rzero(m)
            this%inuse(index) = .true.
            this%n_inuse = this%n_inuse + 1
            return
         end if
      end do
 
      ! all existing matrices in use, we need to expand to add a new one
      index = n_entries + 1
      call this%expand()
      n_entries = n_entries + 1
      n_inuse = n_inuse + 1
      this%inuse(n_entries) = .true.
      call this%entries(n_entries)%init_matrix(nrows, ncols)
      m => this%entries(n_entries)%get_matrix()
 
    end associate
  end subroutine request_matrix
 
 
  subroutine request_field(this, f, index, clear)
    class(scratch_registry_t), target, intent(inout) :: this
    type(field_t), pointer, intent(inout) :: f
    integer, intent(inout) :: index
    logical, intent(in) :: clear
    character(len=10) :: name
 
    if (.not. associated(this%dof)) then
       call neko_error("scratch_registry::request_field: "&
            // "No dofmap assigned to scratch registry.")
    end if
 
    associate(entries => this%entries, n_entries => this%n_entries, &
         n_inuse => this%n_inuse)
 
      do index = 1, this%get_size()
         if (.not. this%inuse(index)) then
 
            if (.not. entries(index)%is_allocated()) then
               write(name, "(A3,I0.3)") "wrk", index
               call entries(index)%init_field(this%dof, trim(name))
               n_entries = n_entries + 1
            else if (entries(index)%get_type() .ne. 'field') then
               cycle
            end if
 
            f => entries(index)%get_field()
            if (clear) call field_rzero(f)
            this%inuse(index) = .true.
            this%n_inuse = this%n_inuse + 1
            return
         end if
      end do
 
      ! all existing fields in use, we need to expand to add a new one
      index = n_entries + 1
      call this%expand()
      n_entries = n_entries + 1
      n_inuse = n_inuse + 1
      this%inuse(n_entries) = .true.
      write (name, "(A3,I0.3)") "wrk", index
      call this%entries(n_entries)%init_field(this%dof, trim(name))
      f => this%entries(n_entries)%get_field()
 
    end associate
  end subroutine request_field
 
  subroutine relinquish_host_array_single(this, index)
    class(scratch_registry_t), target, intent(inout) :: this
    integer, intent(inout) :: index
 
    if (trim(this%entries(index)%get_type()) .ne. 'host_array') then
       call neko_error("scratch_registry::relinquish_host_array_single: " &
            // "Register entry is not a host_array.")
    end if
 
    this%inuse(index) = .false.
    this%n_inuse = this%n_inuse - 1
  end subroutine relinquish_host_array_single
 
  subroutine relinquish_host_array_multiple(this, indices)
    class(scratch_registry_t), target, intent(inout) :: this
    integer, intent(inout) :: indices(:)
    integer :: i
 
    do i = 1, size(indices)
       if (trim(this%entries(indices(i))%get_type()) .ne. 'host_array') then
          call neko_error("scratch_registry::relinquish_host_array_single: " &
               // "Register entry is not a host_array.")
       end if
 
       this%inuse(indices(i)) = .false.
    end do
    this%n_inuse = this%n_inuse - size(indices)
  end subroutine relinquish_host_array_multiple
 
  subroutine relinquish_device_array_single(this, index)
    class(scratch_registry_t), target, intent(inout) :: this
    integer, intent(inout) :: index
 
    if (trim(this%entries(index)%get_type()) .ne. 'device_array') then
       call neko_error("scratch_registry::relinquish_device_array_single: " &
            // "Register entry is not a device_array.")
    end if
 
    this%inuse(index) = .false.
    this%n_inuse = this%n_inuse - 1
  end subroutine relinquish_device_array_single
 
  subroutine relinquish_device_array_multiple(this, indices)
    class(scratch_registry_t), target, intent(inout) :: this
    integer, intent(inout) :: indices(:)
    integer :: i
 
    do i = 1, size(indices)
       if (trim(this%entries(indices(i))%get_type()) .ne. 'device_array') then
          call neko_error("scratch_registry::relinquish_device_array_single: " &
               // "Register entry is not a device_array.")
       end if
 
       this%inuse(indices(i)) = .false.
    end do
    this%n_inuse = this%n_inuse - size(indices)
  end subroutine relinquish_device_array_multiple
 
  subroutine relinquish_vector_single(this, index)
    class(scratch_registry_t), target, intent(inout) :: this
    integer, intent(inout) :: index
 
    if (trim(this%entries(index)%get_type()) .ne. 'vector') then
       call neko_error("scratch_registry::relinquish_vector_single: " &
            // "Register entry is not a vector.")
    end if
 
    this%inuse(index) = .false.
    this%n_inuse = this%n_inuse - 1
  end subroutine relinquish_vector_single
 
  subroutine relinquish_vector_multiple(this, indices)
    class(scratch_registry_t), target, intent(inout) :: this
    integer, intent(inout) :: indices(:)
    integer :: i
 
    do i = 1, size(indices)
       if (trim(this%entries(indices(i))%get_type()) .ne. 'vector') then
          call neko_error("scratch_registry::relinquish_vector_single: " &
               // "Register entry is not a vector.")
       end if
 
       this%inuse(indices(i)) = .false.
    end do
    this%n_inuse = this%n_inuse - size(indices)
  end subroutine relinquish_vector_multiple
 
  subroutine relinquish_matrix_single(this, index)
    class(scratch_registry_t), target, intent(inout) :: this
    integer, intent(inout) :: index
 
    if (trim(this%entries(index)%get_type()) .ne. 'matrix') then
       call neko_error("scratch_registry::relinquish_matrix_single: " &
            // "Register entry is not a matrix.")
    end if
 
    this%inuse(index) = .false.
    this%n_inuse = this%n_inuse - 1
  end subroutine relinquish_matrix_single
 
  subroutine relinquish_matrix_multiple(this, indices)
    class(scratch_registry_t), target, intent(inout) :: this
    integer, intent(inout) :: indices(:)
    integer :: i
 
    do i = 1, size(indices)
       if (trim(this%entries(indices(i))%get_type()) .ne. 'matrix') then
          call neko_error("scratch_registry::relinquish_matrix_single: " &
               // "Register entry is not a matrix.")
       end if
 
       this%inuse(indices(i)) = .false.
    end do
    this%n_inuse = this%n_inuse - size(indices)
  end subroutine relinquish_matrix_multiple
 
  subroutine relinquish_field_single(this, index)
    class(scratch_registry_t), target, intent(inout) :: this
    integer, intent(inout) :: index
 
    if (trim(this%entries(index)%get_type()) .ne. 'field') then
       call neko_error("scratch_registry::relinquish_field_single: " &
            // "Register entry is not a field.")
    end if
 
    this%inuse(index) = .false.
    this%n_inuse = this%n_inuse - 1
  end subroutine relinquish_field_single
 
  subroutine relinquish_field_multiple(this, indices)
    class(scratch_registry_t), target, intent(inout) :: this
    integer, intent(inout) :: indices(:)
    integer :: i
 
    do i = 1, size(indices)
       if (trim(this%entries(indices(i))%get_type()) .ne. 'field') then
          call neko_error("scratch_registry::relinquish_field_single: " &
               // "Register entry is not a field.")
       end if
 
       this%inuse(indices(i)) = .false.
    end do
    this%n_inuse = this%n_inuse - size(indices)
  end subroutine relinquish_field_multiple
 
  subroutine relinquish_single(this, index)
    class(scratch_registry_t), target, intent(inout) :: this
    integer, intent(inout) :: index
 
    this%inuse(index) = .false.
    this%n_inuse = this%n_inuse - 1
  end subroutine relinquish_single
 
  subroutine relinquish_multiple(this, indices)
    class(scratch_registry_t), target, intent(inout) :: this
    integer, intent(inout) :: indices(:)
    integer :: i
 
    do i = 1, size(indices)
       this%inuse(indices(i)) = .false.
    end do
    this%n_inuse = this%n_inuse - size(indices)
  end subroutine relinquish_multiple
 
end module scratch_registry