cbackend.py 23.3 KB
Newer Older
Martin Bauer's avatar
Martin Bauer committed
1
import sympy as sp
Martin Bauer's avatar
Martin Bauer committed
2
3
from collections import namedtuple
from sympy.core import S
4
from typing import Set
5
from sympy.printing.ccode import C89CodePrinter
6

7
from pystencils.cpu.vectorization import vec_any, vec_all
8
9
10
11
12
from pystencils.data_types import (PointerType, VectorType, address_of,
                                   cast_func, create_type, reinterpret_cast_func,
                                   get_type_of_expression,
                                   vector_memory_access)
from pystencils.fast_approximation import fast_division, fast_inv_sqrt, fast_sqrt
13

Martin Bauer's avatar
Martin Bauer committed
14
15
try:
    from sympy.printing.ccode import C99CodePrinter as CCodePrinter
Martin Bauer's avatar
Martin Bauer committed
16
17
except ImportError:
    from sympy.printing.ccode import CCodePrinter  # for sympy versions < 1.1
Martin Bauer's avatar
Martin Bauer committed
18

Martin Bauer's avatar
Martin Bauer committed
19
from pystencils.integer_functions import bitwise_xor, bit_shift_right, bit_shift_left, bitwise_and, \
20
    bitwise_or, modulo_ceil, int_div, int_power_of_2
21
from pystencils.astnodes import Node, KernelFunction
22

23
__all__ = ['generate_c', 'CustomCodeNode', 'PrintNode', 'get_headers', 'CustomSympyPrinter']
24

Martin Bauer's avatar
Martin Bauer committed
25

26
27
KERNCRAFT_NO_TERNARY_MODE = False

Martin Bauer's avatar
Fixes    
Martin Bauer committed
28

29
def generate_c(ast_node: Node, signature_only: bool = False, dialect='c') -> str:
Martin Bauer's avatar
Martin Bauer committed
30
31
32
33
34
35
36
37
38
    """Prints an abstract syntax tree node as C or CUDA code.

    This function does not need to distinguish between C, C++ or CUDA code, it just prints 'C-like' code as encoded
    in the abstract syntax tree (AST). The AST is built differently for C or CUDA by calling different create_kernel
    functions.

    Args:
        ast_node:
        signature_only:
39
        dialect: 'c' or 'cuda'
Martin Bauer's avatar
Martin Bauer committed
40
41
    Returns:
        C-like code for the ast node and its descendants
Martin Bauer's avatar
Martin Bauer committed
42
    """
43
44
45
46
47
48
    global_declarations = get_global_declarations(ast_node)
    for d in global_declarations:
        if hasattr(ast_node, "global_variables"):
            ast_node.global_variables.update(d.symbols_defined)
        else:
            ast_node.global_variables = d.symbols_defined
49
    printer = CBackend(signature_only=signature_only,
50
51
                       vector_instruction_set=ast_node.instruction_set,
                       dialect=dialect)
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
    code = printer(ast_node)
    if not signature_only and isinstance(ast_node, KernelFunction):
        code = "\n" + code
        for declaration in global_declarations:
            code = printer(declaration) + "\n" + code

    return code


def get_global_declarations(ast):
    global_declarations = []

    def visit_node(sub_ast):
        if hasattr(sub_ast, "required_global_declarations"):
            nonlocal global_declarations
            global_declarations += sub_ast.required_global_declarations

        if hasattr(sub_ast, "args"):
            for node in sub_ast.args:
                visit_node(node)

    visit_node(ast)

    return set(global_declarations)
76
77


Martin Bauer's avatar
Martin Bauer committed
78
79
def get_headers(ast_node: Node) -> Set[str]:
    """Return a set of header files, necessary to compile the printed C-like code."""
80
81
    headers = set()

Martin Bauer's avatar
Martin Bauer committed
82
83
84
    if isinstance(ast_node, KernelFunction) and ast_node.instruction_set:
        headers.update(ast_node.instruction_set['headers'])

Martin Bauer's avatar
Martin Bauer committed
85
86
87
    if hasattr(ast_node, 'headers'):
        headers.update(ast_node.headers)
    for a in ast_node.args:
