Flop counting #623

inducer · 2025-10-20T22:07:38Z

Very WIP right now.

majosm · 2025-10-27T20:52:11Z

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 · 2025-11-07T22:53:02Z

@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 · 2025-11-14T17:23:27Z

@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 · 2025-12-04T16:53:35Z

pytato/analysis/__init__.py

            _visited_functions=self._visited_functions)

-    def post_visit(self, expr: Any) -> None:
+    def post_visit(self, expr: ArrayOrNames | FunctionDefinition) -> None:


Would it make sense to define ArrayOrNamesOrFunction?

Yeah, probably. This shows up all over the place, though, so maybe best for a standalone PR?

pytato/reductions.py

pytato/scalar_expr.py

inducer · 2025-12-04T17:02:34Z

pytato/utils.py

+    return (
+        isinstance(expr, Array)
+        and not isinstance(expr, (
+            # FIXME: Is there a nice way to generalize this?


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?

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

Am I allowed to materialize/unmaterialize this expression? (Can't for things like NamedArray and DistributedSendRefHolder.)

Can this expression be lowered to an index lambda? (Can't for things like InputArgumentBase and DistributedRecv.)

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

inducer · 2025-12-04T17:41:02Z

pytato/transform/calls.py

+        return frozenset(expr.bindings.values())
+
+
+class _CallMaterializer(CopyMapper):


I'm not sure I understand this. This appears to enforce a specific view of whether function arguments and/or return values are materialized, when I would prefer that pytato have no opinion in the matter. They could/could not be materialized; they're just expressions, just like everyone else.

Could you explain the intent behind this?

I was trying to look at the flop counting from the perspective of "if I immediately generated code from this DAG". In that sense, any function calls remaining in the DAG would (in theory) be emitted as kernel functions, right? So the inputs/outputs would get materialized?

I removed the call materialization and made the flop counting require function calls to be inlined beforehand for now.

inducer · 2025-12-04T21:28:37Z

pytato/analysis/__init__.py

+        op_name_to_num_flops: Mapping[str, int] | None = None,
+    ) -> ArrayOrScalar:
+    """
+    Count the total number of floating point operations in the DAG *expr*.


Explain more precisely how this interacts with materialization.

Something like 9284664?

inducer · 2025-12-04T21:29:07Z

pytato/analysis/__init__.py

+    Returns a dictionary mapping materialized nodes in DAG *expr* to their floating
+    point operation count.


Explain more precisely how this interacts with materialization.

…yFlopCounter

…ation

…_lambda()

majosm force-pushed the flop-counting branch 6 times, most recently from 31da603 to 83629d9 Compare October 27, 2025 16:31

majosm force-pushed the flop-counting branch from 09a4549 to 7bcb940 Compare November 7, 2025 22:27

majosm force-pushed the flop-counting branch 4 times, most recently from 5236448 to 72b69ac Compare November 14, 2025 17:05

majosm force-pushed the flop-counting branch from 72b69ac to f6d34b8 Compare November 14, 2025 19:08

majosm mentioned this pull request Nov 14, 2025

Tweak MPMS materialization to handle some edge cases #624

Open

majosm force-pushed the flop-counting branch 2 times, most recently from 789e7ac to f077dd1 Compare November 21, 2025 23:11

majosm marked this pull request as ready for review November 21, 2025 23:27

majosm force-pushed the flop-counting branch from f077dd1 to 61aed37 Compare November 25, 2025 14:57

inducer commented Dec 4, 2025

View reviewed changes

majosm force-pushed the flop-counting branch from 57906c5 to 8164074 Compare December 9, 2025 20:38

majosm and others added 8 commits December 15, 2025 10:54

remove Any from WalkMapper

cccc1d0

add ConditionalDependencyMapper

51f8bcf

add flop counting functions

80d97fc

Update baseline

8cd10b1

add note to docs about assumptions when handling conditional expressions

96c574d

change 'pass' -> '...' in ReductionOperation abstract methods

c17fe0c

move op_name_to_num_flops type declaration to FlopCounter class body

d2bf20d

add some details about how flop counts are computed

9284664

majosm added 5 commits December 15, 2025 10:54

clarify how flop counting functions behave w.r.t. DAG functions

e1c0aa6

don't try to count flops for function calls

f142ff1

move own flop counting out of rec and into its own method in _PerEntr…

af81492

…yFlopCounter

change is_materializable -> is_materialized / has_taggable_materializ…

2e9fedd

…ation

use explicit isinstance() check instead of try/except around to_index…

75e5dbe

…_lambda()

majosm force-pushed the flop-counting branch from 8164074 to 75e5dbe Compare December 15, 2025 16:55

remove a couple of FIXMEs

1d73566

majosm mentioned this pull request Dec 15, 2025

Add ArrayOrNamesOrFunctionDef #628

Draft

1 task

		return frozenset(expr.bindings.values())


		class _CallMaterializer(CopyMapper):

		Returns a dictionary mapping materialized nodes in DAG expr to their floating
		point operation count.

Flop counting #623

Are you sure you want to change the base?

Flop counting #623

Uh oh!

Conversation

inducer commented Oct 20, 2025

Uh oh!

majosm commented Oct 27, 2025

Uh oh!

majosm commented Nov 7, 2025

Uh oh!

majosm commented Nov 14, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

majosm Dec 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

majosm Dec 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

majosm Dec 9, 2025 •

edited

Loading

majosm Dec 5, 2025 •

edited

Loading