@@ -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
196203def 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