Merge main

accelforge/frontend/mapping/__init__.py

@@ -18,6 +18,7 @@
     "Temporal",
     "TensorHolder",
     "TensorName",
+    "TextBox",
     "TilePattern",
     "Toll",
 ]

accelforge/frontend/spec.py

@@ -157,10 +157,16 @@ def _spec_eval_expressions(
             st["workload"] = evaluated_workload
 
             if einsum_name is not None:
-                renames = evaluated_workload.einsums[einsum_name].renames
-                st.update(**{k.name: k.source for k in renames})
+                einsum = evaluated_workload.einsums[einsum_name]
+                st.update(**{k.name: k.source for k in einsum.renames})
+                n_computes = evaluated_workload.n_computes(einsum_name)
+                n_outputs = min(
+                    evaluated_workload.get_tensor_size(t) for t in einsum.output_tensor_names
+                )
+                st["einsum_has_reduction"] = n_computes > n_outputs
             else:
                 st.update(evaluated_workload.empty_renames())
+                st["einsum_has_reduction"] = True
 
             if eval_arch:
                 evaluated_arch, st = self.arch._eval_expressions(st)

accelforge/frontend/workload.py

@@ -128,7 +128,7 @@ class TensorAccess(EvalableModel):
     backing_storage_size_scale: float = 1.0
     """ If != 1, then the backing storage size will be scaled by this factor. """
 
-    bits_per_value: int | str | None = None
+    bits_per_value: EvalsTo[int | None] = None
     """ Bits per value for this tensor. """
 
     def model_post_init(self, __context__=None) -> None:
@@ -678,6 +678,7 @@ def _eval_expressions(self, symbol_table: dict[str, Any], *args, **kwargs):
         outputs = self.output_tensor_names
         all_ = inputs | outputs
         persistent = oset(t.name for t in self.tensor_accesses if t.persistent)
+        element_bits = {}
         element_to_child_space = {}
         all_rank_variables = self.rank_variables
         for tensor in self.tensor_names:
@@ -707,6 +708,7 @@ def _eval_expressions(self, symbol_table: dict[str, Any], *args, **kwargs):
             space_type=TensorName,
             child_access_name="rank_variables",
             element_to_child_space=element_to_child_space,
+            element_to_bits_per_value=element_bits,
         )
         kwargs_rank_variables = dict(
             full_space=all_rank_variables,
@@ -809,7 +811,20 @@ def _eval_expressions(self, symbol_table: dict[str, Any], *args, **kwargs):
                     source_field=f"tensor_accesses[{t.name}].bits_per_value",
                 )
             if t.bits_per_value is None:
-                t.bits_per_value = bits_per_value[t.name]
+                t.bits_per_value = eval_expression(
+                    bits_per_value[t.name],
+                    st,
+                    attr_name=f"bits_per_value[{t.name}]",
+                )
+
+        element_bits.update(
+            {t.name: t.bits_per_value for t in evaluated.tensor_accesses}
+        )
+        for r in evaluated.renames:
+            if isinstance(r.source, InvertibleSet) and all(
+                t in element_bits for t in r.source.instance
+            ):
+                r.source.element_to_bits_per_value = element_bits
 
         if symbol_table.get("workload_persistent_tensors", None):
             rename_st_with_evaluated = {**st}
@@ -1336,3 +1351,22 @@ def get_compute_intensity(self, einsum_name: str) -> float:
             self.get_tensor_size(tensor)
             for tensor in self.einsums[einsum_name].tensor_names
         )
+
+    def get_per_tensor_compute_intensity(self) -> dict[TensorName, float]:
+        """
+        Returns the compute intensity of each tensor, defined as the sum of the number
+        of computes of each Einsum that accesses the tensor, divided by the number of
+        elements in the tensor.
+
+        Returns
+        -------
+        dict[TensorName, float]
+            The compute intensity of each tensor in #computes / #tensor elements.
+        """
+        return {
+            tensor: sum(
+                self.n_computes(e.name) for e in self.einsums_with_tensor(tensor)
+            )
+            / self.get_tensor_size(tensor)
+            for tensor in self.tensor_names
+        }

accelforge/util/_setexpressions.py

@@ -26,6 +26,7 @@ class InvertibleSet(BaseModel, Generic[T]):
     space_type: type[T]
     # child_access_name: Optional[str] = None
     element_to_child_space: Optional[dict[str, Any]] = None
+    element_to_bits_per_value: Optional[dict[str, int]] = None
     _bits_per_value: Optional[int] = None
 
     def __init__(self, *args, **kwargs):
@@ -118,6 +119,7 @@ def to_my_space(self, other) -> Union[set, "InvertibleSet"]:
             space_type=self.space_type,
             # child_access_name=self.child_access_name,
             element_to_child_space=self.element_to_child_space,
+            element_to_bits_per_value=self.element_to_bits_per_value,
         )
 
     @staticmethod
@@ -190,6 +192,7 @@ def iter_one_element_sets(self) -> Iterator["InvertibleSet[T]"]:
                 space_type=self.space_type,
                 # child_access_name=self.child_access_name,
                 element_to_child_space=self.element_to_child_space,
+                element_to_bits_per_value=self.element_to_bits_per_value,
             )
 
     @property
