add placeholder class for operator flop counts that aren't specified

Author: majosm

pytato/analysis/__init__.py

@@ -806,13 +806,16 @@ def _get_own_flop_count(self, expr: Array) -> int:
             return 0
         nflops = self.scalar_flop_counter(to_index_lambda(expr).expr)
         if not isinstance(nflops, int):
-            from pytato.scalar_expr import InputGatherer as ScalarInputGatherer
-            var_names: set[str] = set(ScalarInputGatherer()(nflops))
-            var_names.discard("nflops")
-            if var_names:
-                raise UndefinedOpFlopCountError(next(iter(var_names))) from None
+            # Restricting to numerical result here because the flop counters that use
+            # this mapper subsequently multiply the result by things that are
+            # potentially arrays (e.g., shape components), and arrays and scalar
+            # expressions are not interoperable
+            from pytato.scalar_expr import OpFlops, OpFlopsCollector
+            op_flops: frozenset[OpFlops] = OpFlopsCollector()(nflops)
+            if op_flops:
+                raise UndefinedOpFlopCountError(next(iter(op_flops)).op)
             else:
-                raise AssertionError from None
+                raise AssertionError
         return nflops
 
     @override

pytato/scalar_expr.py

@@ -44,6 +44,7 @@
 
 import re
 from collections.abc import Iterable, Mapping, Set as AbstractSet
+from functools import reduce
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -266,8 +267,7 @@ def _get_op_nflops(self, name: str) -> ArithmeticExpression:
         try:
             return self.op_name_to_num_flops[name]
         except KeyError:
-            from pymbolic import var
-            result = var("nflops")(var(name))
+            result = OpFlops(name)
             self.op_name_to_num_flops[name] = result
             return result
 
@@ -480,9 +480,42 @@ class TypeCast(ExpressionBase):
     dtype: np.dtype[Any]
     inner_expr: ScalarExpression
 
+
+@expr_dataclass()
+class OpFlops(prim.AlgebraicLeaf):
+    """
+    Placeholder flop count for an operator.
+
+    .. autoattribute:: op
+    """
+    op: str
+
 # }}}
 
 
+class OpFlopsCollector(CombineMapper[frozenset[OpFlops], []]):
+    """
+    Constructs a :class:`frozenset` containing all instances of
+    :class:`pytato.scalar_expr.OpFlops` found in a scalar expression.
+    """
+    @override
+    def combine(
+            self, values: Iterable[frozenset[OpFlops]]) -> frozenset[OpFlops]:
+        return reduce(
+            lambda x, y: x.union(y),
+            values,
+            cast("frozenset[OpFlops]", frozenset()))
+
+    @override
+    def map_algebraic_leaf(
+            self, expr: prim.AlgebraicLeaf) -> frozenset[OpFlops]:
+        return frozenset([expr]) if isinstance(expr, OpFlops) else frozenset()
+
+    @override
+    def map_constant(self, expr: object) -> frozenset[OpFlops]:
+        return frozenset()
+
+
 class InductionVariableCollector(CombineMapper[AbstractSet[str], []]):
     def combine(self, values: Iterable[AbstractSet[str]]) -> frozenset[str]:
         from functools import reduce