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