heffte_backend_vector.h

/*
    -- heFFTe --
       Univ. of Tennessee, Knoxville
       @date
*/
 
#ifndef HEFFTE_BACKEND_VECTOR_H
#define HEFFTE_BACKEND_VECTOR_H
 
#ifdef Heffte_ENABLE_GPU
 
namespace heffte{
 
namespace gpu {
 
    template<typename scalar_type, typename manipulator>
    class device_vector{
    public:
        using value_type = scalar_type;
        using backend_device = typename manipulator::backend_device;
        using stream_type = typename backend_device::stream_type;
 
        device_vector(size_t num_entries = 0) :
            device(),
            num(num_entries),
            device_data(manipulator::template allocate<scalar_type>(device.stream(), num_entries))
        {}
        device_vector(backend_device const &new_device, size_t num_entries = 0) :
            device(new_device),
            num(num_entries),
            device_data(manipulator::template allocate<scalar_type>(device.stream(), num_entries))
        {}
        device_vector(scalar_type const *begin, scalar_type const *end) :
            device_vector(std::distance(begin, end)){
            manipulator::copy_device_to_device(device.stream(), begin, num, device_data);
        }
        device_vector(backend_device const &new_device, scalar_type const *begin, scalar_type const *end) :
            device_vector(new_device, std::distance(begin, end)){
            manipulator::copy_device_to_device(device.stream(), begin, num, device_data);
        }
 
        device_vector(const device_vector<scalar_type, manipulator>& other) :
            device_vector(other.device, other.num){
            manipulator::copy_device_to_device(device.stream(), other.device_data, num, device_data);
        }
        device_vector(device_vector<scalar_type, manipulator> &&other) :
            device(other.device),
            num(c11_exchange(other.num, 0)),
            device_data(c11_exchange(other.device_data, nullptr))
        {}
 
        device_vector(scalar_type* &&raw_pointer, size_t num_entries) :
            device(),
            num(num_entries),
            device_data(c11_exchange(raw_pointer, nullptr))
        {}
        device_vector(backend_device const &new_device, scalar_type* &&raw_pointer, size_t num_entries) :
            device(new_device),
            num(num_entries),
            device_data(c11_exchange(raw_pointer, nullptr))
        {}
 
        ~device_vector(){ manipulator::free(device.stream(), device_data); }
 
        void operator =(device_vector<scalar_type, manipulator> const &other){
            device_vector<scalar_type, manipulator> temp(other);
            device = temp.device;
            std::swap(num, temp.num);
            std::swap(device_data, temp.device_data);
        }
 
        void operator =(device_vector<scalar_type, manipulator>&& other){
            device_vector<scalar_type, manipulator> temp(std::move(other));
            device = temp.device;
            std::swap(num, temp.num);
            std::swap(device_data, temp.device_data);
        }
 
        scalar_type* data(){ return device_data; }
        const scalar_type* data() const{ return device_data; }
 
        size_t size() const{ return num; }
        bool empty() const{ return (num == 0); }
 
        scalar_type* release(){
            num = 0;
            return c11_exchange(device_data, nullptr);
        }
 
        stream_type device_stream(){ return device.stream(); }
        stream_type device_stream() const{ return device.stream(); }
 
    private:
        backend_device device;
        size_t num;
        scalar_type *device_data;
    };
 
    int device_count();
 
    void device_set(int active_device);
 
    void synchronize_default_stream();
 
}
 
}
 
#endif
 
#endif   /* HEFFTE_BACKEND_VECTOR_H */