Repository: DiamondLightSource/httomo (language: python)

Issue:
Overestimation of max slices for `sino_360_to_180`
Setup:
- `hopper` node
- 1 GPU (Nvidia A100)
- 1 MPI process
- 360 data is `/mnt/gpfs03/scratch/data/imaging/tomography/tmp/testdata/360/112482.nxs`

The following pipeline:
```yaml
- method: standard_tomo
  module_path: httomo.data.hdf.loaders
  parameters:
    name: tomo
    data_path: entry1/tomo_entry/data/data
    image_key_path: entry1/tomo_entry/instrument/detector/image_key
    rotation_angles:
      data_path: /entry1/tomo_entry/data/rotation_angle
- method: find_center_360
  module_path: httomolibgpu.recon.rotation
  parameters:
    ind: mid
    win_width: 10
    side: null
    denoise: true
    norm: false
    use_overlap: false
  id: centering
  side_outputs:
    cor: centre_of_rotation
    overlap: overlap
    side: side
    overlap_position: overlap_position
- method: normalize
  module_path: httomolibgpu.prep.normalize
  parameters:
    cutoff: 10.0
    minus_log: true
    nonnegativity: false
    remove_nans: false
- method: sino_360_to_180
  module_path: httomolibgpu.misc.morph
  parameters:
    overlap: ${{centering.side_outputs.overlap}}
    rotation: right  
```

