Parthenon Enhancement Proposal 1 -- Support Subclassing StateDescriptor

[bprather]

Introduction

The Parthenon "Package" is a powerful way to split up code and data/parameters by function, keeping parameters and fields associated with the functions that handle them. However, current limitations make it difficult for users to extend the definition of a "Package" to better fit their needs. This "Parthenon Enhancement Proposal" (PEP) explains the desired design pattern, and outlines ways to accommodate it in Parthenon if this is desired.

I partially also intend this as a template for future PEPs, and would welcome format feedback. Ideally future submissions will not be so verbose.

Relevant background and current status

A Parthenon package consists of three things:

1. A C++ namespace (by convention, at least)
2. A (function constructing and returning a) StateDescriptor object, which is basically a host for a Parameters object, which is more or less a std::map<std::string,std::any>. A list of StateDescriptor objects is then accessible as a member of any Mesh or MeshBlock object, i.e., from nearly any function in the code
3. Implementations of zero or more common functions, which are assigned to function pointer members of the StateDescriptor object and called at specific times during a Parthenon step (Some of them only if one subclasses a particular Driver)

This is a neat paradigm! It makes nearly arbitrary data & parameters automatically accessible nearly anywhere in a code from the Mesh/MeshBlock objects, while still hierarchically organized. The function pointer or "callback" structure additionally allows implementing new features in a way that can be enabled & disabled cleanly, without a bunch of if (feature_enabled) constructs, which are hard to read and easy to get wrong (e.g. reading options & arrays that don't exist, running an extra op or fix that should have been disabled, etc.)

The full list of current callbacks is (as documented somewhat obliquely in metadata.md):

// Per-package callbacks
void PreFillDerivedBlock(MeshBlockData<Real> *rc);
void PreFillDerivedMesh(MeshData<Real> *rc);
void PostFillDerivedBlock(MeshBlockData<Real> *rc);
void PostFillDerivedMesh(MeshData<Real> *rc);
void FillDerivedBlock(MeshBlockData<Real> *rc);
void FillDerivedMesh(MeshData<Real> *rc);

void PreStepDiagnosticsMesh(SimTime const &simtime, MeshData<Real> *rc);
void PostStepDiagnosticsMesh(SimTime const &simtime, MeshData<Real> *rc);

Real EstimateTimestepBlock(MeshBlockData<Real> *rc);
Real EstimateTimestepMesh(MeshData<Real> *rc);

AmrTag CheckRefinementBlock(MeshBlockData<Real> *rc);

void InitNewlyAllocatedVarsMesh(MeshData<Real> *rc);
void InitNewlyAllocatedVarsBlock(MeshBlockData<Real> *rc);

This list of callbacks is long, yet limiting: only a few of the things the code actually needs to do during a simulation step (notably and understandably not any physics) can be in this list. Even as additions are made, it will not (and should not) cover every part of any given downstream algorithm.

The package pattern has proven a very useful way of splitting up code for handling different physics or interchangeable feature implementations, at least for KHARMA. However, this encapsulation is broken when the MakeTaskList function must repeatedly query which packages are loaded, and special-case on their inclusion to form the task list. There are often a number of functionally similar operations which must be completed by different packages, each for its own variables -- these must simply be listed out and special-cased, with new entries added manually to all Driver implementations.

For example, operations common to many packages for their particular fields, like "compute the primitive variables" or "add any source terms" must be treated like completely different operations in every package and Driver, which leads to code duplication and mistakes, not to mention a lot of cruft checking for package enablement before every possible operation, or on every function call inside a package.

Proposed solution

By writing a subclass of StateDescriptor, a user can add arbitrary function pointers, making a "user package" specific to the case of their code (e.g., conservative finite-volume schemes in curved space, to choose a random example). Extending this pattern is a powerful way to ensure that enabled features are actually enabled, and disabled features disabled, without either writing enablement guards into every package function, or breaking encapsulation in MakeTaskList by querying packages for presence and enablement before every call.

This pattern does not heavily involve changing Parthenon, though I will suggest a few small key additions. Instead, the proposed changes are mostly to support this as a use case, and gradually see it implemented in downstreams and perhaps more complex examples. As the documentation of packages develops, it could be a note in the relevant documentation with a short example snippet implementing the parent function, which just iterates through packages calling the member function.

The main propsed code change, for now, would be to address a pain point when using the package list (Packages_t) with a user package type. With a user type, the list now contains base StateDescriptor objects, such as the "Refinement" package, and user package objects of some subclass type. Iterating through the list for user callbacks therefore becomes cumbersome due to casting, as packages in the list lose their type information. That is, the problem is that the pattern:

for(auto package : pmesh->AllPackages()) {
    MyPackage *my_package = static_cast<MyPackage*>(package.second)
    my_package->MyCallback();
}

will throw errors if there are members of AllPackages which are not of type MyPackage, but instead of the base StateDescriptor type. There are two possible solutions, and I propose implementing one of them in Parthenon as a function pmesh->AllUserPackages(), or better yet pmesh->AllPackagesOfType<Type>().

1. Perhaps there is some clever way to distinguish the underlying types of the pointers within the Packages_t object. I think the list itself erases type information, as package.second always returns a shared_ptr cast to StateDescriptor. However, there may be a list or pattern which preserves the original type information but otherwise behaves similarly. Alternatively, the type information could be carried in a separate list member of Packages_t.
2. If there is not a type-based solution, there's always a manual check. Each Parthenon package could register itself in an internal list, which could be checked to distinguish user packages from Parthenon packages. Note this wouldn't be as flexible as the type-based approach, as it would only allow one type of user package (or would force the user to keep their own lists sorting packages by type). Note that it may be independently useful to keep a list of Parthenon packages separate from user packages, even if they are the same type, as users may not know nor care the full list of loaded Parthenon packages, and should have at least some barrier before unintentionally touching Parthenon packages' parameters in user code: say, when iterating through packages changing a common parameter by name.

Note that another possible solution to this problem would be to eliminate the "Refinement" package, and make the package list entirely user-controlled. This wouldn't be my preferred solution, as I'd like to see parts of Parthenon make use of the package pattern more fully (especially the code for outputs including/especially reductions, and possibly the code for drivers as well).

Corollary considerations

It has proven useful in KHARMA, and I think would be useful generally, to add a compile flag which enables printing exactly what functions are being executed, as they are called. Even in Parthenon's current state, callbacks obfuscate which functions from which packages are actually being called, and this is a significant downside vs an explicit C code, especially when debugging. Parthenon already defines Kokkos profiling regions for many of its own functions, and these can be used with the kokkos-tools kernel logger to print out exactly what the code is doing. Pushing a new region with the package name for every non-trivial callback would make this functionality complete for user code.

Additionally, if we encourage splitting code into many small packages, it would be very useful to provide a way to register package dependencies, and allow Parthenon to load dependent packages automatically. This would ideally mean registering all available packages in a map from package names to initializer functions, then providing a list of dependency names as a part of the StateDescriptor object created by each package.

Implementation

Implementation would consist of adding either Mesh::AllUserPackages(), Mesh::AllPackagesOfType<Type>(), or both functions (with accompanying versions in MeshBlock).

1. Implementing Mesh::AllUserPackages() will additionally require registering all Parthenon packages by name, and returning the package list sans those entries -- I could do this.
2. Implementing Mesh::AllPackagesOfType<Type>() will require either deep incantations to the C++ type system beyond my ken, or somehow adding type info to Packages_t, perhaps as a separate list or map alongside the usual std::map<std::string,shared_ptr<StateDescriptor>>.

I am happy to work on the corollaries and provide prototype PRs if this is accepted, or to workshop them in other issues. I don't personally think they deserve PEPs.

Documentation of packages generally is a separate issue, and I think anyone who's read this PEP can adequately drop in a paragraph or two illustrating the new-callback pattern to the eventual package documentation.

Implementation in downstreams and examples is a separate issue.

[Yurlungur]

Thanks for submitting this! Exciting to see the first PEP ;)

