kokkos-storage.hpp

/*
 * kokkos-storage.hpp
 *
 * Created on: 17 Nov 2022
 * Author: Holger Schmitz
 * Email: holger@notjustphysics.com
 *
 * Copyright 2012-2022 Holger Schmitz
 *
 * This file is part of Schnek.
 *
 * Schnek 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.
 *
 * Schnek 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 Schnek.  If not, see <http://www.gnu.org/licenses/>.
 */

#ifndef SCHNEK_GRID_GRIDSTORAGE_KOKKOSSTORAGE_HPP_
#define SCHNEK_GRID_GRIDSTORAGE_KOKKOSSTORAGE_HPP_

#include "../../config.hpp"

#ifdef SCHNEK_HAVE_KOKKOS

#include "../../macros.hpp"
#include "../array.hpp"
#include "../range.hpp"

#include <memory>
#include <functional>
#include <map>

#include <Kokkos_Core.hpp>

namespace schnek {

    namespace internal {
        template<typename T, size_t rank>
        struct KokkosViewType {
            typedef typename KokkosViewType<T, rank-1>::type* type;
        };

        template<typename T>
        struct KokkosViewType<T, 1> {
            typedef T* type;
        };
    }

    template <
        typename T, 
        size_t rank, 
        class ...ViewProperties
    >
    class KokkosGridStorage
    {
    public:
        typedef Array<int, rank> IndexType;

        typedef Range<int, rank> RangeType;
    private:
        typedef std::function<void(const RangeType&)> UpdaterType;
        typedef std::map<void *, UpdaterType> UpdaterMapType;

        RangeType range;

        IndexType dims;

        Kokkos::View<typename internal::KokkosViewType<T, rank>::type, ViewProperties...> view;

        std::shared_ptr<UpdaterMapType> updaters;
    public:
        KokkosGridStorage();

        KokkosGridStorage(const KokkosGridStorage &);

        KokkosGridStorage(const IndexType &lo, const IndexType &hi);

        KokkosGridStorage(const RangeType &range);

        ~KokkosGridStorage();

        SCHNEK_INLINE const T& get(const IndexType &index) const;

        SCHNEK_INLINE T& get(const IndexType &index);

        SCHNEK_INLINE const IndexType &getLo() const { return this->range.getLo(); }

        SCHNEK_INLINE const IndexType &getHi() const { return this->range.getHi(); }

        SCHNEK_INLINE const RangeType &getRange() const { return this->range; }

        SCHNEK_INLINE const IndexType &getDims() const { return this->dims; }

        SCHNEK_INLINE int getLo(int k) const { return this->range.getLo(k); }

        SCHNEK_INLINE int getHi(int k) const { return this->range.getHi(k); }

        SCHNEK_INLINE int getDims(int k) const { return this->dims[k]; }

        SCHNEK_INLINE int getSize() const { return this->size; }

        void resize(const IndexType &low, const IndexType &high);

        SCHNEK_INLINE ptrdiff_t stride(size_t dim) const;

    private:
        template<std::size_t... I>
        auto createKokkosViewImpl(const IndexType& a, std::index_sequence<I...>)
        {
            Kokkos::View<typename internal::KokkosViewType<T, rank>::type, ViewProperties...> view("schnek", a[I]...);
            return view;
        }
        
        auto createKokkosView(const IndexType& dims)
        {
            return createKokkosViewImpl(dims, std::make_index_sequence<rank>{});
        }

        template<std::size_t... I>
        SCHNEK_INLINE T& getFromViewImpl(const IndexType& pos, std::index_sequence<I...>)
        {
            return view(pos[I]...);
        }

        template<std::size_t... I>
        SCHNEK_INLINE const T& getFromViewImpl(const IndexType& pos, std::index_sequence<I...>) const
        {
            return view(pos[I]...);
        }

        SCHNEK_INLINE T& getFromView(const IndexType& pos)
        {
            return getFromViewImpl(pos, std::make_index_sequence<rank>{});
        }

