accelforge 0.0.1__py3-none-any.whl

tests/distribuffers/spec/logical_arch.yaml

@@ -0,0 +1,32 @@
+arch:
+  nodes:
+  - !Component
+    name: DRAM                 # offchip DRAM is the source of all datatypes
+    class: DRAM                # assume DRAM is large enough to store all the data, so no depth specification needed
+    attributes:
+      width: 64                # width in bits
+
+  - !Container
+    name: PE
+    spatial: {meshX: 1, meshY: 2}
+
+  # registers for the mac unit
+  - !Component
+    name: weight_reg
+    class: reg_storage
+    attributes: {depth: 1, width: 8}
+
+  - !Component
+    name: input_activation_reg
+    class: reg_storage
+    attributes: {depth: 1, width: 8}
+
+  - !Component
+    name: output_activation_reg
+    class: reg_storage
+    attributes: {depth: 1, width: 8}
+
+  - !Component
+    name: mac
+    class: mac_compute
+    attributes: {num_pipline_stages: 2}

tests/distribuffers/spec/physical_arch.yaml

@@ -0,0 +1,69 @@
+arch:
+  name: "test_distributed_spec"
+  nodes:
+  - !Container
+    name: system_arch
+    attributes:
+      # Top-level attributes inherited by all components unless overridden
+      technology: "45nm"
+      cycle_period: 1e-9
+
+  - !Network
+    name: L2
+    topology: Mesh
+    dims:
+      - x
+    constraints:
+      - 0 <= x < 2
+  - !Network
+    name: L1
+    topology: Mesh
+    dims:
+      - x
+      - y
+    constraints:
+      - 0 <= x < 1
+      - -1 <= y <= 1
+    port: |
+      {L2[x] -> L1[x', y'] : x = 1 and x' = 0 and y' = 0}
+
+  - !Component
+    name: DRAM                 # offchip DRAM is the source of all datatypes
+    class: DRAM                # assume DRAM is large enough to store all the data, so no depth specification needed
+    attributes:
+      width: 64                # width in bits
+    network:
+      level: L2
+      placement: "{ DRAM[i] -> L2[i] | i = 1 }"
+
+  - !Container
+    name: PE
+    network:
+      level: L1
+      placement: |
+        {
+          PE[i, j] -> L1[x, y] :
+            j = y and x < 0 and i = -x - 1;
+            j = y and x > 0 and i = x;
+        }
+
+  # registers for the mac unit
+  - !Component
+    name: weight_reg
+    class: reg_storage
+    attributes: {depth: 1, width: 8}
+
+  - !Component
+    name: input_activation_reg
+    class: reg_storage
+    attributes: {depth: 1, width: 8}
+
+  - !Component
+    name: output_activation_reg
+    class: reg_storage
+    attributes: {depth: 1, width: 8}
+
+  - !Component
+    name: mac
+    class: mac_compute
+    attributes: {num_pipline_stages: 2}

tests/distribuffers/test_binding.py

@@ -0,0 +1,48 @@
+from pathlib import Path
+import unittest
+
+from typing import Dict, List
+from islpy import DEFAULT_CONTEXT, Map
+
+import yaml
+from accelforge.frontend._binding import Binding, BindingNode
+
+TESTS_DIR = Path(__file__).parent / "spec" / "binding"
+
+
+class TestBindingMapper(unittest.TestCase):
+    def test_valid_bindings(self):
+        """
+        Tests that the valid bindings translate into the appropriate
+        ISL strings.
+        """
+        specs_file: str = TESTS_DIR / "valid_bindings.yaml"
+        with open(specs_file, mode="r", encoding="utf-8") as f:
+            specs: List = yaml.safe_load(f)
+
+        spec: Dict
+        for spec in specs:
+            binding: Binding = Binding.model_validate(spec["binding"])
+
+            soln: Dict = spec["solution"]
+
+            soln_node: Dict[str, str]
+            binding_node: BindingNode
+            for soln_node, binding_node in zip(soln["nodes"], binding.nodes):
+                isl_relations: Dict[str, Map] = binding_node.isl_relations
+                assert soln_node.keys() == isl_relations.keys(), (
+                    "Not all isl_relations read in properly. Missing \n"
+                    f"{set(soln_node.keys()).difference(isl_relations.keys())} "
+                    "\nfrom isl_relations and \n"
+                    f"{set(isl_relations.keys()).difference(soln_node.keys())} "
+                    "\nfrom solutions.\n"
+                )
+
+                tensor: str
+                for tensor in soln_node:
+                    soln_relation: Map = Map.read_from_str(
+                        DEFAULT_CONTEXT, soln_node[tensor]
+                    )
+                    assert soln_relation.is_equal(
+                        isl_relations[tensor]
+                    ), f"\n{soln_relation} != \n{isl_relations[tensor]}"

