[looptree] Symbolic model can produce latency and energy

Author: Michael Gilbert

pytimeloop/looptree/accesses.py

@@ -1,7 +1,6 @@
-from collections import defaultdict
-from collections.abc import Mapping
 from dataclasses import dataclass
-from numbers import Number
+from functools import reduce
+from operator import mul
 from typing import Optional, overload
 
 from bindings.looptree import TemporalTag, SequentialTag, PipelineTemporalTag
@@ -61,10 +60,12 @@ def summarize_total_and_per_unit_actions(
     reuse_analysis_result
 ):
     result = {}
+    reads_to_parent = reuse_analysis_result.reads_to_parent
+    reads_to_peer = reuse_analysis_result.reads_to_peer
     if isinstance(reuse_analysis_result, IslReuseAnalysisOutput):
         for key, (tags, fill) in reuse_analysis_result.fills.items():
-            read_to_parent = reuse_analysis_result.reads_to_parent[key][1]
-            read_to_peer = reuse_analysis_result.reads_to_peer[key][1]
+            read_to_parent = reads_to_parent[key][1]
+            read_to_peer = reads_to_peer[key][1]
 
             total_fill = get_sum_of_pw_qpolynomial(fill)
             total_read_to_parent = get_sum_of_pw_qpolynomial(read_to_parent)
@@ -89,20 +90,28 @@ def summarize_total_and_per_unit_actions(
                            max_per_unit_read_to_peer)
     elif isinstance(reuse_analysis_result, SummarizedAnalysisOutput):
         for key, (tags, fill) in reuse_analysis_result.fills.items():
-            buffer_id = key[0]
+            buffer_name = key[0]
 
-            read_to_parent = reuse_analysis_result.reads_to_parent[key][1]
-            read_to_peer = reuse_analysis_result.reads_to_peer[key][1]
+            if key in reads_to_parent:
+                read_to_parent = reads_to_parent[key][1]
+            else:
+                read_to_parent = 0
+
+            if key in reads_to_peer:
+                read_to_peer = reads_to_peer[key][1]
+            else:
+                read_to_peer = 0
 
             total_fill = fill
             total_read_to_parent = read_to_parent
             total_read_to_peer = read_to_peer
 
-            fanout = reuse_analysis_result.fanout[buffer_id]
+            fanout = reuse_analysis_result.fanout[buffer_name]
+            total_fanout = reduce(mul, fanout, 1)
 
-            max_per_unit_fill = fill / fanout
-            max_per_unit_read_to_parent = read_to_parent / fanout
-            max_per_unit_read_to_peer = read_to_peer / fanout
+            max_per_unit_fill = fill / total_fanout
+            max_per_unit_read_to_parent = read_to_parent / total_fanout
+            max_per_unit_read_to_peer = read_to_peer / total_fanout
 
             result[key] = (total_fill,
                            total_read_to_parent,

pytimeloop/looptree/energy.py

@@ -7,7 +7,7 @@
 from pytimeloop.looptree.mapping_utilities import *
 
 
-def gather_actions(looptree_results, mapping, workload, bindings, is_path=False):
+def gather_actions(looptree_results, mapping, workload, bindings, is_path=False, use_name=False):
     einsum_name_to_id = workload.einsum_name_to_id()
 
     einsums_with_complete_mapping = \
@@ -23,7 +23,11 @@ def gather_actions(looptree_results, mapping, workload, bindings, is_path=False)
                                                          is_path)
     actions = {}
     for (buf, tensor, einsum), accesses in accesses_stats.items():
-        buf = bindings[buf]
+        if use_name:
+            buf = buf
+        else:
+            buf = bindings[buf]
+
         key = (buf, 'read')
         if key not in actions:
             actions[key] = 0
@@ -65,7 +69,7 @@ def gather_ops(ops, einsums_with_complete_mapping):
             continue
         if isinstance(v, isl.PwQPolynomial):
             total += get_sum_of_pw_qpolynomial(v)
-        elif isinstance(v, Number):
+        elif isinstance(v, Real):
             total += v
         else:
             total += v

pytimeloop/looptree/latency/latency.py

@@ -1,4 +1,5 @@
 from typing import overload
+from sympy import Max
 
 from pytimeloop.isl.singular import get_value_from_singular_qpolynomial
 from pytimeloop.looptree.latency.processors import LATENCY_PROCESSORS
@@ -9,7 +10,7 @@
 from bindings.looptree import SpatialTag
 
 
