DataStream.h

#pragma once

#include <isce3/cuda/core/Event.h>
#include <isce3/cuda/core/Stream.h>
#include <isce3/io/Raster.h>
#include <string>
#include <thrust/host_vector.h>
#include <thrust/system/cuda/experimental/pinned_allocator.h>

namespace isce3 { namespace cuda { namespace io {

// thrust::host_vector whose data buffer uses page-locked memory
template<typename T>
using pinned_host_vector = thrust::host_vector<T,
        thrust::system::cuda::experimental::pinned_allocator<T>>;

// callback for asynchronously writing to files via std::ofstream
class fstreamCallback {
private:
    fstreamCallback() = default;

    fstreamCallback(const std::string * filename,
                    const char * buffer,
                    std::size_t offset,
                    std::size_t count);

    // interface for cudaStreamAddCallback
    // casts *obj* to (fstreamCallback *) and calls write()
    static
    void CUDART_CB cb_write(cudaStream_t, cudaError_t, void * obj);

    // write the contents of *buffer* to file at *filename* via std::ofstream
    void write();

    const std::string * filename;
    const char * buffer;
    std::size_t offset;
    std::size_t count;

    friend class FileDataStream;
};

// callback for asynchronously writing to Raster objects
class RasterCallback {
private:
    RasterCallback() = default;

    RasterCallback(isce3::io::Raster * raster,
                   char * buffer,
                   std::size_t col,
                   std::size_t row,
                   std::size_t width,
                   std::size_t length);

    // interface for cudaStreamAddCallback
    // casts *obj* to (RasterCallback *) and calls setBlock()
    template<typename T>
    static
    void CUDART_CB cb_setBlock(cudaStream_t, cudaError_t, void * obj);

    // forwards members as arguments to Raster::setBlock()
    template<typename T>
    void setBlock();

    isce3::io::Raster * raster;
    char * buffer;
    std::size_t col;
    std::size_t row;
    std::size_t width;
    std::size_t length;

    friend class RasterDataStream;
};

class FileDataStream {
public:
    FileDataStream() = default;

    FileDataStream(const std::string & filename,
                   isce3::cuda::core::Stream stream,
                   std::size_t buffer_size = 0);

    const std::string & filename() const;

    void set_filename(const std::string &);

    isce3::cuda::core::Stream stream() const;

    void set_stream(isce3::cuda::core::Stream stream);

    std::size_t buffer_size() const;

    void resize_buffer(std::size_t buffer_size);

    void load(void * dst, std::size_t offset, std::size_t count);

    void store(const void * src, std::size_t offset, std::size_t count);

private:
    std::string _filename;
    isce3::cuda::core::Stream _stream = 0;
    isce3::cuda::core::Event _mutex;
    pinned_host_vector<char> _buffer;
    fstreamCallback _callback;
};

class RasterDataStream {
public:
    RasterDataStream() = default;

    RasterDataStream(isce3::io::Raster * raster,
                     isce3::cuda::core::Stream stream,
                     std::size_t buffer_size = 0);

    isce3::io::Raster * raster() const;

    void set_raster(isce3::io::Raster *);

    isce3::cuda::core::Stream stream() const;

    void set_stream(isce3::cuda::core::Stream stream);

    std::size_t buffer_size() const;

    void resize_buffer(std::size_t buffer_size);

    template<typename T>
    void load(T * dst,
            std::size_t col,
            std::size_t row,
            std::size_t width,
            std::size_t length);

    template<typename T>
    void store(const T * src,
            std::size_t col,
            std::size_t row,
            std::size_t width,
            std::size_t length);

private:
    isce3::io::Raster * _raster;
    isce3::cuda::core::Stream _stream = 0;
    isce3::cuda::core::Event _mutex;
    pinned_host_vector<char> _buffer;
    RasterCallback _callback;
};

}}}

#define ISCE_CUDA_IO_DATASTREAM_ICC
#include "DataStream.icc"
#undef ISCE_CUDA_IO_DATASTREAM_ICC