produces a CUDA OOM error on the first block being processed by the `sino_360_to_180` method:
```
(/dls/science/users/twi18192/conda-envs/httomo) [twi18192@cs05r-sc-hop01-02 httomo (fourier)]$ mpirun -n 1 python -m httomo run /mnt/gpfs03/scratch/data/imaging/tomography/tmp/testdata/360/112482.nxs /dls/science
/users/twi18192/dls_pipelines/pipelines/bench_recons/bench_recon_gridrec_cpu360.yaml /mnt/gpfs03/scratch/d
ata/imaging/tomography/twi18192/out/
2024-05-31 12:07:42.684 | DEBUG    | httomo.utils:<module>:17 - CuPy is installed
Pipeline has been separated into 2 sections
See the full log file at: /mnt/gpfs03/scratch/data/imaging/tomography/twi18192/out/31-05-2024_12_07_43_output/user.log
Running loader (pattern=projection): standard_tomo...
    Finished loader: standard_tomo (httomo) Took 27448.20ms
Section 0 (pattern=projection) with the following methods:
    data_reducer (httomolib)
    find_center_360 (httomolibgpu)
    normalize (httomolibgpu)
     0%|          | 0/4 [00:07<?, ?block/s]
    25%|##5       | 1/4 [00:32<01:27, 29.10s/block]
    50%|#####     | 2/4 [00:52<00:47, 23.90s/block]
    75%|#######5  | 3/4 [01:07<00:20, 20.93s/block]
    --->The center of rotation is (1583.5791015625, 1950.841796875, 1, 614.158203125)
    Finished processing last block
Section 1 (pattern=sinogram) with the following methods:
    sino_360_to_180 (httomolibgpu)
     0%|          | 0/5 [00:00<?, ?block/s]
Traceback (most recent call last):
  File "/dls/science/users/twi18192/conda-envs/httomo/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/dls/science/users/twi18192/conda-envs/httomo/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/dls/science/users/twi18192/httomo/httomo/__main__.py", line 4, in <module>
    main()
  File "/dls/science/users/twi18192/conda-envs/httomo/lib/python3.10/site-packages/click/core.py", line 1157, in __call__
    return self.main(*args, **kwargs)
  File "/dls/science/users/twi18192/conda-envs/httomo/lib/python3.10/site-packages/click/core.py", line 1078, in main
    rv = self.invoke(ctx)
  File "/dls/science/users/twi18192/conda-envs/httomo/lib/python3.10/site-packages/click/core.py", line 1688, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/dls/science/users/twi18192/conda-envs/httomo/lib/python3.10/site-packages/click/core.py", line 1434, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/dls/science/users/twi18192/conda-envs/httomo/lib/python3.10/site-packages/click/core.py", line 783, in invoke
    return __callback(*args, **kwargs)
  File "/dls/science/users/twi18192/httomo/httomo/cli.py", line 205, in run
    runner.execute()
  File "/dls/science/users/twi18192/httomo/httomo/runner/task_runner.py", line 63, in execute
    self._execute_section(section, i)
  File "/dls/science/users/twi18192/httomo/httomo/runner/task_runner.py", line 131, in _execute_section
    block = self._execute_section_block(section, block)
  File "/dls/science/users/twi18192/httomo/httomo/runner/task_runner.py", line 185, in _execute_section_block
    block = self._execute_method(method, block)
  File "/dls/science/users/twi18192/httomo/httomo/runner/task_runner.py", line 217, in _execute_method
    block = method.execute(block)
  File "/dls/science/users/twi18192/httomo/httomo/method_wrappers/generic.py", line 292, in execute
    block = self._run_method(block, args)
  File "/dls/science/users/twi18192/httomo/httomo/method_wrappers/generic.py", line 302, in _run_method
    ret = self._method(**args)
  File "/dls/science/users/twi18192/httomolibgpu/httomolibgpu/misc/morph.py", line 60, in sino_360_to_180
    return __sino_360_to_180(data, overlap, rotation)
  File "/dls/science/users/twi18192/conda-envs/httomo/lib/python3.10/site-packages/nvtx/nvtx.py", line 101, in inner
    result = func(*args, **kwargs)
  File "/dls/science/users/twi18192/httomolibgpu/httomolibgpu/misc/morph.py", line 100, in __sino_360_to_180
    + (weights * data[n : 2 * n, :, -overlap:])[:, :, ::-1]
  File "cupy/_core/core.pyx", line 1281, in cupy._core.core._ndarray_base.__mul__
  File "cupy/_core/_kernel.pyx", line 1347, in cupy._core._kernel.ufunc.__call__
  File "cupy/_core/_kernel.pyx", line 645, in cupy._core._kernel._get_out_args_from_optionals
  File "cupy/_core/core.pyx", line 2779, in cupy._core.core._ndarray_init
  File "cupy/_core/core.pyx", line 237, in cupy._core.core._ndarray_base._init_fast
  File "cupy/cuda/memory.pyx", line 740, in cupy.cuda.memory.alloc
  File "cupy/cuda/memory.pyx", line 1426, in cupy.cuda.memory.MemoryPool.malloc
  File "cupy/cuda/memory.pyx", line 1447, in cupy.cuda.memory.MemoryPool.malloc
  File "cupy/cuda/memory.pyx", line 1118, in cupy.cuda.memory.SingleDeviceMemoryPool.malloc
  File "cupy/cuda/memory.pyx", line 1139, in cupy.cuda.memory.SingleDeviceMemoryPool._malloc
  File "cupy/cuda/memory.pyx", line 1384, in cupy.cuda.memory.SingleDeviceMemoryPool._try_malloc
  File "cupy/cuda/memory.pyx", line 1387, in cupy.cuda.memory.SingleDeviceMemoryPool._try_malloc
cupy.cuda.memory.OutOfMemoryError: Out of memory allocating 13,007,707,648 bytes (allocated so far: 40,688,911,872 bytes).
--------------------------------------------------------------------------
Primary job  terminated normally, but 1 process returned
a non-zero exit code. Per user-direction, the job has been aborted.
--------------------------------------------------------------------------
--------------------------------------------------------------------------
mpirun detected that one or more processes exited with non-zero status, thus causing
the job to be terminated. The first process to do so was:

  Process name: [[55546,1],0]
  Exit code:    1
--------------------------------------------------------------------------
```

As `sino_360_to_180` is the only method in its section, only its memory estimator is determining the max slices for its section, which implies that its memory estimation must be incorrect.


Discussion / hints:
I found some info in an issue in httomolibgpu that could be helpful for addressing this problem (if the shape of the stitched sinogram is calculateable, then the number of bytes in a single sinogram can be calculated from the stitched sinogram shape + data type).

_Originally posted in https://github.com/DiamondLightSource/httomolibgpu/issues/107#issue-1818248219_

> From playing with the test 360 data, it seems that the output shape of the stitching method is able to be determined in advance by the overlap value that is outputted by the `find_center_360` method and passed to the stitching method, doing something like the following:
>
> ```python
> stitched_sino_width = original_sino_width * 2 - math.ceil(overlap)
> ```
>
> where `original_sino_width` is the width of the original 360 sinogram, and `overlap` is the overlap value produced by
> the `find_center_360` method.
>
> So in principle, I think the correct shape of the output of the stitching method could be returned by its memory
> estimator.
One possible way to achieve the above would be:
- change `output_dims_change: False` to `output_dims_change: True`
- change `- multipliers: [2.2]` to `- multipliers: [None]`
- change `- methods: [direct]` to `- methods: [module]`
- provide a memory estimator function for the method

