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)