batch_ops.F90

Go to the documentation of this file.

!! Copyright (C) 2008 X. Andrade
!!
!! 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_ops_oct_m
  use accel_oct_m
  use batch_oct_m
  use blas_oct_m
  use debug_oct_m
  use iso_c_binding
  use global_oct_m
  use lalg_basic_oct_m
  use math_oct_m
  use messages_oct_m
  use profiling_oct_m
  use types_oct_m
 
  implicit none
 
  private
  public ::                         &
    batch_set_zero,                 &
    batch_axpy,                     &
    batch_axpby,                    &
    batch_scal,                     &
    batch_scal2v,                   &
    batch_xpay,                     &
    batch_set_state,                &
    batch_get_state,                &
    batch_get_points,               &
    batch_set_points,               &
    batch_points_block_size,        &
    batch_pointwise_mul,            &
    batch_mul_mf,                   &
    batch_add_with_map,             &
    batch_copy_with_map,            &
    batch_copy_precision,           &
    dbatch_axpy_function,           &
    zbatch_axpy_function,           &
    dbatch_ax_function_py,          &
    zbatch_ax_function_py,          &
    dbatch_copy_with_map_to_array,  &
    zbatch_copy_with_map_to_array,  &
    batch_split_complex
 
  interface batch_axpy
    module procedure dbatch_axpy_const
    module procedure zbatch_axpy_const
    module procedure sbatch_axpy_const
    module procedure cbatch_axpy_const
    module procedure dbatch_axpy_vec
    module procedure zbatch_axpy_vec
  end interface batch_axpy
 
  interface batch_scal
    module procedure dbatch_scal_const
    module procedure zbatch_scal_const
    module procedure dbatch_scal_vec
    module procedure zbatch_scal_vec
  end interface batch_scal
 
  interface batch_axpby
    module procedure dbatch_axpby
    module procedure zbatch_axpby
  end interface batch_axpby
 
  interface batch_scal2v
    module procedure dbatch_scal2v
    module procedure zbatch_scal2v
  end interface batch_scal2v
 
  interface batch_xpay
    module procedure dbatch_xpay_vec
    module procedure zbatch_xpay_vec
    module procedure dbatch_xpay_const
    module procedure zbatch_xpay_const
  end interface batch_xpay
 
  interface batch_add_with_map
    module procedure batch_add_with_map_cpu
    module procedure batch_add_with_map_accel
  end interface batch_add_with_map
 
  interface batch_copy_with_map
    module procedure batch_copy_with_map_cpu
    module procedure batch_copy_with_map_accel
  end interface batch_copy_with_map
 
  interface batch_set_state
    module procedure dbatch_set_state1
    module procedure zbatch_set_state1
    module procedure sbatch_set_state1
    module procedure cbatch_set_state1
    module procedure dbatch_set_state2
    module procedure zbatch_set_state2
    module procedure sbatch_set_state2
    module procedure cbatch_set_state2
    module procedure dbatch_set_state3
    module procedure zbatch_set_state3
    module procedure sbatch_set_state3
    module procedure cbatch_set_state3
  end interface batch_set_state
 
  interface batch_get_state
    module procedure dbatch_get_state1
    module procedure zbatch_get_state1
    module procedure sbatch_get_state1
    module procedure cbatch_get_state1
    module procedure dbatch_get_state2
    module procedure zbatch_get_state2
    module procedure sbatch_get_state2
    module procedure cbatch_get_state2
    module procedure dbatch_get_state3
    module procedure zbatch_get_state3
    module procedure sbatch_get_state3
    module procedure cbatch_get_state3
  end interface batch_get_state
 
  interface batch_get_points
    module procedure dbatch_get_points
    module procedure zbatch_get_points
    module procedure batch_get_points_accel
  end interface batch_get_points
 
  interface batch_set_points
    module procedure dbatch_set_points
    module procedure zbatch_set_points
    module procedure batch_set_points_accel
  end interface batch_set_points
 
  interface batch_pointwise_mul
    module procedure batch_mul_cj
  end interface batch_pointwise_mul
 
  interface batch_mul_mf
    module procedure dbatch_mul_mf
    module procedure zbatch_mul_mf
  end interface batch_mul_mf
 
 