For reference, the current methods database entry for `sino_360_to_180` is the following: https://github.com/DiamondLightSource/httomo/blob/4429711ab89ed1efe86d9c3be260df25fc7f1999/httomo/methods_database/packages/external/httomolibgpu/httomolibgpu.yaml#L22-L30

Fix patch:
diff --git a/httomo/methods_database/packages/external/httomolibgpu/httomolibgpu.yaml b/httomo/methods_database/packages/external/httomolibgpu/httomolibgpu.yaml
index 86c8d59c1..6852b6294 100644
--- a/httomo/methods_database/packages/external/httomolibgpu/httomolibgpu.yaml
+++ b/httomo/methods_database/packages/external/httomolibgpu/httomolibgpu.yaml
@@ -23,14 +23,14 @@ misc:
   morph:
     sino_360_to_180:
       pattern: sinogram
-      output_dims_change: False
+      output_dims_change: True
       implementation: gpu_cupy
       save_result_default: False
       padding: False
       memory_gpu:
         - datasets: [tomo]
-        - multipliers: [2.2]
-        - methods: [direct]
+        - multipliers: [None]
+        - methods: [module]
     data_resampler:
       pattern: all
       output_dims_change: True
diff --git a/httomo/methods_database/packages/external/httomolibgpu/supporting_funcs/misc/morph.py b/httomo/methods_database/packages/external/httomolibgpu/supporting_funcs/misc/morph.py
index ed3cade3c..687ca6029 100644
--- a/httomo/methods_database/packages/external/httomolibgpu/supporting_funcs/misc/morph.py
+++ b/httomo/methods_database/packages/external/httomolibgpu/supporting_funcs/misc/morph.py
@@ -24,9 +24,13 @@
 from typing import Tuple
 import numpy as np
 
+from httomo.runner.output_ref import OutputRef
+
 __all__ = [
     "_calc_memory_bytes_data_resampler",
     "_calc_output_dim_data_resampler",
+    "_calc_memory_bytes_sino_360_to_180",
+    "_calc_output_dim_sino_360_to_180",
 ]
 
 