Hmm this proposal seems to be entangling two separate things, I think:

• Allowing the user to distinguish between parthenon-added packages and user-added packages
• Allowing the user to extend the StateDescriptor arbitrarily

I am 100% in favor of the former... although I don't know if I have a strong preference between an internal map/list-based approach vs. a type-based approach. The type-based approach makes more sense if the desire is to extend StateDescriptor further. But it does have some advantages even if we don't plan to do that, as getting the compiler to yell at us before we start a big expensive simulation seems desirable.

As far as extending StateDescriptor arbitrarily... what kind of stuff do you want to hide in a more generic one exactly? Can you give an example? I'm not necessarily opposed to allowing this (we obviously do implicitly allow it), but I wonder if the need to do so indicates a bigger gap in the design of packages we should address?

Additionally, if we encourage splitting code into many small packages, it would be very useful to provide a way to register package dependencies, and allow Parthenon to load dependent packages automatically.

This would be a cool feature 👍

It has proven useful in KHARMA, and I think would be useful generally, to add a compile flag which enables printing exactly what functions are being executed, as they are called.

Does this not require some additional decoration by the author of said function? Also I think we have this. If you put your function inside a Kokkos profiling region, the kernel tracker will report it.

[bprather]

Maybe my use case will shed some light here. In KHARMA, for example, I want to add about 8 or 10 callbacks to the Parthenon list:

