| |
| |
| """The row-parallel dgemm used by VBx (diar_exact::npblas::dot_gemm_f64_rows_par) must equal the one-call |
| OpenBLAS product bit for bit, for both VBx call sites' layouts, over the shapes long files produce: rows of C split |
| across threads leave every element's kernel arithmetic unchanged.""" |
| import numpy as np |
| import pytest |
|
|
| SHAPES = [(n, k) for n in (37, 640, 5000, 20774) for k in (2, 3, 7, 40, 70, 513, 1178, 1981)] |
|
|
|
|
| @pytest.mark.parametrize("n,k", SHAPES) |
| @pytest.mark.parametrize("rows", [16, 64, 256]) |
| def test_vbx_gemm_sites(core, n, k, rows): |
| rng = np.random.default_rng(n * 7919 + k) |
| d = 128 |
| |
| gamma = rng.random((n, k)) * rng.random((n, 1)) |
| rho = rng.standard_normal((n, d)) * 3.0 |
| alpha = rng.standard_normal((k, d)) * 0.1 |
| |
| one, par = core.gemm_pair(True, False, k, d, n, gamma.ravel(), k, rho.ravel(), d, rows) |
| assert one.tobytes() == par.tobytes() |
| |
| assert one.tobytes() == gamma.T.dot(rho).ravel().tobytes() |
| |
| one, par = core.gemm_pair(False, True, n, k, d, rho.ravel(), d, alpha.ravel(), d, rows) |
| assert one.tobytes() == par.tobytes() |
| assert one.tobytes() == rho.dot(alpha.T).ravel().tobytes() |
|
|
