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