[RFC, generate_loopy] Propagating tags to lp.Instructions

[kaushikcfd]

In the current design, the tags attached to a materialized array are propagated to the corresponding variable's tags in the loopy kernel. However, no API is present to propagate any tags to the loopy instructions.

Here's what @matthiasdiener and I came up with: Make generate_loopy take in another argument that tells us which Array.tags should make into lp.Instruction.tags. The array expressions that are evaluated/read by a statement contribute toward this propagation. This way we could even capture any tags that were attached to unmaterialized arrays.

class OperationDefiningTag(Tag): pass
class FluxComputingStage(OperationDefiningTag): pass
class QuadratureStage(OperationDefiningTag): pass
class ScalarOpOnBoundary(OperationDefiningTag): pass

def flux(x, y):
     a1 = x[y].tagged(FluxComputingStage())
     a2 = pt.sin(a1).tagged(ScalarOpOnBoundary())
     return 2*a2

def quadrature(D, u):
     a1 = pt.einsum("ij,ej->ei", D, u)
     a2 = 2 * a1.tagged(QuadratureStage())
     return pt.sin(a2)

def face_mass(..) : ...


def main():
    u = ... (DOF values)
    map_ = ... (vol. dof to face dof mapping)
    D = ... (derivative matrix)

    flux_values = flux(u, map_).tagged(ImplStored())
    integral = quadrature(D, u).tagged(ImplStored())

    result = face_mass(flux_value) + integral

    t_unit = pt.generate_loopy(result, insn_tag_t_to_propagate=OperationDefiningTag)

For the above expression, the generated loopy kernel will have 3 statements: one for flux_values, one for integral and one for result. The statement writing flux_values would be tagged with FluxComputingStage and ScalarOpOnBoundary, the statement writing integral would be tagged with QuadratureStage, and the instruction writing result will have zero tags.

[inducer]

I think there are two questions here. The first is what to do about the fact that, to pytato, result and computation are represented as one, whereas in loopy, the two representations drift apart. The second question is that, since we're not going to be tagging each intermediate result, how do tags "propagate" along the DAG.

For the first question, I actually think that maintaining this separation is not super important. If something is tagged X in pytato, we could simply propagate tags to both data and computation X. It's easy to filter later.

For the second question, I expressed some ideas here. Basically, every pytato result has a set of, say, CostAttributionTags. These notionally tag the entire subgraph leading up to that result, up to the point where another CostAttributionTag is encountered, which resets the tag. If there are results that are shared between multiple subexpressions that have cost attribution tags, those sub-DAGs get tagged with both, until another explicit cost attribution tag is hit.