|
|
@@ -1,66 +1,76 @@ |
|
|
|
import pytest |
|
|
|
import numpy as np |
|
|
|
from openblas_wrap import ( |
|
|
|
# level 1 |
|
|
|
dnrm2, ddot, daxpy, |
|
|
|
# level 3 |
|
|
|
dgemm, dsyrk, |
|
|
|
# lapack |
|
|
|
dgesv, # linalg.solve |
|
|
|
dgesdd, dgesdd_lwork, # linalg.svd |
|
|
|
dsyev, dsyev_lwork, # linalg.eigh |
|
|
|
) |
|
|
|
import openblas_wrap as ow |
|
|
|
|
|
|
|
dtype_map = { |
|
|
|
's': np.float32, |
|
|
|
'd': np.float64, |
|
|
|
'c': np.complex64, |
|
|
|
'z': np.complex128, |
|
|
|
'dz': np.complex128, |
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
# ### BLAS level 1 ### |
|
|
|
|
|
|
|
# dnrm2 |
|
|
|
|
|
|
|
dnrm2_sizes = [100, 1000] |
|
|
|
dnrm2_sizes = [100, 200, 400, 600, 800, 1000] |
|
|
|
|
|
|
|
def run_dnrm2(n, x, incx): |
|
|
|
res = dnrm2(x, n, incx=incx) |
|
|
|
def run_dnrm2(n, x, incx, func): |
|
|
|
res = func(x, n, incx=incx) |
|
|
|
return res |
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize('variant', ['d', 'dz']) |
|
|
|
@pytest.mark.parametrize('n', dnrm2_sizes) |
|
|
|
def test_nrm2(benchmark, n): |
|
|
|
def test_nrm2(benchmark, n, variant): |
|
|
|
rndm = np.random.RandomState(1234) |
|
|
|
x = np.array(rndm.uniform(size=(n,)), dtype=float) |
|
|
|
result = benchmark(run_dnrm2, n, x, 1) |
|
|
|
dtyp = dtype_map[variant] |
|
|
|
|
|
|
|
x = np.array(rndm.uniform(size=(n,)), dtype=dtyp) |
|
|
|
nrm2 = ow.get_func('nrm2', variant) |
|
|
|
result = benchmark(run_dnrm2, n, x, 1, nrm2) |
|
|
|
|
|
|
|
|
|
|
|
# ddot |
|
|
|
|
|
|
|
ddot_sizes = [100, 1000] |
|
|
|
ddot_sizes = [100, 200, 400, 600, 800, 1000] |
|
|
|
|
|
|
|
def run_ddot(x, y,): |
|
|
|
res = ddot(x, y) |
|
|
|
def run_ddot(x, y, func): |
|
|
|
res = func(x, y) |
|
|
|
return res |
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize('n', ddot_sizes) |
|
|
|
def test_dot(benchmark, n): |
|
|
|
rndm = np.random.RandomState(1234) |
|
|
|
|
|
|
|
x = np.array(rndm.uniform(size=(n,)), dtype=float) |
|
|
|
y = np.array(rndm.uniform(size=(n,)), dtype=float) |
|
|
|
result = benchmark(run_ddot, x, y) |
|
|
|
dot = ow.get_func('dot', 'd') |
|
|
|
result = benchmark(run_ddot, x, y, dot) |
|
|
|
|
|
|
|
|
|
|
|
# daxpy |
|
|
|
|
|
|
|
daxpy_sizes = [100, 1000] |
|
|
|
daxpy_sizes = [100, 200, 400, 600, 800, 1000] |
|
|
|
|
|
|
|
def run_daxpy(x, y,): |
|
|
|
res = daxpy(x, y, a=2.0) |
|
|
|
def run_daxpy(x, y, func): |
|
|
|
res = func(x, y, a=2.0) |
|
|
|
return res |
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize('variant', ['s', 'd', 'c', 'z']) |
|
|
|
@pytest.mark.parametrize('n', daxpy_sizes) |
|
|
|
def test_daxpy(benchmark, n): |
|
|
|
def test_daxpy(benchmark, n, variant): |
|
|
|
rndm = np.random.RandomState(1234) |
|
|
|
x = np.array(rndm.uniform(size=(n,)), dtype=float) |
|
|
|
y = np.array(rndm.uniform(size=(n,)), dtype=float) |
|
|
|
result = benchmark(run_daxpy, x, y) |
|
|
|
dtyp = dtype_map[variant] |
|
|
|
|
|
|
|
x = np.array(rndm.uniform(size=(n,)), dtype=dtyp) |
|
|
|
y = np.array(rndm.uniform(size=(n,)), dtype=dtyp) |
|
|
|
axpy = ow.get_func('axpy', variant) |
|
|
|
result = benchmark(run_daxpy, x, y, axpy) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@@ -69,40 +79,46 @@ def test_daxpy(benchmark, n): |
|
|
|
|
|
|
|
# dgemm |
|
|
|
|
|
|
|
gemm_sizes = [100, 1000] |
|
|
|
gemm_sizes = [100, 200, 400, 600, 800, 1000] |
|
|
|
|
|
|
|
def run_gemm(a, b, c): |
|
|
|
def run_gemm(a, b, c, func): |
|
|
|
alpha = 1.0 |
|
|
|
res = dgemm(alpha, a, b, c=c, overwrite_c=True) |
|
|
|
res = func(alpha, a, b, c=c, overwrite_c=True) |
|
|
|
return res |
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize('variant', ['s', 'd', 'c', 'z']) |
|
|
|
@pytest.mark.parametrize('n', gemm_sizes) |
|
|
|
def test_gemm(benchmark, n): |
|
|
|
def test_gemm(benchmark, n, variant): |
|
|
|
rndm = np.random.RandomState(1234) |
|
|
|
a = np.array(rndm.uniform(size=(n, n)), dtype=float, order='F') |
|
|
|
b = np.array(rndm.uniform(size=(n, n)), dtype=float, order='F') |
|
|
|
c = np.empty((n, n), dtype=float, order='F') |
|
|
|
result = benchmark(run_gemm, a, b, c) |
|
|
|
dtyp = dtype_map[variant] |
|
|
|
a = np.array(rndm.uniform(size=(n, n)), dtype=dtyp, order='F') |
|
|
|
b = np.array(rndm.uniform(size=(n, n)), dtype=dtyp, order='F') |
|
|
|
c = np.empty((n, n), dtype=dtyp, order='F') |
|
|
|
gemm = ow.get_func('gemm', variant) |
|
|
|
result = benchmark(run_gemm, a, b, c, gemm) |
|
|
|
assert result is c |
|
|
|
|
|
|
|
|
|
|
|
# dsyrk |
|
|
|
|
|
|
|
syrk_sizes = [100, 1000] |
|
|
|
syrk_sizes = [100, 200, 400, 600, 800, 1000] |
|
|
|
|
|
|
|
|
|
|
|
def run_syrk(a, c): |
|
|
|
res = dsyrk(1.0, a, c=c, overwrite_c=True) |
|
|
|
def run_syrk(a, c, func): |
|
|
|
res = func(1.0, a, c=c, overwrite_c=True) |
|
|
|
return res |
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize('variant', ['s', 'd', 'c', 'z']) |
|
|
|
@pytest.mark.parametrize('n', syrk_sizes) |
|
|
|
def test_syrk(benchmark, n): |
|
|
|
def test_syrk(benchmark, n, variant): |
|
|
|
rndm = np.random.RandomState(1234) |
|
|
|
a = np.array(rndm.uniform(size=(n, n)), dtype=float, order='F') |
|
|
|
c = np.empty((n, n), dtype=float, order='F') |
|
|
|
result = benchmark(run_syrk, a, c) |
|
|
|
dtyp = dtype_map[variant] |
|
|
|
a = np.array(rndm.uniform(size=(n, n)), dtype=dtyp, order='F') |
|
|
|
c = np.empty((n, n), dtype=dtyp, order='F') |
|
|
|
syrk = ow.get_func('syrk', variant) |
|
|
|
result = benchmark(run_syrk, a, c, syrk) |
|
|
|
assert result is c |
|
|
|
|
|
|
|
|
|
|
@@ -110,21 +126,25 @@ def test_syrk(benchmark, n): |
|
|
|
|
|
|
|
# linalg.solve |
|
|
|
|
|
|
|
gesv_sizes = [100, 1000] |
|
|
|
gesv_sizes = [100, 200, 400, 600, 800, 1000] |
|
|
|
|
|
|
|
|
|
|
|
def run_gesv(a, b): |
|
|
|
res = dgesv(a, b, overwrite_a=True, overwrite_b=True) |
|
|
|
def run_gesv(a, b, func): |
|
|
|
res = func(a, b, overwrite_a=True, overwrite_b=True) |
|
|
|
return res |
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize('variant', ['s', 'd', 'c', 'z']) |
|
|
|
@pytest.mark.parametrize('n', gesv_sizes) |
|
|
|
def test_gesv(benchmark, n): |
|
|
|
def test_gesv(benchmark, n, variant): |
|
|
|
rndm = np.random.RandomState(1234) |
|
|
|
a = (np.array(rndm.uniform(size=(n, n)), dtype=float, order='F') + |
|
|
|
np.eye(n, order='F')) |
|
|
|
b = np.array(rndm.uniform(size=(n, 1)), order='F') |
|
|
|
lu, piv, x, info = benchmark(run_gesv, a, b) |
|
|
|
dtyp = dtype_map[variant] |
|
|
|
|
|
|
|
a = (np.array(rndm.uniform(size=(n, n)), dtype=dtyp, order='F') + |
|
|
|
np.eye(n, dtype=dtyp, order='F')) |
|
|
|
b = np.array(rndm.uniform(size=(n, 1)), dtype=dtyp, order='F') |
|
|
|
gesv = ow.get_func('gesv', variant) |
|
|
|
lu, piv, x, info = benchmark(run_gesv, a, b, gesv) |
|
|
|
assert lu is a |
|
|
|
assert x is b |
|
|
|
assert info == 0 |
|
|
@@ -135,49 +155,63 @@ def test_gesv(benchmark, n): |
|
|
|
gesdd_sizes = [(100, 5), (1000, 222)] |
|
|
|
|
|
|
|
|
|
|
|
def run_gesdd(a, lwork): |
|
|
|
res = dgesdd(a, lwork=lwork, full_matrices=False, overwrite_a=False) |
|
|
|
def run_gesdd(a, lwork, func): |
|
|
|
res = func(a, lwork=lwork, full_matrices=False, overwrite_a=False) |
|
|
|
return res |
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize('variant', ['s', 'd']) |
|
|
|
@pytest.mark.parametrize('mn', gesdd_sizes) |
|
|
|
def test_gesdd(benchmark, mn): |
|
|
|
def test_gesdd(benchmark, mn, variant): |
|
|
|
m, n = mn |
|
|
|
rndm = np.random.RandomState(1234) |
|
|
|
a = np.array(rndm.uniform(size=(m, n)), dtype=float, order='F') |
|
|
|
dtyp = dtype_map[variant] |
|
|
|
|
|
|
|
a = np.array(rndm.uniform(size=(m, n)), dtype=dtyp, order='F') |
|
|
|
|
|
|
|
gesdd_lwork = ow.get_func('gesdd_lwork', variant) |
|
|
|
|
|
|
|
lwork, info = dgesdd_lwork(m, n) |
|
|
|
lwork, info = gesdd_lwork(m, n) |
|
|
|
lwork = int(lwork) |
|
|
|
assert info == 0 |
|
|
|
|
|
|
|
u, s, vt, info = benchmark(run_gesdd, a, lwork) |
|
|
|
gesdd = ow.get_func('gesdd', variant) |
|
|
|
u, s, vt, info = benchmark(run_gesdd, a, lwork, gesdd) |
|
|
|
|
|
|
|
assert info == 0 |
|
|
|
np.testing.assert_allclose(u @ np.diag(s) @ vt, a, atol=1e-13) |
|
|
|
|
|
|
|
atol = {'s': 1e-5, 'd': 1e-13} |
|
|
|
|
|
|
|
np.testing.assert_allclose(u @ np.diag(s) @ vt, a, atol=atol[variant]) |
|
|
|
|
|
|
|
|
|
|
|
# linalg.eigh |
|
|
|
|
|
|
|
syev_sizes = [50, 200] |
|
|
|
syev_sizes = [50, 64, 128, 200] |
|
|
|
|
|
|
|
|
|
|
|
def run_syev(a, lwork): |
|
|
|
res = dsyev(a, lwork=lwork, overwrite_a=True) |
|
|
|
def run_syev(a, lwork, func): |
|
|
|
res = func(a, lwork=lwork, overwrite_a=True) |
|
|
|
return res |
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize('variant', ['s', 'd']) |
|
|
|
@pytest.mark.parametrize('n', syev_sizes) |
|
|
|
def test_syev(benchmark, n): |
|
|
|
def test_syev(benchmark, n, variant): |
|
|
|
rndm = np.random.RandomState(1234) |
|
|
|
dtyp = dtype_map[variant] |
|
|
|
|
|
|
|
a = rndm.uniform(size=(n, n)) |
|
|
|
a = np.asarray(a + a.T, dtype=float, order='F') |
|
|
|
a = np.asarray(a + a.T, dtype=dtyp, order='F') |
|
|
|
a_ = a.copy() |
|
|
|
|
|
|
|
dsyev_lwork = ow.get_func('syev_lwork', variant) |
|
|
|
lwork, info = dsyev_lwork(n) |
|
|
|
lwork = int(lwork) |
|
|
|
assert info == 0 |
|
|
|
|
|
|
|
w, v, info = benchmark(run_syev, a, lwork) |
|
|
|
syev = ow.get_func('syev', variant) |
|
|
|
w, v, info = benchmark(run_syev, a, lwork, syev) |
|
|
|
|
|
|
|
assert info == 0 |
|
|
|
assert a is v # overwrite_a=True |
|
|
|