Merge branch 'main' into pipeline

Author: gilbertmike

Makefile

@@ -2,7 +2,7 @@ DOCKER_EXE ?= docker
 DOCKER_NAME ?= accelforge
 DOCKER_BUILD ?= ${DOCKER_EXE} buildx build --load --pull
 
-VERSION := 0.1.5
+VERSION := 0.1.6
 
 USER    := timeloopaccelergy
 REPO    := accelforge

accelforge/_accelerated_imports.py

@@ -5,6 +5,7 @@
 if os.environ.get("ACCELFORGE_ACCELERATED_IMPORTS", "0") == "1":
     import cudf as pd
     import cupy as np
+
     pandas = pd
     numpy = np
     # import cupy as scipy
@@ -13,6 +14,7 @@
 else:
     import pandas as pd
     import numpy as np
+
     pandas = pd
     numpy = np
     # import scipy

accelforge/frontend/arch/components.py

@@ -760,6 +760,48 @@ def _eval_tensor2bits(
     return {k2: v for k, v in result.items() for k2 in k}
 
 
+_VALID_DIRECTIONS = {"up", "down", "up_and_down"}
+
+
+def _eval_direction(toeval, symbol_table: dict[str, Any]) -> dict[str, str]:
+    """Evaluate a direction field. If a string, expand to all tensors. If a dict,
+    resolve tensor expression keys."""
+    if isinstance(toeval, str):
+        if toeval not in _VALID_DIRECTIONS:
+            raise EvaluationError(
+                f'Invalid direction: "{toeval}". '
+                f"Must be one of {sorted(_VALID_DIRECTIONS)}."
+            )
+        all_tensors = symbol_table["All"].instance
+        return {t: toeval for t in all_tensors}
+
+    result = {}
+    for key, value in toeval.items():
+        if value not in _VALID_DIRECTIONS:
+            raise EvaluationError(
+                f'Invalid direction for {key}: "{value}". '
+                f"Must be one of {sorted(_VALID_DIRECTIONS)}."
+            )
+        key_evaluated = eval_set_expression(
+            expression=key,
+            symbol_table=symbol_table,
+            expected_space=TensorName,
+            location=f"direction key {key}",
+        ).instance
+        result[key_evaluated] = value
+
+    all_tensors = symbol_table["All"].instance
+    for k in result:
+        all_tensors -= k
+
+    if all_tensors:
+        raise EvaluationError(
+            f"Missing direction for {all_tensors}. Have {result}."
+        )
+
+    return {t: v for k, v in result.items() for t in k}
+
+
 class Tensors(EvalableModel):
     """
     Fields that control which tensor(s) are kept in a :py:class:`~.TensorHolder` and in
@@ -955,6 +997,12 @@ class Memory(TensorHolder, ConcurrentlyBoundable):
     physical units may be flattened into only one logical level.
     """
 
+    skip_initial_output_write: bool = True
+    """
+    If False, the initial value of output tensors will be fetched from above and used to
+    initalize outputs. If True, this initial fetch and fill is skipped.
+    """
+
     def _render_node_shape(self) -> str:
         return "cylinder"
 
@@ -986,10 +1034,15 @@ class Toll(TensorHolder):
     zero.
     """
 
-    direction: Literal["up", "down", "up_and_down"]
+    direction: TryEvalTo[dict]
     """
-    The direction in which data flows through this `Toll`. If "up", then data
-    flows from below `TensorHolder`, through this `Toll` (plus paying
+    The direction in which data flows through this `Toll`. Can be:
+
+    - A string: ``"up"``, ``"down"``, or ``"up_and_down"`` — applies to all tensors.
+    - A dict mapping tensor expressions to direction strings, e.g.
+      ``{input: "down", output: "up"}`` — sets direction per tensor.
+
+    If "up", then data flows from below `TensorHolder`, through this `Toll` (plus paying
     associated costs), and then to the next `TensorHolder` above it. Other data
     movements are assumed to avoid this Toll.
     """
@@ -1000,6 +1053,32 @@ class Toll(TensorHolder):
     def model_post_init(self, __context__=None) -> None:
         self._update_actions(PROCESSING_STAGE_ACTIONS)
 
+    def _eval_expressions(self, *args, **kwargs):
+        if getattr(self, "_evaluated", False):
+            return super()._eval_expressions(*args, **kwargs)
+
+        # Override TensorHolder's _PostCall to also handle direction
+        class MyPostCall(_PostCall):
+            def __call__(self_pc, field, value, evaluated, symbol_table):
+                if field == "bits_per_value_scale":
+                    evaluated = _eval_tensor2bits(
+                        evaluated,
+                        location="bits_per_value_scale",
+                        symbol_table=symbol_table,
+                    )
+                if field == "direction":
+                    evaluated = _eval_direction(
+                        evaluated,
+                        symbol_table=symbol_table,
+                    )
+                return evaluated
+
+        # Skip TensorHolder's _eval_expressions (which adds its own post_calls
+        # for bits_per_value_scale) since we handle it here too
+        return Component._eval_expressions(
+            self, *args, **kwargs, post_calls=(MyPostCall(),)
+        )
+
     def _render_node_shape(self) -> str:
         return "rarrow"
 
@@ -1011,6 +1090,13 @@ class Compute(Component, Leaf, ConcurrentlyBoundable):
     actions: EvalableList[Action] = COMPUTE_ACTIONS
     """ The actions that this `Compute` can perform. """
 
+    skip_initial_output_write: bool = True
+    """
+    If False, the initial value of output tensors will be fetched from above and used to
+    initalize outputs. If True, this initial fetch and fill is skipped.
+    """
+
+
     def model_post_init(self, __context__=None) -> None:
         self._update_actions(COMPUTE_ACTIONS)
 

accelforge/frontend/arch/spatialable.py

@@ -36,11 +36,11 @@ class Spatial(EvalableModel):
     Note: Loops may be removed if they are constrained to only one iteration.
     """
 
-    min_usage: int | float | str = 0.0
+    min_usage: EvalsTo[int | float] = 0.0
     """ The minimum usage of spatial instances, as a value from 0 to 1. A mapping
     is invalid if less than this porportion of this dimension's fanout is utilized.
     Mappers that support it (e.g., FFM) may, if no mappings satisfy this constraint,
-    return the highest-usage mappings.
+    return the highest-usage mappings. These constraints are disabled for copy Einsums.
     """
 
     reuse: TryEvalTo[InvertibleSet[TensorName]] = "Nothing"