From 659319bddac44d2f5df545121f6b316089431ae8 Mon Sep 17 00:00:00 2001
From: Michael Kuron <mkuron@icp.uni-stuttgart.de>
Date: Thu, 16 Jan 2020 11:52:05 +0100
Subject: [PATCH] Correctly print sp.Abs in generated code
---
pystencils/backends/cbackend.py | 6 ++++++
pystencils/cpu/vectorization.py | 4 ++++
2 files changed, 10 insertions(+)
diff --git a/pystencils/backends/cbackend.py b/pystencils/backends/cbackend.py
index 9efca1f53..cf5843130 100644
--- a/pystencils/backends/cbackend.py
+++ b/pystencils/backends/cbackend.py
@@ -353,6 +353,12 @@ class CustomSympyPrinter(CCodePrinter):
result = super(CustomSympyPrinter, self)._print_Piecewise(expr)
return result.replace("\n", "")
+ def _print_Abs(self, expr):
+ if expr.is_integer:
+ return 'abs({0})'.format(self._print(expr.args[0]))
+ else:
+ return 'fabs({0})'.format(self._print(expr.args[0]))
+
def _print_Type(self, node):
return str(node)
diff --git a/pystencils/cpu/vectorization.py b/pystencils/cpu/vectorization.py
index f99c77f2d..12b492cf4 100644
--- a/pystencils/cpu/vectorization.py
+++ b/pystencils/cpu/vectorization.py
@@ -176,6 +176,10 @@ def insert_vector_casts(ast_node):
visit_expr(expr.args[4]))
elif isinstance(expr, cast_func):
return expr
+ elif expr.func is sp.Abs:
+ new_arg = visit_expr(expr.args[0])
+ pw = sp.Piecewise((-1 * new_arg, new_arg < 0), (new_arg, True))
+ return visit_expr(pw)
elif expr.func in handled_functions or isinstance(expr, sp.Rel) or isinstance(expr, BooleanFunction):
new_args = [visit_expr(a) for a in expr.args]
arg_types = [get_type_of_expression(a) for a in new_args]
--
GitLab