tests/frontend/test_mapping_viz.py

@@ -0,0 +1,52 @@
+"""
+Uncommenting the lines with `graph.write_png` will generate PNG pictures of
+LoopTrees so they can be inspected.
+
+TODO: an automated way to check for correctness.
+"""
+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.parent / "examples"
+
+
+class TestModel(unittest.TestCase):
+    def test_default(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)
+        graph = result.data["Total<SEP>mapping"].iloc[0]().render_pydot()
+        # graph.write_png("default.png")
+
+    def test_without_reservations(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)
+        graph = result.data["Total<SEP>mapping"].iloc[0]().render_pydot(with_reservations=False)
+        # graph.write_png("without_reservations.png")
+
+    def test_without_stride(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)
+        graph = result.data["Total<SEP>mapping"].iloc[0]().render_pydot(with_tile_shape=False)
+        # graph.write_png("without_stride.png")

tests/mapper/configs/conv1d/conv1d.mapping.yaml

@@ -0,0 +1,31 @@
+mapping:
+  nodes:
+  - !Storage
+    tensors: [I, W, O]
+    component: MainMemory
+  - !Temporal
+    rank_variable: p
+    # bound is 2
+    tile_pattern:
+      initial_tile_shape: 8
+      stride: 8
+  - !Storage
+    tensors: [I, W, O]
+    component: LocalBuffer
+  - !Temporal
+    rank_variable: p
+    # bound is 8
+    tile_pattern:
+      initial_tile_shape: 1
+      stride: 1
+  - !Spatial
+    rank_variable: r
+    # bound is 3
+    tile_pattern:
+      initial_tile_shape: 1
+      stride: 1
+    component: MAC
+    name: X
+  - !Compute
+    einsum: conv
+    compute: MAC

tests/mapper/configs/conv1d/conv1d.workload.yaml

@@ -0,0 +1,11 @@
+workload:
+  shape:
+    r: 0 <= r < 3
+    p: 0 <= p < 16
+
+  einsums:
+  - name: conv
+    tensor_accesses:
+    - {name: I, projection: [ r, p ]}
+    - {name: W, projection: [ r ]}
+    - {name: O, projection: [ p ], output: True}

tests/mapper/configs/two_conv1d/two_conv1d.expected.yaml