@@ -294,7 +297,7 @@ def eval_set_expression_dict(
     symbol_table: dict[str, InvertibleSet],
     expected_space: type[T],
     location: str,
-    disjoint: bool=True,
+    disjoint: bool = True,
 ) -> list[tuple[str, "frozenset[T]", Any]]:
     """
     Evaluate a dict whose keys are set expressions, returning an ordered list of
@@ -310,7 +313,7 @@ def eval_set_expression_dict(
         )
 
     evaluated: list[tuple[str, Any, Any]] = []
-    
+
     symbol_table = symbol_table.copy()
     symbol_table["Other"] = symbol_table["All"]
 
@@ -324,7 +327,7 @@ def _eval(i):
         ).instance
         symbol_table["Other"] -= ins
         return k, ins, v
-    
+
     eval_order = [i for i in range(len(items)) if i not in others] + others
     for i in eval_order:
         evaluated.append(_eval(i))

accelforge/util/_yaml.py

@@ -8,6 +8,7 @@
 from typing import Callable, List, Dict, Any, OrderedDict, Tuple
 import ruamel.yaml
 import warnings
+from ruamel.yaml.comments import CommentedMap, CommentedSeq
 from ruamel.yaml.error import ReusedAnchorWarning
 from jinja2 import StrictUndefined, Environment, FileSystemLoader, pass_context, nodes
 from jinja2.ext import Extension
@@ -407,8 +408,17 @@ def ordereddict_to_dict(self, dictionary: OrderedDict) -> Dict[str, Any]:
     return self.represent_dict(dictionary)
 
 
+def _update_in_place(values, f):
+    items = values.items() if isinstance(values, dict) else enumerate(values)
+    for k, v in list(items):
+        values[k] = f(v)
+    return values
+
+
 @recursive_mutator_stop
 def recursive_unorder_dict(to_unorder: Dict[str, Any]) -> Dict[str, Any]:
+    if isinstance(to_unorder, (CommentedMap, CommentedSeq)):
+        return _update_in_place(to_unorder, recursive_unorder_dict)
     if isinstance(to_unorder, dict):
         return {k: recursive_unorder_dict(v) for k, v in to_unorder.items()}
     elif isinstance(to_unorder, list):
@@ -418,7 +428,9 @@ def recursive_unorder_dict(to_unorder: Dict[str, Any]) -> Dict[str, Any]:
 
 @recursive_mutator_stop
 def callables2strings(to_convert: Dict[str, Any]) -> Dict[str, Any]:
-    if isinstance(to_convert, dict):
+    if isinstance(to_convert, (CommentedMap, CommentedSeq)):
+        to_convert = _update_in_place(to_convert, callables2strings)
+    elif isinstance(to_convert, dict):
         to_convert = {k: callables2strings(v) for k, v in to_convert.items()}
     elif isinstance(to_convert, list):
         to_convert = [callables2strings(v) for v in to_convert]
@@ -427,6 +439,36 @@ def callables2strings(to_convert: Dict[str, Any]) -> Dict[str, Any]:
     return to_convert
 
 
+def _flow_length(obj) -> float:
+    """
+    Length of obj rendered in flow style on one line, or inf if it can't be (multi-line
+    strings, comments).
+    """
+    ca = getattr(obj, "ca", None)
+    if ca is not None and (ca.comment or ca.items):
+        return float("inf")
+    if isinstance(obj, dict):
+        return 2 + sum(_flow_length(k) + _flow_length(v) + 4 for k, v in obj.items())
+    if isinstance(obj, list):
+        return 2 + sum(_flow_length(v) + 2 for v in obj)
+    s = str(obj)
+    return float("inf") if "\n" in s else len(s)
+
+
+@recursive_mutator_stop
+def compact_flow(obj):
+    """Set flow style on dicts and lists that fit on one line."""
+    if not isinstance(obj, (dict, list)):
+        return obj
+    _update_in_place(obj, compact_flow)
+    if _flow_length(obj) > 100:
+        return obj
+    if not isinstance(obj, (CommentedMap, CommentedSeq)):
+        obj = CommentedMap(obj) if isinstance(obj, dict) else CommentedSeq(obj)
+    obj.fa.set_flow_style()
+    return obj
+
+
 def write_yaml_file(filepath: str, content: Dict[str, Any]) -> None:
     """
     Write YAML content to a file
@@ -451,15 +493,17 @@ def to_yaml_string(content: Dict[str, Any]) -> str:
     with LockAcquirer():
         dumpstream = io.StringIO()
         get_base_yaml().dump(
-            callables2strings(recursive_unorder_dict(content)), stream=dumpstream
+            compact_flow(callables2strings(recursive_unorder_dict(content))),
+            stream=dumpstream,
         )
         return dumpstream.getvalue()
 
 
 def get_base_yaml() -> ruamel.yaml.YAML:
     yaml = ruamel.yaml.YAML(typ="rt")
     # yaml.default_flow_style = None
-    yaml.indent(mapping=4, sequence=4, offset=2)
+    yaml.indent(mapping=2, sequence=2, offset=0)
+    yaml.width = 120
     yaml.preserve_quotes = True
 
     def recursive_mutator_stop(func):