-
Notifications
You must be signed in to change notification settings - Fork 6
/
bench_dispatch.py
90 lines (68 loc) · 2.14 KB
/
bench_dispatch.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
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
"""
Benchmarks for argument dispatching and call overhead of ``@jit`` functions.
"""
import numpy as np
rec_dtype = np.dtype([('a', np.float64),
('b', np.int32),
('c', np.complex64),
])
samples = {
'bool': True,
'int': 100000,
'float': 0.5,
'complex': 0.5 + 1.0j,
'array_1d': np.zeros(10, dtype=np.int64),
'array_3d': np.zeros(20, dtype=np.float64).reshape(2, 2, 5),
'array_records': np.zeros(10, dtype=rec_dtype),
'recarray': np.recarray(10, dtype=rec_dtype),
'tuple': (0.5, 1.0j, ()),
'record': np.empty(1, dtype=rec_dtype)[0],
'bytearray': bytearray(3),
}
def setup():
"""
Precompile jitted functions. This will register many specializations
to choose from.
"""
from numba import jit
global binary, binary_pyobj, unary_default
@jit(nopython=True)
def binary(x, y):
pass
@jit(forceobj=True)
def binary_pyobj(x, y):
pass
@jit(nopython=True)
def unary_default(x=None):
pass
for tp in samples.values():
binary(tp, tp)
binary_pyobj(object(), object())
unary_default()
class NoPythonDispatch:
"""
Time dispatching to a jitted function's specializations based on argument
types.
This stresses two things:
- the typing of arguments (from argument value to typecode)
- the selection of the best specialization amongst all the known ones
"""
# We repeat 1000 times so as to make the overhead of benchmark launching
# negligible.
@classmethod
def generate_benchmarks(cls, names):
for name in names:
def timefunc(self, arg=samples[name]):
func = binary
for i in range(1000):
func(arg, arg)
timefunc.__name__ = "time_dispatch_" + name
setattr(cls, timefunc.__name__, timefunc)
def time_dispatch_defaults(self):
unary_default()
NoPythonDispatch.generate_benchmarks(samples.keys())
class PyObjectDispatch:
def time_dispatch_pyobject(self):
x = object()
for i in range(1000):
binary_pyobj(x, x)