heffte_backend_data_transfer.h

/*
    -- heFFTe --
       Univ. of Tennessee, Knoxville
       @date
*/
 
#ifndef HEFFTE_BACKEND_DATA_TRANSFER_H
#define HEFFTE_BACKEND_DATA_TRANSFER_H
 
#ifdef Heffte_ENABLE_GPU
 
namespace heffte{
 
namespace gpu {
 
    template<typename scalar_type>
    using vector = device_vector<scalar_type, heffte::backend::data_manipulator<heffte::tag::gpu>>;
 
    template<typename manipulator>
    struct device_transfer{
        using backend_device = typename manipulator::backend_device;
        template<typename scalar_type>
        static void copy(device_vector<scalar_type, backend_device> const &source, scalar_type destination[]){
            manipulator::copy_device_to_host(source.stream(), source.data(), source.size(), destination);
        }
        template<typename scalar_type>
        static void copy(scalar_type const source[], device_vector<scalar_type, backend_device> &destination){
            manipulator::copy_device_to_device(destination.stream(), source, destination.size(), destination.data());
        }
 
        template<typename scalar_type>
        static device_vector<scalar_type, manipulator> load(typename backend_device::stream_type stream, scalar_type const *cpu_source, size_t num_entries){
            device_vector<scalar_type, manipulator> result(stream, num_entries);
            manipulator::copy_host_to_device(stream, cpu_source, num_entries, result.data());
            return result;
        }
        template<typename scalar_type>
        static void load(typename backend_device::stream_type stream, scalar_type const *cpu_source, size_t num_entries, scalar_type *gpu_destination){
            manipulator::copy_host_to_device(stream, cpu_source, num_entries, gpu_destination);
        }
        template<typename scalar_type>
        static void load(void*, scalar_type const *cpu_source, size_t num_entries, scalar_type *gpu_destination){
            std::copy_n(cpu_source, num_entries, gpu_destination);
        }
        template<typename scalar_type>
        static device_vector<scalar_type, manipulator> load(scalar_type const *cpu_source, size_t num_entries){
            return load(backend_device().stream(), cpu_source, num_entries);
        }
        template<typename scalar_type>
        static device_vector<scalar_type, manipulator> load(void*, scalar_type const*, size_t){
            return device_vector<scalar_type, manipulator>();
        }
 
        template<typename scalar_type>
        static device_vector<scalar_type, manipulator> load(std::vector<scalar_type> const &cpu_source){
            return load(cpu_source.data(), cpu_source.size());
        }
        template<typename scalar_type>
        static void load(typename backend_device::stream_type stream, std::vector<scalar_type> const &cpu_source, scalar_type gpu_destination[]){
            manipulator::copy_host_to_device(stream, cpu_source.data(), cpu_source.size(), gpu_destination);
        }
        template<typename scalar_type>
        static void load(std::vector<scalar_type> const &cpu_source, scalar_type gpu_destination[]){
            load(backend_device().stream(), cpu_source, gpu_destination);
        }
        template<typename scalar_type>
        static void load(void*, std::vector<scalar_type> const&, scalar_type[]){}
 
        template<typename scalar_type>
        static void load(std::vector<scalar_type> const &a, std::vector<scalar_type> &b){ b = a; }
        template<typename scalar_type>
        static std::vector<scalar_type> unload(std::vector<scalar_type> const &a){ return a; }
        template<typename scalar_type>
        static std::vector<scalar_type> unload(void*, scalar_type const[], size_t){
            return std::vector<scalar_type>();
        }
 
        template<typename scalar_type>
        static std::vector<scalar_type> unload(typename backend_device::stream_type stream, scalar_type const gpu_source[], size_t num_entries){
            std::vector<scalar_type> result(num_entries);
            manipulator::copy_device_to_host(stream, gpu_source, num_entries, result.data());
            return result;
        }
        template<typename scalar_type>
        static std::vector<scalar_type> unload(scalar_type const gpu_source[], size_t num_entries){
            return unload(backend_device().stream(), gpu_source, num_entries);
        }
 
        template<typename scalar_type>
        static void unload(device_vector<scalar_type, manipulator> const &gpu_source, scalar_type *cpu_result){
            manipulator::copy_device_to_host(gpu_source.device_stream(), gpu_source.data(), gpu_source.size(), cpu_result);
        }
        template<typename scalar_type>
        static void unload(typename backend_device::stream_type stream, scalar_type const *gpu_source, size_t num_entries, scalar_type *cpu_result){
            manipulator::copy_device_to_host(stream, gpu_source, num_entries, cpu_result);
        }
        template<typename scalar_type>
        static void unload(void*, scalar_type const *gpu_source, size_t num_entries, scalar_type *cpu_result){
            std::copy_n(gpu_source, num_entries, cpu_result); // void* stream indicates CPU backend
        }
 
        template<typename scalar_type>
        static std::vector<scalar_type> unload(device_vector<scalar_type, manipulator> const &gpu_source){
            std::vector<scalar_type> result(gpu_source.size());
            unload(gpu_source, result.data());
            return result;
        }
        template<typename scalar_type>
        static device_vector<scalar_type, manipulator> capture(scalar_type* &&raw_pointer, size_t num_entries){
            return device_vector<scalar_type, manipulator>(std::forward<scalar_type*>(raw_pointer), num_entries);
        }
    };
 
    using transfer = device_transfer<heffte::backend::data_manipulator<heffte::tag::gpu>>;
 
    int device_count();
 
    void device_set(int active_device);
 
    void synchronize_default_stream();
 
}
 
template<typename scalar_type>
gpu::vector<scalar_type> make_buffer_container(typename gpu::vector<scalar_type>::stream_type stream, size_t size){
    return gpu::vector<scalar_type>(stream, size);
}
 
}
 
#endif
 
#endif