Test updates ; vibe tests

Author: tanner-andrulis

accelforge/frontend/_binding.py

@@ -20,7 +20,7 @@ class Domain(EvalableModel):
     _prefix: str
     name: str
     dims: EvalableList[str]
-    
+
     @property
     def isl_space(self) -> isl.Space:
         return isl.Space.create_from_names(
@@ -36,7 +36,8 @@ class LogicalDomain(Domain):
     """
     Represents the logical architecture domain space of logical dims × tensor ranks.
     """
-    _prefix: str = 'l'
+
+    _prefix: str = "l"
     pass
 
 
@@ -45,7 +46,8 @@ class PhysicalDomain(Domain):
     Represents the logical architecture domain space of physical dims.
     The physical space is defined as the physical architecture dims.
     """
-    _prefix: str = 'p'
+
+    _prefix: str = "p"
     pass
 
 
@@ -109,6 +111,7 @@ class Binding(EvalableModel):
     nodes: EvalableList[BindingNode]
     """Parts of the binding."""
 
+
 '''
 class NetworkOnChip(ParsableModel):
     """A model of a network-on-chip on the physical chip."""

accelforge/model/_looptree/reuse/isl/distributed/bind.py

@@ -1,23 +1,21 @@
 """Applies the binding layer into one that can be used for later analysis,"""
+
 from accelforge.frontend.binding import Binding
 from accelforge.frontend.mapping import Mapping
 from accelforge.frontend.workload import Workload
 
 from accelforge.model._looptree.reuse.isl.mapping_to_isl.analyze_mapping import (
     occupancies_from_mapping,
-    MappingAnalysisResult
+    MappingAnalysisResult,
 )
 
 
-
 def apply_binding(binding: Binding, mapping: Mapping, workload: Workload):
     """
     Given a mapping, apply the mapping from logical components onto physical
     components.
     """
-    map_analysis: MappingAnalysisResult = occupancies_from_mapping(
-        mapping, workload
-    )
+    map_analysis: MappingAnalysisResult = occupancies_from_mapping(mapping, workload)
     for buffet, occ in map_analysis.buffet_to_occupancy.items():
         print(occ)
     print(binding)

accelforge/model/_looptree/reuse/isl/distributed/distributed_buffers.py

@@ -39,7 +39,7 @@ def identify_mesh_casts(
         space, corresponding to the `src` and `dst`, and returns the distance
         between the two points in terms of `hops`, a quantized atomic distance of
         data transmission cost.
-    
+
     Returns
     -------
     { [data] -> [dst -> src] } where { [dst] -> [data] } and { [src] -> [data] } are in
@@ -57,9 +57,11 @@ def identify_mesh_casts(
     data_presence: isl.Map = src_occupancy.reverse()
 
     # { [dst -> data] -> [dst -> src] } where src contains data.
-    fills_to_matches: isl.Map = fill_to_fill.uncurry( # { [[dst -> data] -> dst] -> data }
-                                ).apply_range(data_presence # { [[dst -> data] -> dst] -> src }
-                                ).curry() # { [[dst -> data] -> [dst -> src] }
+    fills_to_matches: isl.Map = (
+        fill_to_fill.uncurry()  # { [[dst -> data] -> dst] -> data }
+        .apply_range(data_presence)  # { [[dst -> data] -> dst] -> src }
+        .curry()
+    )  # { [[dst -> data] -> [dst -> src] }
     if DUMP_ISL_IR:
         logging.info(f"fills_to_matches: {fills_to_matches}")
 
@@ -68,13 +70,19 @@ def identify_mesh_casts(
     fill_min_dist: isl.Map = fills_to_matches.apply_range(dist_fn).lexmin()
     # Isolates the relevant minimal pairs.
     # { [dst -> data] -> [dst -> src] :.dst -> src is minimized distance }
-    minimal_pairs: isl.Map = fill_min_dist.apply_range(
-        # Note: Need to match fill -> min_dist with min_dist -> [fill -> match] as lexmin over
-        # fill and match will minimize distance over the tuple (src, dst, data), but that
-        # overconstrains the optimization as we want to minimize over distance (dst, data)
-        # only for all src.
-        fills_to_matches.range_map().apply_range(dist_fn).reverse()
-    ).range().unwrap()
+    minimal_pairs: isl.Map = (
+        fill_min_dist.apply_range(
+            # Note: Need to match fill -> min_dist with min_dist -> [fill -> match] as lexmin over
+            # fill and match will minimize distance over the tuple (src, dst, data), but that
+            # overconstrains the optimization as we want to minimize over distance (dst, data)
+            # only for all src.
+            fills_to_matches.range_map()
+            .apply_range(dist_fn)
+            .reverse()
+        )
+        .range()
+        .unwrap()
+    )
     if DUMP_ISL_IR:
         logging.info(f"minimal_pairs: {minimal_pairs}")
 
@@ -211,7 +219,7 @@ def _cost_mesh_cast_hypercube(self, mcns: isl.Map) -> int:
         mcns:
             Multicast networks grouped together by [srcs -> data] fulfilling
             [dsts -> data], where there is at least 1 src and 1 dst in each mcn.
-        
+
         Returns
         -------
         The upperbound of doing all the multicasts specified by the multicast

tests/__init__.py

@@ -66,9 +66,7 @@ def test_gamut(self):
             fill: Fill = Fill(dim_tags, test["fill"])
             occ: Occupancy = Occupancy(dim_tags, test["occ"])
             dist_fn: isl.Map = test["dist_fn"]
-            multicast_model: HypercubeMulticastModel = HypercubeMulticastModel(
-                dist_fn
-            )
+            multicast_model: HypercubeMulticastModel = HypercubeMulticastModel(dist_fn)
 
             # Applies the model.
             info: TransferInfo = multicast_model.apply(0, fill, occ)

tests/isl/distributed/test_multicast.py

@@ -1,5 +1,5 @@
 workload:
-  shape:
+  iteration_space_shape:
     r: 0 <= r < 3
     p: 0 <= p < 16
 

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

@@ -1,5 +1,5 @@
 workload:
-  shape:
+  iteration_space_shape:
     ra: 0 <= ra < 3
     pa: 0 <= pa < 16
     rb: 0 <= rb < 3

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

@@ -8,10 +8,11 @@
 
 from accelforge.model._looptree.reuse.isl.isl_functions import *
 
+
 class BasicIslFunctionTests(unittest.TestCase):
     """
     Prima facie checks of `isl_functions` to ensure basic functionality.
     """
 
     def test_project_dim_after(self):
-        pass
+        pass

tests/isl/mapper/test_isl_functions.py

@@ -21,6 +21,7 @@
 
 TEST_CONFIG_PATH: Path = Path(__file__).parent / "configs"
 
+
 class TestMappingToIsl(unittest.TestCase):
     """
     Tests taking a `Mapping` and `Workload` and converting it into relevant isl
@@ -33,9 +34,13 @@ def test_conv1d(self):
         """
         # Loads in the CONV1D Config
         config_path: Path = TEST_CONFIG_PATH / "conv1d"
-        workload: Workload = Workload.from_yaml(config_path / "conv1d.workload.yaml")
+        workload: Workload = Workload.from_yaml(
+            config_path / "conv1d.workload.yaml", top_key="workload"
+        )
 
-        mapping: Mapping = Mapping.from_yaml(config_path / "conv1d.mapping.yaml")
+        mapping: Mapping = Mapping.from_yaml(
+            config_path / "conv1d.mapping.yaml", top_key="mapping"
+        )
         occupancies: MappingAnalysisResult = analyze_mapping.occupancies_from_mapping(
             mapping, workload
         )
@@ -56,10 +61,12 @@ def test_two_conv1d(self):
         # Loads in the CONV1D Config
         config_path: Path = TEST_CONFIG_PATH / "two_conv1d"
         workload: Workload = Workload.from_yaml(
-            config_path / "two_conv1d.workload.yaml"
+            config_path / "two_conv1d.workload.yaml", top_key="workload"
         )
 
-        mapping: Mapping = Mapping.from_yaml(config_path / "two_conv1d.mapping.yaml")
+        mapping: Mapping = Mapping.from_yaml(
+            config_path / "two_conv1d.mapping.yaml", top_key="mapping"
+        )
         occupancies: MappingAnalysisResult = analyze_mapping.occupancies_from_mapping(
             mapping, workload
         )

tests/isl/mapper/test_mapping_to_isl.py

@@ -7,6 +7,7 @@
 
 from ruamel.yaml import YAML
 
+
 def to_isl_maps(obj: str | list | dict) -> dict:
     """
     Given an object, attempt to reduce all strings in tree with isl.Map
@@ -26,10 +27,7 @@ def _to_isl_maps(obj: str | dict | list) -> isl.Map | dict | list:
         if isinstance(obj, str):
             return isl.Map.read_from_str(isl.DEFAULT_CONTEXT, obj)
         if isinstance(obj, dict):
-            return {
-                k: (_to_isl_maps(v) if k != 'type' else v) 
-                for k, v in obj.items() 
-            }
+            return {k: (_to_isl_maps(v) if k != "type" else v) for k, v in obj.items()}
         if isinstance(obj, list):
             return [_to_isl_maps(v) for v in obj]
         return obj