cbackend.py 21 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
from pystencils.fast_approximation import fast_division, fast_sqrt, fast_inv_sqrt

Martin Bauer's avatar
Martin Bauer committed
10
11
try:
    from sympy.printing.ccode import C99CodePrinter as CCodePrinter
Martin Bauer's avatar
Martin Bauer committed
12
13
except ImportError:
    from sympy.printing.ccode import CCodePrinter  # for sympy versions < 1.1
Martin Bauer's avatar
Martin Bauer committed
14

Martin Bauer's avatar
Martin Bauer committed
15
from pystencils.integer_functions import bitwise_xor, bit_shift_right, bit_shift_left, bitwise_and, \
Martin Bauer's avatar
Martin Bauer committed
16
    bitwise_or, modulo_ceil
17
from pystencils.astnodes import Node, KernelFunction
Martin Bauer's avatar
Martin Bauer committed
18
from pystencils.data_types import create_type, PointerType, get_type_of_expression, VectorType, cast_func, \
19
    vector_memory_access, reinterpret_cast_func
20

21
__all__ = ['generate_c', 'CustomCodeNode', 'PrintNode', 'get_headers', 'CustomSympyPrinter']
22

Martin Bauer's avatar
Martin Bauer committed
23

24
def generate_c(ast_node: Node, signature_only: bool = False, dialect='c') -> str:
Martin Bauer's avatar
Martin Bauer committed
25
26
27
28
29
30
31
32
33
    """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:
34
        dialect: 'c' or 'cuda'
Martin Bauer's avatar
Martin Bauer committed
35
36
    Returns:
        C-like code for the ast node and its descendants
Martin Bauer's avatar
Martin Bauer committed
37
    """
38
    printer = CBackend(signature_only=signature_only,
39
40
                       vector_instruction_set=ast_node.instruction_set,
                       dialect=dialect)
Martin Bauer's avatar
Martin Bauer committed
41
    return printer(ast_node)
42
43


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

Martin Bauer's avatar
Martin Bauer committed
48
49
50
    if isinstance(ast_node, KernelFunction) and ast_node.instruction_set:
        headers.update(ast_node.instruction_set['headers'])

Martin Bauer's avatar
Martin Bauer committed
51
52
53
    if hasattr(ast_node, 'headers'):
        headers.update(ast_node.headers)
    for a in ast_node.args:
54
        if isinstance(a, Node):
Martin Bauer's avatar
Martin Bauer committed
55
            headers.update(get_headers(a))
56
57

    return headers
58
59


60
61
62
# --------------------------------------- Backend Specific Nodes -------------------------------------------------------


63
class CustomCodeNode(Node):
Martin Bauer's avatar
Martin Bauer committed
64
    def __init__(self, code, symbols_read, symbols_defined, parent=None):
65
        super(CustomCodeNode, self).__init__(parent=parent)
66
        self._code = "\n" + code
Martin Bauer's avatar
Martin Bauer committed
67
68
        self._symbolsRead = set(symbols_read)
        self._symbolsDefined = set(symbols_defined)
69
        self.headers = []
70

71
    def get_code(self, dialect, vector_instruction_set):
72
73
74
75
76
77
78
        return self._code

    @property
    def args(self):
        return []

    @property
Martin Bauer's avatar
Martin Bauer committed
79
    def symbols_defined(self):
80
81
82
        return self._symbolsDefined

    @property
Martin Bauer's avatar
Martin Bauer committed
83
84
    def undefined_symbols(self):
        return self.symbols_defined - self._symbolsRead
85
86


87
class PrintNode(CustomCodeNode):
Martin Bauer's avatar
Martin Bauer committed
88
89
90
91
    # 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())
92
        self.headers.append("<iostream>")
93
94
95
96


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

97

Martin Bauer's avatar
Martin Bauer committed
98
99
# noinspection PyPep8Naming
class CBackend:
100

101
    def __init__(self, sympy_printer=None,
102
                 signature_only=False, vector_instruction_set=None, dialect='c'):
Martin Bauer's avatar
Martin Bauer committed
103
104
        if sympy_printer is None:
            if vector_instruction_set is not None:
105
                self.sympy_printer = VectorizedCustomSympyPrinter(vector_instruction_set, dialect)
106
            else:
107
                self.sympy_printer = CustomSympyPrinter(dialect)
108
        else:
Martin Bauer's avatar
Martin Bauer committed
109
            self.sympy_printer = sympy_printer
110

111
        self._vector_instruction_set = vector_instruction_set
112
        self._indent = "   "
113
        self._dialect = dialect
Martin Bauer's avatar
Martin Bauer committed
114
        self._signatureOnly = signature_only
115
116

    def __call__(self, node):
Martin Bauer's avatar
Martin Bauer committed
117
        prev_is = VectorType.instruction_set
118
        VectorType.instruction_set = self._vector_instruction_set
119
        result = str(self._print(node))
Martin Bauer's avatar
Martin Bauer committed
120
        VectorType.instruction_set = prev_is
121
        return result
122
123
124

    def _print(self, node):
        for cls in type(node).__mro__:
Martin Bauer's avatar
Martin Bauer committed
125
126
127
128
            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)))
129
130

    def _print_KernelFunction(self, node):
131
        function_arguments = ["%s %s" % (str(s.symbol.dtype), s.symbol.name) for s in node.get_parameters()]
Martin Bauer's avatar
Martin Bauer committed
132
        func_declaration = "FUNC_PREFIX void %s(%s)" % (node.function_name, ", ".join(function_arguments))
133
        if self._signatureOnly:
Martin Bauer's avatar
Martin Bauer committed
134
            return func_declaration
135

136
        body = self._print(node.body)
Martin Bauer's avatar
Martin Bauer committed
137
        return func_declaration + "\n" + body
138
139

    def _print_Block(self, node):
Martin Bauer's avatar
Martin Bauer committed
140
141
        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)))
142
143

    def _print_PragmaBlock(self, node):
Martin Bauer's avatar
Martin Bauer committed
144
        return "%s\n%s" % (node.pragma_line, self._print_Block(node))
145
146

    def _print_LoopOverCoordinate(self, node):
Martin Bauer's avatar
Martin Bauer committed
147
        counter_symbol = node.loop_counter_name
Martin Bauer's avatar
Martin Bauer committed
148
149
150
151
        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)
152

Martin Bauer's avatar
Martin Bauer committed
153
        prefix = "\n".join(node.prefix_lines)
154
155
        if prefix:
            prefix += "\n"
Martin Bauer's avatar
Martin Bauer committed
156
        return "%s%s\n%s" % (prefix, loop_str, self._print(node.body))
157
158

    def _print_SympyAssignment(self, node):
Martin Bauer's avatar
Martin Bauer committed
159
160
        if node.is_declaration:
            data_type = "const " + str(node.lhs.dtype) + " " if node.is_const else str(node.lhs.dtype) + " "
161
162
            return "%s%s = %s;" % (data_type, self.sympy_printer.doprint(node.lhs),
                                   self.sympy_printer.doprint(node.rhs))
163
        else:
Martin Bauer's avatar
Martin Bauer committed
164
            lhs_type = get_type_of_expression(node.lhs)
Martin Bauer's avatar
Martin Bauer committed
165
166
167
168
169
170
            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'

171
172
173
174
175
176
                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

177
178
                return self._vector_instruction_set[instr].format("&" + self.sympy_printer.doprint(node.lhs.args[0]),
                                                                  self.sympy_printer.doprint(rhs)) + ';'
179
            else:
Martin Bauer's avatar
Martin Bauer committed
180
                return "%s = %s;" % (self.sympy_printer.doprint(node.lhs), self.sympy_printer.doprint(node.rhs))
181
182

    def _print_TemporaryMemoryAllocation(self, node):
183
        align = 64
Martin Bauer's avatar
Martin Bauer committed
184
185
186
187
188
189
        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),
190
                           size=self.sympy_printer.doprint(size),
Martin Bauer's avatar
Martin Bauer committed
191
192
                           offset=int(node.offset(align)),
                           align=align)
193
194

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

198
199
    def _print_CustomCodeNode(self, node):
        return node.get_code(self._dialect, self._vector_instruction_set)
200

201
    def _print_Conditional(self, node):
202
203
204
        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
205
206
        condition_expr = self.sympy_printer.doprint(node.condition_expr)
        true_block = self._print_Block(node.true_block)
Martin Bauer's avatar
Martin Bauer committed
207
        result = "if (%s)\n%s " % (condition_expr, true_block)
Martin Bauer's avatar
Martin Bauer committed
208
209
        if node.false_block:
            false_block = self._print_Block(node.false_block)
Martin Bauer's avatar
Martin Bauer committed
210
            result += "else " + false_block
211
212
        return result

213
214
215
216

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


Martin Bauer's avatar
Martin Bauer committed
217
# noinspection PyPep8Naming
218
class CustomSympyPrinter(CCodePrinter):
Martin Bauer's avatar
Martin Bauer committed
219

220
    def __init__(self, dialect):
Martin Bauer's avatar
Martin Bauer committed
221
        super(CustomSympyPrinter, self).__init__()
222
        self._float_type = create_type("float32")
223
        self._dialect = dialect
224
225
226
227
        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
228

229
230
231
    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:
232
            return "(" + self._print(sp.Mul(*[expr.base] * expr.exp, evaluate=False)) + ")"
233
234
        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)))
235
236
237
238
239
        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
240
241
        res = str(expr.evalf().num)
        return res
242
243
244
245
246
247
248
249

    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
250
251
        return result.replace("\n", "")

252
    def _print_Function(self, expr):
253
        infix_functions = {
Martin Bauer's avatar
Martin Bauer committed
254
255
256
257
258
            bitwise_xor: '^',
            bit_shift_right: '>>',
            bit_shift_left: '<<',
            bitwise_or: '|',
            bitwise_and: '&',
Martin Bauer's avatar
Martin Bauer committed
259
        }
Martin Bauer's avatar
Martin Bauer committed
260
261
        if hasattr(expr, 'to_c'):
            return expr.to_c(self._print)
262
263
264
265
        if isinstance(expr, reinterpret_cast_func):
            arg, data_type = expr.args
            return "*((%s)(& %s))" % (PointerType(data_type, restrict=False), self._print(arg))
        elif isinstance(expr, cast_func):
Martin Bauer's avatar
Martin Bauer committed
266
            arg, data_type = expr.args
267
268
269
            if isinstance(arg, sp.Number):
                return self._typed_number(arg, data_type)
            else:
270
271
272
273
274
275
276
277
278
279
280
                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])))
281
282
        elif isinstance(expr, vec_any) or isinstance(expr, vec_all):
            return self._print(expr.args[0])
283
284
285
286
287
        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])))
288
289
        elif expr.func in infix_functions:
            return "(%s %s %s)" % (self._print(expr.args[0]), infix_functions[expr.func], self._print(expr.args[1]))
290
        else:
291
            return super(CustomSympyPrinter, self)._print_Function(expr)
Martin Bauer's avatar
Martin Bauer committed
292

293
294
    def _typed_number(self, number, dtype):
        res = self._print(number)
295
        if dtype.is_float():
296
297
298
299
300
301
302
303
            if dtype == self._float_type:
                if '.' not in res:
                    res += ".0f"
                else:
                    res += "f"
            return res
        else:
            return res
304

305
306
307
    _print_Max = C89CodePrinter._print_Max
    _print_Min = C89CodePrinter._print_Min

308

Martin Bauer's avatar
Martin Bauer committed
309
# noinspection PyPep8Naming
310
311
312
class VectorizedCustomSympyPrinter(CustomSympyPrinter):
    SummandInfo = namedtuple("SummandInfo", ['sign', 'term'])

313
314
    def __init__(self, instruction_set, dialect):
        super(VectorizedCustomSympyPrinter, self).__init__(dialect=dialect)
Martin Bauer's avatar
Martin Bauer committed
315
        self.instruction_set = instruction_set
316

Martin Bauer's avatar
Martin Bauer committed
317
318
319
320
    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)
321
        else:
Martin Bauer's avatar
Martin Bauer committed
322
            assert self.instruction_set['width'] == expr_type.width
323
324
            return None

325
    def _print_Function(self, expr):
326
        if isinstance(expr, vector_memory_access):
Martin Bauer's avatar
Martin Bauer committed
327
328
329
            arg, data_type, aligned, _ = expr.args
            instruction = self.instruction_set['loadA'] if aligned else self.instruction_set['loadU']
            return instruction.format("& " + self._print(arg))
330
        elif isinstance(expr, cast_func):
Martin Bauer's avatar
Martin Bauer committed
331
332
            arg, data_type = expr.args
            if type(data_type) is VectorType:
Martin Bauer's avatar
Martin Bauer committed
333
                return self.instruction_set['makeVec'].format(self._print(arg))
