Complex-number Support for NCCL & Fredholm NCCL #148
Conversation
There was a problem hiding this comment.
@tharittk very good to see support for complex numbers in the NCCL engine 🚀
I just left a few comments here and there and some about your TODO comments.
Only a few things left to do once you address them:
- try to fix the flake8 issue.
- update also https://github.com/PyLops/pylops-mpi/blob/main/docs/source/gpu.rst
- one more thing that we could add to this PR is the
mddtutorial as we have now all component for it to run with NCCL (I think...). Or we could leave it for an independent PR.
pylops_mpi/DistributedArray.py
Outdated
| if deps.nccl_enabled and getattr(self, "base_comm_nccl"): | ||
| return nccl_allgather(self.sub_comm, send_buf, recv_buf) | ||
| if hasattr(send_buf, "shape"): | ||
| send_shapes = self.base_comm.allgather(send_buf.shape) |
There was a problem hiding this comment.
you sure is not self.sub_comm?
There was a problem hiding this comment.
Oh yes, you are right. Strange enough, no test case catches that.
There was a problem hiding this comment.
Interesting, I would think here
the test should fail... the original code (in currentmain branch) looks correct but your previous commit should have called this pylops-mpi/pylops_mpi/DistributedArray.py
Line 378 in 1fdf083
.asarray call... anyways, now it is correct 😄
|
@tharittk great work! I am going to merge this. In a separate PR, we may want to consider adding tests targeting the some of the other methods in the |
2 participants
I introduce the complex number support for NCCL communication here.
Changes made:
_nccl.py- Previously, the number of elements sent is equivalent to the number of elements in thesend_buf. To support the complex number, we can simply send 2x of that size (_nccl_buf_size()). The receiving data layout does not need changes as it will be according to therecv_bufi.e., it will already be complex ifsend_bufis complex.test files - I enabled all the tests. Some were muted earlier due to lack of complex-number support. Now the NCCL and MPI both have 306 tests.
Fredholm1.py- This likely needs some discussion (please see TODO: comment). Particularly, the private function_allgatheris called because Ops do not takebase_comm_ncclas an argument. Also, Fredholm operator needs some array unrolling in_rmatvecTesting: