Attempt to repair device memory fragmentation

Author: mfep

httomolibgpu/prep/stripe.py

@@ -615,6 +615,12 @@ def apply(
         return yl
 
 
+def _repair_memory_fragmentation_if_needed(fragmentation_threshold: float = 0.2):
+    pool = cp.get_default_memory_pool()
+    total = pool.total_bytes()
+    if (total / pool.used_bytes()) - 1 > fragmentation_threshold:
+        pool.free_all_blocks()
+
 def remove_stripe_fw(
     data: cp.ndarray,
     sigma: float = 1,
@@ -710,7 +716,7 @@ def remove_stripe_fw(
         for c in cc:
             mem_stack.free(np.prod(c) * np.float32().itemsize)
         mem_stack.free(np.prod(sli_shape) * np.float32().itemsize)
-        return mem_stack.highwater
+        return int(mem_stack.highwater * 1.1)
 
     sli = cp.zeros(sli_shape, dtype="float32")
     sli[:, 0, (nproj_pad - nproj) // 2 : (nproj_pad + nproj) // 2] = data.swapaxes(0, 1)
@@ -736,12 +742,14 @@ def remove_stripe_fw(
             cc[k][:, 0, 1] = cp.fft.ifft(ifft_in, my, axis=1).real
         del ifft_plan
         del ifft_in
+        _repair_memory_fragmentation_if_needed()
 
     # Wavelet reconstruction.
     for k in range(level)[::-1]:
         shape0 = cc[k][0, 0, 1].shape
         sli = sli[:, :, : shape0[0], : shape0[1]]
         sli = ifm.apply((sli, cc[k]))
+        _repair_memory_fragmentation_if_needed()
 
     data = sli[:, 0, (nproj_pad - nproj) // 2 : (nproj_pad + nproj) // 2, :ni]
     data = data.swapaxes(0, 1)

tests/test_prep/test_stripe.py

@@ -109,12 +109,12 @@ def test_remove_stripe_fw_calc_mem(slices, level, dim_x, wname, ensure_clean_mem
 
 @pytest.mark.parametrize("wname", ["haar", "db4", "sym5", "sym16", "bior4.4"])
 @pytest.mark.parametrize(
-    "slices", [177, 239, 320, 490, 607, 803, 859, 902, 951, 1019, 1074, 1105]
+    "slices", [38, 177, 239, 320, 490, 607, 803, 859, 902, 951, 1019, 1074, 1105]
 )
 @pytest.mark.parametrize("level", [None, 7, 11])
-def test_remove_stripe_fw_calc_mem_big(wname, slices, level, ensure_clean_memory):
-    dim_y = 901
-    dim_x = 1200
+@pytest.mark.parametrize("dims", [(901, 1200), (1801, 2560)])
+def test_remove_stripe_fw_calc_mem_big(wname, slices, level, dims, ensure_clean_memory):
+    dim_y, dim_x = dims
     data_shape = (slices, dim_x, dim_y)
     try:
         estimated_mem_peak = remove_stripe_fw(