From e99e8a932efa97662d72d46aa90a4b67bc1e519b Mon Sep 17 00:00:00 2001
From: Martin Bauer <martin.bauer@fau.de>
Date: Mon, 21 Oct 2019 13:35:25 +0200
Subject: [PATCH] Bugfix in new mask_conditionals (vectorization)
---
pystencils/cpu/vectorization.py | 7 +++++--
1 file changed, 5 insertions(+), 2 deletions(-)
diff --git a/pystencils/cpu/vectorization.py b/pystencils/cpu/vectorization.py
index bfd90ffc4..c5080aa20 100644
--- a/pystencils/cpu/vectorization.py
+++ b/pystencils/cpu/vectorization.py
@@ -144,13 +144,16 @@ def mask_conditionals(loop_body):
def visit_node(node, mask):
if isinstance(node, ast.Conditional):
cond = node.condition_expr
- cond = cond if loop_body.loop_counter_symbol in cond.atoms(sp.Symbol) else True
+ skip = (loop_body.loop_counter_symbol not in cond.atoms(sp.Symbol)) or cond.func in (vec_all, vec_any)
+ cond = True if skip else cond
+
true_mask = sp.And(cond, mask)
visit_node(node.true_block, true_mask)
if node.false_block:
false_mask = sp.And(sp.Not(node.condition_expr), mask)
visit_node(node, false_mask)
- node.condition_expr = vec_any(node.condition_expr)
+ if not skip:
+ node.condition_expr = vec_any(node.condition_expr)
elif isinstance(node, ast.SympyAssignment):
if mask is not True:
s = {ma: vector_memory_access(ma.args[0], ma.args[1], ma.args[2], ma.args[3], sp.And(mask, ma.args[4]))
--
GitLab