accelforge 0.0.1__py3-none-any.whl

tests/test_ffm_make_tile_shapes.py

@@ -0,0 +1,49 @@
+import unittest
+from pathlib import Path
+
+import time
+
+from accelforge.frontend.spec import Spec, Mapping
+from accelforge.mapper.FFM._make_pmappings.contraints.constraints import (
+    MappingConstraints,
+)
+from accelforge.mapper.FFM._make_pmappings.pmapper_job import Job
+from accelforge.mapper.FFM._make_pmappings.make_pmappings import make_tile_shapes
+from accelforge.frontend.mapper.metrics import Metrics
+from accelforge.frontend import arch
+
+
+class TestTileShapeExploration(unittest.TestCase):
+    def test_pmapping(self):
+        PARENT_DIR = Path(__file__).parent
+        spec = Spec.from_yaml(
+            PARENT_DIR / "conv.workload.yaml", PARENT_DIR / "four_level.arch.yaml"
+        )
+        spec = spec.calculate_component_area_energy_latency_leak(area=False)
+
+        mapping = Mapping.from_yaml(PARENT_DIR / "conv_sym.mapping.yaml")
+
+        flattened_arch = spec._get_flattened_architecture()
+        memories_track_all = [
+            m.name for m in flattened_arch if isinstance(m, arch.Memory)
+        ]
+        memories_track_pmappings_only = []
+
+        job = Job(
+            mapping=mapping,
+            constraints=MappingConstraints(),
+            spec=spec,
+            metrics=Metrics.LATENCY,
+            job_id=0,
+            rank_variable_bounds={},
+            memories_track_all=memories_track_all,
+            memories_track_pmappings_only=memories_track_pmappings_only,
+        )
+
+        result, _ = make_tile_shapes(job)
+        data, n_total_pmappings = result
+        self.assertTrue("Total_latency" in data.columns)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_mapper.py

@@ -0,0 +1,100 @@
+import unittest
+from pathlib import Path
+
+from accelforge.frontend.spec import Spec
+from accelforge.mapper import Metrics
+from accelforge.mapper.FFM.main import map_workload_to_arch
+
+
+EXAMPLES_DIR = Path(__file__).parent.parent / "examples"
+
+M_SHAPE = 64
+KN_SHAPE = 64
+
+
+class ActionChecker(unittest.TestCase):
+    def _check_memory_actions_exist(self, spec, memory_names, result):
+        for einsum_name in spec.workload.einsum_names:
+            for memory_name in memory_names:
+                for memory_action in ["read", "write"]:
+                    self.assertTrue(
+                        f"{einsum_name}<SEP>action<SEP>{memory_name}<SEP>{memory_action}" in result.data.columns,
+                        f"{einsum_name}<SEP>action<SEP>{memory_name}<SEP>{memory_action} "
+                        f"not found in {result.data.columns}"
+                    )
+
+
+class TestMapper(ActionChecker, unittest.TestCase):
+    def test_one_matmul(self):
+        spec = Spec.from_yaml(
+            EXAMPLES_DIR / "arches" / "simple.yaml",
+            EXAMPLES_DIR / "workloads" / "matmuls.yaml",
+            jinja_parse_data={"N_EINSUMS": 1, "M": 64, "KN": 64},
+        )
+
+        result = map_workload_to_arch(spec)
+        self._check_memory_actions_exist(spec, ["MainMemory", "GlobalBuffer"], result)
+
+    def test_two_matmuls(self):
+        spec = Spec.from_yaml(
+            EXAMPLES_DIR / "arches" / "simple.yaml",
+            EXAMPLES_DIR / "workloads" / "matmuls.yaml",
+            jinja_parse_data={"N_EINSUMS": 2, "M": 64, "KN": 64},
+        )
+
+        result = map_workload_to_arch(spec)
+        self._check_memory_actions_exist(spec, ["MainMemory", "GlobalBuffer"], result)
+
+
+class TestFanout(ActionChecker):
+    FANOUT = 4
+    """Need to sync this with YAMLs somehow."""
+
+    def _run_with_arch(self, arch_fname: str, n_einsums=1):
+        spec = Spec.from_yaml(
+            EXAMPLES_DIR / "arches" / "fanout_variations" / arch_fname,
+            EXAMPLES_DIR / "workloads" / "matmuls.yaml",
+            jinja_parse_data={"N_EINSUMS": n_einsums, "M": 64, "KN": 64},
+        )
+        spec.mapper.ffm.metrics = Metrics.LATENCY
+
+        result = map_workload_to_arch(spec)
+        self._check_memory_actions_exist(spec, ["MainMemory", "GlobalBuffer"], result)
+        self.assertEqual(
+            result.data["Matmul0<SEP>Total<SEP>latency"].iloc[0],
+            M_SHAPE*KN_SHAPE**2/self.FANOUT
+        )
+        return result
+
+
+class TestMapperFanoutOneMatmul(TestFanout):
+    def test_at_mac(self):
+        self._run_with_arch("at_mac.yaml")
+
+    def test_at_glb(self):
+        self._run_with_arch("at_glb.yaml")
+
+    def test_at_mac_with_fanout_node(self):
+        self._run_with_arch("at_mac_with_fanout_node.yaml")
+
+    def test_at_glb_with_fanout_node(self):
+        self._run_with_arch("at_glb_with_fanout_node.yaml")
+
+
+class TestMapperFanoutTwoMatmuls(TestFanout):
+    def test_at_mac(self):
+        self._run_with_arch("at_mac.yaml", n_einsums=2)
+
+    def test_at_glb(self):
+        self._run_with_arch("at_glb.yaml", n_einsums=2)
+
+    def test_at_mac_with_fanout_node(self):
+        self._run_with_arch("at_mac_with_fanout_node.yaml", n_einsums=2)
+
+    def test_at_glb_with_fanout_node(self):
+        self._run_with_arch("at_glb_with_fanout_node.yaml", n_einsums=2)
+
+
+class TestMapperFanoutConstraints(TestFanout):
+    def test_at_mac_constraints(self):
+        self._run_with_arch("at_mac_with_constraints.yaml")

