Cabana_Grid_GlobalParticleComm.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_GLOBALPARTICLECOMM_HPP
#define CABANA_GRID_GLOBALPARTICLECOMM_HPP
 
#include <Cabana_Grid_LocalMesh.hpp>
#include <Cabana_Grid_Types.hpp>
 
#include <Cabana_Distributor.hpp>
#include <Cabana_Slice.hpp>
 
#include <memory>
#include <stdexcept>
#include <type_traits>
 
namespace Cabana
{
namespace Grid
{
//---------------------------------------------------------------------------//
template <class MemorySpace, class LocalGridType>
class GlobalParticleComm
{
  public:
    static constexpr std::size_t num_space_dim = LocalGridType::num_space_dim;
    using mesh_type = typename LocalGridType::mesh_type;
    using global_grid_type = Cabana::Grid::GlobalGrid<mesh_type>;
    using memory_space = MemorySpace;

    using corner_view_type =
        Kokkos::View<double* [num_space_dim][2], memory_space>;
    using destination_view_type = Kokkos::View<int*, memory_space>;
    using rank_view_type =
        Kokkos::View<int***, Kokkos::LayoutRight, memory_space>;
    using host_rank_view_type =
        Kokkos::View<int***, Kokkos::LayoutRight, Kokkos::HostSpace>;

    GlobalParticleComm( const LocalGridType local_grid )
    {
        auto global_grid = local_grid.globalGrid();
        _destinations = destination_view_type(
            Kokkos::ViewAllocateWithoutInitializing( "global_destination" ),
            0 );
 
        int max_ranks_per_dim = -1;
        for ( std::size_t d = 0; d < num_space_dim; ++d )
        {
            _ranks_per_dim[d] = global_grid.dimNumBlock( d );
            if ( _ranks_per_dim[d] > max_ranks_per_dim )
                max_ranks_per_dim = _ranks_per_dim[d];
        }
        copyRanks( global_grid );
 
        // Purposely using zero-init. Some entries unused in non-cubic
        // decompositions.
        _local_corners =
            corner_view_type( "local_mpi_boundaries", max_ranks_per_dim );
 
        _rank_1d = global_grid.blockId();
        for ( std::size_t d = 0; d < num_space_dim; ++d )
            _rank[d] = global_grid.dimBlockId( d );
 
        auto local_mesh = createLocalMesh<memory_space>( local_grid );
        storeRanks( local_mesh );
 
        // Update local boundaries from all ranks.
        auto comm = global_grid.comm();
        // TODO: Could use subcommunicators instead.
        MPI_Allreduce( MPI_IN_PLACE, _local_corners.data(),
                       _local_corners.size(), MPI_DOUBLE, MPI_SUM, comm );
 
        scaleRanks();
    }

    template <class LocalMeshType>
    void storeRanks( LocalMeshType local_mesh )
    {
        auto local_corners = _local_corners;
        auto rank = _rank;
        auto store_corners = KOKKOS_LAMBDA( const std::size_t d )
        {
            local_corners( rank[d], d, 0 ) =
                local_mesh.lowCorner( Cabana::Grid::Own(), d );
            local_corners( rank[d], d, 1 ) =
                local_mesh.highCorner( Cabana::Grid::Own(), d );
        };
        using exec_space = typename memory_space::execution_space;
        Kokkos::RangePolicy<exec_space> policy( 0, num_space_dim );
        Kokkos::parallel_for( "Cabana::Grid::GlobalParticleComm::storeCorners",
                              policy, store_corners );
        Kokkos::fence();
    }

    void scaleRanks()
    {
        auto scale = getScaling();
 
        auto local_corners = _local_corners;
        auto ranks_per_dim = _ranks_per_dim;
        auto scale_corners = KOKKOS_LAMBDA( const std::size_t d )
        {
            for ( int r = 0; r < ranks_per_dim[d]; ++r )
            {
                local_corners( r, d, 0 ) /= scale[d];
                local_corners( r, d, 1 ) /= scale[d];
            }
        };
        using exec_space = typename memory_space::execution_space;
        Kokkos::RangePolicy<exec_space> policy( 0, num_space_dim );
        Kokkos::parallel_for( "Cabana::Grid::GlobalParticleComm::scaleCorners",
                              policy, scale_corners );
        Kokkos::fence();
    }

    template <std::size_t NSD = num_space_dim>
    std::enable_if_t<3 == NSD, Kokkos::Array<int, num_space_dim>> getScaling()
    {
        Kokkos::Array<int, num_space_dim> scale;
        scale[0] = _ranks_per_dim[1] * _ranks_per_dim[2];
        scale[1] = _ranks_per_dim[0] * _ranks_per_dim[2];
        scale[2] = _ranks_per_dim[0] * _ranks_per_dim[1];
        return scale;
    }

    template <std::size_t NSD = num_space_dim>
    std::enable_if_t<2 == NSD, Kokkos::Array<int, num_space_dim>> getScaling()
    {
        Kokkos::Array<int, num_space_dim> scale;
        scale[0] = _ranks_per_dim[1];
        scale[1] = _ranks_per_dim[0];
        return scale;
    }