88
        if isinstance(a, Node):
Martin Bauer's avatar
Martin Bauer committed
89
            headers.update(get_headers(a))
90
91

    return headers
92
93


94
95
96
# --------------------------------------- Backend Specific Nodes -------------------------------------------------------


97
class CustomCodeNode(Node):
Martin Bauer's avatar
Martin Bauer committed
98
    def __init__(self, code, symbols_read, symbols_defined, parent=None):
99
        super(CustomCodeNode, self).__init__(parent=parent)
100
        self._code = "\n" + code
101
102
        self._symbols_read = set(symbols_read)
        self._symbols_defined = set(symbols_defined)
103
        self.headers = []
104

105
    def get_code(self, dialect, vector_instruction_set):
106
107
108
109
110
111
112
        return self._code

    @property
    def args(self):
        return []

    @property
Martin Bauer's avatar
Martin Bauer committed
113
    def symbols_defined(self):
114
        return self._symbols_defined
115
116

    @property
Martin Bauer's avatar
Martin Bauer committed
117
    def undefined_symbols(self):
118
        return self._symbols_read - self._symbols_defined
119
120


121
class PrintNode(CustomCodeNode):
Martin Bauer's avatar
Martin Bauer committed
122
123
124
125
    # noinspection SpellCheckingInspection
    def __init__(self, symbol_to_print):
        code = '\nstd::cout << "%s  =  " << %s << std::endl; \n' % (symbol_to_print.name, symbol_to_print.name)
        super(PrintNode, self).__init__(code, symbols_read=[symbol_to_print], symbols_defined=set())
126
        self.headers.append("<iostream>")
127
128
129
130


# ------------------------------------------- Printer ------------------------------------------------------------------

131

Martin Bauer's avatar
Martin Bauer committed
132
133
# noinspection PyPep8Naming
class CBackend:
134

Martin Bauer's avatar
Martin Bauer committed
135
    def __init__(self, sympy_printer=None, signature_only=False, vector_instruction_set=None, dialect='c'):
Martin Bauer's avatar
Martin Bauer committed
136
137
        if sympy_printer is None:
            if vector_instruction_set is not None:
138
                self.sympy_printer = VectorizedCustomSympyPrinter(vector_instruction_set, dialect)
139
            else:
140
                self.sympy_printer = CustomSympyPrinter(dialect)
141
        else:
Martin Bauer's avatar
Martin Bauer committed
142
            self.sympy_printer = sympy_printer
143

144
        self._vector_instruction_set = vector_instruction_set
145
        self._indent = "   "
146
        self._dialect = dialect
Martin Bauer's avatar
Martin Bauer committed
147
        self._signatureOnly = signature_only
148
149

    def __call__(self, node):
Martin Bauer's avatar
Martin Bauer committed
150
        prev_is = VectorType.instruction_set
151
        VectorType.instruction_set = self._vector_instruction_set
152
        result = str(self._print(node))
Martin Bauer's avatar
Martin Bauer committed
153
        VectorType.instruction_set = prev_is
154
        return result
155
156
157

    def _print(self, node):
        for cls in type(node).__mro__:
Martin Bauer's avatar
Martin Bauer committed
158
159
160
161
            method_name = "_print_" + cls.__name__
            if hasattr(self, method_name):
                return getattr(self, method_name)(node)
        raise NotImplementedError("CBackend does not support node of type " + str(type(node)))
162
163

    def _print_KernelFunction(self, node):
164
        function_arguments = ["%s %s" % (str(s.symbol.dtype), s.symbol.name) for s in node.get_parameters()]
165
166
167
168
169
170
171
        launch_bounds = ""
        if self._dialect == 'cuda':
            max_threads = node.indexing.max_threads_per_block()
            if max_threads:
                launch_bounds = "__launch_bounds__({}) ".format(max_threads)
        func_declaration = "FUNC_PREFIX %svoid %s(%s)" % (launch_bounds, node.function_name,
                                                          ", ".join(function_arguments))
172
        if self._signatureOnly:
Martin Bauer's avatar
Martin Bauer committed
173
            return func_declaration
174

175
        body = self._print(node.body)
Martin Bauer's avatar
Martin Bauer committed
176
        return func_declaration + "\n" + body
177
178

    def _print_Block(self, node):
Martin Bauer's avatar
Martin Bauer committed
179
180
        block_contents = "\n".join([self._print(child) for child in node.args])
        return "{\n%s\n}" % (self._indent + self._indent.join(block_contents.splitlines(True)))
181
182

    def _print_PragmaBlock(self, node):
Martin Bauer's avatar
Martin Bauer committed
183
        return "%s\n%s" % (node.pragma_line, self._print_Block(node))
184
185

    def _print_LoopOverCoordinate(self, node):
Martin Bauer's avatar
Martin Bauer committed
186
        counter_symbol = node.loop_counter_name
Martin Bauer's avatar
Martin Bauer committed
187
188
189
190
        start = "int %s = %s" % (counter_symbol, self.sympy_printer.doprint(node.start))
        condition = "%s < %s" % (counter_symbol, self.sympy_printer.doprint(node.stop))
        update = "%s += %s" % (counter_symbol, self.sympy_printer.doprint(node.step),)
        loop_str = "for (%s; %s; %s)" % (start, condition, update)
191

Martin Bauer's avatar
Martin Bauer committed
192
        prefix = "\n".join(node.prefix_lines)
193
194
        if prefix:
            prefix += "\n"
Martin Bauer's avatar
Martin Bauer committed
195
        return "%s%s\n%s" % (prefix, loop_str, self._print(node.body))
196
197

    def _print_SympyAssignment(self, node):
Martin Bauer's avatar
Martin Bauer committed
198
199
        if node.is_declaration:
            data_type = "const " + str(node.lhs.dtype) + " " if node.is_const else str(node.lhs.dtype) + " "
200
201
            return "%s%s = %s;" % (data_type, self.sympy_printer.doprint(node.lhs),
                                   self.sympy_printer.doprint(node.rhs))
202
        else:
Martin Bauer's avatar
Martin Bauer committed
203
            lhs_type = get_type_of_expression(node.lhs)
Martin Bauer's avatar
Martin Bauer committed
204
205
206
207
208
209
            if type(lhs_type) is VectorType and isinstance(node.lhs, cast_func):
                arg, data_type, aligned, nontemporal = node.lhs.args
                instr = 'storeU'
                if aligned:
                    instr = 'stream' if nontemporal else 'storeA'

210
211
212
213
214
215
                rhs_type = get_type_of_expression(node.rhs)
                if type(rhs_type) is not VectorType:
                    rhs = cast_func(node.rhs, VectorType(rhs_type))
                else:
                    rhs = node.rhs

216
217
                return self._vector_instruction_set[instr].format("&" + self.sympy_printer.doprint(node.lhs.args[0]),
                                                                  self.sympy_printer.doprint(rhs)) + ';'
218
            else:
Martin Bauer's avatar
Martin Bauer committed
219
                return "%s = %s;" % (self.sympy_printer.doprint(node.lhs), self.sympy_printer.doprint(node.rhs))
220
221

    def _print_TemporaryMemoryAllocation(self, node):
222
        align = 64
Martin Bauer's avatar
Martin Bauer committed
223
224
225
226
227
228
        np_dtype = node.symbol.dtype.base_type.numpy_dtype
        required_size = np_dtype.itemsize * node.size + align
        size = modulo_ceil(required_size, align)
        code = "{dtype} {name}=({dtype})aligned_alloc({align}, {size}) + {offset};"
        return code.format(dtype=node.symbol.dtype,
                           name=self.sympy_printer.doprint(node.symbol.name),
229
                           size=self.sympy_printer.doprint(size),
Martin Bauer's avatar
Martin Bauer committed
230
231
                           offset=int(node.offset(align)),
                           align=align)
232
233

    def _print_TemporaryMemoryFree(self, node):
234
        align = 64
Martin Bauer's avatar
Martin Bauer committed
235
        return "free(%s - %d);" % (self.sympy_printer.doprint(node.symbol.name), node.offset(align))
236

Martin Bauer's avatar
Martin Bauer committed
237
238
239
240
241
242
    def _print_SkipIteration(self, _):
        if self._dialect == 'cuda':
            return "return;"
        else:
            return "continue;"

