batch.F90

Go to the documentation of this file.

!! Copyright (C) 2008 X. Andrade, 2020 S. Ohlmann
!!
!! This program is free software; you can redistribute it and/or modify
!! it under the terms of the GNU General Public License as published by
!! the Free Software Foundation; either version 2, or (at your option)
!! any later version.
!!
!! This program is distributed in the hope that it will be useful,
!! but WITHOUT ANY WARRANTY; without even the implied warranty of
!! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!! GNU General Public License for more details.
!!
!! You should have received a copy of the GNU General Public License
!! along with this program; if not, write to the Free Software
!! Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
!! 02110-1301, USA.
!!
 
#include "global.h"
 
!
module batch_oct_m
  use accel_oct_m
  use allocate_hardware_aware_oct_m
  use blas_oct_m
  use debug_oct_m
  use global_oct_m
  use hardware_oct_m
  use iso_c_binding
  use math_oct_m
  use messages_oct_m
  use mpi_oct_m
  use profiling_oct_m
  use types_oct_m
 
  implicit none
 
  private
  public ::                         &
    batch_t,                        &
    batch_init,                     &
    dbatch_init,                    &
    zbatch_init,                    &
    batch_read_device_to_packed
 
  type batch_t
    private
    integer,                     public :: nst
    integer,                     public :: dim
    integer                             :: np
    integer                             :: ndims
    integer,        allocatable         :: ist_idim_index(:, :)
    
    integer,        allocatable, public :: ist(:)
 
    logical                             :: is_allocated
    logical                             :: own_memory
    !  We also need a linear array with the states in order to calculate derivatives, etc.
    integer,                     public :: nst_linear
 
    integer                             :: status_of
    integer                             :: status_host
    type(type_t)                        :: type_of
    integer                             :: device_buffer_count
    integer                             :: host_buffer_count
    logical                             :: special_memory
    logical                             :: needs_finish_unpack
 
 
    ! unpacked variables; linear variables are pointers with different shapes
    real(real64),    pointer, contiguous,  public :: dff(:, :, :)
    !                                                                 !! indices are (1:np, 1:dim, 1:nst)
    complex(real64), pointer, contiguous,  public :: zff(:, :, :)
    !                                                                 !! indices are (1:np, 1:dim, 1:nst)
    real(real64),    pointer, contiguous,  public :: dff_linear(:, :)
    !                                                                 !! indices are (1:np, 1:nst_linear)
    complex(real64), pointer, contiguous,  public :: zff_linear(:, :)
    !                                                                 !! indices are (1:np, 1:nst_linear)
 
    ! packed variables; only rank-2 arrays due to padding to powers of 2
    real(real64), pointer, contiguous,  public :: dff_pack(:, :)
    !                                                                 !! indices are (1:nst_linear, 1:np)
    complex(real64), pointer, contiguous,  public :: zff_pack(:, :)
    !                                                                 !! indices are (1:nst_linear, 1:np)
 
    integer(int64),                 public :: pack_size(1:2)
    integer(int64),                 public :: pack_size_real(1:2)
 
    type(accel_mem_t),           public :: ff_device
 
  contains
    procedure :: check_compatibility_with => batch_check_compatibility_with 
    procedure :: clone_to => batch_clone_to                                 
    procedure :: clone_to_array => batch_clone_to_array                     
    procedure :: copy_to => batch_copy_to                                   
    procedure :: copy_data_to => batch_copy_data_to                         
    procedure :: do_pack_generic => batch_do_pack_generic
    procedure :: do_pack_target => batch_do_pack_target
    generic :: do_pack => do_pack_generic, do_pack_target                   
    procedure :: do_unpack => batch_do_unpack                               
    procedure :: finish_unpack => batch_finish_unpack                       
    procedure :: end => batch_end                                           
    procedure :: inv_index => batch_inv_index                               
    procedure :: is_packed => batch_is_packed                               
    procedure :: ist_idim_to_linear => batch_ist_idim_to_linear             
    procedure :: linear_to_idim => batch_linear_to_idim                     
    procedure :: linear_to_ist => batch_linear_to_ist                       
    procedure :: pack_total_size => batch_pack_total_size                   
    procedure :: remote_access_start => batch_remote_access_start           
    procedure :: remote_access_stop => batch_remote_access_stop             
    procedure :: status => batch_status                                     
    procedure :: type => batch_type                                         
    procedure :: type_as_int => batch_type_as_integer                       
    procedure, private :: dallocate_unpacked_host => dbatch_allocate_unpacked_host
    
    procedure, private :: zallocate_unpacked_host => zbatch_allocate_unpacked_host
    
    procedure, private :: allocate_unpacked_host => batch_allocate_unpacked_host
    
    procedure, private :: dallocate_packed_host => dbatch_allocate_packed_host
    
    procedure, private :: zallocate_packed_host => zbatch_allocate_packed_host
    
    procedure, private :: allocate_packed_host => batch_allocate_packed_host
    
    procedure, private :: allocate_packed_device => batch_allocate_packed_device
    
    procedure, private :: deallocate_unpacked_host => batch_deallocate_unpacked_host
    
    procedure, private :: deallocate_packed_host => batch_deallocate_packed_host
    
    procedure, private :: deallocate_packed_device => batch_deallocate_packed_device
    
  end type batch_t
 
  !--------------------------------------------------------------
  interface batch_init
    module procedure dbatch_init_with_memory_3
    module procedure zbatch_init_with_memory_3
    module procedure dbatch_init_with_memory_2
    module procedure zbatch_init_with_memory_2
    module procedure dbatch_init_with_memory_1
    module procedure zbatch_init_with_memory_1
  end interface batch_init
 
  integer, public, parameter :: &
    batch_not_packed     = 0,   & !< functions are stored in CPU memory, unpacked order
    batch_packed         = 1,   & 
    batch_device_packed  = 2      
 
  integer, parameter :: gpu_pack_max_buffer_size = 4
 
