diff --git a/pystencils/cpu/vectorization.py b/pystencils/cpu/vectorization.py index ec1137ce26fa37fa4d614b691d4566ce72acaf2a..f4964109456ed2df53c1b9562e25e6defa604efe 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)