From 2d361bc0b2a231b658b4bfdbd4fe7500f17b3083 Mon Sep 17 00:00:00 2001 From: Michael Lange Date: Thu, 14 Nov 2024 09:07:52 +0000 Subject: [PATCH] Parallel: Add missing recursion on remove_block_loops --- loki/transformations/parallel/block_loop.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/loki/transformations/parallel/block_loop.py b/loki/transformations/parallel/block_loop.py index 033e27a8a..1926152b6 100644 --- a/loki/transformations/parallel/block_loop.py +++ b/loki/transformations/parallel/block_loop.py @@ -47,14 +47,16 @@ class RemoveBlockLoopTransformer(Transformer): """ def visit_Loop(self, loop, **kwargs): # pylint: disable=unused-argument + body = self.visit(loop.body, **kwargs) + if not loop.variable == idx: - return loop + return loop._rebuild(body=body) to_remove = tuple( - a for a in FindNodes(ir.Assignment).visit(loop.body) + a for a in FindNodes(ir.Assignment).visit(body) if a.lhs in variables ) - return tuple(n for n in loop.body if n not in to_remove) + return tuple(n for n in body if n not in to_remove) routine.body = RemoveBlockLoopTransformer().visit(routine.body)