private method calculate local_shapes for NCCL

Author: tharittk

pylops_mpi/DistributedArray.py

@@ -340,12 +340,7 @@ def local_shapes(self):
         local_shapes : :obj:`list`
         """
         if deps.nccl_enabled and getattr(self, "base_comm_nccl"):
-            # gather tuple of shapes from every rank and copy from GPU to CPU
-            all_tuples = self._allgather(self.local_shape).get()
-            # NCCL returns the flat array that packs every tuple as 1-dimensional array
-            # unpack each tuple from each rank
-            tuple_len = len(self.local_shape)
-            return [tuple(all_tuples[i : i + tuple_len]) for i in range(0, len(all_tuples), tuple_len)]
+            return self._nccl_local_shapes(self.base_comm_nccl)
         else:
             return self._allgather(self.local_shape)
 
@@ -380,12 +375,7 @@ def asarray(self, masked: bool = False):
             return self.local_array
 
         if deps.nccl_enabled and getattr(self, "base_comm_nccl"):
-            if masked:
-                all_tuples = self._allgather_subcomm(self.local_shape).get()
-                tuple_len = len(self.local_shape)
-                local_shapes = [tuple(all_tuples[i : i + tuple_len]) for i in range(0, len(all_tuples), tuple_len)]
-            else:
-                local_shapes = self.local_shapes
+            local_shapes = self._nccl_local_shapes(self.sub_comm if masked else self.base_comm_nccl)
             return nccl_asarray(self.sub_comm if masked else self.base_comm_nccl,
                                 self.local_array, local_shapes, self.axis)
         else:
@@ -560,6 +550,21 @@ def _recv(self, recv_buf=None, source=0, count=None, tag=None):
             self.base_comm.Recv(buf=recv_buf, source=source, tag=tag)
             return recv_buf
 
+    def _nccl_local_shapes(self, nccl_comm: NcclCommunicatorType):
+        """Get the the list of shapes of every GPU in the communicator
+        """
+        # gather tuple of shapes from every rank within thee communicator and copy from GPU to CPU
+        if nccl_comm == self.sub_comm:
+            all_tuples = self._allgather_subcomm(self.local_shape).get()
+        else:
+            assert (nccl_comm == self.base_comm_nccl)
+            all_tuples = self._allgather(self.local_shape).get()
+        # NCCL returns the flat array that packs every tuple as 1-dimensional array
+        # unpack each tuple from each rank
+        tuple_len = len(self.local_shape)
+        local_shapes = [tuple(all_tuples[i : i + tuple_len]) for i in range(0, len(all_tuples), tuple_len)]
+        return local_shapes
+
     def __neg__(self):
         arr = DistributedArray(global_shape=self.global_shape,
                                base_comm=self.base_comm,