class KHARMAPackage : public StateDescriptor {
    public:
        KHARMAPackage(std::string name) : StateDescriptor(name) {}

        // PHYSICS
        // Recovery of primitive variables from conserved.
        // These can be host-side functions because they are not called from the Uberkernel --
        // rather, they are called on zone center values once per step only.
        // Called by various Flux::*UtoP*
        std::function<void(MeshBlockData<Real>*, IndexDomain, bool)> BlockUtoP = nullptr;
        std::function<void(MeshData<Real>*, IndexDomain, bool)> MeshUtoP = nullptr;

        // Maybe at some point we'll have 
        // Since Flux::prim_to_flux must cover everything, it's not worth splitting now
        //std::function<void(MeshBlockData<Real>*, IndexDomain, bool)> BlockPtoU = nullptr;

        // Source term to add to the conserved variables during each step
        std::function<void(MeshData<Real>*, MeshData<Real>*)> AddSource = nullptr;

        // Source term to apply to primitive variables, needed for some problems in order
        // to control dissipation (Hubble, turbulence).
        // Must be applied over entire domain!
        std::function<void(MeshBlockData<Real>*)> BlockApplyPrimSource = nullptr;

        // Apply any fixes after the initial fluxes are calculated
        std::function<void(MeshData<Real>*)> FixFlux = nullptr;

        // Apply any floors or limiters specific to the package (that is, on the package's variables)
        // Called by Floors::*ApplyFloors
        std::function<void(MeshBlockData<Real>*, IndexDomain)> BlockApplyFloors = nullptr;
        std::function<void(MeshData<Real>*, IndexDomain)> MeshApplyFloors = nullptr;

        // CONVENIENCE
        // Anything to be done after every step is fully complete -- usually reductions or preservation of variables
        std::function<void(Mesh*, ParameterInput*, const SimTime&)> MeshPreStepUserWorkInLoop = nullptr;
        // Anything to be done after every step is fully complete -- usually reductions or preservation of variables
        std::function<void(Mesh*, ParameterInput*, const SimTime&)> MeshPostStepUserWorkInLoop = nullptr;
        std::function<void(MeshBlock*, ParameterInput*)> BlockUserWorkBeforeOutput = nullptr;
};