tests/test_model.py

@@ -0,0 +1,37 @@
+import unittest
+from pathlib import Path
+
+from accelforge.frontend.spec import Spec
+from accelforge.model.main import evaluate_mapping
+from accelforge.util.parallel import set_n_parallel_jobs
+
+set_n_parallel_jobs(1)
+
+
+EXAMPLES_DIR = Path(__file__).parent.parent / "examples"
+
+
+class TestModel(unittest.TestCase):
+    def test_one_matmul(self):
+        spec = Spec.from_yaml(
+            EXAMPLES_DIR / "arches" / "simple.yaml",
+            EXAMPLES_DIR / "workloads" / "matmuls.yaml",
+            EXAMPLES_DIR / "mappings" / "unfused_matmuls_to_simple.yaml",
+            jinja_parse_data={"N_EINSUMS": 1, "M": 64, "KN": 64},
+        )
+
+        result = evaluate_mapping(spec)
+
+    def test_two_matmuls(self):
+        spec = Spec.from_yaml(
+            EXAMPLES_DIR / "arches" / "simple.yaml",
+            EXAMPLES_DIR / "workloads" / "matmuls.yaml",
+            EXAMPLES_DIR / "mappings" / "unfused_matmuls_to_simple.yaml",
+            jinja_parse_data={"N_EINSUMS": 2, "M": 64, "KN": 64},
+        )
+
+        result = evaluate_mapping(spec)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_plotting.py

