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
|
import importlib
import os
import sparse
from utils import benchmark
import numpy as np
import scipy.sparse as sps
LEN = 100000
DENSITY = 0.000001
ITERS = 3
rng = np.random.default_rng(0)
if __name__ == "__main__":
print("SpMv_add Example:\n")
A_sps = sps.random(LEN - 10, LEN, format="csc", density=DENSITY, random_state=rng) * 10
x_sps = rng.random((LEN, 1)) * 10
y_sps = rng.random((LEN - 10, 1)) * 10
# ======= Finch =======
os.environ[sparse._ENV_VAR_NAME] = "Finch"
importlib.reload(sparse)
A = sparse.asarray(A_sps)
x = sparse.asarray(np.array(x_sps, order="C"))
y = sparse.asarray(np.array(y_sps, order="C"))
@sparse.compiled
def spmv_finch(A, x, y):
return sparse.sum(A[:, None, :] * sparse.permute_dims(x, (1, 0))[None, :, :], axis=-1) + y
# Compile
result_finch = spmv_finch(A, x, y)
# Benchmark
benchmark(spmv_finch, args=[A, x, y], info="Finch", iters=ITERS)
# ======= Numba =======
os.environ[sparse._ENV_VAR_NAME] = "Numba"
importlib.reload(sparse)
A = sparse.asarray(A_sps, format="csc")
x = x_sps
y = y_sps
def spmv_numba(A, x, y):
return A @ x + y
# Compile
result_numba = spmv_numba(A, x, y)
assert sparse.nonzero(result_numba)[0].size > 5
# Benchmark
benchmark(spmv_numba, args=[A, x, y], info="Numba", iters=ITERS)
# ======= SciPy =======
def spmv_scipy(A, x, y):
return A @ x + y
A = A_sps
x = x_sps
y = y_sps
result_scipy = spmv_scipy(A, x, y)
# Benchmark
benchmark(spmv_scipy, args=[A, x, y], info="SciPy", iters=ITERS)
np.testing.assert_allclose(result_numba, result_scipy)
np.testing.assert_allclose(result_finch.todense(), result_numba)
np.testing.assert_allclose(result_finch.todense(), result_scipy)
|
