Cabana_Grid_SparseArray.hpp

Go to the documentation of this file.

/****************************************************************************
 * Copyright (c) 2018-2023 by the Cabana authors                            *
 * All rights reserved.                                                     *
 *                                                                          *
 * This file is part of the Cabana library. Cabana is distributed under a   *
 * BSD 3-clause license. For the licensing terms see the LICENSE file in    *
 * the top-level directory.                                                 *
 *                                                                          *
 * SPDX-License-Identifier: BSD-3-Clause                                    *
 ****************************************************************************/

#ifndef CABANA_GRID_SPARSE_ARRAY_HPP
#define CABANA_GRID_SPARSE_ARRAY_HPP
 
#include <Cabana_Grid_SparseIndexSpace.hpp>
#include <Cabana_Grid_SparseLocalGrid.hpp>
#include <Cabana_Grid_Types.hpp>
 
#include <Cabana_AoSoA.hpp>
#include <Cabana_Utils.hpp> // FIXME: remove after next release.
 
#include <Kokkos_Core.hpp>
 
#include <cmath>
#include <memory>
#include <type_traits>
#include <vector>
 
#include <mpi.h>
 
namespace Cabana
{
namespace Grid
{
 
namespace Experimental
{
//---------------------------------------------------------------------------//
template <class DataTypes, class EntityType, class MeshType,
          class SparseMapType>
class SparseArrayLayout
{
  public:
    using mesh_type = MeshType;
    using scalar_type = typename mesh_type::scalar_type;
    // check if mesh_type is SparseMesh
    static_assert( isSparseMesh<MeshType>::value,
                   "[SparesArrayLayout] Support only SparseMesh" );

    using entity_type = EntityType;
    using member_types = DataTypes;

    using sparse_map_type = SparseMapType;
    using value_type = typename sparse_map_type::value_type;
    using key_type = typename sparse_map_type::key_type;
    static constexpr unsigned long long cell_bits_per_tile_dim =
        sparse_map_type::cell_bits_per_tile_dim;
    static constexpr unsigned long long cell_num_per_tile_dim =
        sparse_map_type::cell_num_per_tile_dim;
    static constexpr std::size_t num_space_dim = sparse_map_type::rank;
    using memory_space = typename sparse_map_type::memory_space;

    SparseArrayLayout( const std::shared_ptr<LocalGrid<MeshType>>& local_grid,
                       SparseMapType& sparse_map, const float over_allocation )
        : _over_allocation( over_allocation )
        , _local_grid( local_grid )
        , _map( std::forward<sparse_map_type>( sparse_map ) )
    {
        auto sparse_mesh = _local_grid->globalGrid().globalMesh();
        _cell_size[0] = sparse_mesh.cellSize( 0 );
        _cell_size[1] = sparse_mesh.cellSize( 1 );
        _cell_size[2] = sparse_mesh.cellSize( 2 );
        _global_low_corner[0] = sparse_mesh.lowCorner( 0 );
        _global_low_corner[1] = sparse_mesh.lowCorner( 1 );
        _global_low_corner[2] = sparse_mesh.lowCorner( 2 );
    }

    SparseMapType& sparseMap() { return _map; }

    const std::shared_ptr<LocalGrid<MeshType>> localGrid() const
    {
        return _local_grid;
    }

    inline uint64_t sizeCell() const { return _map.sizeCell(); }

    inline uint64_t sizeTile() const { return _map.sizeTile(); }

    inline uint64_t reservedCellSize( float factor ) const
    {
        return _map.reservedCellSize( factor );
    }

    inline uint64_t arraySize() const { return sizeCell(); }

    inline void clear() { _map.clear(); }