@@ -0,0 +1,72 @@
+import unittest
+from pathlib import Path
+
+from accelforge.frontend.spec import Spec
+from accelforge.model.main import evaluate_mapping
+from accelforge.util.parallel import set_n_parallel_jobs
+from accelforge.plotting.mappings import plot_energy_comparison, plot_energy_breakdown, plot_action_breakdown
+
+set_n_parallel_jobs(1)
+
+
+EXAMPLES_DIR = Path(__file__).parent.parent / "examples"
+
+
+class TestEnergyPlotting(unittest.TestCase):
+    def test_comparison(self):
+        spec = Spec.from_yaml(
+            EXAMPLES_DIR / "arches" / "simple.yaml",
+            EXAMPLES_DIR / "workloads" / "matmuls.yaml",
+            EXAMPLES_DIR / "mappings" / "unfused_matmuls_to_simple.yaml",
+            jinja_parse_data={"N_EINSUMS": 2, "M": 64, "KN": 64},
+        )
+
+        result = evaluate_mapping(spec)
+
+        fig, ax = plot_energy_comparison([result])
+        fig.tight_layout()
+        fig.savefig("total.png", dpi=400, bbox_inches="tight")
+
+    def test_breakdown(self):
+        spec = Spec.from_yaml(
+            EXAMPLES_DIR / "arches" / "simple.yaml",
+            EXAMPLES_DIR / "workloads" / "matmuls.yaml",
+            EXAMPLES_DIR / "mappings" / "unfused_matmuls_to_simple.yaml",
+            jinja_parse_data={"N_EINSUMS": 2, "M": 64, "KN": 64},
+        )
+
+        result = evaluate_mapping(spec)
+
+        fig, axes = plot_energy_breakdown([result, result], ["einsum", "component"])
+        fig.tight_layout()
+        fig.savefig("fig.png", dpi=400, bbox_inches="tight")
+
+    def test_breakdown_stacked(self):
+        spec = Spec.from_yaml(
+            EXAMPLES_DIR / "arches" / "simple.yaml",
+            EXAMPLES_DIR / "workloads" / "matmuls.yaml",
+            EXAMPLES_DIR / "mappings" / "unfused_matmuls_to_simple.yaml",
+            jinja_parse_data={"N_EINSUMS": 2, "M": 64, "KN": 64},
+        )
+
+        result = evaluate_mapping(spec)
+
+        fig, axes = plot_energy_breakdown([result, result], ["einsum", "component"], ["action"])
+        fig.tight_layout()
+        fig.savefig("fig.png", dpi=400, bbox_inches="tight")
+
+
+class TestActionPlotting(unittest.TestCase):
+    def test_breakdown(self):
+        spec = Spec.from_yaml(
+            EXAMPLES_DIR / "arches" / "simple.yaml",
+            EXAMPLES_DIR / "workloads" / "matmuls.yaml",
+            EXAMPLES_DIR / "mappings" / "unfused_matmuls_to_simple.yaml",
+            jinja_parse_data={"N_EINSUMS": 2, "M": 64, "KN": 64},
+        )
+
+        result = evaluate_mapping(spec)
+
+        fig, axes = plot_action_breakdown([result, result], ["einsum", "component"])
+        fig.tight_layout()
+        fig.savefig("fig.png", dpi=400, bbox_inches="tight")

tests/test_processing_stage.py

@@ -0,0 +1,46 @@
+from math import isclose
+import unittest
+from pathlib import Path
+
+from accelforge.frontend import Spec
+from accelforge.frontend.arch import Arch
+from accelforge.frontend.mapping import Mapping
+from accelforge.frontend.workload import Workload
+
+from accelforge.frontend.mapper.metrics import Metrics
+from accelforge.model._looptree.accesses import (
+    isl_buffer_accesses_from_buffet_actions,
+    Accesses,
+)
+from accelforge.model._looptree.energy import gather_actions
+from accelforge.model._looptree.latency import get_latency
+from accelforge.model._looptree.types import Buffet
+from accelforge.model._looptree.reuse.symbolic import (
+    BuffetStats,
+    analyze_reuse_and_add_reservations_to_mapping,
+    Compute,
+)
+from accelforge.mapper.FFM import make_pmappings, join_pmappings
+
+
+PARENT_DIR = Path(__file__).parent
+
+
+class TestProcessingStage(unittest.TestCase):
+    def test_processing_stage(self):
+        spec = Spec.from_yaml(
+            [
+                Path(__file__).parent / "processing_stage.arch.yaml",
+                Path(__file__).parent / "matmul.workload.yaml",
+            ]
+        )
+        spec.mapper.ffm.metrics = Metrics.ENERGY
+        pmappings = make_pmappings(spec)
+        mappings = join_pmappings(spec, pmappings)
+
+        energy = mappings.data.iloc[0]["Total_energy"]
+        self.assertAlmostEqual(energy, 100 * (128 * 64 * 2 + 128 * 128))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_symbolic_model.py

@@ -0,0 +1,248 @@
+from math import isclose
+import unittest
+from pathlib import Path
+
+from accelforge.frontend import Spec
+from accelforge.frontend.mapping import Mapping
+from accelforge.frontend.workload import Workload
+
+from accelforge.mapper.FFM._make_pmappings.pmapper_job import Job
+from accelforge.model._looptree.accesses import (
+    isl_buffer_accesses_from_buffet_actions,
+    Accesses,
+)
+from accelforge.model._looptree.energy import gather_actions
+from accelforge.model._looptree.latency import get_latency
+from accelforge.model._looptree.types import Buffet
+from accelforge.model._looptree.reuse.symbolic import (
+    BuffetStats,
+    analyze_reuse_and_add_reservations_to_mapping,
+    Compute,
+)
+
+
+PARENT_DIR = Path(__file__).parent
+
+
+def make_job(mapping: Mapping, workload: Workload) -> Job:
+    return Job(
+        spec=None,
+        mapping=mapping,
+        workload=workload,
+    )
+
+
+class TestSymbolicModel(unittest.TestCase):
+    def test_q_mapping(self):
+        mapping = Mapping.from_yaml(Path(__file__).parent / "Q_mapping.yaml")
+        workload = Workload.from_yaml(Path(__file__).parent / "mha.yaml")
+
+        result = analyze_reuse_and_add_reservations_to_mapping(
+            make_job(mapping, workload)
+        )
+
+        self.assertAlmostEqual(
+            result.compute_stats[Compute("Q", "MAC")].total_ops, 64.0
+        )
+        self.assertAlmostEqual(
+            result.compute_stats[Compute("Q", "MAC")].max_per_unit_ops, 16.0
+        )
+
+    def test_conv_mapping(self):
+        mapping = Mapping.from_yaml(Path(__file__).parent / "conv.mapping.yaml")
+        workload = Workload.from_yaml(Path(__file__).parent / "conv.workload.yaml")
+
+        result = analyze_reuse_and_add_reservations_to_mapping(
+            make_job(mapping, workload)
+        )
+
+        self.assertAlmostEqual(
+            result.compute_stats[Compute("conv", "MAC")].total_ops, 120.0
+        )
+        self.assertAlmostEqual(
+            result.compute_stats[Compute("conv", "MAC")].max_per_unit_ops, 10.0
+        )
+
+    def test_matmul_mapping(self):
+        mapping = Mapping.from_yaml(Path(__file__).parent / "matmul.mapping.yaml")
+        workload = Workload.from_yaml(Path(__file__).parent / "matmul.workload.yaml")
+
+        result = analyze_reuse_and_add_reservations_to_mapping(
+            make_job(mapping, workload)
+        )
+
+        REF_OCCUPANCY = {"W0": 1, "T0": 128, "T1": 128 * 128}
+        for tensor, ref_occupancy in REF_OCCUPANCY.items():
+            self.assertAlmostEqual(
+                result.buffet_stats[
+                    Buffet(tensor, "Matmul1", "LocalBuffer")
+                ].max_occupancy,
+                ref_occupancy,
+            )
+
+    def test_matmul_spatial(self):
+        mapping = Mapping.from_yaml(PARENT_DIR / "matmul_spatial.mapping.yaml")
+        workload = Workload.from_yaml(PARENT_DIR / "matmul.workload.yaml")
+
+        result = analyze_reuse_and_add_reservations_to_mapping(
+            make_job(mapping, workload)
+        )
+        self.assertAlmostEqual(
+            result.fanout,
+            {
+                ("LocalBuffer", "Matmul1"): {0: 128.0, 1: 4.0},
+                ("MainMemory", "Matmul1"): {},
+            },
+        )
+
+    def test_copy_mapping(self):
+        mapping = Mapping.from_yaml(PARENT_DIR / "copy.mapping.yaml")
+        workload = Workload.from_yaml(PARENT_DIR / "copy.workload.yaml")
+
+        result = analyze_reuse_and_add_reservations_to_mapping(
+            make_job(mapping, workload)
+        )
+
+        self.assertAlmostEqual(
+            result.compute_stats[Compute("copy", "MAC")].total_ops, 0
+        )
+        self.assertAlmostEqual(
+            result.compute_stats[Compute("copy", "MAC")].max_per_unit_ops, 0
+        )
+
+        for tensor in ["O1", "O2", "O3", "O4"]:
+            for memory in [
+                "MainMemory",
+                "GlobalBuffer",
+                "LocalBuffer",
+                "Register",
+                "MAC",
+            ]:
+                buffet = Buffet(level=memory, tensor=tensor, einsum="copy")
+                if buffet not in result.buffet_stats:
+                    continue
+                stats = result.buffet_stats[buffet]
+                self.assertAlmostEqual(stats.net_total_read_actions(), 0)
+                self.assertAlmostEqual(stats.net_max_per_unit_read_actions(), 0)
+                self.assertAlmostEqual(stats.net_total_write_actions(), 0)
+                self.assertAlmostEqual(stats.net_max_per_unit_write_actions(), 0)
+                self.assertAlmostEqual(stats.max_occupancy, 0)
+
+        buffet = Buffet(level="GlobalBuffer", tensor="I", einsum="copy")
+        stats = result.buffet_stats.get(buffet, BuffetStats())
+        self.assertAlmostEqual(stats.net_total_read_actions(), 0)
+        self.assertAlmostEqual(stats.net_max_per_unit_read_actions(), 0)
+        self.assertAlmostEqual(stats.net_total_write_actions(), 0)
+        self.assertAlmostEqual(stats.net_max_per_unit_write_actions(), 0)
+        self.assertAlmostEqual(stats.max_occupancy, 0)
+
+        stats = result.buffet_stats[
+            Buffet(level="LocalBuffer", tensor="I", einsum="copy")
+        ]
+        self.assertAlmostEqual(stats.net_total_read_actions(), 16)
+        self.assertAlmostEqual(stats.net_max_per_unit_read_actions(), 16)
+        self.assertAlmostEqual(stats.net_total_write_actions(), 0)
+        self.assertAlmostEqual(stats.net_max_per_unit_write_actions(), 0)
+        self.assertAlmostEqual(stats.max_occupancy, 2)
+
+        stats = result.buffet_stats[Buffet(level="Register", tensor="I", einsum="copy")]
+        self.assertAlmostEqual(stats.net_total_read_actions(), 0)
+        self.assertAlmostEqual(stats.net_max_per_unit_read_actions(), 0)
+        self.assertAlmostEqual(stats.net_total_write_actions(), 8)
+        self.assertAlmostEqual(stats.net_max_per_unit_write_actions(), 8)
+        self.assertAlmostEqual(stats.max_occupancy, 1)
+
+        stats = result.buffet_stats[
+            Buffet(level="MainMemory", tensor="I", einsum="copy")
+        ]
+        self.assertAlmostEqual(stats.net_total_read_actions(), 8)
+        self.assertAlmostEqual(stats.net_max_per_unit_read_actions(), 8)
+        self.assertAlmostEqual(stats.net_total_write_actions(), 8)
+        self.assertAlmostEqual(stats.net_max_per_unit_write_actions(), 8)
+        self.assertAlmostEqual(stats.max_occupancy, 8)
+
+        stats = result.buffet_stats[Buffet(level="Disk", tensor="I", einsum="copy")]
+        self.assertAlmostEqual(stats.net_total_read_actions(), 0)
+        self.assertAlmostEqual(stats.net_max_per_unit_read_actions(), 0)
+        self.assertAlmostEqual(stats.net_total_write_actions(), 8)
+        self.assertAlmostEqual(stats.net_max_per_unit_write_actions(), 8)
+        self.assertAlmostEqual(stats.max_occupancy, 8)
+
+
+class TestSymbolicAccesses(unittest.TestCase):
+    def test_q_mapping(self):
+        mapping = Mapping.from_yaml(Path(__file__).parent / "Q_mapping.yaml")
+        workload = Workload.from_yaml(Path(__file__).parent / "mha.yaml")
+
+        result = analyze_reuse_and_add_reservations_to_mapping(
+            make_job(mapping, workload)
+        )
+
+        # main_memory_I_accesses = accesses.get_accesses('MainMemory', 'I', 'Q')
+        stats = result.buffet_stats[Buffet(level="MainMemory", tensor="I", einsum="Q")]
+        self.assertAlmostEqual(stats.net_total_read_actions(), 64.0)
+        self.assertAlmostEqual(stats.net_max_per_unit_read_actions(), 64.0)
+        self.assertAlmostEqual(stats.net_total_write_actions(), 0.0)
+        self.assertAlmostEqual(stats.net_max_per_unit_write_actions(), 0.0)
+
+        stats = result.buffet_stats[Buffet(level="LocalBuffer", tensor="I", einsum="Q")]
+        self.assertAlmostEqual(stats.net_total_read_actions(), 64.0)
+        self.assertAlmostEqual(stats.net_max_per_unit_read_actions(), 16.0)
+        self.assertAlmostEqual(stats.net_total_write_actions(), 64.0)
+        self.assertAlmostEqual(stats.net_max_per_unit_write_actions(), 16.0)
+
+        stats = result.buffet_stats[Buffet(level="MainMemory", tensor="Q", einsum="Q")]
+        self.assertAlmostEqual(stats.net_total_read_actions(), 0)
+        self.assertAlmostEqual(stats.net_max_per_unit_read_actions(), 0)
+        self.assertAlmostEqual(stats.net_total_write_actions(), 16.0)
+        self.assertAlmostEqual(stats.net_max_per_unit_write_actions(), 16.0)
+
+        stats = result.buffet_stats[Buffet(level="LocalBuffer", tensor="Q", einsum="Q")]
+        self.assertAlmostEqual(stats.net_total_read_actions(), 64.0)
+        self.assertAlmostEqual(stats.net_max_per_unit_read_actions(), 16.0)
+        self.assertAlmostEqual(stats.net_total_write_actions(), 64.0)
+        self.assertAlmostEqual(stats.net_max_per_unit_write_actions(), 16.0)
+
+
+class TestSymbolicActions(unittest.TestCase):
+    def test_q_mapping(self):
+        mapping = Mapping.from_yaml(Path(__file__).parent / "Q_mapping.yaml")
+        workload = Workload.from_yaml(Path(__file__).parent / "mha.yaml")
+
+        result = analyze_reuse_and_add_reservations_to_mapping(
+            make_job(mapping, workload)
+        )
+        actions = gather_actions(result, None, use_name=True)
+
+        self.assertAlmostEqual(actions[("LocalBuffer", "read")].total, 128.0)
+        self.assertAlmostEqual(actions[("LocalBuffer", "read")].max_per_unit, 32.0)
+        self.assertAlmostEqual(actions[("LocalBuffer", "write")].total, 128.0)
+        self.assertAlmostEqual(actions[("LocalBuffer", "write")].max_per_unit, 32.0)
+
+        self.assertAlmostEqual(actions[("MAC", "compute")].total, 64.0)
+        self.assertAlmostEqual(actions[("MAC", "compute")].max_per_unit, 16.0)
+
+
+class TestSymbolicLatency(unittest.TestCase):
+    def test_q_mapping(self):
+        spec = Spec.from_yaml(
+            [
+                # Path(__file__).parent / 'Q_mapping.yaml',
+                Path(__file__).parent / "mha.yaml",
+                Path(__file__).parent / "four_level.arch.yaml",
+            ]
+        )
+        workload = spec.workload
+        architecture = spec._get_flattened_architecture()
+        mapping = Mapping.from_yaml(Path(__file__).parent / "Q_mapping.yaml")
+
+        result = analyze_reuse_and_add_reservations_to_mapping(
+            make_job(mapping, workload)
+        )
+        overall_latency, _, _ = get_latency(result, mapping, workload, architecture)
+
+        self.assertAlmostEqual(overall_latency, 16.0)
+
+
+if __name__ == "__main__":
+    unittest.main()