Merge branch 'PyLops:main' into complex-support

Author: tharittk

Makefile

@@ -55,6 +55,13 @@ doc:
 	rm -rf source/tutorials && rm -rf build &&\
 	cd .. && sphinx-build -b html docs/source docs/build
 
+doc_nccl:
+	cp tutorials_nccl/* tutorials/
+	cd docs  && rm -rf source/api/generated && rm -rf source/gallery &&\
+	rm -rf source/tutorials && rm -rf source/tutorials && rm -rf build &&\
+	cd .. && sphinx-build -b html docs/source docs/build
+	rm tutorials/*_nccl.py
+
 docupdate:
 	cd docs && make html && cd ..
 

docs/source/conf.py

@@ -2,7 +2,6 @@
 import sys
 import os
 import datetime
-from sphinx_gallery.sorting import ExampleTitleSortKey
 
 # Sphinx needs to be able to import the package to use autodoc and get the version number
 sys.path.insert(0, os.path.abspath("../../pylops_mpi"))
@@ -60,7 +59,7 @@
     # Remove the "Download all examples" button from the top level gallery
     "download_all_examples": False,
     # Sort gallery example by file name instead of number of lines (default)
-    "within_subsection_order": ExampleTitleSortKey,
+    "within_subsection_order": "ExampleTitleSortKey",
     # directory where function granular galleries are stored
     "backreferences_dir": "api/generated/backreferences",
     # Modules for which function level galleries are created.

docs/source/installation.rst

@@ -19,7 +19,7 @@ Additionally, to use the NCCL engine, the following additional
 dependencies are required:
 
 * `CuPy <https://cupy.dev/>`_
-* `NCCL <https://docs.cupy.dev/en/stable/install.html#additional-cuda-libraries>`_
+* `NCCL <https://docs.cupy.dev/en/stable/install.html#additional-cuda-libraries>`__
 
 We highly encourage using the `Anaconda Python distribution <https://www.anaconda.com/download>`_
 or its standalone package manager `Conda <https://docs.conda.io/en/latest/index.html>`_. However,
@@ -120,7 +120,7 @@ If you prefer a ``pip`` installation, we provide the following command
 Note that, differently from the  ``conda`` command, the above **will not** create a virtual environment.
 Make sure you create and activate your environment previously.
 
-Simlarly, if you want to enable `NCCL <https://developer.nvidia.com/nccl>`_ but prefer using pip,
+Similarly, if you want to enable `NCCL <https://developer.nvidia.com/nccl>`_ but prefer using pip,
 you must first check the CUDA version of your system:
 
 .. code-block:: bash

pylops_mpi/DistributedArray.py

@@ -526,7 +526,7 @@ def _send(self, send_buf, dest, count=None, tag=None):
                 count = send_buf.size
             nccl_send(self.base_comm_nccl, send_buf, dest, count)
         else:
-            self.base_comm.Send(send_buf, dest, tag)
+            self.base_comm.send(send_buf, dest, tag)
 
     def _recv(self, recv_buf=None, source=0, count=None, tag=None):
         """ Receive operation
@@ -543,11 +543,8 @@ def _recv(self, recv_buf=None, source=0, count=None, tag=None):
             else:
                 raise ValueError("Using recv with NCCL must also supply receiver buffer ")
         else:
-            # MPI allows a receiver buffer to be optional
-            if recv_buf is None:
-                return self.base_comm.recv(source=source, tag=tag)
-            self.base_comm.Recv(buf=recv_buf, source=source, tag=tag)
-            return recv_buf
+            # MPI allows a receiver buffer to be optional and receives as a Python Object
+            return self.base_comm.recv(source=source, tag=tag)
 
     def _nccl_local_shapes(self, masked: bool):
         """Get the the list of shapes of every GPU in the communicator

tutorials/poststack.py

@@ -277,3 +277,6 @@
     axs[5][2].imshow(minv3d_reg[:, :, 220].T, cmap='gist_rainbow', vmin=m.min(), vmax=m.max())
     axs[5][2].set_title('Regularized Inverted Model iter x-y')
     axs[5][2].axis('tight')
+
+###############################################################################
+# To run this tutorial with our NCCL backend, refer to :ref:`sphx_glr_tutorials_poststack_nccl.py`

tutorials_nccl/poststack_nccl.py

@@ -15,8 +15,6 @@
 from pylops.avo.poststack import PoststackLinearModelling
 
 import pylops_mpi
-import pylops_mpi.utils
-import pylops_mpi.utils._nccl
 
 ###############################################################################
 # NCCL communication can be easily initialized with
@@ -198,5 +196,3 @@
     axs[5][2].imshow(minv3d_reg[:, :, 220].T.get(), cmap='gist_rainbow', vmin=m.min(), vmax=m.max())
     axs[5][2].set_title('Regularized Inverted Model iter x-y')
     axs[5][2].axis('tight')
-
-    plt.savefig("./poststack_inv_nccl.png")