Merge pull request #77 from DiamondLightSource/projection_chunking

Projection chunking

Author: dkazanc

httomo_backends/methods_database/packages/backends/httomolibgpu/supporting_funcs/recon/algorithm.py

@@ -192,7 +192,11 @@ def _calc_memory_bytes_LPRec3d_tomobar(
         detector_pad = 0        
 
     min_mem_usage_filter = False
+    if "min_mem_usage_filter" in kwargs:
+        min_mem_usage_filter = kwargs["min_mem_usage_filter"]
     min_mem_usage_ifft2 = False
+    if "min_mem_usage_ifft2" in kwargs:
+        min_mem_usage_ifft2 = kwargs["min_mem_usage_ifft2"]
 
     angles_tot = non_slice_dims_shape[0]
     DetectorsLengthH_prepad = non_slice_dims_shape[1]
@@ -339,21 +343,21 @@ def add_to_memory_counters(amount, per_slice: bool):
 
     add_to_memory_counters(tmp_p_input_slice, True)
     if min_mem_usage_filter:
-        add_to_memory_counters(rfft_plan_slice_size / 4, False)
-        add_to_memory_counters(irfft_plan_slice_size / 4, False)
-        add_to_memory_counters(padded_tmp_p_input_slice, False)
+        add_to_memory_counters(rfft_plan_slice_size / 4 / projection_chunk_count, False)
+        add_to_memory_counters(irfft_plan_slice_size / 4 / projection_chunk_count, False)
+        add_to_memory_counters(padded_tmp_p_input_slice / projection_chunk_count, False)
 
-        add_to_memory_counters(rfft_result_size, False)
-        add_to_memory_counters(filtered_rfft_result_size, False)
-        add_to_memory_counters(-rfft_result_size, False)
-        add_to_memory_counters(-padded_tmp_p_input_slice, False)
+        add_to_memory_counters(rfft_result_size / projection_chunk_count, False)
+        add_to_memory_counters(filtered_rfft_result_size / projection_chunk_count, False)
+        add_to_memory_counters(-rfft_result_size / projection_chunk_count, False)
+        add_to_memory_counters(-padded_tmp_p_input_slice / projection_chunk_count, False)
 
-        add_to_memory_counters(irfft_scratch_memory_size, False)
-        add_to_memory_counters(-irfft_scratch_memory_size, False)
-        add_to_memory_counters(irfft_result_size, False)
-        add_to_memory_counters(-filtered_rfft_result_size, False)
+        add_to_memory_counters(irfft_scratch_memory_size / projection_chunk_count, False)
+        add_to_memory_counters(-irfft_scratch_memory_size / projection_chunk_count, False)
+        add_to_memory_counters(irfft_result_size / projection_chunk_count, False)
+        add_to_memory_counters(-filtered_rfft_result_size / projection_chunk_count, False)
 
-        add_to_memory_counters(-irfft_result_size, False)
+        add_to_memory_counters(-irfft_result_size / projection_chunk_count, False)
     else:
         add_to_memory_counters(rfft_plan_slice_size / chunk_count / projection_chunk_count * 2, True)
         add_to_memory_counters(irfft_plan_slice_size / chunk_count / projection_chunk_count * 2, True)
@@ -402,8 +406,10 @@ def add_to_memory_counters(amount, per_slice: bool):
     add_to_memory_counters(circular_mask_size, False)
     add_to_memory_counters(after_recon_swapaxis_slice, True)
 
-    if min_mem_usage_ifft2 and min_mem_usage_filter:
-        return (tot_memory_bytes * 1.1 + 30 * 1024 * 1024, fixed_amount)
+    if min_mem_usage_filter and min_mem_usage_ifft2:
+        return (tot_memory_bytes * 1.15, fixed_amount)
+    elif min_mem_usage_filter and not min_mem_usage_ifft2:
+        return (tot_memory_bytes + 60 * 1024 * 1024, fixed_amount)
     else:
         return (tot_memory_bytes * 1.1, fixed_amount)
 

tests/test_httomolibgpu.py

@@ -543,6 +543,7 @@ def test_recon_FBP3d_tomobar_memoryhook(
 
 
 @pytest.mark.cupy
+@pytest.mark.parametrize("min_mem_usage_filter_ifft2", [(False, False), (True, False), (True, True)])
 @pytest.mark.parametrize("padding_detx", [0, 10, 50, 100])
 @pytest.mark.parametrize("projections", [1500, 1801, 2560])
 @pytest.mark.parametrize("detX_size", [2560])
@@ -553,6 +554,7 @@ def test_recon_LPRec3d_tomobar_0_pi_memoryhook(
     detX_size,
     projections,
     projection_angle_range,
+    min_mem_usage_filter_ifft2,
     padding_detx,
     ensure_clean_memory,
 ):
@@ -562,12 +564,14 @@ def test_recon_LPRec3d_tomobar_0_pi_memoryhook(
         projections,
         projection_angle_range,
         padding_detx,
+        min_mem_usage_filter_ifft2,
         ensure_clean_memory,
     )
 
 
 @pytest.mark.full
 @pytest.mark.cupy
+@pytest.mark.parametrize("min_mem_usage_filter_ifft2", [(False, False), (True, False), (True, True)])
 @pytest.mark.parametrize("padding_detx", [0, 10, 50, 100, 800])
 @pytest.mark.parametrize("projections", [1500, 1801, 2560, 3601])
 @pytest.mark.parametrize("detX_size", [2560])
@@ -579,6 +583,7 @@ def test_recon_LPRec3d_tomobar_0_pi_memoryhook_full(
     projections,
     projection_angle_range,
     padding_detx,
+    min_mem_usage_filter_ifft2,
     ensure_clean_memory,
 ):
     __test_recon_LPRec3d_tomobar_memoryhook_common(
@@ -587,12 +592,14 @@ def test_recon_LPRec3d_tomobar_0_pi_memoryhook_full(
         projections,
         projection_angle_range,
         padding_detx,
+        min_mem_usage_filter_ifft2,
         ensure_clean_memory,
     )
 
 
 @pytest.mark.full
 @pytest.mark.cupy
+@pytest.mark.parametrize("min_mem_usage_filter_ifft2", [(False, False), (True, False), (True, True)])
 @pytest.mark.parametrize("padding_detx", [0, 10, 50, 100, 800])
 @pytest.mark.parametrize("projections", [1500, 1801, 2560, 3601])
 @pytest.mark.parametrize("detX_size", [2560])
@@ -606,6 +613,7 @@ def test_recon_LPRec3d_tomobar_memoryhook_full(
     projections,
     projection_angle_range,
     padding_detx,
+    min_mem_usage_filter_ifft2,
     ensure_clean_memory,
 ):
     __test_recon_LPRec3d_tomobar_memoryhook_common(
@@ -614,6 +622,7 @@ def test_recon_LPRec3d_tomobar_memoryhook_full(
         projections,
         projection_angle_range,
         padding_detx,
+        min_mem_usage_filter_ifft2,
         ensure_clean_memory,
     )
 
@@ -624,6 +633,7 @@ def __test_recon_LPRec3d_tomobar_memoryhook_common(
     projections,
     projection_angle_range,
     padding_detx,
+    min_mem_usage_filter_ifft2,
     ensure_clean_memory,
 ):
     angles_number = projections
@@ -634,6 +644,8 @@ def __test_recon_LPRec3d_tomobar_memoryhook_common(
     )
     kwargs["center"] = 1280
     kwargs["detector_pad"] = padding_detx
+    kwargs["min_mem_usage_filter"] = min_mem_usage_filter_ifft2[0]
+    kwargs["min_mem_usage_ifft2"] = min_mem_usage_filter_ifft2[1]
     kwargs["recon_size"] = detX_size
     kwargs["recon_mask_radius"] = 0.8