-
Notifications
You must be signed in to change notification settings - Fork 1.1k
/
Copy pathfastmathpass.py
44 lines (33 loc) · 1.18 KB
/
fastmathpass.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
from llvmlite import ir
from llvmlite.ir.transforms import Visitor, CallVisitor
class FastFloatBinOpVisitor(Visitor):
"""
A pass to add fastmath flag to float-binop instruction if they don't have
any flags.
"""
float_binops = frozenset(['fadd', 'fsub', 'fmul', 'fdiv', 'frem', 'fcmp'])
def __init__(self, flags):
self.flags = flags
def visit_Instruction(self, instr):
if instr.opname in self.float_binops:
if not instr.flags:
for flag in self.flags:
instr.flags.append(flag)
class FastFloatCallVisitor(CallVisitor):
"""
A pass to change all float function calls to use fastmath.
"""
def __init__(self, flags):
self.flags = flags
def visit_Call(self, instr):
# Add to any call that has float/double return type
if instr.type in (ir.FloatType(), ir.DoubleType()):
for flag in self.flags:
instr.fastmath.add(flag)
def rewrite_module(mod, options):
"""
Rewrite the given LLVM module to use fastmath everywhere.
"""
flags = options.flags
FastFloatBinOpVisitor(flags).visit(mod)
FastFloatCallVisitor(flags).visit(mod)