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
|
import importlib
import os
import sparse
from utils import benchmark
import numpy as np
import scipy.sparse as sps
LEN = 100000
DENSITY = 0.00001
ITERS = 3
rng = np.random.default_rng(0)
if __name__ == "__main__":
print("Matmul Example:\n")
a_sps = sps.random(LEN, LEN - 10, format="csr", density=DENSITY, random_state=rng) * 10
a_sps.sum_duplicates()
b_sps = sps.random(LEN - 10, LEN, format="csr", density=DENSITY, random_state=rng) * 10
b_sps.sum_duplicates()
# ======= Finch =======
os.environ[sparse._ENV_VAR_NAME] = "Finch"
importlib.reload(sparse)
a = sparse.asarray(a_sps)
b = sparse.asarray(b_sps)
@sparse.compiled()
def sddmm_finch(a, b):
return a @ b
# Compile & Benchmark
result_finch = benchmark(sddmm_finch, args=[a, b], info="Finch", iters=ITERS)
# ======= Numba =======
os.environ[sparse._ENV_VAR_NAME] = "Numba"
importlib.reload(sparse)
a = sparse.asarray(a_sps)
b = sparse.asarray(b_sps)
def sddmm_numba(a, b):
return a @ b
# Compile & Benchmark
result_numba = benchmark(sddmm_numba, args=[a, b], info="Numba", iters=ITERS)
# ======= SciPy =======
def sddmm_scipy(a, b):
return a @ b
a = a_sps
b = b_sps
# Compile & Benchmark
result_scipy = benchmark(sddmm_scipy, args=[a, b], info="SciPy", iters=ITERS)
# np.testing.assert_allclose(result_numba.todense(), result_scipy.toarray())
# np.testing.assert_allclose(result_finch.todense(), result_numba.todense())
# np.testing.assert_allclose(result_finch.todense(), result_scipy.toarray())
|
