forked from inaos/iron-array-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathexpression-multithreading.py
More file actions
122 lines (98 loc) · 3.27 KB
/
Copy pathexpression-multithreading.py
File metadata and controls
122 lines (98 loc) · 3.27 KB
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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
import iarray as ia
import numpy as np
from time import time
import ctypes
import numexpr as ne
from iarray import udf
from iarray.py2llvm import float64
import numba as nb
# omp = ctypes.CDLL('libiomp5.so')
# omp_set_num_threads = omp.omp_set_num_threads
max_num_threads = 4
nrep = 3
@nb.jit(nopython=True, cache=True, parallel=True)
def poly_numba(x):
y = np.empty(x.shape, x.dtype)
for i in nb.prange(len(x)):
y[i] = (x[i] - 1.35) * (x[i] - 4.45) * (x[i] - 8.5)
return y
# config.THREADING_LAYER = 'omp'
@udf.jit
def poly_udf(x: udf.Array(float64, 1), y: udf.Array(float64, 1)):
n = x.shape[0]
for i in range(n):
y[i] = (x[i] - 1.35) * (x[i] - 4.45) * (x[i] - 8.5)
return 0
# Define array params
shape = [32 * 512 * 1024]
chunkshape = [32 * 1024]
blockshape = [16 * 1024]
dtshape = ia.DTShape(shape)
size = int(np.prod(shape))
bstorage = ia.Storage(chunkshape, blockshape)
pstorage = ia.Storage(plainbuffer=True)
def time_expr(expr, result):
t = []
for _ in range(nrep):
t0 = time()
expr.eval()
t1 = time()
t.append(round(size / 8 / 2 ** 20 / (t1 - t0), 2))
t.remove(max(t))
result.append(np.mean(t))
res = []
for num_threads in range(1, max_num_threads + 1):
print(f"Num. threads: {num_threads}")
# omp_set_num_threads(num_threads)
res_i = []
# Numpy
a1 = np.linspace(0, 10, size).reshape(shape)
t = []
for _ in range(nrep):
t0 = time()
eval("(x - 1.35) * (x - 4.45) * (x - 8.5)", {"x": a1})
t1 = time()
t.append(round(size / 8 / 2 ** 20 / (t1 - t0), 2))
t.remove(max(t))
res_i.append(np.mean(t))
# numba
t = []
for _ in range(nrep):
t0 = time()
poly_numba(a1)
t1 = time()
t.append(round(size / 8 / 2 ** 20 / (t1 - t0), 2))
t.remove(max(t))
res_i.append(np.mean(t))
# Numexpr
t = []
ne.set_num_threads(num_threads)
for _ in range(nrep):
t0 = time()
ne.evaluate("(x - 1.35) * (x - 4.45) * (x - 8.5)", local_dict={"x": a1})
t1 = time()
t.append(round(size / 8 / 2 ** 20 / (t1 - t0), 2))
t.remove(max(t))
res_i.append(np.mean(t))
# Plainbuffer
with ia.config(storage=pstorage, nthreads=num_threads) as cfg:
a1 = ia.linspace(dtshape, 0, 10, cfg=cfg)
expr = ia.expr_from_string("(x - 1.35) * (x - 4.45) * (x - 8.5)", {"x": a1}, cfg=cfg)
time_expr(expr, res_i)
# Superchunk without compression
with ia.config(storage=bstorage, nthreads=num_threads, clevel=0) as cfg:
a1 = ia.linspace(dtshape, 0, 10, cfg=cfg)
expr = ia.expr_from_string("(x - 1.35) * (x - 4.45) * (x - 8.5)", {"x": a1}, cfg=cfg)
time_expr(expr, res_i)
# Superchunk with compression
with ia.config(storage=bstorage, nthreads=num_threads, clevel=9) as cfg:
a1 = ia.linspace(dtshape, 0, 10, cfg=cfg)
expr = ia.expr_from_string("(x - 1.35) * (x - 4.45) * (x - 8.5)", {"x": a1}, cfg=cfg)
time_expr(expr, res_i)
# Superchunk with compression and UDF
with ia.config(storage=bstorage, nthreads=num_threads, clevel=9) as cfg:
a1 = ia.linspace(dtshape, 0, 10, cfg=cfg)
expr = ia.expr_from_udf(poly_udf, [a1], cfg=cfg)
time_expr(expr, res_i)
res.append(res_i)
print(res)