334
        elif expr.func == fast_division:
335
336
            result = self._scalarFallback('_print_Function', expr)
            if not result:
337
338
                result = self.instruction_set['/'].format(self._print(expr.args[0]), self._print(expr.args[1]))
            return result
339
340
341
        elif expr.func == fast_sqrt:
            return "({})".format(self._print(sp.sqrt(expr.args[0])))
        elif expr.func == fast_inv_sqrt:
342
343
344
345
346
347
            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])))
348
349
350
351
352
353
354
355
356
357
358
359
360
        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]))

361
362
        return super(VectorizedCustomSympyPrinter, self)._print_Function(expr)

363
364
365
366
367
    def _print_And(self, expr):
        result = self._scalarFallback('_print_And', expr)
        if result:
            return result

Martin Bauer's avatar
Martin Bauer committed
368
369
370
371
        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
372
            result = self.instruction_set['&'].format(result, item)
373
374
375
376
377
378
379
        return result

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

Martin Bauer's avatar
Martin Bauer committed
380
381
382
383
        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
384
            result = self.instruction_set['|'].format(result, item)
385
386
        return result

387
    def _print_Add(self, expr, order=None):
388
389
390
        result = self._scalarFallback('_print_Add', expr)
        if result:
            return result
391
392
393
394

        summands = []
        for term in expr.args:
            if term.func == sp.Mul:
Martin Bauer's avatar
Martin Bauer committed
395
                sign, t = self._print_Mul(term, inside_add=True)
396
397
398
399
400
401
402
403
404
405
406
407
408
            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
409
            func = self.instruction_set['-'] if summand.sign == -1 else self.instruction_set['+']
410
411
412
            processed = func.format(processed, summand.term)
        return processed

413
    def _print_Pow(self, expr):
414
415
416
        result = self._scalarFallback('_print_Pow', expr)
        if result:
            return result
417

418
419
        one = self.instruction_set['makeVec'].format(1.0)

420
421
        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)) + ")"
422
423
424
425
426
        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))
427
428
429
        elif expr.exp == -0.5:
            root = self.instruction_set['sqrt'].format(self._print(expr.base))
            return self.instruction_set['/'].format(one, root)
430
431
432
        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)))
433
        else:
434
            raise ValueError("Generic exponential not supported: " + str(expr))
435

Martin Bauer's avatar
Martin Bauer committed
436
437
438
439
    def _print_Mul(self, expr, inside_add=False):
        # noinspection PyProtectedMember
        from sympy.core.mul import _keep_coeff

440
441
442
        result = self._scalarFallback('_print_Mul', expr)
        if result:
            return result
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470

        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
471
            result = self.instruction_set['*'].format(result, item)
472
473
474
475

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

Martin Bauer's avatar
Martin Bauer committed
479
        if inside_add:
480
481
482
            return sign, result
        else:
            if sign < 0:
Martin Bauer's avatar
Martin Bauer committed
483
                return self.instruction_set['*'].format(self._print(S.NegativeOne), result)
484
485
486
            else:
                return result

487
    def _print_Relational(self, expr):
488
489
490
        result = self._scalarFallback('_print_Relational', expr)
        if result:
            return result
Martin Bauer's avatar
Martin Bauer committed
491
        return self.instruction_set[expr.rel_op].format(self._print(expr.lhs), self._print(expr.rhs))
492
493

    def _print_Equality(self, expr):
494
495
496
        result = self._scalarFallback('_print_Equality', expr)
        if result:
            return result
Martin Bauer's avatar
Martin Bauer committed
497
        return self.instruction_set['=='].format(self._print(expr.lhs), self._print(expr.rhs))
498
499

    def _print_Piecewise(self, expr):
500
501
502
        result = self._scalarFallback('_print_Piecewise', expr)
        if result:
            return result
503

Martin Bauer's avatar
Martin Bauer committed
504
        if expr.args[-1].cond.args[0] is not sp.sympify(True):
505
506
507
508
509
510
511
512
513
            # 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
514
        for true_expr, condition in reversed(expr.args[:-1]):
Martin Bauer's avatar
Martin Bauer committed
515
            # noinspection SpellCheckingInspection
Martin Bauer's avatar
Martin Bauer committed
516
            result = self.instruction_set['blendv'].format(result, self._print(true_expr), self._print(condition))
517
        return result