@@ -54,3 +58,63 @@ def _calc_memory_bytes_data_resampler(
 
     tot_memory_bytes = input_size + output_size + interpolator
     return (tot_memory_bytes, xi)
+
+
+def _calc_output_dim_sino_360_to_180(
+    non_slice_dims_shape: Tuple[int, int],
+    **kwargs,
+) -> Tuple[int, int]:
+    assert "overlap" in kwargs, "Expected overlap in method parameters"
+    overlap_side_output = kwargs["overlap"]
+    assert isinstance(
+        overlap_side_output, OutputRef
+    ), "Expected overlap to be in an OutputRef"
+    overlap: float = overlap_side_output.value
+
+    original_sino_width = non_slice_dims_shape[1]
+    stitched_sino_width = original_sino_width * 2 - math.ceil(overlap)
+    return non_slice_dims_shape[0] // 2, stitched_sino_width
+
+
+def _calc_memory_bytes_sino_360_to_180(
+    non_slice_dims_shape: Tuple[int, int],
+    dtype: np.dtype,
+    **kwargs,
+) -> Tuple[int, int]:
+    assert "overlap" in kwargs, "Expected overlap in method parameters"
+    overlap_side_output = kwargs["overlap"]
+    assert isinstance(
+        overlap_side_output, OutputRef
+    ), "Expected overlap to be in an OutputRef"
+    overlap: float = overlap_side_output.value
+
+    original_sino_width = non_slice_dims_shape[1]
+    stitched_sino_width = original_sino_width * 2 - math.ceil(overlap)
+    n = non_slice_dims_shape[0] // 2
+    stitched_non_slice_dims = (n, stitched_sino_width)
+
+    input_slice_size = int(np.prod(non_slice_dims_shape)) * dtype.itemsize
+    output_slice_size = int(np.prod(stitched_non_slice_dims)) * dtype.itemsize
+
+    summand_shape: Tuple[int, int] = (n, int(overlap))
+    # Multiplication between a subset of the original data (`float32`) and the 1D weights array
+    # (`float32`) causes a new array to be created that has dtype `float32` (for example,
+    # multiplications like `weights * data[:n, :, -overlap]`
+    summand_size = int(np.prod(summand_shape)) * np.float32().itemsize
+
+    total_memory_bytes = (
+        input_slice_size
+        + output_slice_size
+        # In both the `if` branch and the `else` branch checking the `rotation` variable value,
+        # in total, there are 4 copies of subsets of the `data` array that are made (note that
+        # the expressions below are referencing the `if` branch and are slightly different in
+        # the `else` branch):
+        # 1. fancy indexing: `data[n : 2 * n, :, overlap:][:, :, ::-1]`
+        # 2. multiplication: `weights * data[:n, :, :overlap]`
+        # 3. multiplication: `weights * data[n : 2 * n, :, :overlap]`
+        # 4. fancy indexing (performed on the result of 3):
+        # `(weights * data[n : 2 * n, :, :overlap])[:, :, ::-1]`
+        + 4 * summand_size
+    )
+
+    return total_memory_bytes, 0


Test patch:
diff --git a/tests/test_backends/test_httomolibgpu.py b/tests/test_backends/test_httomolibgpu.py
index 70fa8ecfb..bdb2fc7bb 100644
--- a/tests/test_backends/test_httomolibgpu.py
+++ b/tests/test_backends/test_httomolibgpu.py
@@ -5,7 +5,9 @@
 from mpi4py import MPI
 import numpy as np
 from numpy import uint16, float32
+from unittest import mock
 
+from pytest_mock import MockerFixture
 from numpy.testing import assert_allclose, assert_equal
 import os
 
@@ -16,9 +18,10 @@
 import cupy as cp
 
 from httomo.methods_database.query import get_method_info
+from httomo.runner.output_ref import OutputRef
 
 
-from httomolibgpu.misc.morph import data_resampler
+from httomolibgpu.misc.morph import data_resampler, sino_360_to_180
 from httomolibgpu.prep.normalize import normalize
 from httomolibgpu.prep.phase import paganin_filter_tomopy, paganin_filter_savu
 from httomolibgpu.prep.alignment import distortion_correction_proj_discorpy
@@ -39,6 +42,9 @@
 from httomo.methods_database.packages.external.httomolibgpu.supporting_funcs.misc.rescale import *
 from httomo.methods_database.packages.external.httomolibgpu.supporting_funcs.prep.normalize import *
 
+from tests.testing_utils import make_test_method
+
+
 module_mem_path = "httomo.methods_database.packages.external."
 
 
@@ -588,3 +594,50 @@ def test_rescale_to_int_memoryhook(
     # the resulting percent value should not deviate from max_mem on more than 20%
     assert estimated_memory_mb >= max_mem_mb
     assert percents_relative_maxmem <= 35
+
+
+@pytest.mark.cupy
+@pytest.mark.parametrize("slices", [3, 8, 30, 80])
+@pytest.mark.parametrize("det_x", [600, 2160])
+def test_sino_360_to_180_memoryhook(
+    ensure_clean_memory,
+    mocker: MockerFixture,
+    det_x: int,
+    slices: int,
+):
+    # Use a different overlap value for stitching based on the width of the 360 sinogram
+    overlap = 350 if det_x == 600 else 1950
+    shape = (1801, slices, det_x)
+    data = cp.random.random_sample(shape, dtype=np.float32)
+
+    # Run method to see actual memory usage
+    hook = MaxMemoryHook()
+    with hook:
+        sino_360_to_180(cp.copy(data), overlap)
+
+    # Call memory estimator to estimate memory usage
+    output_ref = OutputRef(
+        method=make_test_method(mocker),
+        mapped_output_name="overlap",
+    )
+    with mock.patch(
+        "httomo.runner.output_ref.OutputRef.value",
+        new_callable=mock.PropertyMock,
+    ) as mock_value_property:
+        mock_value_property.return_value = overlap
+        (estimated_bytes, subtract_bytes) = _calc_memory_bytes_sino_360_to_180(
+            non_slice_dims_shape=(shape[0], shape[2]),
+            dtype=np.float32(),
+            overlap=output_ref,
+        )
+    estimated_bytes *= slices
+
+    max_mem = hook.max_mem - subtract_bytes
+
+    # For the difference between the actual memory usage and estimated memory usage, calculate
+    # that as a percentage of the actual memory used
+    difference = abs(estimated_bytes - max_mem)
+    percentage_difference = round((difference / max_mem) * 100)
+
+    assert estimated_bytes >= max_mem
+    assert percentage_difference <= 35