contains
 
  !--------------------------------------------------------------
  subroutine batch_set_zero(this, np, async)
    class(batch_t),     intent(inout) :: this
    integer, optional,  intent(in)    :: np
    logical, optional,  intent(in)    :: async
 
    integer :: ist_linear, ist, ip, np_
 
    push_sub(batch_set_zero)
 
    assert(not_in_openmp())
 
    call profiling_in("BATCH_SET_ZERO")
 
    select case (this%status())
    case (batch_device_packed)
      np_ = optional_default(np, int(this%pack_size(2), int32))
      assert(np_ <= int(this%pack_size(2), int32))
      call accel_set_buffer_to_zero(this%ff_device, this%type(), (int(this%pack_size(1), int32) * np_), async=async)
 
    case (batch_packed)
      np_ = optional_default(np, int(this%pack_size(2), int32))
      assert(np_ <= int(this%pack_size(2), int32))
      if (this%type() == type_float) then
        !$omp parallel do private(ist) schedule(static)
        do ip = 1, np_
          !$omp simd
          do ist = 1, int(this%pack_size(1), int32)
            this%dff_pack(ist, ip) = m_zero
          end do
        end do
      else if (this%type() == type_float_single) then
        !$omp parallel do private(ist) schedule(static)
        do ip = 1, np_
          !$omp simd
          do ist = 1, int(this%pack_size(1), int32)
            this%sff_pack(ist, ip) = 0.0_real32
          end do
        end do
      else if (this%type() == type_cmplx_single) then
        !$omp parallel do private(ist) schedule(static)
        do ip = 1, np_
          !$omp simd
          do ist = 1, int(this%pack_size(1), int32)
            this%cff_pack(ist, ip) = cmplx(0.0_real32, 0.0_real32, real32)
          end do
        end do
      else
        !$omp parallel do private(ist) schedule(static)
        do ip = 1, np_
          !$omp simd
          do ist = 1, int(this%pack_size(1), int32)
            this%zff_pack(ist, ip) = m_z0
          end do
        end do
      end if
 
    case (batch_not_packed)
      if (this%type() == type_float) then
        np_ = optional_default(np, ubound(this%dff_linear, dim=1))
        assert(np_ <= ubound(this%dff_linear, dim=1))
        do ist_linear = 1, this%nst_linear
          !$omp parallel do schedule(static)
          do ip = 1, np_
            this%dff_linear(ip, ist_linear) = m_zero
          end do
        end do
      else if (this%type() == type_float_single) then
        np_ = optional_default(np, ubound(this%sff_linear, dim=1))
        assert(np_ <= ubound(this%sff_linear, dim=1))
        do ist_linear = 1, this%nst_linear
          !$omp parallel do schedule(static)
          do ip = 1, np_
            this%sff_linear(ip, ist_linear) = 0.0_real32
          end do
        end do
      else if (this%type() == type_cmplx_single) then
        np_ = optional_default(np, ubound(this%cff_linear, dim=1))
        assert(np_ <= ubound(this%cff_linear, dim=1))
        do ist_linear = 1, this%nst_linear
          !$omp parallel do schedule(static)
          do ip = 1, np_
            this%cff_linear(ip, ist_linear) = cmplx(0.0_real32, 0.0_real32, real32)
          end do
        end do
      else
        np_ = optional_default(np, ubound(this%zff_linear, dim=1))
        assert(np_ <= ubound(this%zff_linear, dim=1))
        do ist_linear = 1, this%nst_linear
          !$omp parallel do schedule(static)
          do ip = 1, np_
            this%zff_linear(ip, ist_linear) = m_z0
          end do
        end do
      end if
 
    case default
      message(1) = "batch_set_zero: unknown batch status."
      call messages_fatal(1)
 
    end select
 
    call profiling_out("BATCH_SET_ZERO")
 
    pop_sub(batch_set_zero)
  end subroutine batch_set_zero
 
  ! --------------------------------------------------------------
  !
  subroutine batch_get_points_accel(this, sp, ep, psi, ldpsi1, ldpsi2)
    class(batch_t),      intent(in)    :: this
    integer,             intent(in)    :: sp
    integer,             intent(in)    :: ep
    type(accel_mem_t),   intent(inout) :: psi
    integer,             intent(in)    :: ldpsi1
    integer,             intent(in)    :: ldpsi2
 
    integer :: tsize, ii, it
    type(accel_kernel_t), save :: kernel
    integer, allocatable :: linear_to_ist(:), linear_to_idim(:)
    type(accel_mem_t) :: buff_linear_to_ist, buff_linear_to_idim
 
    push_sub(batch_get_points_accel)
    call profiling_in("GET_POINTS")
 
    select case (this%status())
    case (batch_not_packed, batch_packed)
      call messages_not_implemented('batch_get_points_accel for non-GPU batches')
 
    case (batch_device_packed)
 
      tsize = types_get_size(this%type())/types_get_size(type_float)
      safe_allocate(linear_to_ist(1:this%nst_linear*tsize))
      safe_allocate(linear_to_idim(1:this%nst_linear*tsize))
      do ii = 1, this%nst_linear
        do it = 1, tsize
          linear_to_ist(tsize*(ii-1)+it) = tsize*(this%linear_to_ist(ii) - 1) + it - 1
          linear_to_idim(tsize*(ii-1)+it) = this%linear_to_idim(ii) - 1
        end do
      end do
 
      call accel_create_buffer(buff_linear_to_ist, accel_mem_read_only, type_integer, this%nst_linear*tsize)
      call accel_write_buffer(buff_linear_to_ist, this%nst_linear*tsize, linear_to_ist)
      call accel_create_buffer(buff_linear_to_idim, accel_mem_read_only, type_integer, this%nst_linear*tsize)
      call accel_write_buffer(buff_linear_to_idim, this%nst_linear*tsize, linear_to_idim)
 
      call accel_kernel_start_call(kernel, 'points.cu', 'get_points<double>')
 
      call accel_set_kernel_arg(kernel, 0, sp)
      call accel_set_kernel_arg(kernel, 1, ep)
      call accel_set_kernel_arg(kernel, 2, buff_linear_to_ist)
      call accel_set_kernel_arg(kernel, 3, buff_linear_to_idim)
      call accel_set_kernel_arg(kernel, 4, this%nst_linear*tsize)
      call accel_set_kernel_arg(kernel, 5, this%ff_device)
      call accel_set_kernel_arg(kernel, 6, int(this%pack_size_real(1), int32))
      call accel_set_kernel_arg(kernel, 7, psi)
      call accel_set_kernel_arg(kernel, 8, ldpsi1*tsize)
      call accel_set_kernel_arg(kernel, 9, ldpsi2)
 
      call accel_kernel_run(kernel, (/1_int64, int(ep - sp + 1, int64)/), (/this%pack_size_real(1), 1_int64/))
 
      call accel_free_buffer(buff_linear_to_ist)
      call accel_free_buffer(buff_linear_to_idim)
      safe_deallocate_a(linear_to_ist)
      safe_deallocate_a(linear_to_idim)
 
    end select
 
    call profiling_out("GET_POINTS")
 
    pop_sub(batch_get_points_accel)
  end subroutine batch_get_points_accel
 
  ! --------------------------------------------------------------
  !
  subroutine batch_set_points_accel(this, sp, ep, psi, ldpsi1, ldpsi2)
    class(batch_t),      intent(inout) :: this
    integer,             intent(in)    :: sp
    integer,             intent(in)    :: ep
    type(accel_mem_t),   intent(in)    :: psi
    integer,             intent(in)    :: ldpsi1
    integer,             intent(in)    :: ldpsi2
 
    integer :: tsize, ii, it
    type(accel_kernel_t), save :: kernel
    integer, allocatable :: linear_to_ist(:), linear_to_idim(:)
    type(accel_mem_t) :: buff_linear_to_ist, buff_linear_to_idim
 
    push_sub(batch_set_points_accel)
    call profiling_in("SET_POINTS")
 
    select case (this%status())
    case (batch_not_packed, batch_packed)
      call messages_not_implemented('batch_set_points_accel for non-GPU batches')
 
    case (batch_device_packed)
 
      tsize = types_get_size(this%type())/types_get_size(type_float)
      safe_allocate(linear_to_ist(1:this%nst_linear*tsize))
      safe_allocate(linear_to_idim(1:this%nst_linear*tsize))
      do ii = 1, this%nst_linear
        do it = 1, tsize
          linear_to_ist(tsize*(ii-1)+it) = tsize*(this%linear_to_ist(ii) - 1) + it - 1
          linear_to_idim(tsize*(ii-1)+it) = this%linear_to_idim(ii) - 1
        end do
      end do
 
      call accel_create_buffer(buff_linear_to_ist, accel_mem_read_only, type_integer, this%nst_linear*tsize)
      call accel_write_buffer(buff_linear_to_ist, this%nst_linear*tsize, linear_to_ist)
      call accel_create_buffer(buff_linear_to_idim, accel_mem_read_only, type_integer, this%nst_linear*tsize)
      call accel_write_buffer(buff_linear_to_idim, this%nst_linear*tsize, linear_to_idim)
 
      call accel_kernel_start_call(kernel, 'points.cu', 'set_points<double>')
 
      call accel_set_kernel_arg(kernel, 0, sp)
      call accel_set_kernel_arg(kernel, 1, ep)
      call accel_set_kernel_arg(kernel, 2, buff_linear_to_ist)
      call accel_set_kernel_arg(kernel, 3, buff_linear_to_idim)
      call accel_set_kernel_arg(kernel, 4, this%nst_linear*tsize)
      call accel_set_kernel_arg(kernel, 5, psi)
      call accel_set_kernel_arg(kernel, 6, ldpsi1*tsize)
      call accel_set_kernel_arg(kernel, 7, ldpsi2)
      call accel_set_kernel_arg(kernel, 8, this%ff_device)
      call accel_set_kernel_arg(kernel, 9, int(this%pack_size_real(1), int32))
 
      call accel_kernel_run(kernel, (/1_int64, int(ep - sp + 1, int64)/), (/this%pack_size_real(1), 1_int64/))
 
      call accel_free_buffer(buff_linear_to_ist)
      call accel_free_buffer(buff_linear_to_idim)
      safe_deallocate_a(linear_to_ist)
      safe_deallocate_a(linear_to_idim)
 
    end select
 
    call profiling_out("SET_POINTS")
 
    pop_sub(batch_set_points_accel)
  end subroutine batch_set_points_accel
 
  ! -------------------------
  !
  integer pure function batch_points_block_size() result(block_size)
 
    block_size = 61440
 
  end function batch_points_block_size
 