        SCHNEK_INLINE const T& getFromView(const IndexType& pos) const
        {
            return getFromViewImpl(pos, std::make_index_sequence<rank>{});
        }

        void update(const RangeType& range) {
            for (auto& updater: *updaters) {
                updater.second(range);
            }
        }

        void updateSizeInfo(const RangeType &range) {
            this->range = range;
            for (size_t i = 0; i < rank; ++i)
            {
                dims[i] = range.getHi(i) - range.getLo(i) + 1;
            }
        }
    };

    template<typename T, size_t rank>
    using KokkosDefaultGridStorage = KokkosGridStorage<T, rank>;

    //=================================================================
    //==================== KokkosGridStorage ==========================
    //=================================================================

    template <typename T, size_t rank, class ...ViewProperties>
    KokkosGridStorage<T, rank, ViewProperties...>::KokkosGridStorage() 
        : range{IndexType{0}, IndexType{0}}, 
          dims{0},
          updaters{new UpdaterMapType}
    {
        (*updaters)[this] = [this](const RangeType& range) { this->updateSizeInfo(range); };
    }

    template <typename T, size_t rank, class ...ViewProperties>
    KokkosGridStorage<T, rank, ViewProperties...>::KokkosGridStorage(const KokkosGridStorage &other) 
        : range{other.range}, 
          dims{other.dims}, view{other.view},
          updaters{other.updaters}
    {
        (*updaters)[this] = [this](const RangeType& range) { this->updateSizeInfo(range); };
    }

    template <typename T, size_t rank, class ...ViewProperties>
    KokkosGridStorage<T, rank, ViewProperties...>::KokkosGridStorage(const IndexType &lo, const IndexType &hi) 
        : range{lo, hi},
          updaters{new UpdaterMapType}
    {
        dims = hi - lo + 1;
        view = createKokkosView(dims);
        (*updaters)[this] = [this](const RangeType& range) { this->updateSizeInfo(range); };
    }

    template <typename T, size_t rank, class ...ViewProperties>
    KokkosGridStorage<T, rank, ViewProperties...>::KokkosGridStorage(const RangeType &range) 
        : range{range},
          updaters{new UpdaterMapType}  
    {
        dims = range.getHi() - range.getLo() + 1;
        view = createKokkosView(dims);
        (*updaters)[this] = [this](const RangeType& range) { this->updateSizeInfo(range); };
    }

    template <typename T, size_t rank, class ...ViewProperties>
    KokkosGridStorage<T, rank, ViewProperties...>::~KokkosGridStorage()
    {
        updaters->erase(this);
    }

    template <typename T, size_t rank, class ...ViewProperties>
    SCHNEK_INLINE const T &KokkosGridStorage<T, rank, ViewProperties...>::get(const IndexType &index) const
    {
        IndexType pos;
        for (size_t i=0; i<rank; ++i)
        {
            pos[i] = index[i] - range.getLo(i);
        }
        return getFromView(pos);
    }

    template <typename T, size_t rank, class ...ViewProperties>
    SCHNEK_INLINE T &KokkosGridStorage<T, rank, ViewProperties...>::get(const IndexType &index)
    {       
        IndexType pos;
        for (size_t i=0; i<rank; ++i)
        {
            pos[i] = index[i] - range.getLo(i);
        }
        return getFromView(pos);
    }
    
    template <typename T, size_t rank, class ...ViewProperties>
    void KokkosGridStorage<T, rank, ViewProperties...>::resize(const IndexType &lo, const IndexType &hi)
    {
        IndexType dims = hi - lo + 1;
        this->view = createKokkosView(dims);
        update(RangeType{lo, hi});
    }

    template <typename T, size_t rank, class ...ViewProperties>
    SCHNEK_INLINE ptrdiff_t KokkosGridStorage<T, rank, ViewProperties...>::stride(size_t dim) const
    {
        return this->view.stride(dim);
    }

}


#endif // SCHNEK_HAVE_KOKKOS

#endif // SCHNEK_GRID_GRIDSTORAGE_KOKKOSSTORAGE_HPP_