    template <class GlobalGridType, std::size_t NSD = num_space_dim>
    std::enable_if_t<3 == NSD, void> copyRanks( GlobalGridType global_grid )
    {
        host_rank_view_type global_ranks_host(
            Kokkos::ViewAllocateWithoutInitializing( "ranks_host" ),
            _ranks_per_dim[0], _ranks_per_dim[1], _ranks_per_dim[2] );
        for ( int i = 0; i < _ranks_per_dim[0]; ++i )
            for ( int j = 0; j < _ranks_per_dim[1]; ++j )
                for ( int k = 0; k < _ranks_per_dim[2]; ++k )
                    // Not device accessible (uses MPI), so must be copied.
                    global_ranks_host( i, j, k ) =
                        global_grid.blockRank( i, j, k );
 
        _global_ranks = Kokkos::create_mirror_view_and_copy(
            memory_space(), global_ranks_host );
    }

    template <class GlobalGridType, std::size_t NSD = num_space_dim>
    std::enable_if_t<2 == NSD, void> copyRanks( GlobalGridType global_grid )
    {
        // Storing as 3d for convenience.
        host_rank_view_type global_ranks_host(
            Kokkos::ViewAllocateWithoutInitializing( "ranks_host" ),
            _ranks_per_dim[0], _ranks_per_dim[1], 1 );
        for ( int i = 0; i < _ranks_per_dim[0]; ++i )
            for ( int j = 0; j < _ranks_per_dim[1]; ++j )
                // Not device accessible (uses MPI), so must be copied.
                global_ranks_host( i, j, 0 ) = global_grid.blockRank( i, j );
 
        _global_ranks = Kokkos::create_mirror_view_and_copy(
            memory_space(), global_ranks_host );
    }

    template <class ExecutionSpace, class PositionType>
    void build( ExecutionSpace exec_space, PositionType positions,
                const std::size_t begin, const std::size_t end )
    {
        Kokkos::Profiling::ScopedRegion region(
            "Cabana::Grid::GlobalParticleComm::build" );
 
        static_assert( is_accessible_from<memory_space, ExecutionSpace>{}, "" );
        assert( end >= begin );
        assert( end <= positions.size() );
 
        // Must match the size of all particles, even if some can be ignored in
        // this search.
        Kokkos::resize( _destinations, positions.size() );
        // Start with everything staying on this rank.
        Kokkos::deep_copy( _destinations, _rank_1d );
 
        // Local copies for lambda capture.
        auto local_corners = _local_corners;
        auto ranks_per_dim = _ranks_per_dim;
        auto destinations = _destinations;
        auto global_ranks = _global_ranks;
        auto build_migrate = KOKKOS_LAMBDA( const std::size_t p )
        {
            // This is not num_space_dim because global_ranks is always rank-3
            // out of convenience.
            int ijk[3] = { 0, 0, 0 };
 
            // Find the rank this particle should be moved to.
            for ( std::size_t d = 0; d < num_space_dim; ++d )
            {
                int min = 0;
                int max = ranks_per_dim[d];
 
                // Check if outside the box in this dimension.
                if ( ( positions( p, d ) < local_corners( min, d, 0 ) ) ||
                     ( positions( p, d ) > local_corners( max - 1, d, 1 ) ) )
                    destinations( p ) = -1;
 
                // Do a binary search for this particle in this dimension.
                while ( max - min > 1 )
                {
                    int center = Kokkos::floor( ( max + min ) / 2.0 );
                    if ( positions( p, d ) < local_corners( center, d, 0 ) )
                        max = center;
                    else
                        min = center;
                }
                ijk[d] = min;
            }
            // Keep the destination rank for eventual migration.
            destinations( p ) = global_ranks( ijk[0], ijk[1], ijk[2] );
        };
 
        Kokkos::RangePolicy<ExecutionSpace> policy( exec_space, begin, end );
        Kokkos::parallel_for( "Cabana::Grid::GlobalParticleComm::build", policy,
                              build_migrate );
        Kokkos::fence();
    }

    template <class ExecutionSpace, class PositionType>
    void build( ExecutionSpace exec_space, PositionType positions )
    {
        build( exec_space, positions, 0, positions.size() );
    }

    template <class PositionType>
    void build( PositionType positions )
    {
        using execution_space = typename memory_space::execution_space;
        // TODO: enable views.
        build( execution_space{}, positions, 0, positions.size() );
    }

    template <class AoSoAType>
    void migrate( MPI_Comm comm, AoSoAType& aosoa )
    {
        Cabana::Distributor<memory_space> distributor( comm, _destinations );
        Cabana::migrate( distributor, aosoa );
    }
 
  protected:
    int _rank_1d;
    Kokkos::Array<int, num_space_dim> _rank;
    Kokkos::Array<int, num_space_dim> _ranks_per_dim;
    corner_view_type _local_corners;
    rank_view_type _global_ranks;
    destination_view_type _destinations;
};

template <class MemorySpace, class LocalGridType>
auto createGlobalParticleComm( const LocalGridType& local_grid )
{
    return std::make_shared<GlobalParticleComm<MemorySpace, LocalGridType>>(
        local_grid );
}
 
//---------------------------------------------------------------------------//
 
} // namespace Grid
} // namespace Cabana
 
#endif // end CABANA_GRID_GLOBALPARTICLECOMM_HPP