wip

Author: gilbertmike

accelforge/frontend/renames.py

@@ -50,7 +50,8 @@ def _eval_expressions(self, symbol_table: dict[str, Any], *args, **kwargs):
             )
             raise EvaluationError(
                 f"Expected count is {evaluated.expected_count}, but got "
-                f"{len(evaluated.source)}: {evaluated.source}",
+                f"{len(evaluated.source)}: {evaluated.source}\n"
+                f"Symbol table: {symbol_table}",
                 source_field="source",
             )
         return evaluated, symbol_table

accelforge/frontend/workload.py

@@ -973,7 +973,7 @@ def _canonicalize_einsums_and_adapters(self, data):
             )
 
         def is_einsum(d):
-            return not hasattr(d, "tag") or d.tag == "!Einsum"
+            return _get_tag(d, default="Einsum") == "Einsum"
 
         if has_einsums:
             data["einsums_and_adapters"] = data["einsums"]
@@ -1362,3 +1362,79 @@ 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_adapted_workload(self) -> "Workload":
+        """
+        Return a (deep) copy of the workload that has tensor names mangled
+        via the adapter.
+
+        In general, a tensor will be renamed <name of adapter>__<tensor name>.
+
+        Each adapter is turned into a copy Einsum that reads the original tensor
+        and writes the mangled tensor. Every Einsum that comes after the adapter
+        and accesses the original tensor is rewired to use the mangled name, so
+        the copy sits logically between the tensor's producer (or workload input)
+        and its downstream consumers.
+
+        Returns
+        -------
+        Workload
+            A new workload with adapters lowered into copy Einsums and the
+            affected tensor names mangled.
+        """
+
+        def _projection_for(tensor: TensorName) -> dict[Rank, str]:
+            """Find how `tensor` is accessed so the copy Einsum can mirror it."""
+            for item in self.einsums_and_adapters:
+                if not isinstance(item, Einsum):
+                    continue
+                for ta in item.tensor_accesses:
+                    if ta.name == tensor:
+                        return dict(ta.projection)
+            raise ValueError(
+                f"Adapter references tensor {tensor}, but no Einsum accesses it."
+            )
+
+        new_einsums: list[Einsum] = []
+        # Maps an original tensor name to its mangled name, for every adapter seen
+        # so far. Einsums after an adapter use the mangled name.
+        mangled: dict[TensorName, TensorName] = {}
+
+        for item in self.einsums_and_adapters:
+            if isinstance(item, Einsum):
+                einsum = item.model_copy(deep=True)
+                for ta in einsum.tensor_accesses:
+                    if ta.name in mangled:
+                        ta.name = mangled[ta.name]
+                new_einsums.append(einsum)
+            elif isinstance(item, CopyAdapter):
+                tensor = TensorName(item.tensor)
+                new_name = f"{item.name}__{tensor}"
+                # The copy reads whatever name the tensor currently has (it may
+                # itself have been mangled by an earlier adapter) and writes the
+                # newly mangled name.
+                source = mangled.get(tensor, tensor)
+                projection = _projection_for(tensor)
+                copy_einsum = Einsum(
+                    name=item.name,
+                    tensor_accesses=[
+                        {"name": source, "projection": projection, "output": False},
+                        {"name": new_name, "projection": projection, "output": True},
+                    ],
+                    is_copy_operation=True,
+                )
+                new_einsums.append(copy_einsum)
+                mangled[tensor] = new_name
+            else:
+                raise ValueError(
+                    f"Unsupported adapter type {type(item).__name__} in workload."
+                )
+
+        return Workload(
+            einsums=new_einsums,
+            iteration_space_shape=self.iteration_space_shape,
+            rank_sizes=self.rank_sizes,
+            n_instances=self.n_instances,
+            bits_per_value=self.bits_per_value,
+            persistent_tensors=self.persistent_tensors,
+        )

