Merge remote-tracking branch 'origin/main' into halo_op

Author: astroC86

.github/workflows/build.yml

@@ -16,7 +16,7 @@ jobs:
         os: [ubuntu-latest, macos-latest]
         python-version: ['3.10', '3.11', '3.12', '3.13']
         mpi: ['mpich', 'openmpi', 'intelmpi']
-        rank: ['2', '3', '4']
+        rank: ['2', '4', '9']
         exclude:
           - os: macos-latest
             mpi: 'intelmpi'
@@ -43,4 +43,9 @@ jobs:
       - name: Install pylops-mpi
         run: pip install .
       - name: Testing using pytest-mpi
-        run: mpiexec -n ${{ matrix.rank }} pytest tests/ --with-mpi
+        run: |
+          if [ "${{ matrix.mpi }}" = "openmpi" ]; then
+            mpiexec --mca btl ^openib -n ${{ matrix.rank }} pytest tests/ --with-mpi
+          else
+            mpiexec -n ${{ matrix.rank }} pytest tests/ --with-mpi
+          fi

CHANGELOG.md

@@ -1,3 +1,19 @@
+
+# 0.3.0
+* Added `pylops_mpi.basicoperators.MPIMatrixMult` operator.
+* Added NCCL support to all operators in :mod:`pylops_mpi.basicoperators`, 
+  and  `pylops_mpi.signalprocessing`.
+* Added ``base_comm_nccl`` in constructor of `pylops_mpi.DistributedArray`,
+  to enable NCCL communication backend.
+* Added `pylops_mpi.utils.benchmark` subpackage providing methods
+  to decorate and mark functions / class methods to measure their execution 
+  time.
+* Added `pylops_mpi.utils._nccl` subpackage implementing methods
+  for NCCL communication backend.
+* Added `pylops_mpi.utils.deps` subpackage to safely import ``nccl``
+* Fixed partition in the creation of the output distributed array in 
+  `pylops_mpi.signalprocessing.MPIFredholm1`.
+
 # 0.2.0
 - Added support for using CuPy arrays with PyLops-MPI.
 - Introduced the `pylops_mpi.signalprocessing.MPIFredholm1` and `pylops_mpi.waveeqprocessing.MPIMDC` operators.

Makefile

@@ -1,6 +1,6 @@
 PIP := $(shell command -v pip3 2> /dev/null || command which pip 2> /dev/null)
 PYTHON := $(shell command -v python3 2> /dev/null || command which python 2> /dev/null)
-NUM_PROCESSES = 3
+NUM_PROCESSES = 4
 
 .PHONY: install dev-install dev-install_nccl install_ \
 	conda install_conda_nccl dev-install_conda dev-install_conda_nccl  \
@@ -53,10 +53,11 @@ tests:
 tests_nccl:	
 	mpiexec -n $(NUM_PROCESSES) pytest tests_nccl/ --with-mpi
 
+# sphinx-build does not work well with NCCL
 doc:
 	cd docs  && rm -rf source/api/generated && rm -rf source/gallery &&\
 	rm -rf source/tutorials && rm -rf build &&\
