From 754c776707424bc8155ef7f2494c0a1a90a53ad7 Mon Sep 17 00:00:00 2001
From: Martin Bauer <martin.bauer@fau.de>
Date: Tue, 11 Jun 2019 11:14:31 +0200
Subject: [PATCH] Bugfix - Vectorization of in-place LB update wrong
Block.subs method tried to be too smart:
a = field[..]
b = a + b
was "simplified" incorrectly to
b = field[...] + b
---
pystencils/astnodes.py | 10 ++--------
pystencils_tests/test_vectorization.py | 20 ++++++++++++++++++++
2 files changed, 22 insertions(+), 8 deletions(-)
diff --git a/pystencils/astnodes.py b/pystencils/astnodes.py
index c08c137..b1ed610 100644
--- a/pystencils/astnodes.py
+++ b/pystencils/astnodes.py
@@ -212,6 +212,8 @@ class KernelFunction(Node):
argument_symbols = self._body.undefined_symbols - self.global_variables
parameters = [self.Parameter(symbol, get_fields(symbol)) for symbol in argument_symbols]
+ if hasattr(self, 'indexing'):
+ parameters += [self.Parameter(s, []) for s in self.indexing.symbolic_parameters()]
parameters.sort(key=lambda p: p.symbol.name)
return parameters
@@ -252,14 +254,6 @@ class Block(Node):
return self._nodes
def subs(self, subs_dict) -> None:
- new_args = []
- for a in self.args:
- if isinstance(a, SympyAssignment) and a.is_declaration and a.rhs in subs_dict.keys():
- subs_dict[a.lhs] = subs_dict[a.rhs]
- else:
- new_args.append(a)
- self._nodes = new_args
-
for a in self.args:
a.subs(subs_dict)
diff --git a/pystencils_tests/test_vectorization.py b/pystencils_tests/test_vectorization.py
index 2acded5..2f7d0f9 100644
--- a/pystencils_tests/test_vectorization.py
+++ b/pystencils_tests/test_vectorization.py
@@ -25,6 +25,26 @@ def test_vector_type_propagation():
np.testing.assert_equal(dst[1:-1, 1:-1], 2 * 10.0 + 3)
+def test_inplace_update():
+ shape = (9, 9, 3)
+ arr = np.ones(shape, order='f')
+
+ @ps.kernel
+ def update_rule(s):
+ f = ps.fields("f(3) : [2D]", f=arr)
+ s.tmp0 @= f(0)
+ s.tmp1 @= f(1)
+ s.tmp2 @= f(2)
+ f0, f1, f2 = f(0), f(1), f(2)
+ f0 @= 2 * s.tmp0
+ f1 @= 2 * s.tmp0
+ f2 @= 2 * s.tmp0
+
+ ast = ps.create_kernel(update_rule, cpu_vectorize_info={'instruction_set': 'avx'})
+ kernel = ast.compile()
+ kernel(f=arr)
+ np.testing.assert_equal(arr, 2)
+
def test_vectorization_fixed_size():
configurations = []
# Fixed size - multiple of four
--
GitLab