@@ -0,0 +1,38 @@
+mapping_to_isl:
+  BufferTensorEinsum(buffer='MainMemory', tensor='I', einsum=Compute(type=None, einsum='convA', compute='MAC', component_object=None)): |
+    { MainMemory_spacetime[] -> I[PI] : 0 <= PI < 18 }
+  BufferTensorEinsum(buffer='MainMemory', tensor='WA', einsum=Compute(type=None, einsum='convA', compute='MAC', component_object=None)): |
+    { MainMemory_spacetime[] -> WA[RA] : 0 <= RA < 3 }
+  BufferTensorEinsum(buffer='LocalBuffer', tensor='I', einsum=Compute(type=None, einsum='convA', compute='MAC', component_object=None)): |
+    { LocalBuffer_spacetime[pb0=0] -> I[PI] : 0 <= PI < 12;
+      LocalBuffer_spacetime[pb0=1] -> I[PI] : 10 <= PI < 18 }
+  BufferTensorEinsum(buffer='LocalBuffer', tensor='WA', einsum=Compute(type=None, einsum='convA', compute='MAC', component_object=None)): |
+    { LocalBuffer_spacetime[pb0] -> WA[RA] : 0 <= RA <= 2 and 0 <= pb0 < 2 }
+  BufferTensorEinsum(buffer='LocalBuffer', tensor='A', einsum=Compute(type=None, einsum='convA', compute='MAC', component_object=None)): |
+    { LocalBuffer_spacetime[pb0=0] -> A[PA] : 0 <= PA < 10;
+      LocalBuffer_spacetime[pb0=1] -> A[PA] : 10 <= PA < 16 }
+  BufferTensorEinsum(buffer='MAC', tensor='I', einsum=Compute(type=None, einsum='convA', compute='MAC', component_object=None)): |
+    { MAC_spacetime[pb0=0, 0, pa1, ra0] -> I[PI = pa1 + ra0] : 0 <= pa1 < 10 and 0 <= ra0 < 3;
+      MAC_spacetime[pb0=1, 0, pa1, ra0] -> I[PA = 10 + pa1 + ra0] : 0 <= pa1 < 6 and 0 <= ra0 < 3 }
+  BufferTensorEinsum(buffer='MAC', tensor='WA', einsum=Compute(type=None, einsum='convA', compute='MAC', component_object=None)): |
+    { MAC_spacetime[pb0=0, 0, pa1, ra0] -> WA[RA=ra0] : 0 <= pa1 < 10 and 0 <= ra0 < 3;
+      MAC_spacetime[pb0=1, 0, pa1, ra0] -> WA[RA=ra0] : 0 <= pa1 < 6 and 0 <= ra0 < 3 }
+  BufferTensorEinsum(buffer='MAC', tensor='A', einsum=Compute(type=None, einsum='convA', compute='MAC', component_object=None)): |
+    { MAC_spacetime[pb0=0, 0, pa1, ra0] -> A[PA=pa1] : 0 <= pa1 < 10 and 0 <= ra0 < 3;
+      MAC_spacetime[pb0=1, 0, pa1, ra0] -> A[PA=10 + pa1] : 0 <= pa1 < 6 and 0 <= ra0 < 3 }
+  BufferTensorEinsum(buffer='MainMemory', tensor='WB', einsum=Compute(type=None, einsum='convB', compute='MAC', component_object=None)): |
+    { MainMemory_spacetime[] -> WB[RB] : 0 <= RB < 3 }
+  BufferTensorEinsum(buffer='MainMemory', tensor='B', einsum=Compute(type=None, einsum='convB', compute='MAC', component_object=None)): |
+    { MainMemory_spacetime[] -> B[PB] : 0 <= PB < 16 }
+  BufferTensorEinsum(buffer='LocalBuffer', tensor='A', einsum=Compute(type=None, einsum='convB', compute='MAC', component_object=None)): |
+    { LocalBuffer_spacetime[pb0] -> A[PA] : 0 <= pb0 < 2 and 8pb0 <= PA <= 9 + 8pb0 }
+  BufferTensorEinsum(buffer='LocalBuffer', tensor='B', einsum=Compute(type=None, einsum='convB', compute='MAC', component_object=None)): |
+    { LocalBuffer_spacetime[pb0] -> B[PB] : 0 <= pb0 < 2 and 8pb0 <= PB < 8 + 8pb0 }
+  BufferTensorEinsum(buffer='LocalBuffer', tensor='WB', einsum=Compute(type=None, einsum='convB', compute='MAC', component_object=None)): |
+    { LocalBuffer_spacetime[pb0] -> WB[RB] : 0 <= pb0 < 2 and 0 <= RB <= 2 }
+  BufferTensorEinsum(buffer='MAC', tensor='A', einsum=Compute(type=None, einsum='convB', compute='MAC', component_object=None)): |
+    { MAC_spacetime[pb0, 1, pb2, rb0] -> A[PA = 8pb0 + pb2 + rb0] : 0 <= pb0 <= 1 and 0 <= pb2 <= 7 and 0 <= rb0 <= 2 }
+  BufferTensorEinsum(buffer='MAC', tensor='WB', einsum=Compute(type=None, einsum='convB', compute='MAC', component_object=None)): |
+    { MAC_spacetime[pb0, 1, pb2, rb0] -> WB[RB = rb0] : 0 <= pb0 <= 1 and 0 <= pb2 <= 7 and 0 <= rb0 <= 2 }
+  BufferTensorEinsum(buffer='MAC', tensor='B', einsum=Compute(type=None, einsum='convB', compute='MAC', component_object=None)): |
+    { MAC_spacetime[pb0, 1, pb2, rb0] -> B[PB = 8pb0 + pb2] : 0 <= pb0 <= 1 and 0 <= pb2 <= 7 and 0 <= rb0 <= 2 }