243
244
    def _print_CustomCodeNode(self, node):
        return node.get_code(self._dialect, self._vector_instruction_set)
245

246
    def _print_Conditional(self, node):
247
248
249
        cond_type = get_type_of_expression(node.condition_expr)
        if isinstance(cond_type, VectorType):
            raise ValueError("Problem with Conditional inside vectorized loop - use vec_any or vec_all")
Martin Bauer's avatar
Martin Bauer committed
250
251
        condition_expr = self.sympy_printer.doprint(node.condition_expr)
        true_block = self._print_Block(node.true_block)
Martin Bauer's avatar
Martin Bauer committed
252
        result = "if (%s)\n%s " % (condition_expr, true_block)
Martin Bauer's avatar
Martin Bauer committed
253
254
        if node.false_block:
            false_block = self._print_Block(node.false_block)
Martin Bauer's avatar
Martin Bauer committed
255
            result += "else " + false_block
256
257
        return result

258
259
260
261

# ------------------------------------------ Helper function & classes -------------------------------------------------


Martin Bauer's avatar
Martin Bauer committed
262
# noinspection PyPep8Naming
263
class CustomSympyPrinter(CCodePrinter):
Martin Bauer's avatar
Martin Bauer committed
264

265
    def __init__(self, dialect):
Martin Bauer's avatar
Martin Bauer committed
266
        super(CustomSympyPrinter, self).__init__()
267
        self._float_type = create_type("float32")
268
        self._dialect = dialect
269
270
271
272
        if 'Min' in self.known_functions:
            del self.known_functions['Min']
        if 'Max' in self.known_functions:
            del self.known_functions['Max']
Martin Bauer's avatar
Martin Bauer committed
273

274
275
276
    def _print_Pow(self, expr):
        """Don't use std::pow function, for small integer exponents, write as multiplication"""
        if expr.exp.is_integer and expr.exp.is_number and 0 < expr.exp < 8:
277
            return "(" + self._print(sp.Mul(*[expr.base] * expr.exp, evaluate=False)) + ")"
278
279
        elif expr.exp.is_integer and expr.exp.is_number and - 8 < expr.exp < 0:
            return "1 / ({})".format(self._print(sp.Mul(*[expr.base] * (-expr.exp), evaluate=False)))
280
281
282
283
284
        else:
            return super(CustomSympyPrinter, self)._print_Pow(expr)

    def _print_Rational(self, expr):
        """Evaluate all rationals i.e. print 0.25 instead of 1.0/4.0"""
Martin Bauer's avatar
Martin Bauer committed
285
286
        res = str(expr.evalf().num)
        return res
287
288
289
290
291
292
293
294

    def _print_Equality(self, expr):
        """Equality operator is not printable in default printer"""
        return '((' + self._print(expr.lhs) + ") == (" + self._print(expr.rhs) + '))'

    def _print_Piecewise(self, expr):
        """Print piecewise in one line (remove newlines)"""
        result = super(CustomSympyPrinter, self)._print_Piecewise(expr)
Martin Bauer's avatar
Martin Bauer committed
295
296
        return result.replace("\n", "")

297
    def _print_Function(self, expr):
298
        infix_functions = {
Martin Bauer's avatar
Martin Bauer committed
299
300
301
302
303
            bitwise_xor: '^',
            bit_shift_right: '>>',
            bit_shift_left: '<<',
            bitwise_or: '|',
            bitwise_and: '&',
Martin Bauer's avatar
Martin Bauer committed
304
        }
Martin Bauer's avatar
Martin Bauer committed
305
306
        if hasattr(expr, 'to_c'):
            return expr.to_c(self._print)
307
308
309
        if isinstance(expr, reinterpret_cast_func):
            arg, data_type = expr.args
            return "*((%s)(& %s))" % (PointerType(data_type, restrict=False), self._print(arg))
310
311
312
        elif isinstance(expr, address_of):
            assert len(expr.args) == 1, "address_of must only have one argument"
            return "&(%s)" % self._print(expr.args[0])
313
        elif isinstance(expr, cast_func):
Martin Bauer's avatar
Martin Bauer committed
314
            arg, data_type = expr.args
