[looptree] Refactor: rename LooptreeOutput to IslReuseAnalysisOutput

Author: Michael Gilbert

pytimeloop/fastfusion/fastmodel/fastmodel.py

@@ -5,7 +5,7 @@
 import sympy
 
 from pytimeloop.looptree.equivalent_ranks import EquivalentGroups
-from pytimeloop.looptree.des import LooptreeOutput
+from pytimeloop.looptree.des import IslReuseAnalysisOutput
 
 
 def compile_mapping(mapping,
@@ -58,7 +58,7 @@ def compile_mapping(mapping,
 
     tile_shapes = []
 
-    output = LooptreeOutput()
+    output = IslReuseAnalysisOutput()
 
     latency = 1
     potential_tensor_access_multiplier = defaultdict(lambda: 1)

pytimeloop/fastfusion/mapper/per_einsum_subspaces/subspaces/tile_shape/tile_shape.py

@@ -4,7 +4,7 @@
 
 from .shape_subspace import ShapeSubspace
 
-from pytimeloop.looptree.des import LooptreeOutput
+from pytimeloop.looptree.des import IslReuseAnalysisOutput
 
 
 def explore_tile_shape(
@@ -56,7 +56,7 @@ def explore_tile_shape(
         if only_count:
             continue
 
-        result = LooptreeOutput()
+        result = IslReuseAnalysisOutput()
         result.ops = call_with_arg(compiled_result.ops, shape)
         result.temporal_steps = call_with_arg(compiled_result.temporal_steps, shape)
         result.fanout = call_with_arg(compiled_result.fanout, shape)

pytimeloop/looptree/des.py

@@ -1,34 +1,28 @@
+from dataclasses import dataclass, field
+
 import islpy as isl
 
 import bindings
 
 
-class LooptreeOutput:
-    def __init__(self):
-        self.ops = {}
-        self.fills = {}
-        self.occupancy = {}
-        self.op_occupancy = {}
-        self.reads_to_peer = {}
-        self.reads_to_parent = {}
-        self.temporal_steps = {}
-        self.fanout = {}
-        self.op_intensity = {}
-
-    def __repr__(self):
-        return (
-            f'LooptreeOutput(' +
-            f'ops={self.ops}, ' +
-            f'occupancy={self.occupancy}, ' +
-            f'reads_to_parent={self.reads_to_parent})'
-        )
+@dataclass
+class IslReuseAnalysisOutput:
+    ops: dict = field(default_factory=dict)
+    fills: dict = field(default_factory=dict)
+    occupancy: dict = field(default_factory=dict)
+    op_occupancy: dict = field(default_factory=dict)
+    reads_to_peer: dict = field(default_factory=dict)
+    reads_to_parent: dict = field(default_factory=dict)
+    temporal_steps: dict = field(default_factory=dict)
+    fanout: dict = field(default_factory=dict)
+    op_intensity: dict = field(default_factory=dict)
 
 
 def deserialize_looptree_output(
     looptree_output: bindings.looptree.LooptreeResult,
     isl_ctx: isl.Context
-) -> LooptreeOutput:
-    output = LooptreeOutput()
+) -> IslReuseAnalysisOutput:
+    output = IslReuseAnalysisOutput()
 
     output.ops = {
         k: (dims, isl.PwQPolynomial.read_from_str(isl_ctx, v))

pytimeloop/looptree/latency/latency.py

@@ -6,12 +6,12 @@
     reads_and_writes_from_fill_by_peer
 )
 from pytimeloop.looptree.latency.processors import LATENCY_PROCESSORS
-from pytimeloop.looptree.des import LooptreeOutput
+from pytimeloop.looptree.des import IslReuseAnalysisOutput
 
 from bindings.looptree import SpatialTag
 
 
-def get_latency(looptree_results: LooptreeOutput,
+def get_latency(looptree_results: IslReuseAnalysisOutput,
                 mapping,
                 workload,
                 arch,
@@ -34,7 +34,7 @@ def compute_latency(mapping, temporal_steps, workload):
     ).to_python()
 
 
-def memory_latency(looptree_results: LooptreeOutput,
+def memory_latency(looptree_results: IslReuseAnalysisOutput,
                    arch,
                    mapping,
                    workload,

pytimeloop/looptree/visualization/occupancy.py

@@ -1,9 +1,9 @@
 import matplotlib.pyplot as plt
 
-from pytimeloop.looptree.des import LooptreeOutput
+from pytimeloop.looptree.des import IslReuseAnalysisOutput
 
 
-def plot_occupancy_graph(output: LooptreeOutput, workload):
+def plot_occupancy_graph(output: IslReuseAnalysisOutput, workload):
     einsum_rank_to_shape = {
         einsum: {
             rank: workload.get_rank_shape(rank)

tests/looptree/test_model.py renamed to tests/looptree/test_cpp_reuse_analysis.py

@@ -4,11 +4,10 @@
 
 import islpy as isl
 
-from bindings.config import Config
 from bindings.looptree import *
 from pytimeloop.looptree.des import deserialize_looptree_output
 from .make_model_app import make_model_app
-from tests.util import TEST_TMP_DIR, gather_yaml_configs
+from tests.util import TEST_TMP_DIR
 
 class LooptreeModelAppTest(unittest.TestCase):
     def test_model_with_two_level_mm(self):
@@ -52,7 +51,7 @@ def seconds(ds):
         )
 
 
-class TestLooptreeOutputDeserializer(unittest.TestCase):
+class TestIslReuseAnalysisOutputDeserializer(unittest.TestCase):
     def test_deserializer_with_two_level_mm(self):
         self.check_deserializer(
             Path(__file__).parent.parent / 'test_configs',

tests/test_configs/symbolic-mapping.yaml

@@ -0,0 +1,24 @@
+mapping:
+  type: fused
+  nodes:
+  - type: storage
+    target: 0
+    dspace: [Filter1, Fmap1, Fmap2]
+  - type: storage
+    target: 1
+    dspace: [Filter1]
+  - type: temporal
+    rank: P1
+  - type: storage
+    target: 2
+    dspace: [Fmap2]
+  - type: temporal
+    rank: C1
+  - type: spatial
+    rank: M1
+  - type: storage
+    target: 3
+    dspace: [Fmap1]
+  - type: compute
+    einsum: Fc1
+    target: 4

tests/test_configs/three_level.arch.yaml

@@ -8,7 +8,7 @@ architecture:
   - !Component
     name: MainMemory
     class: DRAM
-    attributes: {width: 256, block_size: 32, word_bits: 8, datawidth: 8}
+    attributes: {depth: 1000000000, width: 256, block_size: 32, word_bits: 8, datawidth: 8}
     required_actions: ['read', 'write']
   - !Component
     name: GlobalBuffer