!------------------------------------------------------------------------------------------------
! Copyright (c) 2021 - 2025, Oleg Shatrov
! All rights reserved.
! This file is part of dtFFT library.

! dtFFT 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 3 of the License, or
! (at your option) any later version.

! dtFFT 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, see <https://www.gnu.org/licenses/>.
!------------------------------------------------------------------------------------------------
#include "dtfft_config.h"
#include "dtfft.f03"
module dtfft_reshape_plan_base
!! Base class for reshape/transpose plans
use iso_c_binding
use iso_fortran_env
use dtfft_abstract_reshape_handle,      only: reshape_container, execute_args
use dtfft_abstract_backend,             only: backend_helper
#ifdef DTFFT_WITH_COMPRESSION
use dtfft_abstract_compressor
#endif
use dtfft_config
#ifdef DTFFT_WITH_ZFP
use dtfft_compressor_zfp,               only: compressor_zfp
#endif
use dtfft_errors
use dtfft_parameters
use dtfft_pencil,                       only: pencil_init, pencil
#ifdef DTFFT_WITH_CUDA
# ifdef NCCL_HAVE_COMMREGISTER
use dtfft_abstract_backend,             only: NCCL_REGISTER_PREALLOC_SIZE
# endif
use dtfft_interface_cuda,               only: load_cuda
use dtfft_interface_cuda_runtime
use dtfft_interface_nvrtc,              only: load_nvrtc
# ifdef DTFFT_WITH_NCCL
use dtfft_interface_nccl
# endif
# ifdef DTFFT_WITH_NVSHMEM
use dtfft_interface_nvshmem
# endif
#endif
use dtfft_reshape_handle_generic,       only: reshape_handle_generic
use dtfft_reshape_handle_datatype,      only: reshape_handle_datatype
use dtfft_utils
#include "_dtfft_cuda.h"
#include "_dtfft_mpi.h"
#include "_dtfft_private.h"
#include "_dtfft_profile.h"
implicit none
private
public :: reshape_plan_base
public :: allocate_plans, destroy_plans
public :: get_aux_bytes_generic
public :: execute_autotune
public :: report_timings

    type :: reshape_plan_base
    !! Base class for reshape/transpose plans
        type(dtfft_backend_t)     :: backend
            !! Backend
        type(backend_helper)      :: helper
            !! Backend helper
        integer(int64)            :: min_buffer_size
            !! Minimal buffer size for transposition
        type(dtfft_platform_t)    :: platform
            !! Platform used for transposition
        type(dtfft_stream_t)      :: stream
            !! CUDA stream
        type(reshape_container), allocatable :: plans(:)
            !! Plans for each reshape operation
        type(string), allocatable :: names(:)
            !! Names of each reshape operation
        integer(int32), allocatable :: colors(:)
    contains
        procedure,  pass(self), non_overridable :: init                 !! Initializes reshape plan base
        procedure,  pass(self), non_overridable :: get_async_active     !! Returns .true. if any of the plans is running asynchronously
        procedure,  pass(self), non_overridable :: get_aux_bytes        !! Returns auxiliary buffer size
        procedure,  pass(self), non_overridable :: is_aux_needed        !! Returns .true. if aux buffer is needed
        procedure,  pass(self), non_overridable :: get_backend          !! Returns backend id
        procedure,  pass(self), non_overridable :: report_backends      !! Reports backends used by each plan. Should be used only for adaptive backend
        procedure,  pass(self), non_overridable :: is_nvshmem
        procedure,  pass(self), non_overridable :: execute              !! Executes reshape
        procedure,  pass(self), non_overridable :: execute_end          !! Finishes asynchronous reshape
        procedure,  pass(self), non_overridable :: mem_alloc            !! Allocates memory
        procedure,  pass(self), non_overridable :: mem_free             !! Frees memory
        procedure,  pass(self)                  :: destroy              !! Destroys reshape plan
#ifdef DTFFT_WITH_COMPRESSION
        procedure,  pass(self)                  :: report_compression   !! Reports compression statistics
#endif
    end type reshape_plan_base

contains

#ifdef DTFFT_WITH_COMPRESSION
    function init(self, platform, effort, backend, autotune_effort, compression_config, base_dtype, base_storage) result(error_code)
#else
    function init(self, platform, effort, backend, autotune_effort) result(error_code)
#endif
    !! Initializes reshape plan base
        class(reshape_plan_base),           intent(inout)   :: self                 !! Reshape class
        type(dtfft_platform_t),             intent(in)      :: platform             !! Execution platform
        type(dtfft_effort_t),               intent(in)      :: effort               !! Effort level
        type(dtfft_backend_t),              intent(in)      :: backend
        type(dtfft_effort_t),               intent(in)      :: autotune_effort
