Flop counting

[inducer]

Very WIP right now.

cc @majosm

[majosm]

ConditionalDependencyMapper addition is unused (I used it in an earlier draft but ended up not needing it). I left it in because it doesn't harm anything and seems potentially useful.

@inducer I think this is ready for a first look, apart from some basedpyright errors that I don't understand. Probably best to go commit by commit.

[majosm]

@inducer I ran into a complication with the flop counting for conditionals. When mirgecom projects something from interior/boundary faces to all faces, meshmode's direct connection code creates an AdvancedIndexInContiguousAxes (code) shaped like the all-faces discretization, and then it follows up with an IndexLambda that zeros out the entries corresponding to elements that aren't part of the input discretization. You can see this for wave below:

Since the flop counting currently evaluates flops for Conditional as max(nflops_then, nflops_else), this is causing the number of flops to be overestimated for this part of the DAG. It's counting the flops (in the case of wave) over 1536 faces instead of the 64 that end up being nonzero. I'm not sure how to deal with this, since it's data-dependent. I could maybe hack in a check for condition being something easily evaluated and then just count the number of true/false cases, but beyond that I don't know how to resolve this. Any thoughts?

[majosm]

@inducer Do you know what's causing these basedpyright errors for rec? Is this another instance of the thing that should be added to the baseline?

Also, about the Conditional flop count handling: if we're going to model it after GPU execution, should it be counted as max(nflops_then, nflops_else) or sum(nflops_then, nflops_else)?

[inducer]

I'm not sure I understand the intent of this.

For example, DistributedRecv is always materialized (it ends up in a buffer after all).

What is this function trying to accomplish?

[majosm]

I'm using this in the flop counting code to tell me a couple of things:

1. Am I allowed to materialize/unmaterialize this expression? (Can't for things like NamedArray and DistributedSendRefHolder.)
2. Can this expression be lowered to an index lambda? (Can't for things like InputArgumentBase and DistributedRecv.)

[majosm]

Is 5e30e95 any better? It still has some isinstance() checks that seem kind of fragile, but I'm not sure how to avoid them.

[majosm]

@inducer This is ready for a look again when you have a chance.

[inducer]

Just the enum thing and then this is good to go from my perspective. 🎉

[inducer]

Let's do a separate node type for nflops.

[inducer]

Suggested change

	raise UndefinedOpFlopCountError(next(iter(var_names))) from None
	raise UndefinedOpFlopCountError(next(iter(var_names)))

[inducer]

Suggested change

	raise AssertionError from None
	raise AssertionError