4
4
Module: test.py
5
5
6
6
This module contains a suite of tests for the `simsimd` package.
7
- It compares various SIMD kernels (like Dot-products, squared Euclidean, and Cosine distances)
8
- with their NumPy or baseline counterparts, testing accuracy for different data types including
7
+ It compares various SIMD kernels (like Dot-products, squared Euclidean, and Cosine distances)
8
+ with their NumPy or baseline counterparts, testing accuracy for different data types including
9
9
floating-point, integer, and complex numbers.
10
10
11
11
The tests cover:
@@ -1421,31 +1421,36 @@ def test_cdist_hamming(ndim, out_dtype, capability):
1421
1421
def test_gil_free_threading ():
1422
1422
"""Test SimSIMD in Python 3.13t free-threaded mode if available."""
1423
1423
import sys
1424
-
1424
+ import sysconfig
1425
+
1425
1426
# Check if we're in a GIL-free environment
1427
+ # https://py-free-threading.github.io/running-gil-disabled/
1426
1428
version = sys .version_info
1427
1429
if version .major == 3 and version .minor >= 13 :
1428
- if sys ._is_gil_enabled ():
1429
- pytest .skip ("GIL is enabled, skipping GIL-free threading test" )
1430
+ is_free_threaded = bool (sysconfig .get_config_var ("Py_GIL_DISABLED" ))
1431
+ if not is_free_threaded :
1432
+ pytest .skip ("Uses non-free-threaded Python, skipping GIL-related tests" )
1433
+ if sys ._is_gil_enabled ():
1434
+ pytest .skip ("GIL is enabled, skipping GIL-related tests" )
1430
1435
else :
1431
- pytest .skip ("Python < 3.13t, skipping GIL-free threading test " )
1432
-
1436
+ pytest .skip ("Python < 3.13t, skipping GIL-related tests " )
1437
+
1433
1438
import multiprocessing
1434
1439
import concurrent .futures
1435
1440
1436
1441
num_threads = multiprocessing .cpu_count ()
1437
1442
vectors_a = np .random .rand (32 * 1024 * num_threads , 1024 ).astype (np .float32 )
1438
1443
vectors_b = np .random .rand (32 * 1024 * num_threads , 1024 ).astype (np .float32 )
1439
1444
distances = np .zeros (vectors_a .shape [0 ], dtype = np .float32 )
1440
-
1445
+
1441
1446
def compute_batch (start_idx , end_idx ) -> float :
1442
1447
"""Compute cosine distances for a batch."""
1443
1448
slice_a = vectors_a [start_idx :end_idx ]
1444
1449
slice_b = vectors_b [start_idx :end_idx ]
1445
1450
slice_distances = distances [start_idx :end_idx ]
1446
- simd .cosine (slice_a , slice_b , out = slice_distances )
1451
+ simd .cosine (slice_a , slice_b , out = slice_distances )
1447
1452
return sum (slice_distances )
1448
-
1453
+
1449
1454
def compute_with_threads (threads : int ) -> float :
1450
1455
"""Compute cosine distances using multiple threads."""
1451
1456
chunk_size = len (vectors_a ) // threads
@@ -1455,36 +1460,37 @@ def compute_with_threads(threads: int) -> float:
1455
1460
start_idx = i * chunk_size
1456
1461
end_idx = (i + 1 ) * chunk_size if i < threads - 1 else len (vectors_a )
1457
1462
futures .append (executor .submit (compute_batch , start_idx , end_idx ))
1458
-
1463
+
1459
1464
total_sum = 0.0
1460
1465
for future in concurrent .futures .as_completed (futures ):
1461
1466
total_sum += future .result ()
1462
-
1467
+
1463
1468
return total_sum
1464
-
1469
+
1465
1470
# Dual-threaded baseline is better than single-threaded,
1466
1471
# as it will include the overhead of thread management.
1467
1472
start_time = time .time ()
1468
1473
baseline_sum = compute_with_threads (2 )
1469
1474
end_time = time .time ()
1470
1475
baseline_duration = end_time - start_time
1471
-
1476
+
1472
1477
# Multi-threaded execution, using all available threads
1473
1478
start_time = time .time ()
1474
1479
multi_sum = compute_with_threads (num_threads )
1475
1480
end_time = time .time ()
1476
1481
multi_duration = end_time - start_time
1477
1482
1478
1483
# Verify results are the same length and reasonable
1479
- assert np .allclose (baseline_sum , multi_sum , atol = SIMSIMD_ATOL , rtol = SIMSIMD_RTOL ), \
1480
- f"Results differ: baseline { baseline_sum } vs multi-threaded { multi_sum } "
1484
+ assert np .allclose (
1485
+ baseline_sum , multi_sum , atol = SIMSIMD_ATOL , rtol = SIMSIMD_RTOL
1486
+ ), f"Results differ: baseline { baseline_sum } vs multi-threaded { multi_sum } "
1481
1487
1482
- # Warn if multi-threaded execution is slower than the baseline
1488
+ # Warn if multi-threaded execution is slower than the baseline
1483
1489
if baseline_duration < multi_duration :
1484
1490
pytest .warns (
1485
1491
UserWarning ,
1486
1492
f"{ num_threads } -threaded execution took longer than 2-threaded baseline: { multi_duration :.2f} s vs { baseline_duration :.2f} s" ,
1487
- )
1493
+ )
1488
1494
1489
1495
1490
1496
if __name__ == "__main__" :
0 commit comments