diff --git a/pystencils/backends/cbackend.py b/pystencils/backends/cbackend.py index 9efca1f53a6870d2b0ce2bdd1da547d16beb6a6c..cf58431309e1ddb0ba963fa808f9db5261448091 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 f99c77f2d95f1d45191726101718beff7c9d3b0e..12b492cf452e27220b3013ac6339aeaf6073c051 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]