-def get_latency(looptree_results: IslReuseAnalysisOutput,
+def get_latency(looptree_results,
                 mapping,
                 workload,
                 arch,
@@ -22,7 +23,7 @@ def get_latency(looptree_results: IslReuseAnalysisOutput,
                                  mapping,
                                  workload,
                                  bindings)
-    overall_latency = max(comp_latency, max(mem_latency.values()))
+    overall_latency = Max(comp_latency, Max(*mem_latency.values()))
     return overall_latency, comp_latency, mem_latency
 
 
@@ -57,7 +58,7 @@ def compute_isl_latency(temporal_steps, mapping, workload):
 
 def compute_summarized_latency(temporal_steps, mapping, workload):
     # TODO: this is only for single-Einsum!!!
-    return sum(value for key, value in temporal_steps)
+    return sum(value for key, value in temporal_steps.items())
 
 
 def _compute_latency(mapping, top_idx: int, temporal_steps, workload):

pytimeloop/looptree/latency/memory.py

@@ -1,4 +1,5 @@
 from collections import defaultdict
+from sympy import Max, Min
 
 from pytimeloop.looptree.accesses import buffer_accesses_from_buffet_actions
 from pytimeloop.looptree.reuse.isl import IslReuseAnalysisOutput
@@ -21,6 +22,8 @@ def memory_latency(
 
     component_to_read_writes = defaultdict(lambda: [None, None])
     for level, component in bindings.items():
+        if isinstance(looptree_results, SummarizedAnalysisOutput):
+            level = component
         read_count = 0
         write_count = 0
         for key, accesses in accesses_stats.items_with_buffer(level):
@@ -47,21 +50,23 @@ def memory_latency(
         # All shared bw for writing
         write_latency = writes / (write_bw + shared_bw)
         read_latency = reads / read_bw
-        if write_latency >= read_latency:
-            component_latency[component] = write_latency
-            continue
+        all_shared_for_write_latency = Max(write_latency, read_latency)
+
         # All shared bw for reading
         write_latency = writes / write_bw
         read_latency = reads / (read_bw + shared_bw)
-        if read_latency >= write_latency:
-            component_latency[component] = read_latency
-            continue
+        all_shared_for_read_latency = Max(write_latency, read_latency)
+
         # Shared bw shared for reading and writing
-        component_latency[component] = (
+        shared_for_read_and_write_latency = (
             (reads + writes)
             / 
             (read_bw + write_bw + shared_bw)
         )
+
+        component_latency[component] = Min(all_shared_for_write_latency,
+                                           all_shared_for_read_latency,
+                                           shared_for_read_and_write_latency)
     return component_latency
 
 

pytimeloop/looptree/reuse/summarized/__init__.py

@@ -1,2 +1,2 @@
-from .symbolic import SummarizedAnalysisOutput
+from .symbolic import SummarizedAnalysisOutput, analyze_reuse
 from .compiler import compile_analysis_result

pytimeloop/looptree/reuse/summarized/symbolic.py

@@ -26,9 +26,11 @@ class SummarizedAnalysisOutput:
     op_intensity: dict = field(default_factory=dict)
 
 
-def analyze_reuse(mapping,
-                  workload: LooptreeWorkload,
-                  analyzer: LooptreeWorkloadDependencyAnalyzer):
+def analyze_reuse(
+    mapping,
+    workload: LooptreeWorkload,
+    analyzer: LooptreeWorkloadDependencyAnalyzer
+) -> SummarizedAnalysisOutput:
     einsum_name_to_id = workload.einsum_name_to_id()
     rank_name_to_id = workload.dimension_name_to_id()
     tensor_name_to_id = workload.data_space_name_to_id()
@@ -73,7 +75,7 @@ def analyze_reuse(mapping,
 
     tile_shapes = []
 
-    output = IslReuseAnalysisOutput()
+    output = SummarizedAnalysisOutput()
 
     latency = 1
     potential_tensor_access_multiplier = defaultdict(lambda: 1)

pytimeloop/looptree/run.py

@@ -4,12 +4,13 @@
 import islpy as isl
 
 from bindings.config import Config
-from bindings.looptree import LooptreeModelApp, LooptreeWorkload
+from bindings.looptree import LooptreeModelApp, LooptreeWorkload, LooptreeWorkloadDependencyAnalyzer
 
 from pytimeloop.file import gather_yaml_configs
 
 from pytimeloop.looptree.capacity import compute_capacity_usage
 from pytimeloop.looptree.reuse.isl.des import deserialize_looptree_output
+from pytimeloop.looptree.reuse.summarized import analyze_reuse
 from pytimeloop.looptree.energy import gather_actions, compute_energy_from_actions
 from pytimeloop.looptree.latency import get_latency
 
@@ -69,3 +70,49 @@ def run_looptree(config_dir, paths, tmp_path, bindings, call_accelergy):
                               actions,
                               mem_latency,
                               capacity_usage=component_capacity_usage)
+
+
+def run_looptree_symbolic(config_dir, paths, tmp_path, bindings, call_accelergy):
+    yaml_str = gather_yaml_configs(config_dir, paths)
+
+    config = Config(yaml_str, 'yaml')
+    workload = LooptreeWorkload.parse_cfg(config.root['problem'])
+    analyzer = LooptreeWorkloadDependencyAnalyzer(workload)
+
+    spec = Specification.from_yaml_files([
+        str(config_dir / p) for p in paths
+    ])
+
+    if call_accelergy:
+        if isinstance(tmp_path, Path):
+            tmp_path = str(tmp_path)
+        call_accelergy_verbose(spec, tmp_path)
+        spec = Specification.from_yaml_files([
+            str(config_dir / p) for p in paths
+        ] + [str(Path(tmp_path) / 'ERT.yaml')])
+
+
+    tile_shapes, result = analyze_reuse(spec.mapping.nodes, workload, analyzer)
+
+    actions = gather_actions(result, spec.mapping, workload, bindings, use_name=True)
+    energy = compute_energy_from_actions(actions, spec.ERT)
+
+    latency, comp_latency, mem_latency = get_latency(result,
+                                                     spec.mapping,
+                                                     workload,
+                                                     spec.architecture,
+                                                     bindings)
+
+    capacity_usage = compute_capacity_usage(spec.mapping.nodes,
+                                            result.occupancy,
+                                            workload)
+    component_capacity_usage = {}
+    for level, component in bindings.items():
+        if level in capacity_usage:
+            component_capacity_usage[component] = capacity_usage[level]
+
+    return LoopTreeStatistics(latency,
+                              energy,
+                              actions,
+                              mem_latency,
+                              capacity_usage=component_capacity_usage)

tests/looptree/reuse_analysis/test_symbolic.py

@@ -29,9 +29,9 @@ def test_model_with_two_level_mm(self):
         P1_tile_shape, C1_tile_shape, M1_tile_shape = tile_shapes
 
         REFERENCE_FILLS = {
-            ('DRAM', 0, 0): (None, 18),
-            ('DRAM', 1, 0): (None, 8),
-            ('DRAM', 2, 0): (None, 36),
+            ('MainMemory', 0, 0): (None, 18),
+            ('MainMemory', 1, 0): (None, 8),
+            ('MainMemory', 2, 0): (None, 36),
             ('GlobalBuffer', 0, 0): (None, 18.0*ceiling(4/M1_tile_shape)),
             ('GlobalBuffer', 1, 0): (None, 8)
         }
@@ -61,9 +61,9 @@ def test_model_with_two_level_mm(self):
         for M1_tile_shape_val in [1, 2, 4]:
 
             REFERENCE_FILLS = {
-                ('DRAM', 0, 0): 18,
-                ('DRAM', 1, 0): 8,
-                ('DRAM', 2, 0): 36,
+                ('MainMemory', 0, 0): 18,
+                ('MainMemory', 1, 0): 8,
+                ('MainMemory', 2, 0): 36,
                 ('GlobalBuffer', 0, 0): 18.0*ceil(4/M1_tile_shape_val),
                 ('GlobalBuffer', 1, 0): 8
             }

tests/looptree/test_run.py

@@ -1,8 +1,12 @@
 from pathlib import Path
+from pprint import pp
 import unittest
 
-from pytimeloop.looptree.run import run_looptree
+from bindings.looptree import LooptreeWorkload, LooptreeWorkloadDependencyAnalyzer
 
+from pytimeloop.looptree.run import run_looptree, run_looptree_symbolic
+
+from tests.load_config_mixin import LoadConfigMixin
 from tests.util import TEST_TMP_DIR
 
 
@@ -40,3 +44,34 @@ def test_fused_sequential(self):
 
         for k, v in stats.energy.items():
             self.assertAlmostEqual(ENERGY_REFS[k], v, 1)
+
+
+class TestLooptreeSymbolic(unittest.TestCase, LoadConfigMixin):
+    def test_two_level_mm(self):
+        BINDINGS = {
+            0: 'MainMemory',
+            1: 'GlobalBuffer',
+            2: 'GlobalBuffer',
+            3: 'GlobalBuffer',
+            4: 'MACC'
+        }
+
+        stats = run_looptree_symbolic(
+            Path(__file__).parent.parent / 'test_configs',
+            [
+                'symbolic-mapping.yaml',
+                'cascaded_mm.workload.yaml',
+                'three_level.arch.yaml'
+            ],
+            TEST_TMP_DIR,
+            BINDINGS,
+            True
+        )
+
+        ACTION_REFS = {
+            ('MainMemory', 'read'): 26,
+            ('MainMemory', 'write'): 36,
+            ('MACC', 'compute'): 72
+        }
+        for key, ref_value in ACTION_REFS.items():
+            self.assertEqual(stats.actions[key], ref_value)

tests/test_configs/symbolic-mapping.yaml

@@ -2,7 +2,7 @@ mapping:
   type: fused
   nodes:
   - type: storage
-    target: DRAM
+    target: MainMemory
     dspace: [Filter1, Fmap1, Fmap2]
   - type: storage
     target: GlobalBuffer