    template <class ExecSpace, class PositionSliceType>
    void registerSparseMap( PositionSliceType& positions,
                            const int particle_num, const int p2g_radius = 1 )
    {
        // get references
        Kokkos::Array<scalar_type, 3> dx_inv = {
            (scalar_type)1.0 / _cell_size[0], (scalar_type)1.0 / _cell_size[1],
            (scalar_type)1.0 / _cell_size[2] };
        auto& map = _map;
        auto& low_corner = _global_low_corner;
        // register sparse map in sparse array layout
        Kokkos::parallel_for(
            "Cabana::Grid::Experimental::SparseArrayLayout::registerSparseMap",
            Kokkos::RangePolicy<ExecSpace>( 0, particle_num ),
            KOKKOS_LAMBDA( const int pid ) {
                scalar_type pos[3] = { positions( pid, 0 ) - low_corner[0],
                                       positions( pid, 1 ) - low_corner[1],
                                       positions( pid, 2 ) - low_corner[2] };
                int grid_base[3] = {
                    static_cast<int>( std::lround( pos[0] * dx_inv[0] ) -
                                      p2g_radius ),
                    static_cast<int>( std::lround( pos[1] * dx_inv[1] ) -
                                      p2g_radius ),
                    static_cast<int>( std::lround( pos[2] * dx_inv[2] ) -
                                      p2g_radius ) };
                // register grids that will have data transfer with the particle
                const int p2g_size = p2g_radius * 2;
                for ( int i = 0; i <= p2g_size; ++i )
                    for ( int j = 0; j <= p2g_size; ++j )
                        for ( int k = 0; k <= p2g_size; ++k )
                        {
                            int cell_id[3] = { grid_base[0] + i,
                                               grid_base[1] + j,
                                               grid_base[2] + k };
                            map.insertCell( cell_id[0], cell_id[1],
                                            cell_id[2] );
                        }
            } );
    }

    KOKKOS_FORCEINLINE_FUNCTION
    value_type queryCell( const int cell_i, const int cell_j,
                          const int cell_k ) const
    {
        return _map.queryCell( cell_i, cell_j, cell_k );
    }

    KOKKOS_FORCEINLINE_FUNCTION
    value_type queryTile( const int cell_i, const int cell_j,
                          const int cell_k ) const
    {
        return _map.queryTile( cell_i, cell_j, cell_k );
    }

    KOKKOS_FORCEINLINE_FUNCTION
    value_type queryTileFromTileId( const int tile_i, const int tile_j,
                                    const int tile_k ) const
    {
        return _map.queryTileFromTileId( tile_i, tile_j, tile_k );
    }

    KOKKOS_FORCEINLINE_FUNCTION
    value_type queryTileFromTileKey( const key_type tile_key ) const
    {
        return _map.queryTileFromTileKey( tile_key );
    }

    KOKKOS_FORCEINLINE_FUNCTION
    value_type cellLocalId( const int cell_i, const int cell_j,
                            const int cell_k ) const
    {
        return _map.cell_local_id( cell_i, cell_j, cell_k );
    }
 
  private:
    float _over_allocation;
    Kokkos::Array<scalar_type, 3> _cell_size;
    Kokkos::Array<scalar_type, 3> _global_low_corner;
    std::shared_ptr<LocalGrid<MeshType>> _local_grid;
    //！ sparse map
    sparse_map_type _map;
}; // end class SparseArrayLayout

template <class>
struct is_sparse_array_layout : public std::false_type
{
};

template <class DataTypes, class EntityType, class MeshType,
          class SparseMapType>
struct is_sparse_array_layout<
    SparseArrayLayout<DataTypes, EntityType, MeshType, SparseMapType>>
    : public std::true_type
{
};

template <class DataTypes, class EntityType, class MeshType,
          class SparseMapType>
struct is_sparse_array_layout<
    const SparseArrayLayout<DataTypes, EntityType, MeshType, SparseMapType>>
    : public std::true_type
{
};
 
//---------------------------------------------------------------------------//
// Array layout creation.
//---------------------------------------------------------------------------//
template <class DataTypes, class EntityType, class MeshType,
          class SparseMapType>
auto createSparseArrayLayout(
    const std::shared_ptr<LocalGrid<MeshType>>& local_grid,
    SparseMapType& sparse_map, EntityType, const float over_allocation = 1.01f )
{
    return std::make_shared<
        SparseArrayLayout<DataTypes, EntityType, MeshType, SparseMapType>>(
        local_grid, sparse_map, over_allocation );
}
 
//---------------------------------------------------------------------------//
template <class DataTypes, class MemorySpace, class EntityType, class MeshType,
          class SparseMapType>
class SparseArray
{
  public:
    using sparse_array_type = SparseArray<DataTypes, MemorySpace, EntityType,
                                          MeshType, SparseMapType>;
    using memory_space = MemorySpace;
    static_assert( Kokkos::is_memory_space<MemorySpace>() );
    using execution_space = typename memory_space::execution_space;
    using size_type = typename memory_space::size_type;
    using entity_type = EntityType;
    using mesh_type = MeshType;
    using sparse_map_type = SparseMapType;
    static constexpr std::size_t num_space_dim = mesh_type::num_space_dim;
    static constexpr unsigned long long cell_bits_per_tile =
        sparse_map_type::cell_bits_per_tile;
    static constexpr unsigned long long cell_mask_per_tile =
        sparse_map_type::cell_mask_per_tile;
 
