SVector.h

/*
* (C) Copyright 2013 ECMWF.
*
* This software is licensed under the terms of the Apache Licence Version 2.0
* which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
* In applying this licence, ECMWF does not waive the privileges and immunities
* granted to it by virtue of its status as an intergovernmental organisation nor
* does it submit to any jurisdiction.
*/

#pragma once

#include <algorithm>
#include <cassert>
#include <cstddef>

#include "atlas/library/config.h"
#include "atlas/util/Allocate.h"
#ifndef __CUDA_ARCH__
#include "atlas/runtime/Exception.h"
#endif

namespace atlas {
namespace array {

//------------------------------------------------------------------------------

template <typename T>
class SVector {
public:
    ATLAS_HOST_DEVICE
    SVector() : data_( nullptr ), size_( 0 ), externally_allocated_( false ) {}

    ATLAS_HOST_DEVICE
    SVector( const T* data, const idx_t size ) : data_( data ), size_( size ), externally_allocated_( true ) {}

#ifdef __CUDACC__
    // Note that this does not copy!!! It is mainly intended to be passed to a CUDA kernel which requires value semantics for this class
    ATLAS_HOST_DEVICE
    SVector( SVector const& other ) : data_( other.data_ ), size_( other.size_ ), externally_allocated_( true ) {}
#endif

    ATLAS_HOST_DEVICE
    SVector( SVector&& other ) :
        data_( other.data_ ), size_( other.size_ ), externally_allocated_( other.externally_allocated_ ) {
        other.data_                 = nullptr;
        other.size_                 = 0;
        other.externally_allocated_ = true;
    }

    ATLAS_HOST_DEVICE
    SVector& operator=( SVector&& other ) {
        data_                       = other.data_;
        size_                       = other.size_;
        externally_allocated_       = other.externally_allocated_;
        other.data_                 = nullptr;
        other.size_                 = 0;
        other.externally_allocated_ = true;
        return *this;
    }

    ATLAS_HOST_DEVICE
    SVector( T* data, idx_t size ) : data_( data ), size_( size ), externally_allocated_( true ) {}

    SVector( idx_t N ) : data_( nullptr ), size_( N ), externally_allocated_( false ) { allocate( data_, N ); }

    ATLAS_HOST_DEVICE
    ~SVector() { clear(); }

    ATLAS_HOST_DEVICE
    void clear() {
        if ( data_ && !externally_allocated_ ) {
#ifndef __CUDA_ARCH__
            deallocate( data_, size_ );
#endif
        }
        data_                 = nullptr;
        size_                 = 0;
        externally_allocated_ = false;
    }

    void insert( idx_t pos, idx_t dimsize ) {
        T* data = nullptr;
        allocate( data, size_ + dimsize );
        for ( idx_t c = 0; c < pos; ++c ) {
            data[c] = std::move( data_[c] );
        }
        for ( idx_t c = pos; c < size_; ++c ) {
            data[c + dimsize] = std::move( data_[c] );
        }
        deallocate( data_, size_ );
        data_ = data;
        size_ += dimsize;
    }

    size_t footprint() const { return sizeof( T ) * size_; }

    ATLAS_HOST_DEVICE
    T* data() { return data_; }

    ATLAS_HOST_DEVICE
    T const* data() const { return data_; }

    ATLAS_HOST_DEVICE
    T& operator()( const idx_t idx ) {
        assert( data_ && idx < size_ );
        return data_[idx];
    }
    ATLAS_HOST_DEVICE
    T const& operator()( const idx_t idx ) const {
        assert( data_ && idx < size_ );
        return data_[idx];
    }

    ATLAS_HOST_DEVICE
    T& operator[]( const idx_t idx ) {
        assert( data_ && idx < size_ );
        return data_[idx];
    }
    ATLAS_HOST_DEVICE
    T const& operator[]( const idx_t idx ) const {
        assert( data_ && idx < size_ );
        return data_[idx];
    }

    ATLAS_HOST_DEVICE
    idx_t size() const { return size_; }

    void resize_impl( idx_t N ) {
        if ( N == size_ )
            return;

        T* d_ = nullptr;
        allocate( d_, N );
        for ( idx_t c = 0; c < std::min( size_, N ); ++c ) {
            d_[c] = std::move( data_[c] );
        }
        deallocate( data_, size_ );
        data_ = d_;
    }

    void resize( idx_t N ) {
#ifndef __CUDA_ARCH__
        ATLAS_ASSERT( not externally_allocated_, "Cannot resize externally allocated (or wrapped) data" );
#endif
        resize_impl( N );
        size_ = N;
    }

private:
    static void allocate( T*& ptr, idx_t size ) {
        if ( size > 0 ) {
            util::allocate_managedmem( ptr, size );
            for ( idx_t c = 0; c < size; ++c ) {
                new ( ptr + c ) T();
            }
        }
    }
    static void deallocate( T*& ptr, idx_t size ) {
        if ( ptr ) {
            for ( idx_t c = 0; c < size; ++c ) {
                ptr[c].~T();
            }
            util::delete_managedmem( ptr );
        }
    }

private:
    T* data_;
    idx_t size_;
    bool externally_allocated_;
};

//------------------------------------------------------------------------------

}  // namespace array
}  // namespace atlas