315
316
317
            if isinstance(arg, sp.Number):
                return self._typed_number(arg, data_type)
            else:
318
319
320
321
322
323
324
325
326
327
328
                return "((%s)(%s))" % (data_type, self._print(arg))
        elif isinstance(expr, fast_division):
            if self._dialect == "cuda":
                return "__fdividef(%s, %s)" % tuple(self._print(a) for a in expr.args)
            else:
                return "({})".format(self._print(expr.args[0] / expr.args[1]))
        elif isinstance(expr, fast_sqrt):
            if self._dialect == "cuda":
                return "__fsqrt_rn(%s)" % tuple(self._print(a) for a in expr.args)
            else:
                return "({})".format(self._print(sp.sqrt(expr.args[0])))
329
330
        elif isinstance(expr, vec_any) or isinstance(expr, vec_all):
            return self._print(expr.args[0])
331
332
333
334
335
        elif isinstance(expr, fast_inv_sqrt):
            if self._dialect == "cuda":
                return "__frsqrt_rn(%s)" % tuple(self._print(a) for a in expr.args)
            else:
                return "({})".format(self._print(1 / sp.sqrt(expr.args[0])))
336
337
        elif expr.func in infix_functions:
            return "(%s %s %s)" % (self._print(expr.args[0]), infix_functions[expr.func], self._print(expr.args[1]))
338
339
340
341
        elif expr.func == int_power_of_2:
            return "(1 << (%s))" % (self._print(expr.args[0]))
        elif expr.func == int_div:
            return "((%s) / (%s))" % (self._print(expr.args[0]), self._print(expr.args[1]))
342
        else:
343
            return super(CustomSympyPrinter, self)._print_Function(expr)
Martin Bauer's avatar
Martin Bauer committed
344

345
346
    def _typed_number(self, number, dtype):
        res = self._print(number)
347
        if dtype.is_float():
348
349
350
351
352
353
354
355
            if dtype == self._float_type:
                if '.' not in res:
                    res += ".0f"
                else:
                    res += "f"
            return res
        else:
            return res
356

357
358
359
    _print_Max = C89CodePrinter._print_Max
    _print_Min = C89CodePrinter._print_Min

360

Martin Bauer's avatar
Martin Bauer committed
361
# noinspection PyPep8Naming
362
363
364
class VectorizedCustomSympyPrinter(CustomSympyPrinter):
    SummandInfo = namedtuple("SummandInfo", ['sign', 'term'])

365
366
    def __init__(self, instruction_set, dialect):
        super(VectorizedCustomSympyPrinter, self).__init__(dialect=dialect)
Martin Bauer's avatar
Martin Bauer committed
367
        self.instruction_set = instruction_set
368

Martin Bauer's avatar
Martin Bauer committed
369
370
371
372
    def _scalarFallback(self, func_name, expr, *args, **kwargs):
        expr_type = get_type_of_expression(expr)
        if type(expr_type) is not VectorType:
            return getattr(super(VectorizedCustomSympyPrinter, self), func_name)(expr, *args, **kwargs)
373
        else:
Martin Bauer's avatar
Martin Bauer committed
374
            assert self.instruction_set['width'] == expr_type.width
375
376
            return None

377
    def _print_Function(self, expr):
378
        if isinstance(expr, vector_memory_access):
Martin Bauer's avatar
Martin Bauer committed
379
380
381
            arg, data_type, aligned, _ = expr.args
            instruction = self.instruction_set['loadA'] if aligned else self.instruction_set['loadU']
            return instruction.format("& " + self._print(arg))
382
        elif isinstance(expr, cast_func):
Martin Bauer's avatar
Martin Bauer committed
383
384
            arg, data_type = expr.args
            if type(data_type) is VectorType:
Martin Bauer's avatar
Martin Bauer committed
385
                return self.instruction_set['makeVec'].format(self._print(arg))
386
        elif expr.func == fast_division:
387
388
            result = self._scalarFallback('_print_Function', expr)
            if not result:
389
390
                result = self.instruction_set['/'].format(self._print(expr.args[0]), self._print(expr.args[1]))
            return result