tests/mapper/configs/two_conv1d/two_conv1d.mapping.yaml

@@ -0,0 +1,54 @@
+mapping:
+  nodes:
+  - !Storage
+    tensors: [I, WA, WB, B]
+    component: MainMemory
+  - !Temporal
+    rank_variable: pb
+    # bound is 2
+    tile_pattern:
+      initial_tile_shape: 8
+      stride: 8
+  - !Storage
+    tensors: [I, WA, A, WB, B]
+    component: LocalBuffer
+  - !Sequential
+    nodes:
+    - !Nested
+      nodes:
+      - !Temporal
+        rank_variable: pa
+        # bound is 8
+        tile_pattern:
+          initial_tile_shape: 1
+          stride: 1
+      - !Spatial
+        rank_variable: ra
+        # bound is 3
+        tile_pattern:
+          initial_tile_shape: 1
+          stride: 1
+        component: MAC
+        name: X
+      - !Compute
+        einsum: convA
+        compute: MAC
+    - !Nested
+      nodes:
+      - !Temporal
+        rank_variable: pb
+        # bound is 8
+        tile_pattern:
+          initial_tile_shape: 1
+          stride: 1
+      - !Spatial
+        rank_variable: rb
+        # bound is 3
+        tile_pattern:
+          initial_tile_shape: 1
+          stride: 1
+        component: MAC
+        name: X
+      - !Compute
+        einsum: convB
+        compute: MAC

tests/mapper/configs/two_conv1d/two_conv1d.workload.yaml

@@ -0,0 +1,19 @@
+workload:
+  shape:
+    ra: 0 <= ra < 3
+    pa: 0 <= pa < 16
+    rb: 0 <= rb < 3
+    pb: 0 <= pb < 16
+
+  einsums:
+  - name: convA
+    tensor_accesses:
+    - {name: I, projection: { PI: pa+ra }}
+    - {name: WA, projection: [ ra ]}
+    - {name: A, projection: [ pa ], output: True}
+
+  - name: convB
+    tensor_accesses:
+    - {name: A, projection: { PA: pb+rb }}
+    - {name: WB, projection: [ rb ]}
+    - {name: B, projection: [ pb ], output: True}

tests/mapper/test_mapping_to_isl.py

