Dispatcher
Kernels in kamayan will dispatch into functions, often based on compile time template parameters. This allows for modularity in the code/solver capabilities without code duplication, and
parthenon::par_for_inner(
member, 0, nrecon - 1, ib.s - 1, ib.e + 1, [&](const int var, const int i) {
auto stencil = SubPack<Axis::IAXIS>(pack_recon, b, var, k, j, i);
Reconstruct<reconstruction_traits>(stencil, vM(var, i), vP(var, i));
});
member.team_barrier();
parthenon::par_for_inner(member, ib.s, ib.e + 1, [&](const int i) {
// riemann solve
auto vL = MakeScratchIndexer(pack_recon, vP, b, i - 1);
auto vR = MakeScratchIndexer(pack_recon, vM, b, i);
auto pack_indexer = SubPack(pack_flux, b, k, j, i);
if constexpr (hydro_traits::MHD == Mhd::ct) {
vL(MAGC(0)) = pack_indexer(TopologicalElement::F1, MAG());
vR(MAGC(0)) = pack_indexer(TopologicalElement::F1, MAG());
}
RiemannFlux<TE::F1, riemann, hydro_traits>(pack_indexer, vL, vR);
});
In the above snippet the Reconstruct and RiemannFlux functions
will dispatch to the correct methods depending on the template
parameters recon & riemann.
Having these template parameters makes for concise code in the kernels, but as the number of parameters used in a function increases the combinatorial combinations that need to be explicitly instantiated, if we want the choices to be runtime decisions. Kamayan provides a dispatcher that simplifies the dispatching of templated functions from an arbitrary number of runtime parameters.
// functor with template parameters that we want to launch from
// runtime values
struct MyFunctor {
// put together the allowed values to instantiate from, to be checked
// against the runtime values
using options = OptTypeList<OptList<Foo, Foo::a, Foo::b>, OptList<Bar, Bar::d, Bar::e>,
OptList<Baz, Baz::f, Baz::g>>;
// return type of our functor
using value = void;
// function we want to dispatch, needs to be named like this
template <Foo FOO, Bar BAR, Baz BAZ>
value dispatch(int foo, int bar, int baz) const {
EXPECT_EQ(foo_func<FOO>(), foo);
EXPECT_EQ(bar_func<BAR>(), bar);
EXPECT_EQ(baz_func<BAZ>(), baz);
}
};
Such functions can be provided as a functor that has two using decls:
-
options-- defines aOptTypeListthat describes how the functor can be built from non-type template parameterenum classs, and their possible runtime values. -
value-- the return type of the function to be dispatched
The dispatcher will template and call the contained dispatch function,
forwarding any arguments that it is passed.
// will execute MyFunctor().template dispatch<foo, bar, baz>(foo_v, bar_v, baz_v)
Dispatcher<MyFunctor>(PARTHENON_AUTO_LABEL, foo, bar, baz).execute(foo_v, bar_v, baz_v);
Options
enum classs that can be used as template parameters are declared with the
POLYMORPHIC_PARM macro defined in dispatcher/options.hpp.
// define our enums, the macro will take care of
// specializing OptInfo for each type that can give
// debug output when dispatcher fails
POLYMORPHIC_PARM(Foo, a, b);
POLYMORPHIC_PARM(Bar, d, e);
POLYMORPHIC_PARM(Baz, f, g);
Note
POLYMORPHIC_PARMs must be declared in the kamayan namespace in order
to specialize the OptInfo_t that is used to inform the dispatcher
of some debugging information.
Only enums that have been declared like this can be used in the options list
used by the dispatcher. The options OptTypeList is a list of types, one for
each template parameter in the dispatch function, that are used to
instantiate the template parameters. It can be helpful to mentally map the
OptList types onto a switch statement for the template parameter.
// OptList<Foo, Foo::a, Foo::b> roughly corresponds to
template<typename... Args>
void Dispatcher<MyFunctor>(Foo foo_v, Args ...&&args) {
switch (foo_v) {
case Foo::a :
MyFunctor.template <Foo::a>(std::forward<Args>(args)...);
break;
case Foo::b :
MyFunctor.template <Foo::b>(std::forward<Args>(args)...);
break;
}
}
Composite Types
Another common pattern with the non-type template parameters we have discussed
so far is to use them to build more complex type traits as composite options
built from our POLYMORPHIC_PARMs.
// we can also template our dispatches on types that are templated on
// multiple enum options
template <Foo f, Bar b>
struct CompositeOption {
static constexpr auto foo = f;
static constexpr auto bar = b;
};
struct MyCompositeFunctor {
// composite types need to provide the factory rather than an OptList. Composite types
// always come first
using options = OptTypeList<CompositeFactory, OptList<Baz, Baz::f, Baz::g>>;
using value = void;
template <typename Composite, Baz BAZ>
value dispatch(int foo, int bar, int baz) const {
constexpr auto FOO = Composite::foo;
constexpr auto BAR = Composite::bar;
EXPECT_EQ(foo_func<FOO>(), foo);
EXPECT_EQ(bar_func<BAR>(), bar);
EXPECT_EQ(baz_func<BAZ>(), baz);
}
};
The functors now look very similar to the ones we looked at before. The main
difference is that we need to provide a factory type, CompositeFactory,
that informs the dispatcher how to build our composite type from
the provided runtime options.
// In order to have composite type options we need to define
// a factory that can tell dispatch how to build our type
struct CompositeFactory : OptionFactory {
using options = OptTypeList<OptList<Foo, Foo::a, Foo::b>, OptList<Bar, Bar::d, Bar::e>>;
template <Foo f, Bar b>
using composite = CompositeOption<f, b>;
using type = CompositeFactory;
};
The factory interface is very similar to the previous functor definition, with
three using decl definitions
options-- andOptTypeListlike in the functors to enumerate the runtime options and their allowed values.composite-- a templatedusing declthat aliases the type the factory producestype-- always points to itself
Additionally the factory needs to inherit from the OptionFactory type to inform
the dispatcher that this is a composite type.
Dispatcher<MyCompositeFunctor>(PARTHENON_AUTO_LABEL, Foo::a, Bar::d, Baz::f)
.execute(1, 0, 1);
Using Config for dispatch
Instead of passing each individual runtime template parameter value to the dispatcher
the global kamayan Config may also be used. The dispatcher will instead pull the
runtime values directly from POLYMORPHIC_PARMS that have been registered to the Config
struct PrepareConserved_impl {
using options = OptTypeList<HydroFactory>;
using value = TaskStatus;
template <typename hydro_traits>
requires(NonTypeTemplateSpecialization<hydro_traits, HydroTraits>)
value dispatch(MeshData *md) {
auto pack = grid::GetPack(typename hydro_traits::ConsPrim(), md);
const int nblocks = pack.GetNBlocks();
auto ib = md->GetBoundsI(IndexDomain::interior);
auto jb = md->GetBoundsJ(IndexDomain::interior);
auto kb = md->GetBoundsK(IndexDomain::interior);
const auto ndim = md->GetNDim();
parthenon::par_for(
PARTHENON_AUTO_LABEL, 0, nblocks - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) {
capture(ndim);
// also need to average the face-fields if doing constrained transport
if constexpr (hydro_traits::MHD == Mhd::ct) {
using te = TopologicalElement;
if (ndim > 1) {
pack(b, MAGC(0), k, j, i) = 0.5 * (pack(b, te::F1, MAG(), k, j, i + 1) +
pack(b, te::F1, MAG(), k, j, i));
pack(b, MAGC(1), k, j, i) = 0.5 * (pack(b, te::F2, MAG(), k, j + 1, i) +
pack(b, te::F2, MAG(), k, j, i));
}
if (ndim > 2) {
pack(b, MAGC(2), k, j, i) = 0.5 * (pack(b, te::F3, MAG(), k + 1, j, i) +
pack(b, te::F3, MAG(), k, j, i));
}
}
auto U = SubPack(pack, b, k, j, i);
Prim2Cons<hydro_traits>(U, U);
});
return TaskStatus::complete;
}
};
TaskStatus PrepareConserved(MeshData *md) {
auto cfg = GetConfig(md);
return Dispatcher<PrepareConserved_impl>(PARTHENON_AUTO_LABEL, cfg.get()).execute(md);
}