! -------------------------
  subroutine batch_mul_cj(np, xx, yy, zz, conjugate_yy)
    integer,           intent(in)    :: np
    class(batch_t),    intent(in)    :: xx
    class(batch_t),    intent(in)    :: yy
    class(batch_t),    intent(inout) :: zz
    logical, optional, intent(in)    :: conjugate_yy
 
    integer :: ii, ip
    logical :: conj_yy
    integer(int64), dimension(3) :: gsizes, bsizes
    type(accel_kernel_t), save, target :: kernel_zmul_conj
    type(accel_kernel_t), save, target :: kernel_zmul
    type(accel_kernel_t), save, target :: kernel_dmul
    type(accel_kernel_t), pointer :: kernel
    push_sub(batch_mul_cj)
 
    call xx%check_compatibility_with(yy)
    call xx%check_compatibility_with(zz)
 
    conj_yy = optional_default(conjugate_yy, .true.)
 
    select case (xx%status())
    case (batch_not_packed)
      if (xx%type() == type_cmplx) then
        if (conj_yy) then
          !$omp parallel private(ii, ip)
          do ii = 1, xx%nst_linear
            !$omp do schedule(static)
            do ip = 1, np
              zz%zff_linear(ip, ii) = xx%zff_linear(ip, ii)*conjg(yy%zff_linear(ip, ii))
            end do
            !$omp end do
          end do
          !$omp end parallel
        else
          !$omp parallel private(ii, ip)
          do ii = 1, xx%nst_linear
            !$omp do schedule(static)
            do ip = 1, np
              zz%zff_linear(ip, ii) = xx%zff_linear(ip, ii)*yy%zff_linear(ip, ii)
            end do
            !$omp end do
          end do
          !$omp end parallel
        end if
      else
        !$omp parallel private(ii, ip)
        do ii = 1, xx%nst_linear
          !$omp do schedule(static)
          do ip = 1, np
            zz%dff_linear(ip, ii) = xx%dff_linear(ip, ii)*yy%dff_linear(ip, ii)
          end do
          !$omp end do
        end do
        !$omp end parallel
      end if
 
    case (batch_packed)
      if (xx%type() == type_cmplx) then
        if (conj_yy) then
          !$omp parallel do private(ii)
          do ip = 1, np
            !$omp simd
            do ii = 1, xx%nst_linear
              zz%zff_pack(ii, ip) = xx%zff_pack(ii, ip)*conjg(yy%zff_pack(ii, ip))
            end do
          end do
        else
          !$omp parallel do private(ii)
          do ip = 1, np
            !$omp simd
            do ii = 1, xx%nst_linear
              zz%zff_pack(ii, ip) = xx%zff_pack(ii, ip)*yy%zff_pack(ii, ip)
            end do
          end do
        end if
      else
        !$omp parallel do private(ii)
        do ip = 1, np
          !$omp simd
          do ii = 1, xx%nst_linear
            zz%dff_pack(ii, ip) = xx%dff_pack(ii, ip)*yy%dff_pack(ii, ip)
          end do
        end do
      end if
 
    case (batch_device_packed)
      if (xx%type() == type_cmplx) then
        if (conj_yy) then
          kernel => kernel_zmul_conj
          call accel_kernel_start_call(kernel, 'batch_mul.cu', 'mul_batch<complex<double>, Conjugation::Yes>')
        else
          kernel => kernel_zmul
          call accel_kernel_start_call(kernel, 'batch_mul.cu', 'mul_batch<complex<double>, Conjugation::No>')
        end if
 
        call accel_set_kernel_arg(kernel, 0, np)
        call accel_set_kernel_arg(kernel, 1, xx%ff_device)
        call accel_set_kernel_arg(kernel, 2, log2(int(xx%pack_size(1), int32)))
        call accel_set_kernel_arg(kernel, 3, yy%ff_device)
        call accel_set_kernel_arg(kernel, 4, log2(int(yy%pack_size(1), int32)))
        call accel_set_kernel_arg(kernel, 5, zz%ff_device)
        call accel_set_kernel_arg(kernel, 6, log2(int(zz%pack_size(1), int32)))
 
        call accel_grid_size_extend_dim(int(np, int64), xx%pack_size(1), gsizes, bsizes, kernel)
      else
        kernel => kernel_dmul
        call accel_kernel_start_call(kernel, 'batch_mul.cu', 'mul_batch<double>')
 
        call accel_set_kernel_arg(kernel, 0, np)
        call accel_set_kernel_arg(kernel, 1, xx%ff_device)
        call accel_set_kernel_arg(kernel, 2, log2(int(xx%pack_size_real(1), int32)))
        call accel_set_kernel_arg(kernel, 3, yy%ff_device)
        call accel_set_kernel_arg(kernel, 4, log2(int(yy%pack_size_real(1), int32)))
        call accel_set_kernel_arg(kernel, 5, zz%ff_device)
        call accel_set_kernel_arg(kernel, 6, log2(int(zz%pack_size_real(1), int32)))
 
        call accel_grid_size_extend_dim(int(np, int64), xx%pack_size_real(1), gsizes, bsizes, kernel)
      end if
 
      call accel_kernel_run(kernel, gsizes, bsizes)
    end select
 
    pop_sub(batch_mul_cj)
  end subroutine batch_mul_cj
 
