Merge pull request #731 from DiamondLightSource/peak-reslice-estimator

Add CPU estimation of largest reslice in pipeline

Author: yousefmoazzam

httomo/cli.py

@@ -562,17 +562,18 @@ def estimate_cpu_memory(in_data_file: Path, pipeline_file: Path, nprocs: int) ->
 
     section_memory_peak = 0
     for idx in range(len(sections)):
+        mem, previewed_shape = estimate_section_memory(
+            nprocs,
+            0,
+            None,
+            dtype,
+            previewed_shape,
+            sections,
+            idx,
+        )
         section_memory_peak = max(
             section_memory_peak,
-            estimate_section_memory(
-                nprocs,
-                0,
-                None,
-                dtype,
-                previewed_shape,
-                sections,
-                idx,
-            ),
+            mem,
         )
 
     return section_memory_peak

httomo/runner/dataset_store_backing.py

@@ -52,7 +52,7 @@ def calculate_section_output_chunk_shape(
     for method in section.methods:
         if method.memory_gpu is None:
             continue
-        output_non_slice_dims = method.calculate_output_dims(input_non_slice_dims)
+        output_non_slice_dims = method.calculate_output_dims(output_non_slice_dims)
 
     output_chunk_shape = list(output_non_slice_dims)
     output_chunk_shape.insert(slicing_dim, chunk_shape[slicing_dim])
@@ -74,7 +74,7 @@ def estimate_section_memory(
     sections: List[Section],
     section_idx: int,
     consider_pinned_memory_pool: bool = False,
-) -> int:
+) -> tuple[int, tuple[int, int, int]]:
     # Get chunk shape created by reader of section `n` (the current section) that will account
     # for padding. This chunk shape is based on the chunk shape written by the writer of
     # section `n - 1` (the previous section)
@@ -150,12 +150,26 @@ def estimate_section_memory(
         # See https://github.com/cupy/cupy/issues/9813
         cupy_transfer_overhead = total_mem
 
+    # Calculate the shape of the global data of the output of the section to pass back to the
+    # caller.
+    #
+    # NOTE: The caller will only use this if performing CPU memory estimation of the entire
+    # pipeline not at runtime of htttomo (ie, the `memory-check` CLI command). If performing
+    # CPU memory estimation of a section during runtime of httomo, the caller won't use this
+    # value.
+    current_section_slicing_dim = _get_slicing_dim(sections[section_idx].pattern) - 1
+    output_global_shape = list(output_chunk_shape)
+    output_global_shape[current_section_slicing_dim] = (
+        output_global_shape[current_section_slicing_dim] * nprocs
+    )
+
     return (
         padded_input_chunk_bytes
         + output_chunk_bytes
         + reslice_bytes
         + cupy_pinned_cpu_pool_memory
-        + cupy_transfer_overhead
+        + cupy_transfer_overhead,
+        output_global_shape,
     )
 
 
@@ -167,7 +181,7 @@ def determine_store_backing(
     global_shape: Tuple[int, int, int],
     section_idx: int,
 ) -> DataSetStoreBacking:
-    section_memory = estimate_section_memory(
+    section_memory, _ = estimate_section_memory(
         comm.size,
         comm.rank,
         comm.allgather,