    // AoSoA related types
    using member_types = DataTypes;
    static constexpr int vector_length = sparse_map_type::cell_num_per_tile;
    using aosoa_type = Cabana::AoSoA<member_types, memory_space, vector_length>;
    using soa_type = Cabana::SoA<member_types, vector_length>;
    using tuple_type = Cabana::Tuple<member_types>;

    using array_layout = SparseArrayLayout<member_types, entity_type, mesh_type,
                                           sparse_map_type>;
    SparseArray( const std::string label, array_layout layout )
        : _layout( layout )
        , _data( label )
    {
    }
 
    // ------------------------------------------------------------------------
    // AoSoA-related interfaces
    aosoa_type& aosoa() { return _data; }
    std::string label() const { return _data.label(); }
    array_layout& layout() { return _layout; }
    const array_layout& layout() const { return _layout; }

    inline void resize( const size_type n ) { _data.resize( n ); }

    inline void resize() { resize( _layout.sizeCell() ); }

    inline void reserve( const size_type n ) { _data.reserve( n ); }

    inline void reserveFromMap( const double factor = 1.2 )
    {
        reserve( _layout.reservedCellSize( factor ) );
    }

    inline void shrinkToFit() { _data.shrinkToFit(); }
    inline void clear()
    {
        resize( 0 );
        _layout.clear();
    };

    KOKKOS_FUNCTION
    size_type capacity() { return _data.capacity(); }
    KOKKOS_FUNCTION
    size_type size() const { return _data.size(); }
    KOKKOS_FUNCTION
    bool empty() const { return ( size() == 0 ); }
    KOKKOS_FORCEINLINE_FUNCTION
    size_type numSoA() const { return _data.numSoA(); }
    KOKKOS_FORCEINLINE_FUNCTION
    size_type arraySize( const size_type s ) const
    {
        return _data.arraySize( s );
    }

    template <class PositionSliceType>
    void registerSparseGrid( PositionSliceType& positions, int particle_num,
                             const int p2g_radius = 1 )
    {
        _layout.template registerSparseMap<execution_space, PositionSliceType>(
            positions, particle_num, p2g_radius );
        this->resize( _layout.sparseMap().sizeCell() );
    }
 
    // ------------------------------------------------------------------------
    KOKKOS_FORCEINLINE_FUNCTION
    soa_type& accessTile( const int tile_i, const int tile_j,
                          const int tile_k ) const
    {
        auto tile_id = _layout.queryTileFromTileId( tile_i, tile_j, tile_k );
        return _data.access( tile_id );
    }

    template <typename Value>
    KOKKOS_FORCEINLINE_FUNCTION soa_type&
    accessTile( const Value tile_id ) const
    {
        return _data.access( tile_id );
    }

    KOKKOS_FORCEINLINE_FUNCTION
    soa_type& accessTileFromCell( const int cell_i, const int cell_j,
                                  const int cell_k ) const
    {
        auto tile_id = _layout.queryTile( cell_i, cell_j, cell_k );
        return _data.access( tile_id );
    }
 
    // ------------------------------------------------------------------------
    template <typename Key>
    KOKKOS_FORCEINLINE_FUNCTION tuple_type
    getTuple( const Key tile_key, const int cell_local_id ) const
    {
        auto tile_id = _layout.queryTileFromTileKey( tile_key );
        return _data.getTuple( ( tile_id << cell_bits_per_tile ) |
                               ( cell_local_id & cell_mask_per_tile ) );
    }
 
    // ------------------------------------------------------------------------
    template <std::size_t M, typename... Indices>
    KOKKOS_FORCEINLINE_FUNCTION
        typename soa_type::template member_reference_type<M>
        get( const Kokkos::Array<int, 3> cell_ijk, Indices&&... ids ) const
    {
        auto& soa = accessTileFromCell( cell_ijk[0], cell_ijk[1], cell_ijk[2] );
        auto array_index =
            _layout.cellLocalId( cell_ijk[0], cell_ijk[1], cell_ijk[2] );
        return Cabana::get<M>( soa, array_index, ids... );
    }

    template <std::size_t M>
    KOKKOS_FORCEINLINE_FUNCTION
        typename soa_type::template member_reference_type<M>
        get( const Kokkos::Array<int, 3> cell_ijk ) const
    {
        auto& soa = accessTileFromCell( cell_ijk[0], cell_ijk[1], cell_ijk[2] );
        auto array_index =
            _layout.cellLocalId( cell_ijk[0], cell_ijk[1], cell_ijk[2] );
        return Cabana::get<M>( soa, array_index );
    }

