Merge branch 'main' of github.com:Accelergy-Project/timeloop-python

Author: tanner-andrulis

pytimeloop/fastfusion/fastmodel/fastmodel.py

@@ -133,11 +133,13 @@ def compile_mapping(mapping,
         elif node['type'] == 'storage':
             target = node['target']
             tensor_names = node['dspace']
+            exploits_reuse = 'exploits_reuse' not in node or node['exploits_reuse']
             for tensor_id in tensor_names:
                 if isinstance(tensor_id, int):
                     tensor_id = tensor_id
                 else:
                     tensor_id = tensor_name_to_id[tensor_id]
+
                 if tensor_id not in tensors:
                     continue
 
@@ -152,6 +154,10 @@ def compile_mapping(mapping,
 
                 output.occupancy[(target, tensor_id)] = tensor_size[tensor_id]
 
+                if not exploits_reuse:
+                    actual_tensor_access_multiplier[tensor_id] = \
+                        potential_tensor_access_multiplier[tensor_id]
+
                 output.fills[(target, tensor_id, einsum_id)] = (
                     None,
                     (
@@ -173,6 +179,8 @@ def compile_mapping(mapping,
 
                 actual_tensor_access_multiplier[tensor_id] *= \
                     fill_multicast_factor[tensor_id]
+                potential_tensor_access_multiplier[tensor_id] *= \
+                    fill_multicast_factor[tensor_id]
                 fill_multicast_factor[tensor_id] = 1
 
                 if target not in fanout:

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

@@ -66,8 +66,11 @@ def add_sequential(self, idx=None):
     def add_pipeline(self):
         self.mapping.append({"type": "pipeline"})
 
-    def add_storage(self, target, dspaces, idx=None):
-        node = {"type": "storage", "target": target, "dspace": dspaces}
+    def add_storage(self, target, dspaces, idx=None, exploits_reuse=False):
+        node = {"type": "storage",
+                "target": target,
+                "dspace": dspaces,
+                "exploits_reuse": exploits_reuse}
         if idx is None:
             self.mapping.append(node)
         else: