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
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 committed Nov 27, 2024
1 parent e06308d commit 8961aa5
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 8961aa5

Please sign in to comment.