[looptree] Temporary symbolic reuse analysis compiler

Author: Michael Gilbert

pytimeloop/looptree/reuse/summarized/__init__.py

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

pytimeloop/looptree/reuse/summarized/compiler.py

@@ -12,3 +12,17 @@ def lambdify(d, tile_shapes):
             k: sympy.lambdify(tile_shapes, v)
             for k, v in d.items()
         }
+
+
+def compile_analysis_result(result, tile_shapes):
+    lambdify_with_tile_shapes = lambda x: lambdify(x, tile_shapes)
+
+    result.ops = lambdify_with_tile_shapes(result.ops)
+    result.temporal_steps = lambdify_with_tile_shapes(result.temporal_steps)
+    result.fanout = lambdify_with_tile_shapes(result.fanout)
+    result.occupancy = lambdify_with_tile_shapes(result.occupancy)
+    result.fills = lambdify_with_tile_shapes(result.fills)
+    result.reads_to_parent = lambdify_with_tile_shapes(result.reads_to_parent)
+    result.op_intensity = lambdify_with_tile_shapes(result.op_intensity)
+
+    return result

pytimeloop/looptree/reuse/summarized/symbolic.py

@@ -232,15 +232,3 @@ def analyze_reuse(mapping,
     output.fanout = fanout
 
     return tile_shapes, output
-
-    lambdify_with_tile_shapes = lambda x: lambdify(x, tile_shapes)
-
-    output.ops = lambdify_with_tile_shapes(output.ops)
-    output.temporal_steps = lambdify_with_tile_shapes(output.temporal_steps)
-    output.fanout = lambdify_with_tile_shapes(output.fanout)
-    output.occupancy = lambdify_with_tile_shapes(output.occupancy)
-    output.fills = lambdify_with_tile_shapes(output.fills)
-    output.reads_to_parent = lambdify_with_tile_shapes(output.reads_to_parent)
-    output.op_intensity = lambdify_with_tile_shapes(output.op_intensity)
-
-    return tile_shapes, output

tests/looptree/test_symbolic_reuse_analysis.py renamed to tests/looptree/reuse_analysis/test_symbolic.py

@@ -1,11 +1,14 @@
 import unittest
 from sympy import ceiling
+from math import ceil
+from pprint import pp
 
 from bindings.looptree import LooptreeWorkload, LooptreeWorkloadDependencyAnalyzer
 
 from tests.load_config_mixin import LoadConfigMixin
 
 from pytimeloop.looptree.reuse.summarized.symbolic import analyze_reuse
+from pytimeloop.looptree.reuse.summarized import compile_analysis_result
 
 
 class TestSymbolicReuseAnalysis(unittest.TestCase, LoadConfigMixin):
@@ -37,3 +40,37 @@ def test_model_with_two_level_mm(self):
             self.assertEqual(result.fills[key],
                              ref_value,
                              f'fills for {key} do not match')
+
+
+class TestCompiler(unittest.TestCase, LoadConfigMixin):
+    def test_model_with_two_level_mm(self):
+        config, spec = self.load_config([
+            'symbolic-mapping.yaml',
+            'cascaded_mm.workload.yaml',
+            'three_level.arch.yaml'
+        ])
+
+        mapping = spec['mapping']['nodes']
+        workload = LooptreeWorkload.parse_cfg(config.root['problem'])
+        analyzer = LooptreeWorkloadDependencyAnalyzer(workload)
+
+        tile_shapes, result = analyze_reuse(mapping, workload, analyzer)
+        self.assertEqual(len(tile_shapes), 3)
+        compiled_result = compile_analysis_result(result, tile_shapes)
+
+        for M1_tile_shape_val in [1, 2, 4]:
+
+            REFERENCE_FILLS = {
+                ('DRAM', 0, 0): 18,
+                ('DRAM', 1, 0): 8,
+                ('DRAM', 2, 0): 36,
+                ('GlobalBuffer', 0, 0): 18.0*ceil(4/M1_tile_shape_val),
+                ('GlobalBuffer', 1, 0): 8
+            }
+
+            for key, ref_value in REFERENCE_FILLS.items():
+                self.assertEqual(
+                    compiled_result.fills[key][1](1, 1, M1_tile_shape_val),
+                    ref_value,
+                    f'fills for {key} do not match'
+                )