Tile matching bugfix

Author: tanner-andrulis

pytimeloop/fastfusion/mapper/per_einsum_subspaces/snowcat.py

@@ -63,7 +63,7 @@ def glb_storage(mapping, unfused_tensors):
                                     explore_uneven=True,
                                     add_split_at_tensors=set(),
                                     must_have_terminal_storage=True,
-                                    apply_lrp_after_loop_idx=last_fused_loop_idx+1)
+                                    apply_lrp_after_loop_idx=last_fused_loop_idx)
 
     def tile_shape_optimization(mapping):
         for partial_mapping in infer_smallest_tile_shape(mapping,

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

@@ -78,7 +78,7 @@ def make_storage(
                 any_irrelevant_loop = any_irrelevant_loop or not is_relevant
 
                 auto_lower = i > apply_lrp_after_loop_idx
-                auto_lower = False
+                # auto_lower = False
 
                 if not auto_lower or (last_is_relevant and not is_relevant):
                     tensor_choices.append(i)
@@ -89,7 +89,7 @@ def make_storage(
                     break
 
         auto_lower = len(mapping) > apply_lrp_after_loop_idx
-        auto_lower = False
+        # auto_lower = False
 
         # Lowest possible storage node
         if (not auto_lower or last_is_relevant) and not (tensor_must_be_fully_reused and any_irrelevant_loop):

pytimeloop/fastfusion/mapper/simexplore.py

@@ -154,7 +154,7 @@ def fuse_sims(
                         print(s)
             elif DO_PRINT:
                 for a in left[k]:
-                    print(f"\tNo match for {k} |||| {a.tiling_str()}")
+                    print(f"\tNo match for {a.tiling}")
 
         if not combined:
             print(f'No valid combinations found.')
@@ -173,7 +173,7 @@ def fuse_sims(
             for k in right:
                 if k not in left:
                     for b in right[k]:
-                        print(f"\tREVERSE: No match for {k} |||| {b.tiling_str()}")
+                        print(f"\tREVERSE: No match for {b.tiling}")
 
         left = combined
         print(f"Number of buckets: {len(left)}")

pytimeloop/fastfusion/sim.py

@@ -137,6 +137,11 @@ class Tiling:
     loops: tuple[Loop, ...]
     tensors: fzs[TensorStorage]
     tags: fzs[Any] = fzs()
+    
+    def __post_init__(self):
+        assert isinstance(self.tensors, frozenset)
+        assert isinstance(self.loops, tuple)
+        assert isinstance(self.tags, frozenset)
 
     @cached_property
     def tensor_names(self) -> set[str]:
@@ -172,7 +177,7 @@ def clear_dead_tensors(
             if keep_loops
             else self.loops[: self.shared_loop_index(live_tensors) + 1]
         )
-        tensors = tuple(t for t in self.tensors if t.tensor_id in live_tensors)
+        tensors = frozenset(t for t in self.tensors if t.tensor_id in live_tensors)
         return Tiling(loops, tensors, self.tags)
 
     def __lt__(self, other):