contains
 
  !--------------------------------------------------------------
  !
  subroutine batch_end(this, copy)
    class(batch_t),          intent(inout) :: this
    logical,       optional, intent(in)    :: copy
 
    push_sub(batch_end)
 
    if (this%own_memory .and. this%is_packed()) then
      !deallocate directly to avoid unnecessary copies
      if (this%status() == batch_device_packed) then
        call this%deallocate_packed_device()
      end if
      if (this%status() == batch_packed .or. this%status_host == batch_packed) then
        call this%deallocate_packed_host()
      end if
      this%status_of = batch_not_packed
      this%status_host = batch_not_packed
      this%host_buffer_count = 0
      this%device_buffer_count = 0
    end if
    if (this%status() == batch_device_packed) call this%do_unpack(copy, force = .true.)
    if (this%status() == batch_packed) call this%do_unpack(copy, force = .true.)
 
    if (this%is_allocated) then
      call this%deallocate_unpacked_host()
    end if
 
    safe_deallocate_a(this%ist_idim_index)
    safe_deallocate_a(this%ist)
 
    pop_sub(batch_end)
  end subroutine batch_end
 
  !--------------------------------------------------------------
  !
  subroutine batch_deallocate_unpacked_host(this)
    class(batch_t),  intent(inout) :: this
 
    push_sub(batch_deallocate_unpacked_host)
 
    this%is_allocated = .false.
 
    if (this%special_memory) then
      if (associated(this%dff)) then
        call deallocate_hardware_aware(c_loc(this%dff(1,1,1)), int(this%np, int64)*this%dim*this%nst*8)
      end if
      if (associated(this%zff)) then
        call deallocate_hardware_aware(c_loc(this%zff(1,1,1)), int(this%np, int64)*this%dim*this%nst*16)
      end if
    else
      safe_deallocate_p(this%dff)
      safe_deallocate_p(this%zff)
    end if
    nullify(this%dff)
    nullify(this%dff_linear)
    nullify(this%zff)
    nullify(this%zff_linear)
 
    pop_sub(batch_deallocate_unpacked_host)
  end subroutine batch_deallocate_unpacked_host
 
  !--------------------------------------------------------------
  !
  subroutine batch_deallocate_packed_host(this)
    class(batch_t),  intent(inout) :: this
 
    push_sub(batch_deallocate_packed_host)
 
    if (this%special_memory) then
      if (associated(this%dff_pack)) then
        call deallocate_hardware_aware(c_loc(this%dff_pack(1,1)), int(this%pack_size(1), int64)*this%pack_size(2)*8)
      end if
      if (associated(this%zff_pack)) then
        call deallocate_hardware_aware(c_loc(this%zff_pack(1,1)), int(this%pack_size(1), int64)*this%pack_size(2)*16)
      end if
    else
      safe_deallocate_p(this%dff_pack)
      safe_deallocate_p(this%zff_pack)
    end if
    nullify(this%dff_pack)
    nullify(this%zff_pack)
 
    pop_sub(batch_deallocate_packed_host)
  end subroutine batch_deallocate_packed_host
 
  !--------------------------------------------------------------
  !
  subroutine batch_deallocate_packed_device(this)
    class(batch_t),  intent(inout) :: this
 
    push_sub(batch_deallocate_packed_device)
 
    call accel_free_buffer(this%ff_device)
 
    pop_sub(batch_deallocate_packed_device)
  end subroutine batch_deallocate_packed_device
 
  !--------------------------------------------------------------
  !
  subroutine batch_allocate_unpacked_host(this)
    class(batch_t),  intent(inout) :: this
 
    push_sub(batch_allocate_unpacked_host)
 
    if (this%type() == type_float) then
      call this%dallocate_unpacked_host()
    else if (this%type() == type_cmplx) then
      call this%zallocate_unpacked_host()
    end if
 
    pop_sub(batch_allocate_unpacked_host)
  end subroutine batch_allocate_unpacked_host
 
  !--------------------------------------------------------------
  !
  subroutine batch_allocate_packed_host(this)
    class(batch_t),  intent(inout) :: this
 
    push_sub(batch_allocate_packed_host)
 
    if (this%type() == type_float) then
      call this%dallocate_packed_host()
    else if (this%type() == type_cmplx) then
      call this%zallocate_packed_host()
    end if
 
    pop_sub(batch_allocate_packed_host)
  end subroutine batch_allocate_packed_host
 
  !--------------------------------------------------------------
  !
  subroutine batch_allocate_packed_device(this)
    class(batch_t),  intent(inout) :: this
 
    push_sub(batch_allocate_packed_device)
 
    call accel_create_buffer(this%ff_device, accel_mem_read_write, this%type(), &
      product(this%pack_size))
 
    pop_sub(batch_allocate_packed_device)
  end subroutine batch_allocate_packed_device
 
  !--------------------------------------------------------------
  !
  subroutine batch_init_empty (this, dim, nst, np)
    type(batch_t), intent(out)   :: this
    integer,       intent(in)    :: dim
    integer,       intent(in)    :: nst
    integer,       intent(in)    :: np
 
    push_sub(batch_init_empty)
 
    this%is_allocated = .false.
    this%own_memory = .false.
    this%special_memory = .false.
    this%needs_finish_unpack = .false.
    this%nst = nst
    this%dim = dim
    this%type_of = type_none
 
    this%nst_linear = nst*dim
 
    this%np = np
    this%device_buffer_count = 0
    this%host_buffer_count = 0
    this%status_of = batch_not_packed
    this%status_host = batch_not_packed
 
    this%ndims = 2
    safe_allocate(this%ist_idim_index(1:this%nst_linear, 1:this%ndims))
    safe_allocate(this%ist(1:this%nst))
 
    nullify(this%dff, this%zff, this%dff_linear, this%zff_linear)
    nullify(this%dff_pack, this%zff_pack)
 
    pop_sub(batch_init_empty)
  end subroutine batch_init_empty
 
  !--------------------------------------------------------------
  !
  subroutine batch_clone_to(this, dest, pack, copy_data, new_np, special, dest_type)
    class(batch_t),              intent(in)    :: this
    class(batch_t), allocatable, intent(out)   :: dest
    logical,        optional,    intent(in)    :: pack
    logical,        optional,    intent(in)    :: copy_data
    integer,        optional,    intent(in)    :: new_np
    logical,        optional,    intent(in)    :: special
    type(type_t),   optional,    intent(in)    :: dest_type
 
    push_sub(batch_clone_to)
 
    if (.not. allocated(dest)) then
      safe_allocate_type(batch_t, dest)
    else
      message(1) = "Internal error: destination batch in batch_clone_to has been previously allocated."
      call messages_fatal(1)
    end if
 
    call this%copy_to(dest, pack, copy_data, new_np, special, dest_type)
 
    pop_sub(batch_clone_to)
  end subroutine batch_clone_to
 
  !--------------------------------------------------------------
 
  subroutine batch_clone_to_array(this, dest, n_batches, pack, copy_data, new_np, special, dest_type)
    class(batch_t),              intent(in)    :: this
    class(batch_t), allocatable, intent(out)   :: dest(:)
    integer,                     intent(in)    :: n_batches
    logical,        optional,    intent(in)    :: pack
    logical,        optional,    intent(in)    :: copy_data
    integer,        optional,    intent(in)    :: new_np
    logical,        optional,    intent(in)    :: special
    type(type_t),   optional,    intent(in)    :: dest_type
 
    integer :: ib
 
    push_sub(batch_clone_to_array)
 
    if (.not. allocated(dest)) then
      safe_allocate_type_array(batch_t, dest, (1:n_batches))
    else
      message(1) = "Internal error: destination batch in batch_clone_to_array has been previously allocated."
      call messages_fatal(1)
    end if
 
    do ib = 1, n_batches
      call this%copy_to(dest(ib), pack, copy_data, new_np, special, dest_type)
    end do
 
    pop_sub(batch_clone_to_array)
  end subroutine batch_clone_to_array
 
  !--------------------------------------------------------------
  subroutine batch_copy_to(this, dest, pack, copy_data, new_np, special, dest_type)
    class(batch_t),          intent(in)    :: this
    class(batch_t),          intent(out)   :: dest
    logical,       optional, intent(in)    :: pack
    logical,       optional, intent(in)    :: copy_data
    integer,       optional, intent(in)    :: new_np
    logical,       optional, intent(in)    :: special
    type(type_t),  optional, intent(in)    :: dest_type
 
    logical :: host_packed, special_
    integer :: np_
    type(type_t) :: type_
 
    push_sub(batch_copy_to)
 
    np_ = optional_default(new_np, this%np)
 
    host_packed = this%host_buffer_count > 0
    ! use special memory here only for batches not on the GPU to avoid allocating
    ! pinned memory for temporary batches because that leads to a severe performance
    ! decrease for GPU runs (up to 20x)
    if (present(special)) then
      special_ = this%special_memory
    else
      special_ = this%special_memory .and. .not. this%device_buffer_count > 0
    end if
 
    if (present(dest_type)) then
      type_ = dest_type
    else
      type_ = this%type()
    end if
 
    if (type_ == type_float) then
      call dbatch_init(dest, this%dim, 1, this%nst, np_, packed=host_packed, special=special_)
    else if (type_ == type_cmplx) then
      call zbatch_init(dest, this%dim, 1, this%nst, np_, packed=host_packed, special=special_)
    else
      message(1) = "Internal error: unknown batch type in batch_copy_to."
      call messages_fatal(1)
    end if
 
    if (this%status() /= dest%status() .and. optional_default(pack, this%is_packed())) call dest%do_pack(copy = .false.)
 
    dest%ist_idim_index(1:this%nst_linear, 1:this%ndims) = this%ist_idim_index(1:this%nst_linear, 1:this%ndims)
    dest%ist(1:this%nst) = this%ist(1:this%nst)
 
    if (optional_default(copy_data, .false.)) then
      assert(np_ == this%np)
      call this%copy_data_to(min(this%np, np_), dest)
    end if
 
    pop_sub(batch_copy_to)
  end subroutine batch_copy_to
 
  ! ----------------------------------------------------
  type(type_t) pure function batch_type(this) result(btype)
    class(batch_t),      intent(in)    :: this
 
    btype = this%type_of
 
  end function batch_type
 
  ! ----------------------------------------------------
  integer pure function batch_type_as_integer(this) result(itype)
    class(batch_t),      intent(in)    :: this
 
    type(type_t) :: btype
 
    itype = 0
    btype = this%type()
    if (btype == type_float) itype = 1
    if (btype == type_cmplx) itype = 2
 
  end function batch_type_as_integer
 
  ! ----------------------------------------------------
  integer pure function batch_status(this) result(bstatus)
    class(batch_t),      intent(in)    :: this
 
    bstatus = this%status_of
  end function batch_status
 
  ! ----------------------------------------------------
 
  logical pure function batch_is_packed(this) result(in_buffer)
    class(batch_t),      intent(in)    :: this
 
    in_buffer = (this%device_buffer_count > 0) .or. (this%host_buffer_count > 0)
  end function batch_is_packed
 
  ! ----------------------------------------------------
 
  integer(int64) function batch_pack_total_size(this) result(size)
    class(batch_t),      intent(inout) :: this
 
    size = this%np
    if (accel_is_enabled()) size = accel_padded_size(size)
    size = size*pad_pow2(this%nst_linear)*types_get_size(this%type())
 
  end function batch_pack_total_size
 
  ! ----------------------------------------------------
 
  subroutine batch_do_pack_generic(this, copy, async)
    class(batch_t),      intent(inout) :: this
    logical,   optional, intent(in)    :: copy
    logical,   optional, intent(in)    :: async
 
    integer               :: source, target
 
    ! no push_sub, called too frequently
 
    ! get source and target states for this batch
    source = this%status()
    select case (source)
    case (batch_not_packed, batch_packed)
      if (accel_is_enabled()) then
        target = batch_device_packed
      else
        target = batch_packed
      end if
    case (batch_device_packed)
      target = batch_device_packed
    end select
 
    call this%do_pack(target, copy, async)
 
  end subroutine batch_do_pack_generic
 
  subroutine batch_do_pack_target(this, target, copy, async, cpu_only)
    class(batch_t),      intent(inout) :: this
    integer,             intent(in)    :: target
    logical,   optional, intent(in)    :: copy
    logical,   optional, intent(in)    :: async
    logical,   optional, intent(in)    :: cpu_only
 
    logical               :: copy_
    logical               :: async_
    integer               :: source
 
    ! no push_sub, called too frequently
 
    call profiling_in("BATCH_DO_PACK")
 
    copy_ = optional_default(copy, .true.)
 
    async_ = optional_default(async, .false.)
 
    ! get source state for this batch
    source = this%status()
 
    ! only do something if target is different from source
    if (source /= target) then
      select case (target)
      case (batch_device_packed)
        call this%allocate_packed_device()
        this%status_of = batch_device_packed ! Note that this%status_host remains untouched.
 
        if (copy_) then
          select case (source)
          case (batch_not_packed)
            ! copy from unpacked host array to device
            call batch_write_unpacked_to_device(this)
          case (batch_packed)
            ! copy from packed host array to device
            call batch_write_packed_to_device(this, async_)
          end select
        end if
      case (batch_packed)
        call this%allocate_packed_host()
        this%status_of = batch_packed
        this%status_host = batch_packed
 
        if (copy_) then
          if (this%type() == type_float) then
            call dbatch_pack_copy(this)
          else if (this%type() == type_cmplx) then
            call zbatch_pack_copy(this)
          end if
        end if
        if (this%own_memory) call this%deallocate_unpacked_host()
      case (batch_not_packed)
        call messages_not_implemented("Error: batch_do_pack called with BATCH_NOT_PACKED as target.")
      end select
    end if
 
    select case (target)
    case (batch_device_packed)
      this%device_buffer_count = this%device_buffer_count + 1
    case (batch_packed)
      this%host_buffer_count = this%host_buffer_count + 1
    end select
 
    call profiling_out("BATCH_DO_PACK")
  end subroutine batch_do_pack_target
 
 
  ! ----------------------------------------------------
  subroutine batch_do_unpack(this, copy, force, async)
    class(batch_t),     intent(inout) :: this
    logical, optional,  intent(in)    :: copy
    logical, optional,  intent(in)    :: force
    logical, optional,  intent(in)    :: async
 
    logical :: copy_, force_, async_
    integer               :: source, target
 
    push_sub(batch_do_unpack)
 
    call profiling_in("BATCH_DO_UNPACK")
 
    copy_ = optional_default(copy, .true.)
 
    force_ = optional_default(force, .false.)
 
    async_ = optional_default(async, .false.)
 
    ! get source and target states for this batch
    source = this%status()
    select case (source)
    case (batch_not_packed)
      target = source
    case (batch_packed)
      target = batch_not_packed
    case (batch_device_packed)
      target = this%status_host
    end select
 
    ! only do something if target is different from source
    if (source /= target) then
      select case (source)
      case (batch_packed)
        if (this%host_buffer_count == 1 .or. force_) then
          if (this%own_memory) call this%allocate_unpacked_host()
          ! unpack from packed_host to unpacked_host
          if (copy_ .or. this%own_memory) then
            if (this%type() == type_float) then
              call dbatch_unpack_copy(this)
            else if (this%type() == type_cmplx) then
              call zbatch_unpack_copy(this)
            end if
          end if
          call this%deallocate_packed_host()
          this%status_host = target
          this%status_of = target
          this%host_buffer_count = 1
        end if
        this%host_buffer_count = this%host_buffer_count - 1
      case (batch_device_packed)
        if (this%device_buffer_count == 1 .or. force_) then
          if (copy_) then
            select case (target)
              ! unpack from packed_device to unpacked_host
            case (batch_not_packed)
              call batch_read_device_to_unpacked(this)
              ! unpack from packed_device to packed_host
            case (batch_packed)
              call batch_read_device_to_packed(this, async_)
            end select
          end if
          if (async_) then
            this%needs_finish_unpack = .true.
          else
            call this%deallocate_packed_device()
          end if
          this%status_of = target
          this%device_buffer_count = 1
        end if
        this%device_buffer_count = this%device_buffer_count - 1
      end select
    end if
 
    call profiling_out("BATCH_DO_UNPACK")
 
    pop_sub(batch_do_unpack)
  end subroutine batch_do_unpack
 
  ! ----------------------------------------------------
  subroutine batch_finish_unpack(this)
    class(batch_t),      intent(inout)  :: this
 
    push_sub(batch_finish_unpack)
    if (this%needs_finish_unpack) then
      call accel_finish()
      call this%deallocate_packed_device()
      this%needs_finish_unpack = .false.
    end if
    pop_sub(batch_finish_unpack)
  end subroutine batch_finish_unpack
 
  ! ----------------------------------------------------
 
  subroutine batch_write_unpacked_to_device(this)
    class(batch_t),      intent(inout)  :: this
 
    integer :: ist, ist2
    integer(int64) :: unroll, bsize, gsize
    type(accel_mem_t) :: tmp
    type(accel_kernel_t), pointer :: kernel
 
    push_sub(batch_write_unpacked_to_device)
 
    call profiling_in("BATCH_WRT_UNPACK_ACCEL")
    if (this%nst_linear == 1) then
      ! we can copy directly
      if (this%type() == type_float) then
        call accel_write_buffer(this%ff_device, ubound(this%dff_linear, dim=1), this%dff_linear(:, 1))
      else if (this%type() == type_cmplx) then
        call accel_write_buffer(this%ff_device, ubound(this%zff_linear, dim=1), this%zff_linear(:, 1))
      else
        assert(.false.)
      end if
 
    else
      ! we copy to a temporary array and then we re-arrange data
 
      if (this%type() == type_float) then
        kernel => dpack
      else
        kernel => zpack
      end if
 
      unroll = min(int(gpu_pack_max_buffer_size, int64), this%pack_size(1))
 
      call accel_create_buffer(tmp, accel_mem_read_only, this%type(), unroll*this%pack_size(2))
 
      do ist = 1, this%nst_linear, int(unroll, int32)
 
        ! copy a number 'unroll' of states to the buffer
        do ist2 = ist, min(ist + int(unroll, int32) - 1, this%nst_linear)
 
          if (this%type() == type_float) then
            call accel_write_buffer(tmp, ubound(this%dff_linear, dim=1, kind=int64), this%dff_linear(:, ist2), &
              offset = (ist2 - ist)*this%pack_size(2))
          else
            call accel_write_buffer(tmp, ubound(this%zff_linear, dim=1, kind=int64), this%zff_linear(:, ist2), &
              offset = (ist2 - ist)*this%pack_size(2))
          end if
        end do
 
        ! now call a kernel to rearrange the data
        call accel_set_kernel_arg(kernel, 0, int(this%pack_size(1), int32))
        call accel_set_kernel_arg(kernel, 1, this%np)
        call accel_set_kernel_arg(kernel, 2, this%nst_linear)
        call accel_set_kernel_arg(kernel, 3, ist - 1)
        call accel_set_kernel_arg(kernel, 4, tmp)
        call accel_set_kernel_arg(kernel, 5, this%ff_device)
 
        call profiling_in("GPU_PACK")
 
        ! Compute the grid size
        bsize = accel_max_block_size()/unroll
        call accel_grid_size(this%pack_size(2), bsize, gsize)
 
        call accel_kernel_run(kernel, (/gsize, 1_int64/), (/bsize, unroll/))
 
        if (this%type() == type_float) then
          call profiling_count_transfers(unroll*this%pack_size(2), m_one)
        else
          call profiling_count_transfers(unroll*this%pack_size(2), m_zi)
        end if
 
        call accel_finish()
        call profiling_out("GPU_PACK")
 
      end do
 
      call accel_free_buffer(tmp)
 
    end if
 
    call profiling_out("BATCH_WRT_UNPACK_ACCEL")
    pop_sub(batch_write_unpacked_to_device)
  end subroutine batch_write_unpacked_to_device
 
  ! ------------------------------------------------------------------
 
  subroutine batch_read_device_to_unpacked(this)
    class(batch_t),      intent(inout) :: this
 
    integer :: ist, ist2
    integer(int64) :: unroll, bsize, gsize
    type(accel_mem_t) :: tmp
    type(accel_kernel_t), pointer :: kernel
 
    push_sub(batch_read_device_to_unpacked)
    call profiling_in("BATCH_READ_UNPACKED_ACCEL")
 
    if (this%nst_linear == 1) then
      ! we can copy directly
      if (this%type() == type_float) then
        call accel_read_buffer(this%ff_device, ubound(this%dff_linear, dim=1), this%dff_linear(:, 1))
      else
        call accel_read_buffer(this%ff_device, ubound(this%zff_linear, dim=1), this%zff_linear(:, 1))
      end if
    else
 
      unroll = min(int(gpu_pack_max_buffer_size, int64), this%pack_size(1))
 
      ! we use a kernel to move to a temporary array and then we read
      call accel_create_buffer(tmp, accel_mem_write_only, this%type(), unroll*this%pack_size(2))
 
      if (this%type() == type_float) then
        kernel => dunpack
      else
        kernel => zunpack
      end if
 
      do ist = 1, this%nst_linear, int(unroll, int32)
        call accel_set_kernel_arg(kernel, 0, int(this%pack_size(1), int32))
        call accel_set_kernel_arg(kernel, 1, this%np)
        call accel_set_kernel_arg(kernel, 2, this%nst_linear)
        call accel_set_kernel_arg(kernel, 3, ist - 1)
        call accel_set_kernel_arg(kernel, 4, this%ff_device)
        call accel_set_kernel_arg(kernel, 5, tmp)
 
        call profiling_in("GPU_UNPACK")
 
        ! Compute the grid size
        bsize = accel_max_block_size()/unroll
        call accel_grid_size(this%pack_size(2), bsize, gsize)
 
        call accel_kernel_run(kernel, (/1_int64, gsize/), (/unroll, bsize/))
 
        if (this%type() == type_float) then
          call profiling_count_transfers(unroll*this%pack_size(2), m_one)
        else
          call profiling_count_transfers(unroll*this%pack_size(2), m_zi)
        end if
 
        call accel_finish()
        call profiling_out("GPU_UNPACK")
 
        ! copy a number 'unroll' of states from the buffer
        do ist2 = ist, min(ist + int(unroll, int32) - 1, this%nst_linear)
 
          if (this%type() == type_float) then
            call accel_read_buffer(tmp, ubound(this%dff_linear, dim=1, kind=int64), this%dff_linear(:, ist2), &
              offset = (ist2 - ist)*this%pack_size(2))
          else
            call accel_read_buffer(tmp, ubound(this%zff_linear, dim=1, kind=int64), this%zff_linear(:, ist2), &
              offset = (ist2 - ist)*this%pack_size(2))
          end if
        end do
 
      end do
 
      call accel_free_buffer(tmp)
    end if
 
    call profiling_out("BATCH_READ_UNPACKED_ACCEL")
    pop_sub(batch_read_device_to_unpacked)
  end subroutine batch_read_device_to_unpacked
 
  ! ------------------------------------------------------------------
  subroutine batch_write_packed_to_device(this, async)
    class(batch_t),      intent(inout)  :: this
    logical,   optional, intent(in)     :: async
 
 
    push_sub(batch_write_packed_to_device)
 
    call profiling_in("BATCH_WRITE_PACKED_ACCEL")
    if (this%type() == type_float) then
      call accel_write_buffer(this%ff_device, this%pack_size(1), this%pack_size(2), this%dff_pack, async=async)
    else
      call accel_write_buffer(this%ff_device, this%pack_size(1), this%pack_size(2), this%zff_pack, async=async)
    end if
    call profiling_out("BATCH_WRITE_PACKED_ACCEL")
 
    pop_sub(batch_write_packed_to_device)
  end subroutine batch_write_packed_to_device
 
  ! ------------------------------------------------------------------
  subroutine batch_read_device_to_packed(this, async)
    class(batch_t),      intent(inout) :: this
    logical,   optional, intent(in)    :: async
 
 
    push_sub(batch_read_device_to_packed)
 
    call profiling_in("BATCH_READ_PACKED_ACCEL")
    if (this%type() == type_float) then
      call accel_read_buffer(this%ff_device, this%pack_size(1), this%pack_size(2), this%dff_pack, async=async)
    else
      call accel_read_buffer(this%ff_device, this%pack_size(1), this%pack_size(2), this%zff_pack, async=async)
    end if
    call profiling_out("BATCH_READ_PACKED_ACCEL")
 
    pop_sub(batch_read_device_to_packed)
  end subroutine batch_read_device_to_packed
 
  ! ------------------------------------------------------
  integer function batch_inv_index(this, cind) result(index)
    class(batch_t),     intent(in)    :: this
    integer,            intent(in)    :: cind(:)
 
    do index = 1, this%nst_linear
      if (all(cind(1:this%ndims) == this%ist_idim_index(index, 1:this%ndims))) exit
    end do
 
    assert(index <= this%nst_linear)
 
  end function batch_inv_index
 
  ! ------------------------------------------------------
  !
  integer pure function batch_ist_idim_to_linear(this, cind) result(index)
    class(batch_t),     intent(in)    :: this
    integer,            intent(in)    :: cind(:)
 
    if (ubound(cind, dim = 1) == 1) then
      index = cind(1)
    else
      index = (cind(1) - 1)*this%dim + cind(2)
    end if
 
  end function batch_ist_idim_to_linear
 
  ! ------------------------------------------------------
  !
  integer pure function batch_linear_to_ist(this, linear_index) result(ist)
    class(batch_t),     intent(in)    :: this
    integer,            intent(in)    :: linear_index
 
    ist = this%ist_idim_index(linear_index, 1)
 
  end function batch_linear_to_ist
 
  ! ------------------------------------------------------
  !
  integer pure function batch_linear_to_idim(this, linear_index) result(idim)
    class(batch_t),     intent(in)    :: this
    integer,            intent(in)    :: linear_index
 
    idim = this%ist_idim_index(linear_index, 2)
 
  end function batch_linear_to_idim
 
  ! ------------------------------------------------------
  !
  subroutine batch_remote_access_start(this, mpi_grp, rma_win)
    class(batch_t),  intent(inout) :: this
    type(mpi_grp_t), intent(in)    :: mpi_grp
    type(mpi_win),   intent(out)   :: rma_win
 
    push_sub(batch_remote_access_start)
 
    if (mpi_grp%size > 1) then
 
      assert(.not. accel_is_enabled())
 
      call this%do_pack()
 
      assert(int(product(this%pack_size)*types_get_size(this%type())) > 0) ! In case of overflow, this becomes negative
      if (this%type() == type_cmplx) then
#ifdef HAVE_MPI
        call mpi_win_create(this%zff_pack(1, 1), int(product(this%pack_size)*types_get_size(this%type()), mpi_address_kind), &
          types_get_size(this%type()), mpi_info_null, mpi_grp%comm, rma_win)
#endif
      else if (this%type() == type_float) then
#ifdef HAVE_MPI
        call mpi_win_create(this%dff_pack(1, 1), int(product(this%pack_size)*types_get_size(this%type()), mpi_address_kind), &
          types_get_size(this%type()), mpi_info_null, mpi_grp%comm, rma_win)
#endif
      else
        message(1) = "Internal error: unknown batch type in batch_remote_access_start."
        call messages_fatal(1)
      end if
 
    else
      rma_win = mpi_win_null
    end if
 
    pop_sub(batch_remote_access_start)
  end subroutine batch_remote_access_start
 
  ! ------------------------------------------------------
  !
  subroutine batch_remote_access_stop(this, rma_win)
    class(batch_t),     intent(inout) :: this
    type(mpi_win),      intent(inout) :: rma_win
 
    push_sub(batch_remote_access_stop)
 
    if (rma_win /= mpi_win_null) then
#ifdef HAVE_MPI
      call mpi_win_free(rma_win)
#endif
      call this%do_unpack()
    end if
 
    pop_sub(batch_remote_access_stop)
  end subroutine batch_remote_access_stop
 
  ! --------------------------------------------------------------
  !
  subroutine batch_copy_data_to(this, np, dest, async)
    class(batch_t),    intent(in)    :: this
    integer,           intent(in)    :: np
    class(batch_t),    intent(inout) :: dest
    logical, optional, intent(in)    :: async
 
    integer(int64), dimension(3) :: gsizes, bsizes
    integer :: ist, ip
 
    push_sub(batch_copy_data_to)
    call profiling_in("BATCH_COPY_DATA_TO")
 
    ! this routine can be used to copy data between batches of different type
    call this%check_compatibility_with(dest, type_check=.false.)
 
    if (this%type() == dest%type()) then
      select case (this%status())
      case (batch_device_packed)
        call accel_set_kernel_arg(kernel_copy, 0, np)
        call accel_set_kernel_arg(kernel_copy, 1, this%ff_device)
        call accel_set_kernel_arg(kernel_copy, 2, log2(int(this%pack_size_real(1), int32)))
        call accel_set_kernel_arg(kernel_copy, 3, dest%ff_device)
        call accel_set_kernel_arg(kernel_copy, 4, log2(int(dest%pack_size_real(1), int32)))
 
        ! Compute the grid (extend to another dimensions if the size of the problem is too big)
        call accel_grid_size_extend_dim(int(np, int64), dest%pack_size_real(1), gsizes, bsizes, kernel_copy)
 
        call accel_kernel_run(kernel_copy, gsizes, bsizes)
 
        if(.not. optional_default(async, .false.)) call accel_finish()
 
      case (batch_packed)
        if (np*this%pack_size(1) > huge(0_int32)) then
          ! BLAS cannot handle 8-byte integers, so we need a special version here
          do ip = 1, np
            if (dest%type() == type_float) then
              call blas_copy(int(this%pack_size(1), int32), this%dff_pack(1, ip), 1, dest%dff_pack(1, ip), 1)
            else
              call blas_copy(int(this%pack_size(1), int32), this%zff_pack(1, ip), 1, dest%zff_pack(1, ip), 1)
            end if
          end do
        else
          if (dest%type() == type_float) then
            call blas_copy(int(this%pack_size(1)*np, int32), this%dff_pack(1, 1), 1, dest%dff_pack(1, 1), 1)
          else
            call blas_copy(int(this%pack_size(1)*np, int32), this%zff_pack(1, 1), 1, dest%zff_pack(1, 1), 1)
          end if
        end if
 
      case (batch_not_packed)
        do ist = 1, dest%nst_linear
          if (dest%type() == type_cmplx) then
            call blas_copy(np, this%zff_linear(1, ist), 1, dest%zff_linear(1, ist), 1)
          else
            call blas_copy(np, this%dff_linear(1, ist), 1, dest%dff_linear(1, ist), 1)
          end if
        end do
 
      end select
    else if (this%type() == type_cmplx) then
      ! copy complex -> real
      select case (this%status())
      case (batch_device_packed)
        call accel_set_kernel_arg(kernel_copy_complex_to_real, 0, np)
        call accel_set_kernel_arg(kernel_copy_complex_to_real, 1, this%ff_device)
        call accel_set_kernel_arg(kernel_copy_complex_to_real, 2, log2(int(this%pack_size_real(1), int32)))
        call accel_set_kernel_arg(kernel_copy_complex_to_real, 3, dest%ff_device)
        call accel_set_kernel_arg(kernel_copy_complex_to_real, 4, log2(int(dest%pack_size_real(1), int32)))
 
        ! Compute the grid (extend to another dimensions if the size of the problem is too big)
        call accel_grid_size_extend_dim(int(np, int64), dest%pack_size_real(1), gsizes, bsizes, kernel_copy_complex_to_real)
 
        call accel_kernel_run(kernel_copy_complex_to_real, gsizes, bsizes)
 
        if(.not. optional_default(async, .false.)) call accel_finish()
 
      case (batch_packed)
        !$omp parallel do private(ist, ip)
        do ip = 1, np
          !$omp simd
          do ist = 1, dest%nst_linear
            dest%dff_pack(ist, ip) = real(this%zff_pack(ist, ip), real64)
          end do
        end do
        !$omp end parallel do
      case (batch_not_packed)
        !$omp parallel private(ist, ip)
        do ist = 1, dest%nst_linear
          !$omp do
          do ip = 1, np
            dest%dff_linear(ip, ist) = real(this%zff_linear(ip, ist), real64)
          end do
          !$omp end do nowait
        end do
        !$omp end parallel
      end select
    else if (this%type() == type_float) then
      ! copy real -> complex
      select case (this%status())
      case (batch_device_packed)
        call accel_set_kernel_arg(kernel_copy_real_to_complex, 0, np)
        call accel_set_kernel_arg(kernel_copy_real_to_complex, 1, this%ff_device)
        call accel_set_kernel_arg(kernel_copy_real_to_complex, 2, log2(int(this%pack_size_real(1), int32)))
        call accel_set_kernel_arg(kernel_copy_real_to_complex, 3, dest%ff_device)
        call accel_set_kernel_arg(kernel_copy_real_to_complex, 4, log2(int(dest%pack_size_real(1), int32)))
 
        ! Compute the grid (extend to another dimensions if the size of the problem is too big)
        call accel_grid_size_extend_dim(int(np, int64), this%pack_size_real(1), gsizes, bsizes, kernel_copy_real_to_complex)
 
        call accel_kernel_run(kernel_copy_real_to_complex, gsizes, bsizes)
 
        if(.not. optional_default(async, .false.)) call accel_finish()
 
      case (batch_packed)
        !$omp parallel do private(ist, ip)
        do ip = 1, np
          !$omp simd
          do ist = 1, dest%nst_linear
            dest%zff_pack(ist, ip) = cmplx(this%dff_pack(ist, ip), m_zero, real64)
          end do
        end do
        !$omp end parallel do
      case (batch_not_packed)
        !$omp parallel private(ist, ip)
        do ist = 1, dest%nst_linear
          !$omp do
          do ip = 1, np
            dest%zff_linear(ip, ist) = cmplx(this%dff_linear(ip, ist), m_zero, real64)
          end do
          !$omp end do nowait
        end do
        !$omp end parallel
      end select
    else
      message(1) = "Error! This should not happen."
      call messages_fatal(1)
    end if
 
    call profiling_out("BATCH_COPY_DATA_TO")
    pop_sub(batch_copy_data_to)
  end subroutine batch_copy_data_to
 
  ! --------------------------------------------------------------
  !
  subroutine batch_check_compatibility_with(this, target, only_check_dim, type_check)
    class(batch_t),    intent(in) :: this
    class(batch_t),    intent(in) :: target
    logical, optional, intent(in) :: only_check_dim
    logical, optional, intent(in) :: type_check
 
    push_sub(batch_check_compatibility_with)
 
    if (optional_default(type_check, .true.)) then
      assert(this%type() == target%type())
    end if
    if (.not. optional_default(only_check_dim, .false.)) then
      assert(this%nst_linear == target%nst_linear)
    end if
    assert(this%status() == target%status())
    assert(this%dim == target%dim)
 
    pop_sub(batch_check_compatibility_with)
 
  end subroutine batch_check_compatibility_with
 
