From 695560af83cf5cdf23f80294e13fea9889d98155 Mon Sep 17 00:00:00 2001 From: Stuart Slattery Date: Thu, 24 Jun 2021 11:14:05 -0400 Subject: [PATCH] add multiple index space grid operator apply --- src/Picasso_GridOperator.hpp | 106 ++++++++++++++++++++++++++++---- unit_test/tstGridOperator2d.hpp | 44 +++++++++++++ unit_test/tstGridOperator3d.hpp | 40 ++++++++++++ 3 files changed, 179 insertions(+), 11 deletions(-) diff --git a/src/Picasso_GridOperator.hpp b/src/Picasso_GridOperator.hpp index 3101371..e49ef23 100644 --- a/src/Picasso_GridOperator.hpp +++ b/src/Picasso_GridOperator.hpp @@ -398,14 +398,14 @@ class GridOperator // // The functor is given a ParticleView allowing the kernel to read and // write particle data. - template - void apply( FieldLocation::Particle, const ExecutionSpace& exec_space, - const FieldManager& fm, const ParticleList_t& pl, - const WorkTag&, const Func& func ) const + void apply( const FieldLocation::Particle& location, + const ExecutionSpace& exec_space, const FieldManager& fm, + const ParticleList& pl, const WorkTag&, + const Func& func ) const { - applyImpl( fm, exec_space, FieldLocation::Particle(), pl, - func ); + applyImpl( fm, exec_space, location, pl, func ); } // Apply the operator in a loop over particles. Functor does not have a @@ -416,12 +416,13 @@ class GridOperator // // The functor is given a ParticleView allowing the kernel to read and // write particle data. - template - void apply( FieldLocation::Particle, const ExecutionSpace& exec_space, - const FieldManager& fm, const ParticleList_t& pl, + template + void apply( const FieldLocation::Particle& location, + const ExecutionSpace& exec_space, const FieldManager& fm, + const ParticleList& pl, const Func& func ) const { - applyImpl( fm, exec_space, FieldLocation::Particle(), pl, func ); + applyImpl( fm, exec_space, location, pl, func ); } // Apply the operator in a loop over the owned entities of the given @@ -440,6 +441,7 @@ class GridOperator // Apply the operator in a loop over the owned entities of the given // type. Functor does not have a work tag. + // // Functor signature: // func( local_mesh, gather_deps, scatter_deps, local_deps, i, j, k ) template @@ -449,6 +451,35 @@ class GridOperator applyImpl( fm, exec_space, location, func ); } + // Apply the operator in a fused loop over a list of provided index + // spaces. A work tag specifies the functor instance to use. + // + // Functor signature: + // func( local_mesh, gather_deps, scatter_deps, local_deps, s, i, j, k ) + template + void apply( const Kokkos::Array, + NumSpace>& spaces, + const ExecutionSpace& exec_space, const FieldManager& fm, + const WorkTag&, const Func& func ) const + { + applyImpl( fm, exec_space, spaces, func ); + } + + // Apply the operator in a fused loop over a list of provided index + // spaces. Functor does not have a work tag. + // + // Functor signature: + // func( local_mesh, gather_deps, scatter_deps, local_deps, s, i, j, k ) + template + void apply( const Kokkos::Array, + NumSpace>& spaces, + const ExecutionSpace& exec_space, const FieldManager& fm, + const Func& func ) const + { + applyImpl( fm, exec_space, spaces, func ); + } + public: // Manage field dependencies and apply the operator. template @@ -669,7 +700,7 @@ class GridOperator { // Apply kernel to each entity. The user functor gets a local mesh for // geometric operations, gather, scatter, and local dependencies for - // field operations (all of which may be empty), and the local ijk + // field operations (all of which may be empty), and the local ij // index of the entity they are currently working on. Cajita::grid_parallel_for( "operator_apply", exec_space, *( _mesh->localGrid() ), @@ -680,6 +711,59 @@ class GridOperator } ); } + // Apply the operator in a loop over an array of the given index space. 3D + // specialization. + template + void applyOp( const LocalMesh& local_mesh, const GatherFields& gather_deps, + const ScatterFields& scatter_deps, + const LocalFields& local_deps, + const ExecutionSpace& exec_space, + const Kokkos::Array, NumSpace>& spaces, + const Func& func ) const + { + // Apply kernel to each entity. The user functor gets a local mesh for + // geometric operations, gather, scatter, and local dependencies for + // field operations (all of which may be empty), the local id, s, of + // the space currently being operated on, and the local ijk + // index of the entity they are currently working on. + Cajita::grid_parallel_for( + "operator_apply", exec_space, spaces, + KOKKOS_LAMBDA( const int s, const int i, const int j, + const int k ) { + functorTagDispatch( func, local_mesh, gather_deps, + scatter_deps, local_deps, s, i, j, + k ); + } ); + } + + // Apply the operator in a loop over an array of the given index space. 2D + // specialization. + template + void applyOp( const LocalMesh& local_mesh, const GatherFields& gather_deps, + const ScatterFields& scatter_deps, + const LocalFields& local_deps, + const ExecutionSpace& exec_space, + const Kokkos::Array, NumSpace>& spaces, + const Func& func ) const + { + // Apply kernel to each entity. The user functor gets a local mesh for + // geometric operations, gather, scatter, and local dependencies for + // field operations (all of which may be empty), the local id, s, of + // the space currently being operated on, and the local ij + // index of the entity they are currently working on. + Cajita::grid_parallel_for( + "operator_apply", exec_space, spaces, + KOKKOS_LAMBDA( const int s, const int i, const int j ) { + functorTagDispatch( func, local_mesh, gather_deps, + scatter_deps, local_deps, s, i, + j ); + } ); + } + private: std::shared_ptr _mesh; std::shared_ptr> _gather_halo; diff --git a/unit_test/tstGridOperator2d.hpp b/unit_test/tstGridOperator2d.hpp index 424d723..5cdb9df 100644 --- a/unit_test/tstGridOperator2d.hpp +++ b/unit_test/tstGridOperator2d.hpp @@ -115,6 +115,25 @@ struct GridFunc } }; +//---------------------------------------------------------------------------// +// Grid operation. +struct FusedFunc +{ + template + KOKKOS_INLINE_FUNCTION void + operator()( const LocalMeshType& local_mesh, + const GatherDependencies& gather_deps, + const ScatterDependencies& scatter_deps, + const LocalDependencies& local_deps, const int, const int i, + const int j ) const + { + GridFunc gf; + gf( GridFunc::Tag{}, local_mesh, gather_deps, scatter_deps, local_deps, + i, j ); + } +}; + //---------------------------------------------------------------------------// void gatherScatterTest() { @@ -186,6 +205,31 @@ void gatherScatterTest() EXPECT_EQ( foo_out_host( i, j, d ), 4.0 + i + j ); EXPECT_EQ( bar_out_host( i, j, 0 ), 6.0 + i + j ); } ); + + // Apply the fused operator. + FusedFunc fused_func; + Kokkos::Array, 1> spaces = { + fm->mesh()->localGrid()->indexSpace( Cajita::Own{}, Cajita::Cell{}, + Cajita::Local{} ) }; + grid_op->apply( spaces, TEST_EXECSPACE(), *fm, fused_func ); + + // Check the grid results. + foo_out_host = Kokkos::create_mirror_view_and_copy( + Kokkos::HostSpace(), fm->view( FieldLocation::Cell(), FooOut() ) ); + bar_out_host = Kokkos::create_mirror_view_and_copy( + Kokkos::HostSpace(), fm->view( FieldLocation::Cell(), BarOut() ) ); + Kokkos::deep_copy( foo_out_host, + fm->view( FieldLocation::Cell(), FooOut() ) ); + Kokkos::deep_copy( bar_out_host, + fm->view( FieldLocation::Cell(), BarOut() ) ); + Cajita::grid_parallel_for( + "check_grid_out", Kokkos::DefaultHostExecutionSpace(), + *( fm->mesh()->localGrid() ), Cajita::Own(), Cajita::Cell(), + KOKKOS_LAMBDA( const int i, const int j ) { + for ( int d = 0; d < 2; ++d ) + EXPECT_EQ( foo_out_host( i, j, d ), 4.0 + i + j ); + EXPECT_EQ( bar_out_host( i, j, 0 ), 6.0 + i + j ); + } ); } //---------------------------------------------------------------------------// diff --git a/unit_test/tstGridOperator3d.hpp b/unit_test/tstGridOperator3d.hpp index 00c9756..c4a5950 100644 --- a/unit_test/tstGridOperator3d.hpp +++ b/unit_test/tstGridOperator3d.hpp @@ -158,6 +158,25 @@ struct GridFunc } }; +//---------------------------------------------------------------------------// +// Grid operation. +struct FusedFunc +{ + template + KOKKOS_INLINE_FUNCTION void + operator()( const LocalMeshType& local_mesh, + const GatherDependencies& gather_deps, + const ScatterDependencies& scatter_deps, + const LocalDependencies& local_deps, const int, const int i, + const int j, const int k ) const + { + GridFunc gf; + gf( GridFunc::Tag{}, local_mesh, gather_deps, scatter_deps, local_deps, + i, j, k ); + } +}; + //---------------------------------------------------------------------------// void gatherScatterTest() { @@ -277,6 +296,27 @@ void gatherScatterTest() EXPECT_EQ( foo_out_host( i, j, k, d ), 4.0 + i + j + k ); EXPECT_EQ( bar_out_host( i, j, k, 0 ), 6.0 + i + j + k ); } ); + + // Apply the fused operator. + FusedFunc fused_func; + Kokkos::Array, 1> spaces = { + fm->mesh()->localGrid()->indexSpace( Cajita::Own{}, Cajita::Cell{}, + Cajita::Local{} ) }; + grid_op->apply( spaces, TEST_EXECSPACE(), *fm, fused_func ); + + // Check the grid results. + Kokkos::deep_copy( foo_out_host, + fm->view( FieldLocation::Cell(), FooOut() ) ); + Kokkos::deep_copy( bar_out_host, + fm->view( FieldLocation::Cell(), BarOut() ) ); + Cajita::grid_parallel_for( + "check_grid_out", Kokkos::DefaultHostExecutionSpace(), + *( mesh->localGrid() ), Cajita::Own(), Cajita::Cell(), + KOKKOS_LAMBDA( const int i, const int j, const int k ) { + for ( int d = 0; d < 3; ++d ) + EXPECT_EQ( foo_out_host( i, j, k, d ), 4.0 + i + j + k ); + EXPECT_EQ( bar_out_host( i, j, k, 0 ), 6.0 + i + j + k ); + } ); } //---------------------------------------------------------------------------//