From c0d959cb429a285d1e0715e10afeef33aab81b8f Mon Sep 17 00:00:00 2001
From: Daniel Bauer <daniel.j.bauer@fau.de>
Date: Thu, 9 Nov 2023 08:19:44 +0100
Subject: [PATCH] apply vector substitutions in all ast nodes

---
 pystencils/cpu/vectorization.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/pystencils/cpu/vectorization.py b/pystencils/cpu/vectorization.py
index ec1137ce2..f49641094 100644
--- a/pystencils/cpu/vectorization.py
+++ b/pystencils/cpu/vectorization.py
@@ -429,8 +429,12 @@ def insert_vector_casts(ast_node, instruction_set, loop_counter_symbol, default_
                 arg.condition_expr = visit_expr(arg.condition_expr, default_type)
                 visit_node(arg, substitution_dict, default_type)
             elif isinstance(arg, ast.ForLoop): # ForLoop is intended not to be vectorized but to set up small finite loops within kernels
-                    visit_node(arg.body, substitution_dict, default_type)
+                visit_node(arg.body, substitution_dict, default_type)
             else:
                 visit_node(arg, substitution_dict, default_type)
+        node = fast_subs(
+            node, substitution_dict,
+            skip=lambda e: isinstance(e, (ast.SympyAssignment, ast.Conditional, ast.ForLoop, ast.ResolvedFieldAccess))
+        )
 
     visit_node(ast_node, {}, default_float_type)
-- 
GitLab