Update fastfusion; Update test configs

pytimeloop/fastfusion/mapper/__init__.py

@@ -1 +0,0 @@
-from .mapper import *

pytimeloop/fastfusion/mapper/level_mapper/top_level.py

@@ -2,9 +2,10 @@
 from collections import defaultdict
 from functools import reduce
 from itertools import permutations, product
-from more_itertools import powerset
 from operator import mul
 
+from more_itertools import powerset
+
 import pandas as pd
 
 from pytimeloop.fastfusion.mapper.shape_subspace import ShapeSubspace

pytimeloop/fastfusion/mapper/mapper.py

@@ -4,7 +4,10 @@
 from ruamel.yaml import YAML
 yaml = YAML(typ='safe')
 
-from bindings.looptree import LooptreeWorkload
+from bindings.looptree import (
+    LooptreeWorkload,
+    LooptreeWorkloadDependencyAnalyzer
+)
 
 from pytimeloop.fastfusion.pareto import OpData, Pareto
 
@@ -14,14 +17,42 @@
 from .stepped_model import Stats, SteppedModel, SteppedModelState
 
 
-def mapper(config,
-           name_of_einsum_to_eval,
-           fusable_tensors,
-           neighbors,
-           workload,
-           analyzer,
-           spec,
-           tmp_path):
+def mapper(config, spec, tmp_path):
+    workload = LooptreeWorkload.parse_cfg(config.root['problem'])
+    analyzer = LooptreeWorkloadDependencyAnalyzer(workload)
+
+    adj_list = get_neighbors(workload)
+
+    einsum_name_to_id = workload.einsum_name_to_id()
+    einsum_to_initial_data = {}
+    for einsum_name, einsum_id in einsum_name_to_id.items():
+        tensors = (
+            workload.tensors_read_by_einsum(einsum_id)
+            |
+            workload.tensors_written_by_einsum(einsum_id)
+        )
+        fusable_tensors = tensors & get_intermediate_tensors(workload)
+
+        data = layer_mapper(config,
+                            einsum_name,
+                            fusable_tensors,
+                            adj_list[einsum_id],
+                            workload,
+                            analyzer,
+                            spec,
+                            tmp_path)
+        einsum_to_initial_data[einsum_name] = data
+    return einsum_to_initial_data
+
+
+def layer_mapper(config,
+                 name_of_einsum_to_eval,
+                 fusable_tensors,
+                 neighbors,
+                 workload,
+                 analyzer,
+                 spec,
+                 tmp_path):
     einsum_name_to_id = workload.einsum_name_to_id()
     id_of_einsum_to_eval = einsum_name_to_id[name_of_einsum_to_eval]
 

pytimeloop/fastfusion/mapper/shape_subspace.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Mapping, Sequence
 
 from combinatorics.integer import integer_factorizations_to_n_parts
 

pytimeloop/fastfusion/mapper/stepped_model.py

@@ -1,23 +1,14 @@
 from dataclasses import dataclass
 from pathlib import Path
-from time import time
 
 from ruamel.yaml import YAML
 yaml = YAML(typ='safe')
 
-import islpy as isl
-
-from bindings.config import Config
-from bindings.looptree import LooptreeModelApp
-
-from pytimeloop.looptree.des import deserialize_looptree_output
 from pytimeloop.looptree.energy import gather_actions, compute_energy_from_actions
-from pytimeloop.looptree.latency import compute_latency
 from pytimeloop.looptree.fastmodel import run_fastmodel
 
 from pytimeloop.timeloopfe.v4 import Ert
 from pytimeloop.timeloopfe.common.backend_calls import call_accelergy_verbose
-from pytimeloop.timeloopfe.common.version_transpilers import v4_to_v3
 
 
 @dataclass

tests/fastfusion/test_mapper.py

@@ -1,8 +1,8 @@
 import unittest
 
-from bindings.looptree import LooptreeWorkload, LooptreeWorkloadDependencyAnalyzer
+from bindings.looptree import LooptreeWorkload
 
-from pytimeloop.fastfusion.mapper import mapper, get_neighbors, get_intermediate_tensors
+from pytimeloop.fastfusion.mapper.mapper import mapper
 from pytimeloop.fastfusion.mapper.shape_subspace import ShapeSubspace
 
 from tests.load_config_mixin import LoadConfigMixin
@@ -15,31 +15,11 @@ def test_mapper(self):
             'cascaded_mm.workload.yaml',
             'three_level.arch.yaml'
         ])
-
-        workload = LooptreeWorkload.parse_cfg(config.root['problem'])
-        analyzer = LooptreeWorkloadDependencyAnalyzer(workload)
-        NAME_OF_EINSUM_TO_MAP = 'Fc1'
-        einsum_name_to_id = workload.einsum_name_to_id()
-        id_of_einsum_to_eval = einsum_name_to_id[NAME_OF_EINSUM_TO_MAP]
-        tensors = (
-            workload.tensors_read_by_einsum(id_of_einsum_to_eval)
-            |
-            workload.tensors_written_by_einsum(id_of_einsum_to_eval)
-        )
-        fusable_tensors = tensors & get_intermediate_tensors(workload)
-
-        adj_list = get_neighbors(workload)
-
+        print('start')
 
         result = mapper(config,
-                        NAME_OF_EINSUM_TO_MAP,
-                        fusable_tensors,
-                        adj_list[id_of_einsum_to_eval],
-                        workload,
-                        analyzer,
                         spec,
                         TEST_TMP_DIR)
-        print(result)
 
 
 class TestShapeSubspace(LoadConfigMixin, unittest.TestCase):

tests/looptree/make_model_app.py