#ifdef DTFFT_WITH_COMPRESSION
        type(dtfft_compression_config_t),   intent(in)      :: compression_config   !! Compression configuration
        TYPE_MPI_DATATYPE,                  intent(in)      :: base_dtype           !! Base datatype
        integer(int64),                     intent(in)      :: base_storage         !! Number of bytes to store single element of base datatype
#endif
        integer(int32)                                      :: error_code           !! Error code

        error_code = DTFFT_SUCCESS
        self%stream = NULL_STREAM
        self%backend = backend
        self%platform = platform

        if ( self%backend == DTFFT_BACKEND_ADAPTIVE .and. effort%val < autotune_effort%val ) then
            error_code = DTFFT_ERROR_INVALID_BACKEND
            return
        endif

        if ( platform == DTFFT_PLATFORM_HOST ) then
            if ( .not. get_conf_datatype_enabled() .and. .not. get_conf_mpi_enabled() .and. effort%val >= autotune_effort%val) then
                error_code = DTFFT_ERROR_BACKENDS_DISABLED
                return
            endif

            if ( (is_backend_nccl(self%backend) .or. is_backend_nvshmem(self%backend)) .and. effort%val < autotune_effort%val) then
                error_code = DTFFT_ERROR_INVALID_PLATFORM_BACKEND
                return
            endif
#ifdef DTFFT_WITH_CUDA
        else
            if ( .not.get_conf_mpi_enabled() .and. .not.get_conf_nccl_enabled() .and. .not.get_conf_nvshmem_enabled() .and. effort%val >= autotune_effort%val) then
                error_code = DTFFT_ERROR_BACKENDS_DISABLED
                return
            endif

            ! TODO: fix this
            if ( self%backend == DTFFT_BACKEND_ADAPTIVE ) then
                error_code = DTFFT_ERROR_INVALID_PLATFORM_BACKEND
                return
            endif

            CHECK_CALL( load_cuda(), error_code )
            CHECK_CALL( load_nvrtc(), error_code )
            self%stream = get_conf_stream()
#endif
        endif
#ifdef DTFFT_WITH_COMPRESSION

        if ( (effort%val < autotune_effort%val .and. is_backend_compressed(backend)) .or. (effort%val >=autotune_effort%val .and. get_conf_compression_enabled()) ) then
            CHECK_CALL( check_compression_config(compression_config), error_code )

            block
                class(abstract_compressor), allocatable :: compressor

                if ( compression_config%compression_lib == DTFFT_COMPRESSION_LIB_ZFP ) then
# ifdef DTFFT_WITH_ZFP
                    allocate( compressor_zfp :: compressor )
# else
                    INTERNAL_ERROR("ZFP compression requested but dtFFT is not built with ZFP support")
# endif
                endif

                ! Checking compression parameters by creating temporary compressor
                CHECK_CALL( compressor%create(2_int8, compression_config, platform, base_dtype, base_storage, DIMS_PERMUTE_NONE), error_code )
                call compressor%destroy()
                deallocate(compressor)
            endblock
        endif
#else
        if ( is_backend_compressed(self%backend) .and. effort%val < autotune_effort%val ) then
            error_code = DTFFT_ERROR_INVALID_BACKEND
            return
        endif
#endif
    end function init

    subroutine execute(self, in, out, reshape_type, exec_type, aux, error_code)
    !! Executes reshape plan
        class(reshape_plan_base),     intent(inout) :: self           !! Transposition class
        type(c_ptr),                  intent(in)    :: in             !! Incoming buffer
        type(c_ptr),                  intent(in)    :: out            !! Resulting buffer
        integer(int32),               intent(in)    :: reshape_type   !! Type of reshape to execute
        type(c_ptr),                  intent(in)    :: aux            !! Optional auxiliary buffer
        type(async_exec_t),           intent(in)    :: exec_type      !! Type of execution (sync/async)
        integer(int32),    optional,  intent(out)   :: error_code     !! Error code
        type(execute_args)    :: kwargs     !! Additional arguments for execution
        integer(int32)        :: ierr       !! Error code
        logical :: is_pipe, is_fused
        type(dtfft_backend_t) :: backend

        backend = self%plans(reshape_type)%p%get_backend()
        is_pipe = is_backend_pipelined(backend)
        is_fused = is_backend_fused(backend) .or. is_backend_compressed(backend)