391
392
393
        elif expr.func == fast_sqrt:
            return "({})".format(self._print(sp.sqrt(expr.args[0])))
        elif expr.func == fast_inv_sqrt:
394
395
396
397
398
399
            result = self._scalarFallback('_print_Function', expr)
            if not result:
                if self.instruction_set['rsqrt']:
                    return self.instruction_set['rsqrt'].format(self._print(expr.args[0]))
                else:
                    return "({})".format(self._print(1 / sp.sqrt(expr.args[0])))
400
401
402
403
404
405
406
407
408
409
410
411
412
        elif isinstance(expr, vec_any):
            expr_type = get_type_of_expression(expr.args[0])
            if type(expr_type) is not VectorType:
                return self._print(expr.args[0])
            else:
                return self.instruction_set['any'].format(self._print(expr.args[0]))
        elif isinstance(expr, vec_all):
            expr_type = get_type_of_expression(expr.args[0])
            if type(expr_type) is not VectorType:
                return self._print(expr.args[0])
            else:
                return self.instruction_set['all'].format(self._print(expr.args[0]))

413
414
        return super(VectorizedCustomSympyPrinter, self)._print_Function(expr)

415
416
417
418
419
    def _print_And(self, expr):
        result = self._scalarFallback('_print_And', expr)
        if result:
            return result

Martin Bauer's avatar
Martin Bauer committed
420
421
422
423
        arg_strings = [self._print(a) for a in expr.args]
        assert len(arg_strings) > 0
        result = arg_strings[0]
        for item in arg_strings[1:]:
Martin Bauer's avatar
Martin Bauer committed
424
            result = self.instruction_set['&'].format(result, item)
425
426
427
428
429
430
431
        return result

    def _print_Or(self, expr):
        result = self._scalarFallback('_print_Or', expr)
        if result:
            return result

Martin Bauer's avatar
Martin Bauer committed
432
433
434
435
        arg_strings = [self._print(a) for a in expr.args]
        assert len(arg_strings) > 0
        result = arg_strings[0]
        for item in arg_strings[1:]:
Martin Bauer's avatar
Martin Bauer committed
436
            result = self.instruction_set['|'].format(result, item)
437
438
        return result

439
    def _print_Add(self, expr, order=None):
440
441
442
        result = self._scalarFallback('_print_Add', expr)
        if result:
            return result
443
444
445
446

        summands = []
        for term in expr.args:
            if term.func == sp.Mul:
Martin Bauer's avatar
Martin Bauer committed
447
                sign, t = self._print_Mul(term, inside_add=True)
448
449
450
451
452
453
454
455
456
457
458
459
460
            else:
                t = self._print(term)
                sign = 1
            summands.append(self.SummandInfo(sign, t))
        # Use positive terms first
        summands.sort(key=lambda e: e.sign, reverse=True)
        # if no positive term exists, prepend a zero
        if summands[0].sign == -1:
            summands.insert(0, self.SummandInfo(1, "0"))

        assert len(summands) >= 2
        processed = summands[0].term
        for summand in summands[1:]:
Martin Bauer's avatar
Martin Bauer committed
461
            func = self.instruction_set['-'] if summand.sign == -1 else self.instruction_set['+']
462
463
464
            processed = func.format(processed, summand.term)
        return processed

465
    def _print_Pow(self, expr):
466
467
468
        result = self._scalarFallback('_print_Pow', expr)
        if result:
            return result
469

470
471
        one = self.instruction_set['makeVec'].format(1.0)

472
473
        if expr.exp.is_integer and expr.exp.is_number and 0 < expr.exp < 8:
            return "(" + self._print(sp.Mul(*[expr.base] * expr.exp, evaluate=False)) + ")"
474
475
476
477
478
        elif expr.exp == -1:
            one = self.instruction_set['makeVec'].format(1.0)
            return self.instruction_set['/'].format(one, self._print(expr.base))
        elif expr.exp == 0.5:
            return self.instruction_set['sqrt'].format(self._print(expr.base))
479
480
481
        elif expr.exp == -0.5:
            root = self.instruction_set['sqrt'].format(self._print(expr.base))
            return self.instruction_set['/'].format(one, root)