@@ -0,0 +1,90 @@
+"""
+Some non-exhaustive tests to make sure core functionality of the non-fusing and
+fusing portions of ISL tiling is not violated under changes.
+"""
+
+from pathlib import Path
+from pprint import pformat
+import unittest
+
+import islpy as isl
+
+from accelforge.frontend.workload import Workload
+from accelforge.frontend.mapping import Mapping
+
+from accelforge.model._looptree.reuse.isl.mapping_to_isl import analyze_mapping
+from accelforge.model._looptree.reuse.isl.mapping_to_isl.types import (
+    MappingAnalysisResult,
+)
+
+from .util import TEST_CONFIG_PATH, load_solutions
+
+
+class TestMappingToIsl(unittest.TestCase):
+    """
+    Tests taking a `Mapping` and `Workload` and converting it into relevant isl
+    objects.
+    """
+
+    def test_conv1d(self):
+        """
+        Non-fusing one 1-dimensional convolution test.
+        """
+        # Loads in the CONV1D Config
+        config_path: Path = TEST_CONFIG_PATH / "conv1d"
+        workload: Workload = Workload.from_yaml(config_path / "conv1d.workload.yaml")
+
+        mapping: Mapping = Mapping.from_yaml(config_path / "conv1d.mapping.yaml")
+        occupancies: MappingAnalysisResult = analyze_mapping.occupancies_from_mapping(
+            mapping, workload
+        )
+
+        for buffer, occupancy in occupancies.buffet_to_occupancy.items():
+            if buffer == list(occupancies.buffet_to_occupancy.keys())[-1]:
+                soln: isl.Map = isl.Map.read_from_str(
+                    isl.DEFAULT_CONTEXT,
+                    "{ MAC_spacetime[P1, P0, R] -> O[P=8*P1 + P0] : "
+                    "0 <= R < 3 and 0 <= P1 < 2 and 0 <= P0 < 8}",
+                )
+                assert occupancy.map_ == soln
+
+    def test_two_conv1d(self):
+        """
+        Fusing two 1-dimensional convolutions test.
+        """
+        # Loads in the CONV1D Config
+        config_path: Path = TEST_CONFIG_PATH / "two_conv1d"
+        workload: Workload = Workload.from_yaml(
+            config_path / "two_conv1d.workload.yaml"
+        )
+
+        mapping: Mapping = Mapping.from_yaml(config_path / "two_conv1d.mapping.yaml")
+        occupancies: MappingAnalysisResult = analyze_mapping.occupancies_from_mapping(
+            mapping, workload
+        )
+        solns: dict = load_solutions(config_path / "two_conv1d.expected.yaml")[
+            "mapping_to_isl"
+        ]
+
+        errors: list = []
+        buffers_seen: set = set()
+        for buffer, occupancy in occupancies.buffet_to_occupancy.items():
+            try:
+                soln = solns[repr(buffer)]
+                assert (
+                    occupancy.map_ == soln
+                ), f"{buffer} should hold:\n{soln}\ninstead holds:\n{occupancy.map_}"
+            except (AssertionError, KeyError) as e:
+                errors.append(e)
+            buffers_seen.add(repr(buffer))
+
+        buffers_known: set = set(solns.keys())
+        try:
+            assert buffers_seen == buffers_known, (
+                f"Buffers Missing: {pformat(buffers_seen - buffers_known)}\n"
+                f"Buffers Extra: {pformat(buffers_known - buffers_seen)}"
+            )
+        except AssertionError as e:
+            errors.append(e)
+
+        assert len(errors) == 0, pformat(errors)

tests/mapper/test_spatial_reuse_analysis.py