#ifdef DTFFT_DEBUG
        if ( is_same_ptr(in, out) .or. is_same_ptr(in, aux) .or. is_same_ptr(out, aux) ) then
            INTERNAL_ERROR("reshape_plan_base: wrong pointers for reshape "//self%names(reshape_type)%raw)
        endif
        if ( is_null_ptr(in) .or. is_null_ptr(out) ) then
            INTERNAL_ERROR("reshape_plan_base: null pointer detected for reshape "//self%names(reshape_type)%raw)
        endif
        if ( is_null_ptr(aux) .and. (is_pipe .or. is_fused) ) then
            INTERNAL_ERROR("reshape_plan_base: aux null pointer detected for reshape "//self%names(reshape_type)%raw)
        endif
#endif
        ! if ( reshape_type == DTFFT_RESHAPE_Z_PENCILS_TO_BRICKS%val ) then
        !     ! select type ( plan => self%plans(reshape_type)%p )
        !     ! class is ( reshape_handle_datatype )
        !     block
        !         integer :: comm_rank

        !         call MPI_Comm_rank(MPI_COMM_WORLD, comm_rank, ierr)
        !         if (comm_rank == 0) print*,'comm_rank = ',comm_rank, 'executing', transfer(in, int64), '->', transfer(out, int64), 'aux = ', transfer(aux, int64)
        !     endblock
        !     ! endselect
        ! endif

        REGION_BEGIN(self%names(reshape_type)%raw, self%colors(reshape_type))
        kwargs%exec_type = exec_type
        kwargs%stream = self%stream
        kwargs%p1 = aux
        call self%plans(reshape_type)%p%execute(in, out, kwargs, ierr)
        if( present( error_code ) ) error_code = ierr
        REGION_END(self%names(reshape_type)%raw)
    end subroutine execute

    subroutine execute_end(self, in, out, reshape_type, aux, error_code)
    !! Finishes asynchronous reshape
        class(reshape_plan_base),   intent(inout) :: self           !! Reshape class
        type(c_ptr),                intent(in)    :: in             !! Incoming buffer
        type(c_ptr),                intent(in)    :: out            !! Resulting buffer
        integer(int32),             intent(in)    :: reshape_type   !! Type of reshape to execute
        type(c_ptr),                intent(in)    :: aux            !! Optional auxiliary buffer
        integer(int32),             intent(out)   :: error_code     !! Error code
        ! real(real32),   pointer :: pin(:)   !! Source buffer
        ! real(real32),   pointer :: pout(:)  !! Destination buffer
        ! real(real32),   pointer :: paux(:)  !! Aux buffer
        type(execute_args)      :: kwargs   !! Additional arguments for execution

        REGION_BEGIN(self%names(reshape_type)%raw//' end', self%colors(reshape_type))
        kwargs%p1 = in
        kwargs%p2 = out
        kwargs%p3 = aux
        kwargs%stream = self%stream
        call self%plans(reshape_type)%p%execute_end(kwargs, error_code)
        REGION_END(self%names(reshape_type)%raw//' end')
    end subroutine execute_end

    logical function get_async_active(self)
    !! Returns .true. if any of the plans is running asynchronously
        class(reshape_plan_base),   intent(in)    :: self           !! Reshape class
        integer(int32)  :: i

        get_async_active = .false.
        do i = lbound(self%plans, dim=1), ubound(self%plans, dim=1)
            if ( allocated( self%plans(i)%p ) ) then
                get_async_active = get_async_active .or. self%plans(i)%p%get_async_active()
            endif
        enddo
    end function get_async_active

    pure integer(int64) function get_aux_bytes(self)
    !! Returns maximum auxiliary memory size needed by reshape plan
        class(reshape_plan_base), intent(in)    :: self             !! Reshape class
        get_aux_bytes = 0_int64
        if ( .not. allocated( self%plans ) ) return
        get_aux_bytes = get_aux_bytes_generic(self%plans)
    end function get_aux_bytes

    pure logical function is_aux_needed(self)
    !! Returns true if aux is needed. false otherwise
        class(reshape_plan_base), intent(in)    :: self             !! Reshape class
        is_aux_needed = self%get_aux_bytes() > 0
    end function is_aux_needed

    type(dtfft_backend_t) function get_backend(self)
    !! Returns plan backend
        class(reshape_plan_base), intent(in)    :: self             !! Reshape class
        get_backend = self%backend
    end function get_backend

    subroutine destroy(self)
    !! Destroys reshape plan
        class(reshape_plan_base),   intent(inout) :: self           !! Reshape class
        integer(int32) :: i

        if ( allocated( self%plans ) ) then
            call destroy_plans(self%plans)
            deallocate( self%plans )
        endif
        call self%helper%destroy()
        if ( allocated(self%names) ) then
            do i = lbound(self%names, dim=1), ubound(self%names, dim=1)
                call self%names(i)%destroy()
            enddo
            deallocate(self%names)
        endif
        if ( allocated(self%colors) ) deallocate(self%colors)
    end subroutine destroy

    subroutine report_backends(self)
    !! Reports backends used by each plan. Should be used only for adaptive backend
        class(reshape_plan_base),   intent(in)  :: self         !! Reshape class
        integer(int32) :: i
        type(dtfft_backend_t) :: backend

        if ( .not. self%backend == DTFFT_BACKEND_ADAPTIVE ) then
            INTERNAL_ERROR("report_backends: .not. self%backend == DTFFT_BACKEND_ADAPTIVE")
        endif

        do i = lbound(self%plans, dim=1), ubound(self%plans, dim=1)
            if ( allocated( self%plans(i)%p ) ) then
                backend = self%plans(i)%p%get_backend()
                if ( backend == BACKEND_DUMMY ) cycle
                WRITE_REPORT("      "//self%names(i)%raw//" :  "//dtfft_get_backend_string(backend))
            endif
        enddo
    end subroutine report_backends

    pure logical function is_nvshmem(self)
        class(reshape_plan_base),   intent(in)  :: self         !! Reshape class

        is_nvshmem = is_nvshmem_generic(self%plans, self%backend)
    end function is_nvshmem

    pure logical function is_nvshmem_generic(plans, backend)
        type(reshape_container),    intent(in)  :: plans(:)          !! Reshape handles
        type(dtfft_backend_t),      intent(in)  :: backend
        integer(int32)  :: i

        is_nvshmem_generic = .false.
        if ( is_backend_nvshmem( backend ) ) then
            is_nvshmem_generic = .true.
            return
        endif
        if ( .not. backend == DTFFT_BACKEND_ADAPTIVE ) return
        do i = 1, size(plans)
            if ( allocated( plans(i)%p ) ) then
                is_nvshmem_generic = is_nvshmem_generic .or. is_backend_nvshmem( plans(i)%p%get_backend() )
            endif
        enddo
    end function is_nvshmem_generic

#ifdef DTFFT_WITH_COMPRESSION
    subroutine report_compression(self)
    !! Reports compression statistics
        class(reshape_plan_base),   intent(in)  :: self         !! Reshape class
        integer(int32) :: i

        do i = lbound(self%plans, dim=1), ubound(self%plans, dim=1)
            if ( allocated( self%plans(i)%p ) ) then
                call self%plans(i)%p%report_compression(self%names(i)%raw)
            endif
        enddo
    end subroutine report_compression
#endif

    subroutine mem_alloc(self, comm, alloc_bytes, ptr, error_code)
    !! Allocates memory based on selected backend
        class(reshape_plan_base),   intent(inout) :: self               !! Reshape class
        TYPE_MPI_COMM,              intent(in)    :: comm               !! MPI communicator
        integer(int64),             intent(in)    :: alloc_bytes        !! Number of bytes to allocate
        type(c_ptr),                intent(out)   :: ptr                !! Pointer to the allocated memory
        integer(int32),             intent(out)   :: error_code         !! Error code
        type(dtfft_backend_t) :: backend

        backend = self%backend; if ( self%is_nvshmem() ) backend = DTFFT_BACKEND_CUFFTMP
        call alloc_mem(self%platform, self%helper, backend, comm, alloc_bytes, ptr, error_code)
    end subroutine mem_alloc

    subroutine mem_free(self, ptr, error_code)
    !! Frees memory allocated with mem_alloc
        class(reshape_plan_base),   intent(inout) :: self               !! Reshape class
        type(c_ptr),                intent(in)    :: ptr                !! Pointer to the memory to free
        integer(int32),             intent(out)   :: error_code         !! Error code
        type(dtfft_backend_t) :: backend

        backend = self%backend; if ( self%is_nvshmem() ) backend = DTFFT_BACKEND_CUFFTMP
        call free_mem(self%platform, self%helper, backend, ptr, error_code)
    end subroutine mem_free

    subroutine alloc_mem(platform, helper, backend, comm, alloc_bytes, ptr, error_code)
    !! Allocates memory based on ``backend``
        type(dtfft_platform_t), intent(in)    :: platform           !! Execution platform
        type(backend_helper),   intent(inout) :: helper             !! Backend helper
        type(dtfft_backend_t),  intent(in)    :: backend            !! Backend to allocate memory for
        TYPE_MPI_COMM,          intent(in)    :: comm               !! MPI communicator
        integer(int64),         intent(in)    :: alloc_bytes        !! Number of bytes to allocate
        type(c_ptr),            intent(out)   :: ptr                !! Pointer to the allocated memory
        integer(int32),         intent(out)   :: error_code         !! Error code
#ifdef DTFFT_WITH_CUDA
        integer(int64)  :: free_mem_avail, total_mem_avail
#endif

        error_code = DTFFT_SUCCESS
        if ( platform == DTFFT_PLATFORM_HOST ) then
            ptr = mem_alloc_host(alloc_bytes)
#ifdef DTFFT_WITH_CUDA
        else
            CUDA_CALL( cudaMemGetInfo(free_mem_avail, total_mem_avail) )
# ifdef DTFFT_DEBUG
            block
            integer(int64) :: min_mem, max_mem, min_free_mem, max_free_mem
            integer(int32) :: mpi_err

                call MPI_Allreduce(alloc_bytes, max_mem, 1, MPI_INTEGER8, MPI_MAX, comm, mpi_err)
                call MPI_Allreduce(alloc_bytes, min_mem, 1, MPI_INTEGER8, MPI_MIN, comm, mpi_err)
                call MPI_Allreduce(free_mem_avail, max_free_mem, 1, MPI_INTEGER8, MPI_MAX, comm, mpi_err)
                call MPI_Allreduce(free_mem_avail, min_free_mem, 1, MPI_INTEGER8, MPI_MIN, comm, mpi_err)
                WRITE_DEBUG("Trying to allocate "//to_str(min_mem)//"/"//to_str(max_mem)//" (min/max) bytes for backend: '"//dtfft_get_backend_string(backend)//"'")
                WRITE_DEBUG("Free memory available: "//to_str(min_free_mem)//"/"//to_str(max_free_mem)//" (min/max) bytes")
            endblock
# endif
            if ( alloc_bytes > free_mem_avail ) then
                error_code = DTFFT_ERROR_ALLOC_FAILED
                return
            endif
            if ( is_backend_nccl(backend) ) then
# ifdef DTFFT_WITH_NCCL
#   ifdef NCCL_HAVE_MEMALLOC
                error_code = ncclMemAlloc(ptr, alloc_bytes)
#   else
                error_code = cudaMalloc(ptr, alloc_bytes)
#   endif
#   ifdef NCCL_HAVE_COMMREGISTER
                if ( error_code == cudaSuccess .and. helper%should_register ) then
                    block
                    type(c_ptr), allocatable :: temp(:,:)
                    type(c_ptr) :: handle

                        if ( size(helper%nccl_register, dim=2) == helper%nccl_register_size ) then
                            allocate( temp(2, helper%nccl_register_size + NCCL_REGISTER_PREALLOC_SIZE) )
                            temp(2, 1:helper%nccl_register_size) = helper%nccl_register(2, 1:helper%nccl_register_size)
                            deallocate( helper%nccl_register )
                            call move_alloc(temp, helper%nccl_register)
                        endif
                        helper%nccl_register_size = helper%nccl_register_size + 1

                        NCCL_CALL( ncclCommRegister(helper%nccl_comm, ptr, alloc_bytes, handle) )
                        helper%nccl_register(1, helper%nccl_register_size) = ptr
                        helper%nccl_register(2, helper%nccl_register_size) = handle
                        WRITE_DEBUG("Registered pointer "//to_str(transfer(ptr, int64)))
                    endblock
                endif
#   endif
# else
            INTERNAL_ERROR("not DTFFT_WITH_NCCL")
# endif
            else if ( is_backend_nvshmem(backend) ) then
# ifdef DTFFT_WITH_NVSHMEM
                block
                integer(int64)  :: max_alloc_bytes
                integer(int32)  :: mpi_err

                    call MPI_Allreduce(alloc_bytes, max_alloc_bytes, 1, MPI_INTEGER8, MPI_MAX, comm, mpi_err)
                    ptr = nvshmem_malloc(max_alloc_bytes)
                    if ( is_null_ptr(ptr) ) error_code = DTFFT_ERROR_ALLOC_FAILED
                endblock
# else
        INTERNAL_ERROR("not DTFFT_WITH_NVSHMEM")
# endif
            else
                error_code = cudaMalloc(ptr, alloc_bytes)
            endif
#endif
        endif
        if ( error_code /= DTFFT_SUCCESS ) error_code = DTFFT_ERROR_ALLOC_FAILED
    end subroutine alloc_mem

    subroutine free_mem(platform, helper, backend, ptr, error_code)
    !! Frees memory based on ``backend``
        type(dtfft_platform_t),         intent(in)    :: platform       !! Execution platform
        type(backend_helper),           intent(inout) :: helper         !! Backend helper
        type(dtfft_backend_t),          intent(in)    :: backend        !! Backend to free memory for
        type(c_ptr),                    intent(in)    :: ptr            !! Pointer to the memory to free
        integer(int32),                 intent(out)   :: error_code     !! Error code

        error_code = DTFFT_SUCCESS
        if ( platform == DTFFT_PLATFORM_HOST ) then
            call mem_free_host(ptr)
#ifdef DTFFT_WITH_CUDA
        else
            if ( is_backend_nccl(backend) ) then
# ifdef NCCL_HAVE_COMMREGISTER
                if ( helper%should_register ) then
                    block
                    integer(int32) :: i

                        do i = 1, size(helper%nccl_register, dim=2)
                            if ( .not. is_same_ptr(ptr, helper%nccl_register(1, i)) ) cycle
                            NCCL_CALL( ncclCommDeregister(helper%nccl_comm, helper%nccl_register(2, i)) )
                            helper%nccl_register(1, i) = c_null_ptr
                            helper%nccl_register(2, i) = c_null_ptr
                            helper%nccl_register_size = helper%nccl_register_size - 1
                            WRITE_DEBUG("Pointer "//to_str(transfer(ptr, int64))//" has been removed from registry")
                        enddo
                    endblock
                endif
# endif
# ifdef DTFFT_WITH_NCCL
#   ifdef NCCL_HAVE_MEMALLOC
                error_code = ncclMemFree(ptr)
#   else
                error_code = cudaFree(ptr)
#   endif
# else
                INTERNAL_ERROR("not DTFFT_WITH_NCCL")
# endif
            else if ( is_backend_nvshmem(backend) ) then
# ifdef DTFFT_WITH_NVSHMEM
                call nvshmem_free(ptr)
# else
                INTERNAL_ERROR("not DTFFT_WITH_NVSHMEM")
# endif
            else
                error_code = cudaFree(ptr)
            endif
#endif
        endif
        if ( error_code /= DTFFT_SUCCESS ) error_code = DTFFT_ERROR_FREE_FAILED
    end subroutine free_mem

    subroutine allocate_plans(plans, backend, backends)
    !! Allocates array of plans
        type(reshape_container),            intent(inout) :: plans(:)       !! Plans to allocate
        type(dtfft_backend_t),  optional,   intent(in)    :: backend        !! Backend to use
        type(dtfft_backend_t),  optional,   intent(in)    :: backends(:)    !! Backends to use
        integer(int32) :: i
        type(dtfft_backend_t) :: backend_

#ifdef DTFFT_DEBUG
        if ( .not.present(backend) .and. .not.present(backends) ) then
            INTERNAL_ERROR("allocate_plans: .not.present(backend) .and. .not.present(backends)")
        endif
#endif

        do i = 1, size(plans)
            if ( present(backend) ) then
                backend_ = backend
            else
                backend_ = backends(i)
            endif
            if ( backend_ == DTFFT_BACKEND_MPI_DATATYPE ) then
                allocate( reshape_handle_datatype :: plans(i)%p )
            else
                allocate( reshape_handle_generic :: plans(i)%p )
            endif
        enddo
    end subroutine allocate_plans

    pure integer(int64) function get_aux_bytes_generic(plans)
    !! Returns maximum auxiliary memory size needed by handles array
        type(reshape_container),   intent(in)  :: plans(:)          !! Reshape handles
        integer(int32)  :: i

        get_aux_bytes_generic = 0_int64
        do i = 1, size(plans)
            if ( allocated( plans(i)%p ) ) then
                get_aux_bytes_generic = max( get_aux_bytes_generic, plans(i)%p%get_aux_bytes() )
            endif
        enddo
    end function get_aux_bytes_generic

    subroutine destroy_plans(plans)
    !! Destroys array of plans
        type(reshape_container),    intent(inout) :: plans(:)       !! Plans to destroy
        integer(int32) :: i

        do i = 1, size(plans)
            if( allocated(plans(i)%p) ) then
                call plans(i)%p%destroy()
                deallocate(plans(i)%p)
            endif
        enddo
    end subroutine destroy_plans

    function execute_autotune(plans, comm, backend, platform, helper, stream, buffer_size, report_space_count) result(execution_time)
    !! Executes autotuning for given plans and returns execution time in milliseconds
        type(reshape_container),    intent(inout) :: plans(:)       !! Allocated plans
        TYPE_MPI_COMM,              intent(in)    :: comm           !! Communicator
        type(dtfft_backend_t),      intent(in)    :: backend        !! Backend to use
        type(dtfft_platform_t),     intent(in)    :: platform       !! Platform used
        type(backend_helper),       intent(inout) :: helper         !! Helper to use
        type(dtfft_stream_t),       intent(in)    :: stream         !! Stream to use
        integer(int64),             intent(in)    :: buffer_size        !!
        integer(int32), optional,   intent(in)    :: report_space_count
        real(real32)                              :: execution_time !! Execution time
        type(execute_args) :: execute_kwargs
#ifdef DTFFT_WITH_CUDA
        type(cudaEvent) :: timer_start, timer_stop
#endif
        type(c_ptr) :: in, out, aux
        ! real(real32), pointer :: pin(:)   !! Source buffer
        ! real(real32), pointer :: pout(:)  !! Destination buffer
        ! real(real32), pointer :: paux(:)  !! Auxiliary buffer
        integer(int32) :: i, iter, ierr
        integer(int64) :: float_buffer_size
        ! character(len=:), allocatable :: testing_phase
        logical :: is_aux_alloc
        integer(int32) :: n_warmup_iters, n_iters
        real(real64) :: ts, te


        n_warmup_iters = get_conf_measure_warmup_iters()
        n_iters = get_conf_measure_iters()

#ifdef DTFFT_WITH_CUDA
        if ( platform == DTFFT_PLATFORM_CUDA ) then
        CUDA_CALL( cudaEventCreate(timer_start) )
        CUDA_CALL( cudaEventCreate(timer_stop) )
        endif
#endif
        float_buffer_size = buffer_size / FLOAT_STORAGE_SIZE

        call alloc_mem(platform, helper, backend, comm, buffer_size, in, ierr); DTFFT_CHECK(ierr)
        call alloc_mem(platform, helper, backend, comm, buffer_size, out, ierr); DTFFT_CHECK(ierr)

        ! call c_f_pointer(in, pin, [float_buffer_size])
        ! call c_f_pointer(out, pout, [float_buffer_size])

        ! testing_phase = "Testing backend "//dtfft_get_backend_string(backend)
        ! REGION_BEGIN(testing_phase, COLOR_AUTOTUNE2)
        ! WRITE_INFO(testing_phase)

        call alloc_and_set_aux(platform, helper, backend, comm, aux, plans, is_aux_alloc)
        execute_kwargs%p1 = aux
        ! if( is_aux_alloc ) then
        !     execute_kwargs%p1 => paux
        ! else
        !     execute_kwargs%p1 => pout
        ! endif
        execute_kwargs%exec_type = EXEC_BLOCKING
        execute_kwargs%stream = stream

        ! REGION_BEGIN("Warmup", COLOR_TRANSPOSE)
        do iter = 1, n_warmup_iters
            do i = 1, size(plans)
                call plans(i)%p%execute(in, out, execute_kwargs, ierr)
            enddo
        enddo
#ifdef DTFFT_WITH_CUDA
        if ( platform == DTFFT_PLATFORM_CUDA ) then
            CUDA_CALL( cudaStreamSynchronize(stream) )
        endif
#endif
        ! REGION_END("Warmup")

        ! REGION_BEGIN("Measure", COLOR_EXECUTE)

        call MPI_Barrier(comm, ierr)
        if ( platform == DTFFT_PLATFORM_HOST ) then
            ts = MPI_Wtime()
#ifdef DTFFT_WITH_CUDA
        else
            CUDA_CALL( cudaEventRecord(timer_start, stream) )
#endif
        endif
        do iter = 1, n_iters
            do i = 1, size(plans)
                call plans(i)%p%execute(in, out, execute_kwargs, ierr)
            enddo
        enddo
        if ( platform == DTFFT_PLATFORM_HOST ) then
            te = MPI_Wtime()
            execution_time = real(te - ts, real32) * 1000._real32
#ifdef DTFFT_WITH_CUDA
        else
            CUDA_CALL( cudaEventRecord(timer_stop, stream) )
            CUDA_CALL( cudaEventSynchronize(timer_stop) )
            CUDA_CALL( cudaEventElapsedTime(execution_time, timer_start, timer_stop) )
#endif
        endif
        ! REGION_END("Measure")

        ! REGION_END(testing_phase)

        execution_time = report_timings(comm, execution_time, n_iters, report_space_count)

#ifdef DTFFT_WITH_CUDA
        if ( platform == DTFFT_PLATFORM_CUDA ) then
            CUDA_CALL( cudaEventDestroy(timer_start) )
            CUDA_CALL( cudaEventDestroy(timer_stop) )
        endif
#endif

        call free_mem(platform, helper, backend, in, ierr)
        call free_mem(platform, helper, backend, out, ierr)
        if ( is_aux_alloc ) then
            call free_mem(platform, helper, backend, aux, ierr)
        endif

        ! if ( platform == DTFFT_PLATFORM_CUDA .and. is_backend_mpi(backend)) then
        !     call try_free_mpi_handles(comm, platform, helper)
        ! endif
    end function execute_autotune

    function report_timings(comm, elapsed_time, n_iters, space_count) result(max_time)
    !! Aggregates and writing timings to stdout
        TYPE_MPI_COMM,            intent(in)  :: comm
        real(real32),             intent(in)  :: elapsed_time
        integer(int32),           intent(in)  :: n_iters
        integer(int32), optional, intent(in)  :: space_count
        real(real32)                          :: max_time
        real(real32) :: execution_time, min_time, avg_time
        integer(int32) :: ierr, comm_size, space_count_

        execution_time = elapsed_time / real(n_iters, real32)
        space_count_ = 2; if ( present(space_count) ) space_count_ = space_count

        call MPI_Allreduce(execution_time, min_time, 1, MPI_REAL4, MPI_MIN, comm, ierr)
        call MPI_Allreduce(execution_time, max_time, 1, MPI_REAL4, MPI_MAX, comm, ierr)
        call MPI_Allreduce(execution_time, avg_time, 1, MPI_REAL4, MPI_SUM, comm, ierr)
        call MPI_Comm_size(comm, comm_size, ierr)

        avg_time = avg_time / real(comm_size, real32)

        WRITE_INFO(repeat(" ", space_count_)//"max: "//to_str(real(max_time, real64))//" [ms]")
        WRITE_INFO(repeat(" ", space_count_)//"min: "//to_str(real(min_time, real64))//" [ms]")
        WRITE_INFO(repeat(" ", space_count_)//"avg: "//to_str(real(avg_time, real64))//" [ms]")
    end function report_timings

    subroutine alloc_and_set_aux(platform, helper, backend, cart_comm, aux, plans, is_aux_alloc)
    !! Allocates auxiliary memory according to the backend and sets it to the plans
        type(dtfft_platform_t),       intent(in)                :: platform
        type(backend_helper),         intent(inout)             :: helper       !! Backend helper
        type(dtfft_backend_t),        intent(in)                :: backend      !! GPU backend
        TYPE_MPI_COMM,                intent(in)                :: cart_comm    !! Cartesian communicator
        type(c_ptr),                  intent(inout)             :: aux          !! Allocatable auxiliary memory
        ! real(real32),     pointer,    intent(inout)             :: paux(:)      !! Pointer to auxiliary memory
        type(reshape_container),      intent(in)                :: plans(:)
        logical                                                 :: is_aux_alloc !! Is auxiliary memory allocated
        integer(int64) :: max_work_size_local, max_work_size_global
        integer(int32)  :: mpi_ierr
        integer(int32) :: alloc_ierr
        type(dtfft_backend_t) :: backend_

        max_work_size_local = get_aux_bytes_generic(plans)
        call MPI_Allreduce(max_work_size_local, max_work_size_global, 1, MPI_INTEGER8, MPI_MAX, cart_comm, mpi_ierr)

        is_aux_alloc = .false.
        ! paux => null()
        if ( max_work_size_global > 0 ) then
            backend_ = backend; if ( is_nvshmem_generic(plans, backend_) ) backend_ = DTFFT_BACKEND_CUFFTMP
            call alloc_mem(platform, helper, backend_, cart_comm, max_work_size_global, aux, alloc_ierr)
            DTFFT_CHECK(alloc_ierr)
            ! call c_f_pointer(aux, paux, [max_work_size_global / 4_int64])
            is_aux_alloc = .true.
        endif
    end subroutine alloc_and_set_aux

! subroutine try_free_mpi_handles(comm, platform, helper)
! TYPE_MPI_COMM,          intent(in)      :: comm
! type(dtfft_platform_t), intent(in)      :: platform
! type(backend_helper),   intent(inout)   :: helper
! #ifdef DTFFT_WITH_CUDA
! integer(int64) :: alloc_size
! integer(int32) :: base_size, comm_size, ierr
! type(c_ptr) :: ptr1, ptr2, ptr3
! integer(int64)  :: free_mem_before, free_mem_after, total_mem_avail
! real(real32), pointer, contiguous :: rptr1(:), rptr2(:), rptr3(:)
! #endif

!     if ( platform == DTFFT_PLATFORM_HOST ) return
! #ifdef DTFFT_WITH_CUDA

!     call MPI_Comm_size(comm, comm_size, ierr)

!     base_size = 1024 * 1024 * 30    ! 10Mb per process
!     alloc_size = int(base_size, int64) * comm_size

!     CUDA_CALL( cudaMemGetInfo(free_mem_before, total_mem_avail) )

!     call alloc_mem(platform, helper, DTFFT_BACKEND_MPI_P2P, comm, alloc_size, ptr1, ierr);  call c_f_pointer(ptr1, rptr1, [1])
!     call alloc_mem(platform, helper, DTFFT_BACKEND_MPI_P2P, comm, alloc_size, ptr2, ierr);  call c_f_pointer(ptr2, rptr2, [1])
!     call alloc_mem(platform, helper, DTFFT_BACKEND_MPI_P2P, comm, alloc_size, ptr3, ierr);  call c_f_pointer(ptr3, rptr3, [1])

!     call MPI_Alltoall(rptr1, base_size, MPI_BYTE, rptr2, base_size, MPI_BYTE, comm, ierr)
!     call MPI_Alltoall(rptr1, base_size, MPI_BYTE, rptr3, base_size, MPI_BYTE, comm, ierr)
!     call MPI_Alltoall(rptr2, base_size, MPI_BYTE, rptr3, base_size, MPI_BYTE, comm, ierr)

!     call free_mem(platform, helper, DTFFT_BACKEND_MPI_P2P, ptr1, ierr)
!     call free_mem(platform, helper, DTFFT_BACKEND_MPI_P2P, ptr2, ierr)
!     call free_mem(platform, helper, DTFFT_BACKEND_MPI_P2P, ptr3, ierr)

!     CUDA_CALL( cudaMemGetInfo(free_mem_after, total_mem_avail) )

!     WRITE_INFO("Tried to release internal MPI handles. Free mem before: "//to_str(free_mem_before)//", after = "//to_str(free_mem_after))

! #endif
! end subroutine try_free_mpi_handles

end module dtfft_reshape_plan_base