From 4a790e1c48f32c07fc4058de9b20734bcea9cca0 Mon Sep 17 00:00:00 2001
From: Nils Kohl <nils.kohl@fau.de>
Date: Mon, 3 Jun 2024 21:21:08 +0200
Subject: [PATCH] Removing the replacement of loop nodes with blocks for loops
of length one in 5f3be7b3 disabled safety mechanism that ensures that the new
cutting point is actually in the loop bounds. It's the 'with_conditional'
parameters. Reintroduced it in this commit - now both can be combined.
---
pystencils/transformations.py | 16 ++++++++++++----
1 file changed, 12 insertions(+), 4 deletions(-)
diff --git a/pystencils/transformations.py b/pystencils/transformations.py
index 2d8210a1b..8ac7b7230 100644
--- a/pystencils/transformations.py
+++ b/pystencils/transformations.py
@@ -901,14 +901,22 @@ def cut_loop(loop_node, cutting_points, with_conditional: bool = False, replace_
new_start = loop_node.start
cutting_points = list(cutting_points) + [loop_node.stop]
for new_end in cutting_points:
- if replace_loops_with_length_one and (new_end - new_start == 1):
+ if new_end - new_start == 1:
new_body = deepcopy(loop_node.body)
- new_body.subs({loop_node.loop_counter_symbol: new_start})
+ if replace_loops_with_length_one:
+ new_body.subs({loop_node.loop_counter_symbol: new_start})
if with_conditional:
conditional_expr = sp.And(sp.Ge(new_start, loop_node.start), sp.Le(new_start, loop_node.stop))
- new_loops.append(ast.Conditional(conditional_expr, new_body))
+ new_body_wrapped = ast.Block(ast.Conditional(conditional_expr, new_body))
else:
- new_loops.append(new_body)
+ new_body_wrapped = new_body
+ if not replace_loops_with_length_one:
+ new_loop = ast.LoopOverCoordinate(
+ new_body_wrapped, loop_node.coordinate_to_loop_over,
+ new_start, new_end, loop_node.step)
+ else:
+ new_loop = new_body_wrapped
+ new_loops.append(new_loop)
elif new_end - new_start == 0:
pass
else:
--
GitLab