-	cd .. && sphinx-build -b html docs/source docs/build
+	cd .. && NCCL_PYLOPS_MPI=0 sphinx-build -b html docs/source docs/build
 
 doc_cupy:
 	cp tutorials_cupy/* tutorials/

docs/source/api/index.rst

@@ -118,9 +118,11 @@ Utils
     local_split
 
 
-.. currentmodule:: pylops_mpi.utils.dottest
+.. currentmodule:: pylops_mpi.utils
 
 .. autosummary::
    :toctree: generated/
 
-    dottest
+    dottest
+    benchmark
+    mark

docs/source/benchmarking.rst

@@ -0,0 +1,42 @@
+.. _benchmarkutility:
+
+Benchmarking
+============
+
+PyLops-MPI users can convenienly benchmark the performance of their code with a simple decorator.
+:py:func:`pylops_mpi.utils.benchmark` and :py:func:`pylops_mpi.utils.mark` support various
+function calling patterns that may arise when benchmarking distributed code.
+
+- :py:func:`pylops_mpi.utils.benchmark` is a **decorator** used to time the execution of entire functions.
+- :py:func:`pylops_mpi.utils.mark` is a **function** used inside decorated functions to insert fine-grained time measurements.
+
+.. note::
+   This benchmark utility is enabled by default i.e., if the user decorates the function with :py:func:`@benchmark`, the function will go through
+   the time measurements, adding overheads. Users can turn off the benchmark while leaving the decorator in-place with
+
+   .. code-block:: bash
+
+      >> export BENCH_PYLOPS_MPI=0
+
+The usage can be as simple as:
+
+.. code-block:: python
+
+   @benchmark
+   def function_to_time():
+       # Your computation
+
+The result will print out to the standard output.
+For fine-grained time measurements, :py:func:`pylops_mpi.utils.mark` can be inserted in the code region of benchmarked functions:
+
+.. code-block:: python
+
+   @benchmark
+   def funtion_to_time():
+       # You computation that you may want to ignore it in benchmark
+       mark("Begin Region")
+       # You computation
+       mark("Finish Region")
+
+You can also nest benchmarked functions to track execution times across layers of function calls with the output being correctly formatted.
+Additionally, the result can also be exported to the text file. For completed and runnable examples, visit :ref:`sphx_glr_tutorials_benchmarking.py`

docs/source/changelog.rst

@@ -3,6 +3,27 @@
 Changelog
 =========
 
+
+Version 0.3.0
+-------------
+
+*Released on: 05/08/2025*
+
+* Added :class:`pylops_mpi.basicoperators.MPIMatrixMult` operator.
+* Added NCCL support to all operators in :mod:`pylops_mpi.basicoperators`,
+  and  :mod:`pylops_mpi.signalprocessing`.
+* Added ``base_comm_nccl`` in constructor of :class:`pylops_mpi.DistributedArray`,
+  to enable NCCL communication backend.
+* Added :class:`pylops_mpi.utils.benchmark` subpackage providing methods
+  to decorate and mark functions / class methods to measure their execution 
+  time.
+* Added :class:`pylops_mpi.utils._nccl` subpackage implementing methods
+  for NCCL communication backend.
+* Added :class:`pylops_mpi.utils.deps` subpackage to safely import ``nccl``
+* Fixed partition in the creation of the output distributed array in 
+  :class:`pylops_mpi.signalprocessing.MPIFredholm1`.
+  
+
 Version 0.2.0
 -------------
 
@@ -14,6 +35,7 @@ Version 0.2.0
 * Added a dottest function to perform dot tests on PyLops-MPI operators.
 * Created a tutorial for Multi-Dimensional Deconvolution (MDD).
 
+
 Version 0.1.0
 -------------
 

docs/source/index.rst

@@ -76,6 +76,7 @@ class and implementing the ``_matvec`` and ``_rmatvec``.
     self
     installation.rst
     gpu.rst
+    benchmarking.rst
 
 .. toctree::
    :maxdepth: 2

pylops_mpi/DistributedArray.py

@@ -452,7 +452,7 @@ def _check_local_shapes(self, local_shapes):
             elif self.partition is Partition.SCATTER:
                 local_shape = local_shapes[self.rank]
                 # Check if local shape sum up to global shape and other dimensions align with global shape
-                if self._allreduce(local_shape[self.axis]) != self.global_shape[self.axis] or \
+                if self.base_comm.allreduce(local_shape[self.axis]) != self.global_shape[self.axis] or \
                         not np.array_equal(np.delete(local_shape, self.axis), np.delete(self.global_shape, self.axis)):
                     raise ValueError(f"Local shapes don't align with the global shape;"
                                      f"{local_shapes} != {self.global_shape}")

pylops_mpi/basicoperators/Laplacian.py

@@ -1,10 +1,10 @@
 from typing import Tuple
 import numpy as np
-from numpy.core.multiarray import normalize_axis_index
 from mpi4py import MPI
 
 from pylops.utils.typing import DTypeLike, InputDimsLike
 from pylops.basicoperators import SecondDerivative
+from pylops.utils.backend import get_normalize_axis_index
 
 from pylops_mpi import DistributedArray, MPILinearOperator, Partition
 from pylops_mpi.DistributedArray import local_split
@@ -75,7 +75,7 @@ def __init__(self, dims: InputDimsLike,
                  base_comm: MPI.Comm = MPI.COMM_WORLD,
                  dtype: DTypeLike = np.float64):
         self.dims = dims
-        axes = tuple(normalize_axis_index(ax, len(dims)) for ax in axes)
+        axes = tuple(get_normalize_axis_index()(ax, len(dims)) for ax in axes)
         if not (len(axes) == len(weights) == len(sampling)):
             raise ValueError("axes, weights, and sampling have different size")
         self.axes = axes

pylops_mpi/utils/__init__.py

@@ -1,4 +1,5 @@
 # isort: skip_file
 
+from .benchmark import *
 from .dottest import *
 from .deps import *