Skip to content

multiple index space grid operator apply #35

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sslattery wants to merge 1 commit into
base: master
Choose a base branch
from
Open

multiple index space grid operator apply #35

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 95 additions & 11 deletions src/Picasso_GridOperator.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -398,14 +398,14 @@ class GridOperator
//
// The functor is given a ParticleView allowing the kernel to read and
// write particle data.
template <class ExecutionSpace, class ParticleList_t, class WorkTag,
template <class ExecutionSpace, class... ParticleFields, class WorkTag,
class Func>
void apply( FieldLocation::Particle, const ExecutionSpace& exec_space,
const FieldManager<Mesh>& fm, const ParticleList_t& pl,
const WorkTag&, const Func& func ) const
void apply( const FieldLocation::Particle& location,
const ExecutionSpace& exec_space, const FieldManager<Mesh>& fm,
const ParticleList<Mesh, ParticleFields...>& pl, const WorkTag&,
const Func& func ) const
{
applyImpl<WorkTag>( fm, exec_space, FieldLocation::Particle(), pl,
func );
applyImpl<WorkTag>( fm, exec_space, location, pl, func );
}

// Apply the operator in a loop over particles. Functor does not have a
Expand All @@ -416,12 +416,13 @@ class GridOperator
//
// The functor is given a ParticleView allowing the kernel to read and
// write particle data.
template <class ExecutionSpace, class ParticleList_t, class Func>
void apply( FieldLocation::Particle, const ExecutionSpace& exec_space,
const FieldManager<Mesh>& fm, const ParticleList_t& pl,
template <class ExecutionSpace, class... ParticleFields, class Func>
void apply( const FieldLocation::Particle& location,
const ExecutionSpace& exec_space, const FieldManager<Mesh>& fm,
const ParticleList<Mesh, ParticleFields...>& pl,
const Func& func ) const
{
applyImpl<void>( fm, exec_space, FieldLocation::Particle(), pl, func );
applyImpl<void>( fm, exec_space, location, pl, func );
}

// Apply the operator in a loop over the owned entities of the given
Expand All @@ -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 <class ExecutionSpace, class Location, class Func>
Expand All @@ -449,6 +451,35 @@ class GridOperator
applyImpl<void>( 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 <std::size_t NumSpace, class ExecutionSpace, class WorkTag,
class Func>
void apply( const Kokkos::Array<Cajita::IndexSpace<Mesh::num_space_dim>,
NumSpace>& spaces,
const ExecutionSpace& exec_space, const FieldManager<Mesh>& fm,
const WorkTag&, const Func& func ) const
{
applyImpl<WorkTag>( 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 <std::size_t NumSpace, class ExecutionSpace, class Func>
void apply( const Kokkos::Array<Cajita::IndexSpace<Mesh::num_space_dim>,
NumSpace>& spaces,
const ExecutionSpace& exec_space, const FieldManager<Mesh>& fm,
const Func& func ) const
{
applyImpl<void>( fm, exec_space, spaces, func );
}

public:
// Manage field dependencies and apply the operator.
template <class WorkTag, class ExecutionSpace, class... Args>
Expand Down Expand Up @@ -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() ),
Expand All @@ -680,6 +711,59 @@ class GridOperator
} );
}

// Apply the operator in a loop over an array of the given index space. 3D
// specialization.
template <class WorkTag, class LocalMesh, class GatherFields,
class ScatterFields, class LocalFields, class ExecutionSpace,
std::size_t NumSpace, class Func>
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<Cajita::IndexSpace<3>, 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<WorkTag>( 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 <class WorkTag, class LocalMesh, class GatherFields,
class ScatterFields, class LocalFields, class ExecutionSpace,
std::size_t NumSpace, class Func>
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<Cajita::IndexSpace<2>, 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<WorkTag>( func, local_mesh, gather_deps,
scatter_deps, local_deps, s, i,
j );
} );
}

private:
std::shared_ptr<Mesh> _mesh;
std::shared_ptr<Cajita::Halo<memory_space>> _gather_halo;
Expand Down
44 changes: 44 additions & 0 deletions unit_test/tstGridOperator2d.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,25 @@ struct GridFunc
}
};

//---------------------------------------------------------------------------//
// Grid operation.
struct FusedFunc
{
template <class LocalMeshType, class GatherDependencies,
class ScatterDependencies, class LocalDependencies>
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()
{
Expand Down Expand Up @@ -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<Cajita::IndexSpace<2>, 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 );
} );
}

//---------------------------------------------------------------------------//
Expand Down
40 changes: 40 additions & 0 deletions unit_test/tstGridOperator3d.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,25 @@ struct GridFunc
}
};

//---------------------------------------------------------------------------//
// Grid operation.
struct FusedFunc
{
template <class LocalMeshType, class GatherDependencies,
class ScatterDependencies, class LocalDependencies>
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()
{
Expand Down Expand Up @@ -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<Cajita::IndexSpace<3>, 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 );
} );
}

//---------------------------------------------------------------------------//
Expand Down