Skip to content

Commit 11f10da

Browse files
committed
Redo loopy loop priorities
The previous solution to avoiding loop interchanges did not account for having multiple loop nests.
1 parent 796c0e8 commit 11f10da

File tree

1 file changed

+10
-4
lines changed

1 file changed

+10
-4
lines changed

tsfc/loopy.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,7 @@ def __init__(self, target=None):
126126
self.gem_to_pymbolic = {} # gem node -> pymbolic variable
127127
self.name_gen = UniqueNameGenerator()
128128
self.target = target
129+
self.loop_priorities = set() # used to avoid disadvantageous loop interchanges
129130

130131
def fetch_multiindex(self, multiindex):
131132
indices = []
@@ -191,6 +192,12 @@ def active_inames(self):
191192
# Return all active indices
192193
return frozenset([i.name for i in self.active_indices.values()])
193194

195+
def save_loop_ordering(self):
196+
"""Save the active loops to prevent loop reordering."""
197+
priority = tuple(map(str, self.active_indices.values()))
198+
if len(priority) > 1:
199+
self.loop_priorities.add(priority)
200+
194201

195202
@contextmanager
196203
def active_indices(mapping, ctx):
@@ -199,6 +206,7 @@ def active_indices(mapping, ctx):
199206
:arg ctx: code generation context.
200207
:returns: new code generation context."""
201208
ctx.active_indices.update(mapping)
209+
ctx.save_loop_ordering()
202210
yield ctx
203211
for key in mapping:
204212
ctx.active_indices.pop(key)
@@ -261,12 +269,10 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name
261269
seq_dependencies=True,
262270
silenced_warnings=["summing_if_branches_ops"],
263271
lang_version=(2018, 2),
264-
preambles=preamble
272+
preambles=preamble,
273+
loop_priority=frozenset(ctx.loop_priorities),
265274
)
266275

267-
# Prevent loopy interchange by loopy
268-
knl = lp.prioritize_loops(knl, ",".join(ctx.index_extent.keys()))
269-
270276
return knl, event_name
271277

272278

0 commit comments

Comments
 (0)