Skip to content
This repository has been archived by the owner on Dec 6, 2024. It is now read-only.

Commit

Permalink
Redo loopy loop priorities (#333)
Browse files Browse the repository at this point in the history
The previous solution to avoiding loop interchanges did not account for
having multiple loop nests.
  • Loading branch information
connorjward authored Nov 28, 2024
1 parent 796c0e8 commit 7cdfe10
Showing 1 changed file with 10 additions and 4 deletions.
14 changes: 10 additions & 4 deletions tsfc/loopy.py
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,7 @@ def __init__(self, target=None):
self.gem_to_pymbolic = {} # gem node -> pymbolic variable
self.name_gen = UniqueNameGenerator()
self.target = target
self.loop_priorities = set() # used to avoid disadvantageous loop interchanges

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

def save_loop_ordering(self):
"""Save the active loops to prevent loop reordering."""
priority = tuple(map(str, self.active_indices.values()))
if len(priority) > 1:
self.loop_priorities.add(priority)


@contextmanager
def active_indices(mapping, ctx):
Expand All @@ -199,6 +206,7 @@ def active_indices(mapping, ctx):
:arg ctx: code generation context.
:returns: new code generation context."""
ctx.active_indices.update(mapping)
ctx.save_loop_ordering()
yield ctx
for key in mapping:
ctx.active_indices.pop(key)
Expand Down Expand Up @@ -261,12 +269,10 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name
seq_dependencies=True,
silenced_warnings=["summing_if_branches_ops"],
lang_version=(2018, 2),
preambles=preamble
preambles=preamble,
loop_priority=frozenset(ctx.loop_priorities),
)

# Prevent loopy interchange by loopy
knl = lp.prioritize_loops(knl, ",".join(ctx.index_extent.keys()))

return knl, event_name


Expand Down

0 comments on commit 7cdfe10

Please sign in to comment.