! -------------------------
  subroutine batch_add_with_map_cpu(np, map, xx, yy, zz)
    integer,           intent(in)    :: np
    integer,           intent(in)    :: map(:)
    class(batch_t),    intent(in)    :: xx
    class(batch_t),    intent(in)    :: yy
    class(batch_t),    intent(inout) :: zz
    type(accel_mem_t) :: buff_map
 
    push_sub(batch_add_with_map_cpu)
 
    if (xx%status() /= batch_device_packed) then
      if (xx%type() == type_float) then
        call dbatch_add_with_map(np, map, xx, yy, zz)
      else
        call zbatch_add_with_map(np, map, xx, yy, zz)
      end if
    else
      ! copy map to GPU if not already there
      call accel_create_buffer(buff_map, accel_mem_read_only, type_integer, np)
      call accel_write_buffer(buff_map, np, map)
      call batch_add_with_map_accel(np, buff_map, xx, yy, zz)
      call accel_free_buffer(buff_map)
    end if
 
    pop_sub(batch_add_with_map_cpu)
  end subroutine batch_add_with_map_cpu
 
! -------------------------
  subroutine batch_add_with_map_accel(np, map, xx, yy, zz)
    integer,            intent(in)    :: np
    class(accel_mem_t), intent(in)    :: map
    class(batch_t),     intent(in)    :: xx
    class(batch_t),     intent(in)    :: yy
    class(batch_t),     intent(inout) :: zz
 
    type(accel_kernel_t), save :: kernel
    integer(int64), dimension(3) :: gsizes, bsizes
 
    push_sub(batch_add_with_map_accel)
 
    call accel_kernel_start_call(kernel, 'copy.cu', 'add_with_map<double>')
 
    call accel_set_kernel_arg(kernel, 0, np)
    call accel_set_kernel_arg(kernel, 1, map)
    call accel_set_kernel_arg(kernel, 2, xx%ff_device)
    call accel_set_kernel_arg(kernel, 3, log2(int(xx%pack_size_real(1), int32)))
    call accel_set_kernel_arg(kernel, 4, yy%ff_device)
    call accel_set_kernel_arg(kernel, 5, log2(int(yy%pack_size_real(1), int32)))
    call accel_set_kernel_arg(kernel, 6, zz%ff_device)
    call accel_set_kernel_arg(kernel, 7, log2(int(zz%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), xx%pack_size_real(1), gsizes, bsizes, kernel)
 
    call accel_kernel_run(kernel, gsizes, bsizes)
 
    pop_sub(batch_add_with_map_accel)
  end subroutine batch_add_with_map_accel
 
! -------------------------
  subroutine batch_copy_with_map_cpu(np, map, xx, yy)
    integer,           intent(in)    :: np
    integer,           intent(in)    :: map(:)
    class(batch_t),    intent(in)    :: xx
    class(batch_t),    intent(inout) :: yy
    type(accel_mem_t) :: buff_map
 
    push_sub(batch_copy_with_map_cpu)
 
    if (xx%status() /= batch_device_packed) then
      if (xx%type() == type_float) then
        call dbatch_copy_with_map(np, map, xx, yy)
      else
        call zbatch_copy_with_map(np, map, xx, yy)
      end if
    else
      ! copy map to GPU if not already there
      call accel_create_buffer(buff_map, accel_mem_read_only, type_integer, np)
      call accel_write_buffer(buff_map, np, map)
      call batch_copy_with_map_accel(np, buff_map, xx, yy)
      call accel_free_buffer(buff_map)
    end if
 
    pop_sub(batch_copy_with_map_cpu)
  end subroutine batch_copy_with_map_cpu
 
! -------------------------
  subroutine batch_copy_with_map_accel(np, map, xx, yy)
    integer,            intent(in)    :: np
    class(accel_mem_t), intent(in)    :: map
    class(batch_t),     intent(in)    :: xx
    class(batch_t),     intent(inout) :: yy
 
    type(accel_kernel_t), save :: kernel
    integer(int64), dimension(3) :: gsizes, bsizes
 
    push_sub(batch_copy_with_map_accel)
 
    call accel_kernel_start_call(kernel, 'copy.cu', 'copy_with_map<double>')
 
    ! execute only if map has at least one element
    if (np > 0) then
      call accel_set_kernel_arg(kernel, 0, np)
      call accel_set_kernel_arg(kernel, 1, map)
      call accel_set_kernel_arg(kernel, 2, xx%ff_device)
      call accel_set_kernel_arg(kernel, 3, log2(int(xx%pack_size_real(1), int32)))
      call accel_set_kernel_arg(kernel, 4, yy%ff_device)
      call accel_set_kernel_arg(kernel, 5, log2(int(yy%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), xx%pack_size_real(1), gsizes, bsizes, kernel)
 
      call accel_kernel_run(kernel, gsizes, bsizes)
    end if
 
    pop_sub(batch_copy_with_map_accel)
  end subroutine batch_copy_with_map_accel
 
  ! -------------------------
  !
  subroutine batch_split_complex(np, xx, yy, zz)
    integer,           intent(in)    :: np
    class(batch_t),    intent(in)    :: xx
    class(batch_t),    intent(inout) :: yy
    class(batch_t),    intent(inout) :: zz
 
    integer :: ist_linear, ip
    type(accel_kernel_t), save :: kernel
    integer(int64), dimension(3) :: gsizes, bsizes
 
    push_sub(batch_split_complex)
 
    assert(xx%type() == type_cmplx)
    assert(yy%type() == type_float)
    assert(zz%type() == type_float)
    assert(xx%status() == yy%status())
    assert(xx%status() == zz%status())
 
    select case (xx%status())
    case (batch_not_packed)
      do ist_linear = 1, xx%nst_linear
        !$omp parallel do schedule(static)
        do ip = 1, np
          yy%dff_linear(ip, ist_linear) = real(xx%zff_linear(ip, ist_linear), real64)
          zz%dff_linear(ip, ist_linear) = aimag(xx%zff_linear(ip, ist_linear))
        end do
      end do
    case (batch_packed)
      !$omp parallel do private(ist_linear) schedule(static)
      do ip = 1, np
        do ist_linear = 1, xx%nst_linear
          yy%dff_pack(ist_linear, ip) = real(xx%zff_pack(ist_linear, ip), real64)
          zz%dff_pack(ist_linear, ip) = aimag(xx%zff_pack(ist_linear, ip))
        end do
      end do
    case (batch_device_packed)
      call accel_kernel_start_call(kernel, 'split.cu', 'split_complex<double>')
 
      call accel_set_kernel_arg(kernel, 0, int(xx%pack_size(2), int32))
      call accel_set_kernel_arg(kernel, 1, xx%ff_device)
      call accel_set_kernel_arg(kernel, 2, log2(int(xx%pack_size(1), int32)))
      call accel_set_kernel_arg(kernel, 3, yy%ff_device)
      call accel_set_kernel_arg(kernel, 4, log2(int(yy%pack_size(1), int32)))
      call accel_set_kernel_arg(kernel, 5, zz%ff_device)
      call accel_set_kernel_arg(kernel, 6, log2(int(zz%pack_size(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), xx%pack_size(1), gsizes, bsizes, kernel)
 
      call accel_kernel_run(kernel, gsizes, bsizes)
    end select
 
    pop_sub(batch_split_complex)
  end subroutine batch_split_complex
 
  ! --------------------------------------------------------------
  subroutine batch_copy_precision(src, dst, np)
    class(batch_t),    intent(in)    :: src
    class(batch_t),    intent(inout) :: dst
    integer,           intent(in)    :: np
 
    integer :: ist, ip
    integer(int64) :: gsizes(3), bsizes(3)
    logical :: is_cplx, to_single
    character(len=47) :: kernel_name
    type(accel_kernel_t), save, target :: kernels(1:2, 1:2)
    type(accel_kernel_t), pointer :: kernel
 
    push_sub(batch_copy_precision)
 
    assert(src%nst_linear == dst%nst_linear)
    assert(src%status() == dst%status())
    assert(type_is_complex(src%type()) .eqv. type_is_complex(dst%type()))
    assert(type_is_single(src%type()) .neqv. type_is_single(dst%type()))
 
    is_cplx = type_is_complex(src%type())
    to_single = type_is_single(dst%type())
 
    select case (src%status())
    case (batch_device_packed)
      if (is_cplx) then
        kernel_name = merge('precision_copy<complex<double>, complex<float>>', &
          'precision_copy<complex<float>, complex<double>>', to_single)
      else
        kernel_name = merge('precision_copy<double, float>', 'precision_copy<float, double>', to_single)
      end if
      kernel => kernels(merge(2, 1, is_cplx), merge(2, 1, to_single))
      call accel_kernel_start_call(kernel, 'precision_copy.cu', trim(kernel_name))
      call accel_set_kernel_arg(kernel, 0, np)
      call accel_set_kernel_arg(kernel, 1, src%nst_linear)
      call accel_set_kernel_arg(kernel, 2, src%ff_device)
      call accel_set_kernel_arg(kernel, 3, log2(int(src%pack_size(1), int32)))
      call accel_set_kernel_arg(kernel, 4, dst%ff_device)
      call accel_set_kernel_arg(kernel, 5, log2(int(dst%pack_size(1), int32)))
      call accel_grid_size_extend_dim(int(np, int64), dst%pack_size(1), gsizes, bsizes, kernel)
      call accel_kernel_run(kernel, gsizes, bsizes)
      call accel_finish()
 
    case (batch_packed)
      if (to_single) then
        if (is_cplx) then
          !$omp parallel do private(ist) schedule(static)
          do ip = 1, np
            do ist = 1, dst%nst_linear
              dst%cff_pack(ist, ip) = cmplx(src%zff_pack(ist, ip), kind=real32)
            end do
          end do
        else
          !$omp parallel do private(ist) schedule(static)
          do ip = 1, np
            do ist = 1, dst%nst_linear
              dst%sff_pack(ist, ip) = real(src%dff_pack(ist, ip), kind=real32)
            end do
          end do
        end if
      else
        if (is_cplx) then
          !$omp parallel do private(ist) schedule(static)
          do ip = 1, np
            do ist = 1, dst%nst_linear
              dst%zff_pack(ist, ip) = src%cff_pack(ist, ip)
            end do
          end do
        else
          !$omp parallel do private(ist) schedule(static)
          do ip = 1, np
            do ist = 1, dst%nst_linear
              dst%dff_pack(ist, ip) = src%sff_pack(ist, ip)
            end do
          end do
        end if
      end if
 
    case (batch_not_packed)
      if (to_single) then
        if (is_cplx) then
          !$omp parallel do private(ip) schedule(static)
          do ist = 1, dst%nst_linear
            do ip = 1, np
              dst%cff_linear(ip, ist) = cmplx(src%zff_linear(ip, ist), kind=real32)
            end do
          end do
        else
          !$omp parallel do private(ip) schedule(static)
          do ist = 1, dst%nst_linear
            do ip = 1, np
              dst%sff_linear(ip, ist) = real(src%dff_linear(ip, ist), kind=real32)
            end do
          end do
        end if
      else
        if (is_cplx) then
          !$omp parallel do private(ip) schedule(static)
          do ist = 1, dst%nst_linear
            do ip = 1, np
              dst%zff_linear(ip, ist) = src%cff_linear(ip, ist)
            end do
          end do
        else
          !$omp parallel do private(ip) schedule(static)
          do ist = 1, dst%nst_linear
            do ip = 1, np
              dst%dff_linear(ip, ist) = src%sff_linear(ip, ist)
            end do
          end do
        end if
      end if
 
    end select
 
    pop_sub(batch_copy_precision)
  end subroutine batch_copy_precision
 
 
#undef SPECIALIZED
#include "undef.F90"
#include "real.F90"
#include "batch_ops_inc.F90"
 
#include "undef.F90"
#include "complex.F90"
#include "batch_ops_inc.F90"
 
! Specialized kernels for nst_linear=2
#define SPECIALIZED 2
#include "undef.F90"
#include "real.F90"
#include "batch_ops_inc.F90"
 
#include "undef.F90"
#include "complex.F90"
#include "batch_ops_inc.F90"
#undef SPECIALIZED
 
! Specialized kernels for nst_linear=3
#define SPECIALIZED 3
#include "undef.F90"
#include "real.F90"
#include "batch_ops_inc.F90"
 
#include "undef.F90"
#include "complex.F90"
#include "batch_ops_inc.F90"
#undef SPECIALIZED
 
! Specialized kernels for nst_linear=4
#define SPECIALIZED 4
#include "undef.F90"
#include "real.F90"
#include "batch_ops_inc.F90"
 
#include "undef.F90"
#include "complex.F90"
#include "batch_ops_inc.F90"
#undef SPECIALIZED
 
! Specialized kernels for nst_linear=6
#define SPECIALIZED 6
#include "undef.F90"
#include "real.F90"
#include "batch_ops_inc.F90"
 
#include "undef.F90"
#include "complex.F90"
#include "batch_ops_inc.F90"
#undef SPECIALIZED
 
! Specialized kernels for nst_linear=8
#define SPECIALIZED 8
#include "undef.F90"
#include "real.F90"
#include "batch_ops_inc.F90"
 
#include "undef.F90"
#include "complex.F90"
#include "batch_ops_inc.F90"
#undef SPECIALIZED
 
! Single-precision batch operations: batch_ops_inc.F90 skips everything that needs a
! single-precision lalg wrapper and keeps the subset the exchange path uses.
#include "undef.F90"
#include "real32.F90"
#include "batch_ops_inc.F90"
 
#include "undef.F90"
#include "complex32.F90"
#include "batch_ops_inc.F90"
 
end module batch_ops_oct_m
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: