[frontend] Adapter spec proposal that is backwards compatible

Author: gilbertmike

accelforge/frontend/workload.py

@@ -826,6 +826,23 @@ def _eval_expressions(self, symbol_table: dict[str, Any], *args, **kwargs):
         return evaluated, symbol_table
 
 
+class Adapter(EvalableModel):
+    """
+    An adapter between or before (in the case of shared inputs) two or more
+    Einsums.
+    """
+    name: EvalsTo[str]
+
+    adapter_class: EvalsTo[str]
+    """
+    The class of the adapter. Supported classes are:
+    - copy
+    - collective
+    """
+
+    tensor: EvalsTo[str]
+
+
 class Workload(EvalableModel):
     """
     The workload specification as a cascade of Einsums, with each Einsum being a
@@ -874,6 +891,11 @@ class Workload(EvalableModel):
     matching tensors as persistent. Example: "weight" or "~(Outputs | Intermediates)".
     """
 
+    adapters: EvalableList[Adapter] = EvalableList()
+    """
+    Adapters for shared tensors.
+    """
+
     def _for_einsum(self, einsum_name: EinsumName) -> "Workload":
         """Return a copy of the workload with only the Einsum with the given name."""
         new = self.model_copy(deep=False)
@@ -1312,3 +1334,9 @@ 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_adapter_for(self, tensor_name: TensorName) -> Adapter:
+        for adapter in self.adapters:
+            if adapter.tensor == tensor_name:
+                return adapter
+        return ValueError(f"No adapter found for tensor {tensor_name}")