11# Copyright ExplosionAI GmbH, released under BSD.
2- from math import sqrt , floor
3-
4- from hypothesis import given , assume
5- from hypothesis .strategies import integers
6-
72import numpy as np
8- from numpy .testing import assert_allclose
93import pytest
4+ from hypothesis import given
5+ from hypothesis import strategies as st
6+ from numpy .testing import assert_allclose
107
11- from blis_tests_common import ndarrays
8+ from blis_tests_common import dtypes_tols , ndarrays , run_threaded
129from blis .py import gemm
1310
1411
15- def _stretch_matrix (data , m , n ):
16- ratio = sqrt (len (data ) / (m * n ))
17- m = int (floor (m * ratio ))
18- n = int (floor (n * ratio ))
19- data = np .ascontiguousarray (data [: m * n ], dtype = data .dtype )
20- return data .reshape ((m , n )), m , n
21-
22-
23- def _reshape_for_gemm (
24- A , B , a_rows , a_cols , out_cols , dtype , trans_a = False , trans_b = False
25- ):
26- A , a_rows , a_cols = _stretch_matrix (A , a_rows , a_cols )
27- if len (B ) < a_cols or a_cols < 1 :
28- return (None , None , None )
29- b_cols = int (floor (len (B ) / a_cols ))
30- B = np .ascontiguousarray (B .flatten ()[: a_cols * b_cols ], dtype = dtype )
31- B = B .reshape ((a_cols , b_cols ))
32- out_cols = B .shape [1 ]
33- C = np .zeros (shape = (A .shape [0 ], B .shape [1 ]), dtype = dtype )
34- if trans_a :
35- A = np .ascontiguousarray (A .T , dtype = dtype )
36- return A , B , C
12+ @st .composite
13+ def gemm_arrays (draw ):
14+ """Draw two NumPy arrays with shapes (l, m) and (m, n) of the same dtype"""
15+ max_size = 100
16+ l = draw (st .integers (min_value = 1 , max_value = max_size )) # noqa: E741
17+ n = draw (st .integers (min_value = 1 , max_value = max_size ))
18+ m = draw (st .integers (min_value = 1 , max_value = max_size // max (l , n )))
19+ dtype , atol , rtol = draw (dtypes_tols )
20+ A = draw (ndarrays ((l , m ), lo = - 100.0 , hi = 100.0 , dtype = dtype ))
21+ B = draw (ndarrays ((m , n ), lo = - 100.0 , hi = 100.0 , dtype = dtype ))
22+ return A , B , atol , rtol
3723
3824
3925def test_incompatible_shape ():
@@ -47,41 +33,25 @@ def test_incompatible_shape():
4733 gemm (np .zeros ((3 , 2 )), np .zeros ((3 , 2 )), trans1 = True , trans2 = True )
4834
4935
50- @given (
51- ndarrays (min_len = 10 , max_len = 100 , min_val = - 100.0 , max_val = 100.0 , dtype = "float64" ),
52- ndarrays (min_len = 10 , max_len = 100 , min_val = - 100.0 , max_val = 100.0 , dtype = "float64" ),
53- integers (min_value = 2 , max_value = 1000 ),
54- integers (min_value = 2 , max_value = 1000 ),
55- integers (min_value = 2 , max_value = 1000 ),
56- )
57- def test_memoryview_double_notrans (A , B , a_rows , a_cols , out_cols ):
58- A , B , C = _reshape_for_gemm (A , B , a_rows , a_cols , out_cols , "float64" )
59- assume (A is not None )
60- assume (B is not None )
61- assume (C is not None )
62- assume (A .size >= 1 )
63- assume (B .size >= 1 )
64- assume (C .size >= 1 )
65- gemm (A , B , out = C )
36+ @given (gemm_arrays ())
37+ def test_memoryview_notrans (arrays ):
38+ A , B , atol , rtol = arrays
6639 numpy_result = A .dot (B )
67- assert_allclose (numpy_result , C , atol = 1e-3 , rtol = 1e-3 )
40+ C = np .zeros_like (numpy_result ) # (l, n)
41+ gemm (A , B , out = C )
42+ assert_allclose (C , numpy_result , atol = atol , rtol = rtol )
6843
6944
70- @given (
71- ndarrays (min_len = 10 , max_len = 100 , min_val = - 100.0 , max_val = 100.0 , dtype = "float32" ),
72- ndarrays (min_len = 10 , max_len = 100 , min_val = - 100.0 , max_val = 100.0 , dtype = "float32" ),
73- integers (min_value = 2 , max_value = 1000 ),
74- integers (min_value = 2 , max_value = 1000 ),
75- integers (min_value = 2 , max_value = 1000 ),
76- )
77- def test_memoryview_float_notrans (A , B , a_rows , a_cols , out_cols ):
78- A , B , C = _reshape_for_gemm (A , B , a_rows , a_cols , out_cols , dtype = "float32" )
79- assume (A is not None )
80- assume (B is not None )
81- assume (C is not None )
82- assume (A .size >= 1 )
83- assume (B .size >= 1 )
84- assume (C .size >= 1 )
85- gemm (A , B , out = C )
45+ @given (gemm_arrays ())
46+ @pytest .mark .thread_unsafe (reason = "Uses run_threaded" )
47+ def test_threads_share_input (arrays ):
48+ """Test when multiple threads share the same input arrays."""
49+ A , B , atol , rtol = arrays
8650 numpy_result = A .dot (B )
87- assert_allclose (numpy_result , C , atol = 1e-3 , rtol = 1e-3 )
51+
52+ def f ():
53+ C = np .zeros_like (numpy_result ) # (l, n)
54+ gemm (A , B , out = C )
55+ assert_allclose (C , numpy_result , atol = atol , rtol = rtol )
56+
57+ run_threaded (f )
0 commit comments