add note to docs about assumptions when handling conditional expressions

Author: majosm

pytato/analysis/__init__.py

@@ -1016,11 +1016,20 @@ def get_default_op_name_to_num_flops() -> dict[str, int]:
         "max": 1}
 
 
+# FIXME: Should the cost of "If" be the max of the two branches, or the sum?
 def get_num_flops(
         expr: ArrayOrNames,
         op_name_to_num_flops: Mapping[str, int] | None = None,
     ) -> ArrayOrScalar:
-    """Count the total number of floating point operations in the DAG *expr*."""
+    """
+    Count the total number of floating point operations in the DAG *expr*.
+
+    .. note::
+
+        For arrays whose index lambda form contains :class:`pymbolic.primitives.If`,
+        this function assumes a SIMT-like model of computation in which the per-entry
+        cost is the maximum(??? FIXME) of the costs of the two branches.
+    """
     from pytato.codegen import normalize_outputs
     expr = normalize_outputs(expr)
     expr = _normalize_materialization(expr)
@@ -1036,13 +1045,20 @@ def get_num_flops(
         + sum(fc.call_to_nflops.values()))
 
 
+# FIXME: Should the cost of "If" be the max of the two branches, or the sum?
 def get_materialized_node_flop_counts(
         expr: ArrayOrNames,
         op_name_to_num_flops: Mapping[str, int] | None = None,
     ) -> dict[Array, ArrayOrScalar]:
     """
     Returns a dictionary mapping materialized nodes in DAG *expr* to their floating
     point operation count.
+
+    .. note::
+
+        For arrays whose index lambda form contains :class:`pymbolic.primitives.If`,
+        this function assumes a SIMT-like model of computation in which the per-entry
+        cost is the maximum(??? FIXME) of the costs of the two branches.
     """
     from pytato.codegen import normalize_outputs
     expr = normalize_outputs(expr)
@@ -1063,6 +1079,7 @@ class UnmaterializedNodeFlopCounts:
     nflops_if_materialized: ArrayOrScalar
 
 
+# FIXME: Should the cost of "If" be the max of the two branches, or the sum?
 def get_unmaterialized_node_flop_counts(
         expr: ArrayOrNames,
         op_name_to_num_flops: Mapping[str, int] | None = None,
@@ -1071,6 +1088,12 @@ def get_unmaterialized_node_flop_counts(
     Returns a dictionary mapping unmaterialized nodes in DAG *expr* to a
     :class:`UnmaterializedNodeFlopCounts` containing floating-point operation count
     information.
+
+    .. note::
+
+        For arrays whose index lambda form contains :class:`pymbolic.primitives.If`,
+        this function assumes a SIMT-like model of computation in which the per-entry
+        cost is the maximum(??? FIXME) of the costs of the two branches.
     """
     from pytato.codegen import normalize_outputs
     expr = normalize_outputs(expr)