accelforge/mapper/FFM/_make_adapters/copy.py

@@ -0,0 +1,24 @@
+from accelforge import Spec
+from accelforge.frontend.mapper.metrics import Metrics
+from accelforge.mapper.FFM.pmappings import MultiEinsumPmappings
+from accelforge.mapper.FFM.mappings import Mappings
+import accelforge.mapper.FFM._make_pmappings.make_pmappings as pmapper
+from accelforge.frontend.workload import EinsumName
+from accelforge.util._frozenset import oset
+
+
+def make_copy_adapter(spec: Spec) -> MultiEinsumPmappings:
+    """
+    Return a MultiEinsumPmappings that simply allows two pmappings to be
+    compatible iff they are already compatible.
+    """
+    return MultiEinsumPmappings(
+        spec=spec,
+        einsum2pmappings={},
+        pmapping_objects={},
+        einsum2jobs={},
+        can_combine_multiple_runs=True,
+        einsums_with_pmappings_generated=oset(),
+        flattened_arches={},
+        evaluated_specs={},
+    )

accelforge/preprocessor/identity.py

@@ -0,0 +1,47 @@
+from accelforge.frontend import Spec, EinsumName
+
+
+def add_adapter_between_einsums(
+    spec: Spec, einsum_a: EinsumName, einsum_b: EinsumName, adapter
+):
+    """Insert an adapter between Einsums in a spec, modifying the intermediate tensor names."""
+    workload = spec.workload
+    e_a = workload.einsums[einsum_a]
+    e_b = workload.einsums[einsum_b]
+
+    a2b = e_a.output_tensor_names & e_b.input_tensor_names
+    b2a = e_b.output_tensor_names & e_a.input_tensor_names
+
+    if a2b and not b2a:
+        producer, consumer = e_a, e_b
+        intermediate = next(iter(a2b))
+    elif b2a and not a2b:
+        producer, consumer = e_b, e_a
+        intermediate = next(iter(b2a))
+    else:
+        raise ValueError(
+            f"Cannot insert adapter: expected exactly one intermediate tensor "
+            f"flowing between {einsum_a} and {einsum_b}, found a->b={a2b}, b->a={b2a}"
+        )
+
+    new_name = f"{intermediate}_{adapter.name}"
+
+    for ta in consumer.tensor_accesses:
+        if ta.name == intermediate:
+            ta.name = new_name
+
+    adapter_inputs = [ta for ta in adapter.tensor_accesses if not ta.output]
+    adapter_outputs = [ta for ta in adapter.tensor_accesses if ta.output]
+    if len(adapter_inputs) != 1 or len(adapter_outputs) != 1:
+        raise ValueError(
+            f"Adapter {adapter.name} must have exactly one input and one output "
+            f"tensor access, found {len(adapter_inputs)} inputs and "
+            f"{len(adapter_outputs)} outputs"
+        )
+    adapter_inputs[0].name = intermediate
+    adapter_outputs[0].name = new_name
+
+    producer_idx = next(
+        i for i, e in enumerate(workload.einsums) if e.name == producer.name
+    )
+    workload.einsums.insert(producer_idx + 1, adapter)

tests/input_files/adapters/gpt3_6.7B.yaml

@@ -35,9 +35,12 @@ workload:
   - !Copy
     name: copy_I
     tensor: I
-  - "V[b, m, h, e] = I[b, m, d] * WV[h, e, d]"
-  - "K[b, m, h, e] = I[b, m, d] * WK[h, e, d]"
-  - "Q[b, m, h, e] = I[b, m, d] * WQ[h, e, d]"
+  - einsum: "V[b, m, h, e] = I[b, m, d] * WV[h, e, d]"
+    renames: {input: I}
+  - einsum: "K[b, m, h, e] = I[b, m, d] * WK[h, e, d]"
+    renames: {input: I}
+  - einsum: "Q[b, m, h, e] = I[b, m, d] * WQ[h, e, d]"
+    renames: {input: I}
 
   - einsum: "QK[b, m, p, h] = Q[b, m, h, e] * K[b, M: p, h, e]"
     renames: {input: Q}

