Update mem estimator and tests for very large sizes

Author: mfep

httomolibgpu/prep/stripe.py

@@ -337,8 +337,7 @@ def _conv_transpose2d(
     out_shape = [b, ci, hi, wi]
     if mem_stack:
         # The trouble here is that we allocate more than the returned size
-        out_actual_bytes = np.prod(out_shape) * np.float32().itemsize
-        mem_stack.malloc(out_actual_bytes)
+        mem_stack.malloc((np.prod(out_shape) + w.size) * np.float32().itemsize)
         if pad != 0:
             new_out_shape = [
                 out_shape[0],
@@ -347,7 +346,7 @@ def _conv_transpose2d(
                 out_shape[3] - 2 * pad[1],
             ]
             mem_stack.malloc(np.prod(new_out_shape) * np.float32().itemsize)
-            mem_stack.free(np.prod(out_shape) * np.float32().itemsize)
+            mem_stack.free((np.prod(out_shape) + w.size) * np.float32().itemsize)
             out_shape = new_out_shape
         return out_shape
 
@@ -673,7 +672,7 @@ def remove_stripe_fw(
             mem_stack.free(np.prod(c) * np.float32().itemsize)
         mem_stack.malloc(np.prod(data) * 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)

tests/test_prep/test_stripe.py

@@ -86,8 +86,8 @@ def ensure_clean_memory():
 
 
 @pytest.mark.parametrize("wname", ["haar", "db4", "sym5", "sym16", "bior4.4"])
-@pytest.mark.parametrize("slices", [55, 80])
-@pytest.mark.parametrize("level", [None, 1, 3, 7, 11])
+@pytest.mark.parametrize("slices", [3, 7, 32, 61, 109, 120, 150])
+@pytest.mark.parametrize("level", [None, 1, 3, 11])
 @pytest.mark.parametrize("dim_x", [128, 140])
 def test_remove_stripe_fw_calc_mem(slices, level, dim_x, wname, ensure_clean_memory):
     dim_y = 159
@@ -105,31 +105,32 @@ def test_remove_stripe_fw_calc_mem(slices, level, dim_x, wname, ensure_clean_mem
     assert hook.max_mem == 0
 
     assert actual_mem_peak * 0.99 <= estimated_mem_peak
-    assert estimated_mem_peak <= actual_mem_peak * 1.2
+    assert estimated_mem_peak <= actual_mem_peak * 1.3
 
 
-@pytest.mark.parametrize("wname", ['db4', 'sym16'])
+@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])
-def test_remove_stripe_fw_calc_mem_big(wname, slices, ensure_clean_memory):
+@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
     data_shape = (slices, dim_x, dim_y)
     hook = MaxMemoryHook()
     with hook:
-        estimated_mem_peak = remove_stripe_fw(data_shape, wname=wname, calc_peak_gpu_mem=True)
+        estimated_mem_peak = remove_stripe_fw(data_shape, wname=wname, level=level, calc_peak_gpu_mem=True)
     assert hook.max_mem == 0
     av_mem = cp.cuda.Device().mem_info[0]
-    if av_mem < estimated_mem_peak * 1.1:
+    if av_mem < estimated_mem_peak:
         pytest.skip("Not enough GPU memory to run this test")
 
     hook = MaxMemoryHook()
     with hook:
         data = cp.random.random_sample(data_shape, dtype=np.float32)
-        remove_stripe_fw(data, wname=wname)
+        remove_stripe_fw(data, wname=wname, level=level)
     actual_mem_peak = hook.max_mem
 
     assert actual_mem_peak * 0.99 <= estimated_mem_peak
-    assert estimated_mem_peak <= actual_mem_peak * 1.2
+    assert estimated_mem_peak <= actual_mem_peak * 1.3
 
 
 @pytest.mark.parametrize("angles", [180, 181])