482
483
484
        elif expr.exp.is_integer and expr.exp.is_number and - 8 < expr.exp < 0:
            return self.instruction_set['/'].format(one,
                                                    self._print(sp.Mul(*[expr.base] * (-expr.exp), evaluate=False)))
485
        else:
486
            raise ValueError("Generic exponential not supported: " + str(expr))
487

Martin Bauer's avatar
Martin Bauer committed
488
489
490
491
    def _print_Mul(self, expr, inside_add=False):
        # noinspection PyProtectedMember
        from sympy.core.mul import _keep_coeff

492
493
494
        result = self._scalarFallback('_print_Mul', expr)
        if result:
            return result
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522

        c, e = expr.as_coeff_Mul()
        if c < 0:
            expr = _keep_coeff(-c, e)
            sign = -1
        else:
            sign = 1

        a = []  # items in the numerator
        b = []  # items that are in the denominator (if any)

        # Gather args for numerator/denominator
        for item in expr.as_ordered_factors():
            if item.is_commutative and item.is_Pow and item.exp.is_Rational and item.exp.is_negative:
                if item.exp != -1:
                    b.append(sp.Pow(item.base, -item.exp, evaluate=False))
                else:
                    b.append(sp.Pow(item.base, -item.exp))
            else:
                a.append(item)

        a = a or [S.One]

        a_str = [self._print(x) for x in a]
        b_str = [self._print(x) for x in b]

        result = a_str[0]
        for item in a_str[1:]:
Martin Bauer's avatar
Martin Bauer committed
523
            result = self.instruction_set['*'].format(result, item)
524
525
526
527

        if len(b) > 0:
            denominator_str = b_str[0]
            for item in b_str[1:]:
Martin Bauer's avatar
Martin Bauer committed
528
529
                denominator_str = self.instruction_set['*'].format(denominator_str, item)
            result = self.instruction_set['/'].format(result, denominator_str)
530

Martin Bauer's avatar
Martin Bauer committed
531
        if inside_add:
532
533
534
            return sign, result
        else:
            if sign < 0:
Martin Bauer's avatar
Martin Bauer committed
535
                return self.instruction_set['*'].format(self._print(S.NegativeOne), result)
536
537
538
            else:
                return result

539
    def _print_Relational(self, expr):
540
541
542
        result = self._scalarFallback('_print_Relational', expr)
        if result:
            return result
Martin Bauer's avatar
Martin Bauer committed
543
        return self.instruction_set[expr.rel_op].format(self._print(expr.lhs), self._print(expr.rhs))
544
545

    def _print_Equality(self, expr):
546
547
548
        result = self._scalarFallback('_print_Equality', expr)
        if result:
            return result
Martin Bauer's avatar
Martin Bauer committed
549
        return self.instruction_set['=='].format(self._print(expr.lhs), self._print(expr.rhs))
550
551

    def _print_Piecewise(self, expr):
552
553
554
        result = self._scalarFallback('_print_Piecewise', expr)
        if result:
            return result
555

Martin Bauer's avatar
Martin Bauer committed
556
        if expr.args[-1].cond.args[0] is not sp.sympify(True):
557
558
559
560
561
562
563
564
565
            # We need the last conditional to be a True, otherwise the resulting
            # function may not return a result.
            raise ValueError("All Piecewise expressions must contain an "
                             "(expr, True) statement to be used as a default "
                             "condition. Without one, the generated "
                             "expression may not evaluate to anything under "
                             "some condition.")

        result = self._print(expr.args[-1][0])
Martin Bauer's avatar
Martin Bauer committed
566
        for true_expr, condition in reversed(expr.args[:-1]):
567
            if isinstance(condition, cast_func) and get_type_of_expression(condition.args[0]) == create_type("bool"):
568
569
570
571
572
                if not KERNCRAFT_NO_TERNARY_MODE:
                    result = "(({}) ? ({}) : ({}))".format(self._print(condition.args[0]), self._print(true_expr),
                                                           result)
                else:
                    print("Warning - skipping ternary op")
573
574
575
            else:
                # noinspection SpellCheckingInspection
                result = self.instruction_set['blendv'].format(result, self._print(true_expr), self._print(condition))
576
        return result