@@ -11,7 +11,7 @@
 def make_model_app(config_dir: Path, paths: list, tmp_path, call_accelergy: bool = True):
     yaml_str = gather_yaml_configs(config_dir, paths)
     config = Config(yaml_str, 'yaml')
-    model = LooptreeModelApp(config, str(tmp_path), 'looptree-model')
+    model = LooptreeModelApp(config)
 
     spec = Specification.from_yaml_files([
         str(config_dir / p) for p in paths

tests/looptree/test_accesses.py

@@ -12,14 +12,19 @@
 from pytimeloop.looptree.accesses import *
 
 from tests.util import TEST_TMP_DIR, gather_yaml_configs
+from .make_model_app import make_model_app
 
 
 class TestLoopTreeAccess(unittest.TestCase):
     def test_accesses_with_two_level_mm_fused(self):
         self.maxDiff = None
         self.check_accesses(
             Path(__file__).parent.parent / 'test_configs',
-            ['looptree-test-fused.yaml'],
+            [
+                'looptree-test-fused.yaml',
+                'cascaded_mm.workload.yaml',
+                'three_level.arch.yaml'
+            ],
             TEST_TMP_DIR,
             {
                 (0, 'Fmap1', 'Fc1'): 18,
@@ -45,7 +50,11 @@ def test_accesses_with_two_level_mm_fused(self):
     def test_accesses_with_two_level_mm_unfused(self):
         self.check_accesses(
             Path(__file__).parent.parent / 'test_configs',
-            ['looptree-test-unfused.yaml'],
+            [
+                'looptree-test-unfused.yaml',
+                'cascaded_mm.workload.yaml',
+                'three_level.arch.yaml'
+            ],
             TEST_TMP_DIR,
             {
                 (0, 'Fmap1', 'Fc1'): 18,
@@ -85,9 +94,9 @@ def make_model_app(config_dir, paths, tmp_path):
     def check_accesses(
         self, config_dir, paths, tmp_path, read_refs, write_refs
     ):
-        model, config, workload = self.make_model_app(config_dir,
-                                                      paths,
-                                                      tmp_path)
+        model, config, workload = make_model_app(config_dir,
+                                                 paths,
+                                                 tmp_path)
 
         result = model.run()
         result = deserialize_looptree_output(result, isl.DEFAULT_CONTEXT)
@@ -104,6 +113,7 @@ def check_accesses(
 
         reads, writes = reads_and_writes_from_fill_by_peer(
             result.fills_by_peer,
+            config['mapping'],
             workload
         )
         reads = get_total_accesses(reads)

tests/looptree/test_capacity.py

@@ -14,7 +14,11 @@ class TestCapacityAggregators(unittest.TestCase):
     def test_capacity_analysis(self):
         model, spec, workload = make_model_app(
             Path(__file__).parent.parent / 'test_configs',
-            ['looptree-test-fused.yaml'],
+            [
+                'looptree-test-fused.yaml',
+                'cascaded_mm.workload.yaml',
+                'three_level.arch.yaml'
+            ],
             TEST_TMP_DIR,
             False
         )

tests/looptree/test_fastmodel.py

@@ -12,7 +12,9 @@
 class TestLooptreeFastModel(LoadConfigMixin, unittest.TestCase):
     def test_with_fused(self):
         config, spec = self.load_config([
-            'looptree-test-fused.yaml'
+            'looptree-test-fused.yaml',
+            'cascaded_mm.workload.yaml',
+            'three_level.arch.yaml'
         ])
 
         BINDINGS = {

tests/looptree/test_model.py

@@ -14,22 +14,19 @@ class LooptreeModelAppTest(unittest.TestCase):
     def test_model_with_two_level_mm(self):
         self.check_model_app(
             Path(__file__).parent.parent / 'test_configs',
-            ['looptree-test-fused.yaml'],
+            [
+                'looptree-test-fused.yaml',
+                'cascaded_mm.workload.yaml',
+                'three_level.arch.yaml'
+            ],
             TEST_TMP_DIR
         )
 
-    @staticmethod
-    def make_model_app(config_dir, paths, tmp_path):
-        yaml_str = gather_yaml_configs(config_dir, paths)
-        config = Config(yaml_str, 'yaml')
-        model = LooptreeModelApp(config, str(tmp_path), 'looptree-model')
-        return model
-
     def check_model_app(
         self, config_dir, paths, tmp_path
     ):
         self.maxDiff = None
-        model = self.make_model_app(config_dir, paths, tmp_path)
+        model, _, _ = make_model_app(config_dir, paths, tmp_path)
         result = model.run()
 
         def compare_dim(d, dtype):
@@ -59,19 +56,16 @@ class TestLooptreeOutputDeserializer(unittest.TestCase):
     def test_deserializer_with_two_level_mm(self):
         self.check_deserializer(
             Path(__file__).parent.parent / 'test_configs',
-            ['looptree-test-fused.yaml'],
+            [
+                'looptree-test-fused.yaml',
+                'cascaded_mm.workload.yaml',
+                'three_level.arch.yaml'
+            ],
             TEST_TMP_DIR
         )
 
-    @staticmethod
-    def make_model_app(config_dir, paths, tmp_path):
-        yaml_str = gather_yaml_configs(config_dir, paths)
-        return LooptreeModelApp(Config(yaml_str, 'yaml'),
-                                str(tmp_path),
-                                'looptree-model')
-
     def check_deserializer(self, config_dir, paths, tmp_path):
-        model = self.make_model_app(config_dir, paths, tmp_path)
+        model, _, _ = make_model_app(config_dir, paths, tmp_path)
         result = model.run()
 
         des_result = deserialize_looptree_output(result, isl.DEFAULT_CONTEXT)