avoid potential CUDA out of bounds memory access in test case

rapidsai · Mar 5, 2025 · 60d1212 · 60d1212
1 parent a2c0a05
commit 60d1212
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/python/cucim/src/cucim/skimage/filters/tests/test_thresholding.py b/python/cucim/src/cucim/skimage/filters/tests/test_thresholding.py
@@ -637,6 +637,12 @@ def test_mean():
 @pytest.mark.parametrize("kwargs", [{}, {"nbins": 300000}])
 @pytest.mark.parametrize("dtype", [cp.uint8, cp.int16, cp.float16, cp.float32])
 def test_triangle_uniform_images(dtype, kwargs):
+    if dtype == cp.float16:
+        # Avoid potential NaN bin edges if nbins exceeds float16 range.
+        # (NaN in bin edges can cause out of bounds CUDA memory access in
+        # `cp.histogram`)
+        if "nbins" in kwargs:
+            kwargs["nbins"] = min(kwargs["nbins"], np.finfo(np.float16).max)
     assert threshold_triangle(cp.zeros((10, 10), dtype=dtype), **kwargs) == 0
     assert threshold_triangle(cp.ones((10, 10), dtype=dtype), **kwargs) == 1
     assert threshold_triangle(cp.full((10, 10), 2, dtype=dtype), **kwargs) == 2