@@ -0,0 +1,67 @@
+"""
+Tests the spatial ISL models.
+"""
+
+import unittest
+
+import islpy as isl
+
+from accelforge.frontend.mapping import MappingNode
+from accelforge.model._looptree.reuse.isl.mapping_to_isl.types import (
+    Fill,
+    Occupancy,
+    SpatialTag,
+    TemporalTag,
+)
+from accelforge.model._looptree.reuse.isl.spatial import (
+    SimpleLinkTransferModel,
+    TransferInfo,
+)
+
+
+class TestSimpleLinkTransferModel(unittest.TestCase):
+    """
+    Test the simple link transfer model. Adapted from the original test at the
+    following link:
+    https://github.com/NVlabs/timeloop/blob/master/src/unit-test/test-simple-link-transfer.cpp
+    """
+
+    def test_simple_link_transfer_model_sandbox(self):
+        """Independent sanity check of `SimpleLinkTransferModel`."""
+        buffer: MappingNode = MappingNode()
+        fill: Fill = Fill(
+            [TemporalTag(), SpatialTag(0, buffer), SpatialTag(1, buffer)],  # type: ignore
+            isl.Map.read_from_str(
+                isl.DEFAULT_CONTEXT,
+                "{ spacetime[t, x, y] -> data[t+x+y] : 0 <= x < 2 and 0 <= y < 2 and 0 <= t < 2}",
+            ),
+        )
+        occ: Occupancy = Occupancy(fill.tags, fill.map_)
+        link_transfer_model: SimpleLinkTransferModel = SimpleLinkTransferModel()
+        info: TransferInfo = link_transfer_model.apply(buffer, fill, occ)
+
+        assert info.fulfilled_fill.map_ == (
+            fulfilled_soln := isl.Map.read_from_str(
+                isl.DEFAULT_CONTEXT,
+                "{ spacetime[t = 1, x, y = 0] -> data[1 + x] : 0 <= x <= 1; "
+                "  spacetime[t = 1, x = 0, y] -> data[1 + y] : 0 <= y <= 1 }",
+            )
+        ), (
+            "`fulfilled_fill` and solution mismatch\n"
+            "--------------------------------------\n"
+            f"solution:\n{fulfilled_soln}\n"
+            f"received:\n{info.fulfilled_fill.map_}\n"
+        )
+
+        assert info.unfulfilled_fill.map_ == (
+            unfulfilled_soln := isl.Map.read_from_str(
+                isl.DEFAULT_CONTEXT,
+                "{ spacetime[t = 0, x, y] -> data[x + y] : 0 <= x <= 1 and 0 <= y <= 1; "
+                "  spacetime[t = 1, x = 1, y = 1] -> data[3] }",
+            )
+        ), (
+            "`unfulfilled_fill` and solution mismatch\n"
+            "----------------------------------------\n"
+            f"solution:\n{unfulfilled_soln}\n"
+            f"received:\n{info.unfulfilled_fill.map_}\n"
+        )

tests/mapper/test_temporal_reuse_analysis.py

@@ -0,0 +1,56 @@
+"""
+Tests the temporal reuse analysis in ISL.
+"""
+
+import unittest
+
+import islpy as isl
+
+from accelforge.model._looptree.reuse.isl.mapping_to_isl.types import (
+    Occupancy,
+    SpatialTag,
+    TemporalTag,
+)
+from accelforge.model._looptree.reuse.isl.temporal import (
+    TemporalReuse,
+    analyze_temporal_reuse,
+)
+
+
+class TestTemporalReuseAnalysis(unittest.TestCase):
+    """
+    Tests the temporal reuse analysis in ISL.
+    """
+
+    def test_multiple_loop_reuse_sandbox(self):
+        """
+        Tests a very basic temporal reuse analysis case without needing the
+        `occupancies_from_mapping` function.
+        """
+        # Loads in the CONV1D Config
+        occ: Occupancy = Occupancy(
+            [TemporalTag(), SpatialTag(0, None), TemporalTag()],
+            isl.Map.read_from_str(
+                isl.DEFAULT_CONTEXT,
+                (
+                    "{ generic_iteration[t1, x, t0] -> tensor[d] : "
+                    "t0 <= d < t0+3 and 0 <= t1 < 2 and 0 <= x < 2 and 0 <= t0 < 2 }"
+                ),
+            ).coalesce(),
+        )
+
+        result: TemporalReuse = analyze_temporal_reuse(occ, True, True)
+        soln: isl.Map = isl.Map.read_from_str(
+            isl.DEFAULT_CONTEXT,
+            "{ generic_iteration[t1, x, t0] -> tensor[d] : "
+            "0 <= x < 2 and "
+            "(((t1 = 0) and (t0 = 0) and (0 <= d < 3)) or "
+            " ((t1 = 0) and (t0 = 1) and (d = 3)) or "
+            " ((t1 = 1) and (t0 = 0) and (d = 0)) or "
+            " ((t1 = 1) and (t0 = 1) and (d = 3))"
+            ")}",
+        ).coalesce()
+        fill: isl.Map = result.fill.map_
+        eff_occ: isl.Map = result.effective_occupancy.map_
+        print(eff_occ)
+        assert soln.is_equal(fill), f"Expected:\n{soln}\nReceived:\n{fill}"