I then create a new namespace Packages with the parent functions that iterate through packages calling these members. Except for the type issues this raises in Packages_t as described, I think this is a really elegant way to describe generally what each physics module of a particular code can or should do. KHARMA certainly looks much cleaner with all of these operations (and boundaries, but that's for another time) moved into callbacks, vs. using a lot of similarly-named functions which bear no formal relation to one another. A few of these functions, namely the delegations of code-wide callbacks at the end, probably deserve inclusion in Parthenon. But things like UtoP are probably not desired Parthenon features, let alone naming conventions.

I should say, upon reflection adding callbacks doesn't require subclassing StateDescriptor, since we could just add a std::map<std::string, std::function> to Parameters. However, I like the ability to choose my calling conventions, and the nice strict typing of members vs generating runtime errors when a member can't be found. I don't see how further loading down Parameters is a better alternative to subclassing for the case of adding callbacks.

Re: print flags: If Kokkos regions allow printing during runtime on entry and exit, that kills two birds with one stone, and KHARMA should definitely adopt regions too. What I was proposing, though, was definitely more involved than the existing regions in Parthenon: placing a print flag (or region markers) around each package.second->DoAThing(rc) so that Parthenon will print out which package & function of user code it is entering, as it does so. Then, when I as a user write a dumb segfault, I know which function it was in before even needing to muck about with backtrace hooks.

[Yurlungur]

Ok I see. Is the use-case here that you have multiple packages that define, e.g., a con2prim because they work on different variables? And you want to loop over all of them?

That makes sense to me and I'm not opposed to supporting that. One related thing that's been on my mind has been pgens and related things. But maybe that should wait for now.

Re: print flags: If Kokkos regions allow printing during runtime on entry and exit, that kills two birds with one stone, and KHARMA should definitely adopt regions too. What I was proposing, though, was definitely more involved than the existing regions in Parthenon: placing a print flag (or region markers) around each package.second->DoAThing(rc) so that Parthenon will print out which package & function of user code it is entering, as it does so. Then, when I as a user write a dumb segfault, I know which function it was in before even needing to muck about with backtrace hooks.

I see. I think profiling regions basically gets you that, since the kernel logger will tell you everything so long as you add regions where you need them, and you get profiling for free. That said, pushing and popping profiling regions is a little manual. So maybe we should think about a way to do that more automatically.

[bprather]

Yeah, that's the idea. Incidentally, I'm imagining down the line there could be a "ProblemDescriptor" subclass which takes away (forcibly nulls) most of the callbacks in StateDescriptor, and allows only a Parameters object and a couple of basic function pointers including the default per-block pgen. (Incidentally, making problems a package subclass would also let them state their package dependencies if/when the dependency stuff gets written)

I think having Parthenon set a Kokkos profiling region per (nontrivial) call into a user function would be fantastic, and save the users a lot of Kokkos::push_region(). If users need to split regions further they can, but that's an excellent start. I'll amend the PEP to propose this as the corollary.

Re: approaches, I'm increasingly thinking we should do both. There's the (probably very common) case where a user has not subclassed StateDescriptor, yet still wants a list of only their own packages to iterate. Then, there's the (also hopefully common) case that a user has multiple types of packages, distinguishing a FluidPackage type from an OutputPackage or a ProblemPackage, for example.

If a registry and type list seems like too much, we could always subclass StateDescriptor into ParthenonPackage, and reserve that type for Parthenon-specific packages. That would allow implementing both solutions with just the type trick.

[bprather]

So, I feel a bit silly. It turns out the "deep type incantation" is:

if (MyPackage *my_package = dynamic_cast<MyPackage*> package.second) {
    my_package->MyCallback();
}

Coupled with a virtual destructor in StateDescriptor:

virtual ~StateDescriptor() = default;

No errors, no mess. Marking the destructor virtual makes StateDescriptor polymorphic to allow use of the runtime dynamic_cast call, which just returns nullptr if the cast cannot be made. Types are actually fun and easy, don't let anyone tell you otherwise.

This makes AllPackagesOfType<Type>() trivial to implement with just the virtual destructor change. Personally, I would still implement AllUserPackages by writing a trivial class ParthenonPackage : public StateDescriptor, declaring "Refinement" to be of this type, and making AllUserPackages return everything that can't be cast to ParthenonPackage. Open to ideas if this seems too intrusive, and/or especially alternate naming. I personally find the name StateDescriptor completely obtuse, so it's hard to work around.

[pgrete]

👍 for the first PEP!

In general, I see the use/need for this kind of pattern -- especially with respect to ProblemGenerators as packages so that all kind of source terms can cleanly be incorporated.
In AthenaPK we currently still working with (additional) manual callbacks, but we're getting to a point where calling multiple source terms becomes more and more intrusive (in the sense of requiring code changes in places like the driver that are in principle unrelated to the problem setup itself).

So, I feel a bit silly. It turns out the "deep type incantation" is:

if (MyPackage *my_package = dynamic_cast<MyPackage*> package.second) {
    my_package->MyCallback();
}

This looks like a very clean pattern to me.

This makes AllPackagesOfType<Type>() trivial to implement with just the virtual destructor change. Personally, I would still implement AllUserPackages by writing a trivial class ParthenonPackage : public StateDescriptor, declaring "Refinement" to be of this type, and making AllUserPackages return everything that can't be cast to ParthenonPackage. Open to ideas if this seems too intrusive, and/or especially alternate naming.

Just to be clear: How would that translate to the calls within Parthenon itself? Would all places where we call we iterate over packages be split into iterations over Parthenon and User Packages?

And regarding the snippet above with the if ... dynamic_cast, this would only live in downstream codes as Parthenon itself would not know about MyPackage?
So in other words, in your downstrea task list would have a function call your callback function of all custom packages, correct?

I personally find the name StateDescriptor completely obtuse, so it's hard to work around.

Agreed (in fact, agreed for a long time #397).
We discussed to rename StateDescriptor to Package some time ago and it looks like I didn't do my assigned task :D (#568).
Any objection going forward with this (independent of this PEP)?