    template <std::size_t M, typename... Indices>
    KOKKOS_FORCEINLINE_FUNCTION
        typename soa_type::template member_reference_type<M>
        get( const Kokkos::Array<int, 3> tile_ijk,
             const Kokkos::Array<int, 3> local_cell_ijk,
             Indices&&... ids ) const
    {
        auto& soa = accessTile( tile_ijk[0], tile_ijk[1], tile_ijk[2] );
        auto array_index = _layout.cellLocalId(
            local_cell_ijk[0], local_cell_ijk[1], local_cell_ijk[2] );
        return Cabana::get<M>( soa, array_index, ids... );
    }

    template <std::size_t M>
    KOKKOS_FORCEINLINE_FUNCTION
        typename soa_type::template member_reference_type<M>
        get( const Kokkos::Array<int, 3> tile_ijk,
             const Kokkos::Array<int, 3> local_cell_ijk ) const
    {
        auto& soa = accessTile( tile_ijk[0], tile_ijk[1], tile_ijk[2] );
        auto array_index = _layout.cellLocalId(
            local_cell_ijk[0], local_cell_ijk[1], local_cell_ijk[2] );
        return Cabana::get<M>( soa, array_index );
    }

    template <std::size_t M, typename... Indices>
    KOKKOS_FORCEINLINE_FUNCTION
        typename soa_type::template member_reference_type<M>
        get( const int tile_id, const Kokkos::Array<int, 3> local_cell_ijk,
             Indices&&... ids ) const
    {
        auto& soa = _data.access( tile_id );
        auto array_index = _layout.cellLocalId(
            local_cell_ijk[0], local_cell_ijk[1], local_cell_ijk[2] );
        return Cabana::get<M>( soa, array_index, ids... );
    }

    template <std::size_t M>
    KOKKOS_FORCEINLINE_FUNCTION
        typename soa_type::template member_reference_type<M>
        get( const int tile_id,
             const Kokkos::Array<int, 3> local_cell_ijk ) const
    {
        auto& soa = _data.access( tile_id );
        auto array_index = _layout.cellLocalId(
            local_cell_ijk[0], local_cell_ijk[1], local_cell_ijk[2] );
        return Cabana::get<M>( soa, array_index );
    }

    template <std::size_t M, typename... Indices>
    KOKKOS_FORCEINLINE_FUNCTION
        typename soa_type::template member_reference_type<M>
        get( const int tile_id, const int cell_id, Indices&&... ids ) const
    {
        auto& soa = _data.access( tile_id );
        return Cabana::get<M>( soa, cell_id, ids... );
    }

    template <std::size_t M>
    KOKKOS_FORCEINLINE_FUNCTION
        typename soa_type::template member_reference_type<M>
        get( const int tile_id, const int cell_id ) const
    {
        auto& soa = _data.access( tile_id );
        return Cabana::get<M>( soa, cell_id );
    }
 
  private:
    array_layout _layout;
    aosoa_type _data;
 
}; // end class SparseArray
 
//---------------------------------------------------------------------------//
// Scatic type checker.
//---------------------------------------------------------------------------//
template <class>
struct is_sparse_array : public std::false_type
{
};

template <class DataTypes, class MemorySpace, class EntityType, class MeshType,
          class SparseMapType>
struct is_sparse_array<
    SparseArray<DataTypes, MemorySpace, EntityType, MeshType, SparseMapType>>
    : public std::true_type
{
};

template <class DataTypes, class MemorySpace, class EntityType, class MeshType,
          class SparseMapType>
struct is_sparse_array<const SparseArray<DataTypes, MemorySpace, EntityType,
                                         MeshType, SparseMapType>>
    : public std::true_type
{
};
 
//---------------------------------------------------------------------------//
// Array creation.
//---------------------------------------------------------------------------//
template <class MemorySpace, class DataTypes, class EntityType, class MeshType,
          class SparseMapType>
auto createSparseArray(
    const std::string label,
    SparseArrayLayout<DataTypes, EntityType, MeshType, SparseMapType>& layout )
{
    return std::make_shared<SparseArray<DataTypes, MemorySpace, EntityType,
                                        MeshType, SparseMapType>>( label,
                                                                   layout );
}
 
} // namespace Experimental
} // namespace Grid
} // namespace Cabana
 
#endif // CABANA_GRID_SPARSE_ARRAY_HPP