tests/test_adapter.py

@@ -10,5 +10,63 @@ class TestParsing(unittest.TestCase):
     def test_gpt3(self):
         spec = af.Spec.from_yaml(INPUT_FILES_DIR / "gpt3_6.7B.yaml")
         self.assertEqual(
-            spec.workload.einsum_names, ["V", "K", "Q", "QK_softmax", "Z", "FFA", "FFB"]
+            spec.workload.einsum_names,
+            ["V", "K", "Q", "QK", "QK_softmax", "AV", "Z", "FFA", "FFB"],
         )
+
+
+class TestMangling(unittest.TestCase):
+    def setUp(self):
+        self.spec = af.Spec.from_yaml(INPUT_FILES_DIR / "gpt3_6.7B.yaml")
+        self.workload = self.spec.workload
+        self.adapted = self.workload.get_adapted_workload()
+
+    def test_gpt3(self):
+        self.assertIn("copy_I__I", self.adapted.einsums["Q"].input_tensor_names)
+
+    def test_all_consumers_of_adapted_tensor_are_mangled(self):
+        # Every Einsum that read the original input I should now read the mangled
+        # name instead, and should no longer reference I directly.
+        for name in ["V", "K", "Q"]:
+            inputs = self.adapted.einsums[name].input_tensor_names
+            self.assertIn("copy_I__I", inputs)
+            self.assertNotIn("I", inputs)
+
+    def test_copy_einsum_is_inserted(self):
+        # The adapter is lowered into a copy Einsum named after the adapter.
+        self.assertIn("copy_I", self.adapted.einsum_names)
+        copy_einsum = self.adapted.einsums["copy_I"]
+        self.assertTrue(copy_einsum.is_copy_operation)
+        self.assertEqual(copy_einsum.input_tensor_names, {"I"})
+        self.assertEqual(copy_einsum.output_tensor_names, {"copy_I__I"})
+
+    def test_copy_einsum_mirrors_original_projection(self):
+        # The copy reads/writes the same ranks the original tensor was accessed by.
+        copy_einsum = self.adapted.einsums["copy_I"]
+        src = next(t for t in copy_einsum.tensor_accesses if t.name == "I")
+        dst = next(t for t in copy_einsum.tensor_accesses if t.name == "copy_I__I")
+        self.assertEqual(set(src.ranks), {"B", "M", "D"})
+        self.assertEqual(set(dst.ranks), {"B", "M", "D"})
+
+    def test_original_tensor_only_remains_on_copy(self):
+        # After adapting, the original I is produced/consumed only by the copy
+        # Einsum; downstream Einsums use the mangled name.
+        einsums_with_I = [e.name for e in self.adapted.einsums_with_tensor("I")]
+        self.assertEqual(einsums_with_I, ["copy_I"])
+
+    def test_downstream_einsums_unaffected(self):
+        # Tensors unrelated to the adapter keep their names.
+        qk_inputs = self.adapted.einsums["QK"].input_tensor_names
+        self.assertEqual(qk_inputs, {"Q", "K"})
+
+    def test_einsum_order_preserved(self):
+        self.assertEqual(
+            self.adapted.einsum_names,
+            ["copy_I", "V", "K", "Q", "QK", "QK_softmax", "AV", "Z", "FFA", "FFB"],
+        )
+
+    def test_original_workload_unchanged(self):
+        # get_adapted_workload returns a copy; the source workload is untouched.
+        self.assertIn("I", self.workload.einsums["Q"].input_tensor_names)
+        self.assertNotIn("copy_I__I", self.workload.einsums["Q"].input_tensor_names)
+        self.assertNotIn("copy_I", self.workload.einsum_names)