Skip to content

Commit 92bd55e

Browse files
author
Michael Gilbert
committed
[fastfusion] wip exploits_reuse flag
1 parent 830a109 commit 92bd55e

File tree

2 files changed

+11
-2
lines changed

2 files changed

+11
-2
lines changed

pytimeloop/fastfusion/fastmodel/fastmodel.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,11 +133,13 @@ def compile_mapping(mapping,
133133
elif node['type'] == 'storage':
134134
target = node['target']
135135
tensor_names = node['dspace']
136+
exploits_reuse = 'exploits_reuse' not in node or node['exploits_reuse']
136137
for tensor_id in tensor_names:
137138
if isinstance(tensor_id, int):
138139
tensor_id = tensor_id
139140
else:
140141
tensor_id = tensor_name_to_id[tensor_id]
142+
141143
if tensor_id not in tensors:
142144
continue
143145

@@ -152,6 +154,10 @@ def compile_mapping(mapping,
152154

153155
output.occupancy[(target, tensor_id)] = tensor_size[tensor_id]
154156

157+
if not exploits_reuse:
158+
actual_tensor_access_multiplier[tensor_id] = \
159+
potential_tensor_access_multiplier[tensor_id]
160+
155161
output.fills[(target, tensor_id, einsum_id)] = (
156162
None,
157163
(

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,8 +66,11 @@ def add_sequential(self, idx=None):
6666
def add_pipeline(self):
6767
self.mapping.append({"type": "pipeline"})
6868

69-
def add_storage(self, target, dspaces, idx=None):
70-
node = {"type": "storage", "target": target, "dspace": dspaces}
69+
def add_storage(self, target, dspaces, idx=None, exploits_reuse=False):
70+
node = {"type": "storage",
71+
"target": target,
72+
"dspace": dspaces,
73+
"exploits_reuse": exploits_reuse}
7174
if idx is None:
7275
self.mapping.append(node)
7376
else:

0 commit comments

Comments
 (0)