tests/mapper/util.py

@@ -0,0 +1,58 @@
+"""
+Utility functions common to testing the isl mapper functions.
+"""
+
+from pathlib import Path
+import islpy as isl
+
+from ruamel.yaml import YAML
+
+TEST_CONFIG_PATH: Path = Path(__file__).parent / "configs"
+
+
+def to_isl_maps(obj: str | list | dict) -> dict:
+    """
+    Given an object, attempt to reduce all strings in tree with isl.Map
+
+    Parameters
+    ----------
+    obj:
+        A DAG which can be explored and contains isl.Map strings within it.
+
+    Returns
+    -------
+    `obj` but all strings are converted to isl.Map.
+    """
+
+    def _to_isl_maps(obj: str | dict | list) -> isl.Map | dict | list:
+        """Recursively convert string ISL maps to isl.Map; leave others alone."""
+        if isinstance(obj, str):
+            return isl.Map.read_from_str(isl.DEFAULT_CONTEXT, obj)
+        if isinstance(obj, dict):
+            return {k: _to_isl_maps(v) for k, v in obj.items()}
+        if isinstance(obj, list):
+            return [_to_isl_maps(v) for v in obj]
+        return obj
+
+    return _to_isl_maps(obj)  # type: ignore
+
+
+def load_solutions(path: Path) -> dict:
+    """
+    Loads in a dictionary with the isl solutions to a workload problem.
+
+    Parameters
+    ----------
+    path:
+        The path to the solutions.
+
+    Returns
+    -------
+    A dictionary relating Python-based keys generated by the mapper (e.g.,
+    `BufferTensorEinsum` to their corresponding isl.Map.)
+    """
+    # Load expected solutions (YAML file with string ISL maps)
+    yaml: YAML = YAML(typ="safe")
+
+    with open(path, "r", encoding="utf-8") as f:
+        return to_isl_maps(yaml.load(f))

tests/matmul.mapping.yaml

@@ -0,0 +1,29 @@
+mapping:
+  nodes:
+  - !Storage
+    tensors: [T1, T0, W0]
+    component: MainMemory
+  - !Temporal
+    rank_variable: m
+    tile_shape: 128
+  - !Storage
+    tensors: [T1]
+    component: LocalBuffer
+  - !Temporal
+    rank_variable: n0
+    tile_shape: 1
+  - !Storage
+    tensors: [T0]
+    component: LocalBuffer
+  - !Temporal
+    rank_variable: n1
+    tile_shape: 1
+  - !Storage
+    tensors: [W0]
+    component: LocalBuffer
+  - !Temporal
+    rank_variable: m
+    tile_shape: 1
+  - !Compute
+    einsum: Matmul1
+    compute: MAC

tests/matmul.workload.yaml

@@ -0,0 +1,12 @@
+workload:
+  shape:
+    m: 0 <= m < 128
+    n0: 0 <= n0 < 64
+    n1: 0 <= n1 < 128
+
+  einsums:
+  - name: Matmul1
+    tensor_accesses:
+    - {name: T0, projection: [m, n0]}
+    - {name: W0, projection: [n0, n1]}
+    - {name: T1, projection: [m, n1], output: True}