!--------------------------------------------------------------
  subroutine batch_build_indices(this, st_start, st_end)
    class(batch_t), intent(inout) :: this
    integer,        intent(in)    :: st_start
    integer,        intent(in)    :: st_end
 
    integer :: idim, ii, ist
 
    push_sub(batch_build_indices)
 
    do ist = st_start, st_end
      ! now we also populate the linear array
      do idim = 1, this%dim
        ii = this%dim*(ist - st_start) + idim
        this%ist_idim_index(ii, 1) = ist
        this%ist_idim_index(ii, 2) = idim
      end do
      this%ist(ist - st_start + 1) = ist
    end do
 
    ! compute packed sizes
    this%pack_size(1) = pad_pow2(this%nst_linear)
    this%pack_size(2) = this%np
    if (accel_is_enabled()) this%pack_size(2) = accel_padded_size(this%pack_size(2))
 
    this%pack_size_real = this%pack_size
    if (type_is_complex(this%type())) this%pack_size_real(1) = 2*this%pack_size_real(1)
 
    pop_sub(batch_build_indices)
  end subroutine batch_build_indices
 
 
#include "real.F90"
#include "batch_inc.F90"
#include "undef.F90"
 
#include "complex.F90"
#include "batch_inc.F90"
#include "undef.F90"
 
end module batch_oct_m
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: