NCCL Complex wrapper

[nikopj]

The NCCL backend in distributed utils does not support complex values see issue. Can we add a conveinent wrapper in the NCCLEXT to support broadcast, reduce, etc., likely via using reim and complex? I'm happy to get started on it but would like some feedback on a preferred location and so on.

[avik-pal]

There are 3 possible solutions:

1. for the cases with LuxCUDADevice (

   Lux.jl/ext/LuxMPINCCLExt.jl

   Lines 38 to 42 in 6823845

   	function DistributedUtils.__bcast!(
   	backend::NCCLBackend, sendrecvbuf, ::LuxCUDADevice; root=0)
   	NCCL.Broadcast!(sendrecvbuf, backend.comm; root)
   	return sendrecvbuf
   	end

   ), do the reim and complex if the buffer type is ::AbstractArray{<:Complex}.
2. Alternatively, you can detect complex numbers the same way as point 1 and forward the call to MPI similar to

   Lux.jl/ext/LuxMPINCCLExt.jl

   Lines 44 to 47 in 6823845

   	function DistributedUtils.__bcast!(
   	backend::NCCLBackend, sendrecvbuf, dev::AbstractLuxDevice; root=0)
   	return DistributedUtils.__bcast!(backend.mpi_backend, sendrecvbuf, dev; root)
   	end

   . If MPI is CUDA-aware then no device to host copy is performed and MPI AFAIK supports Complex numbers.
3. Add directly to NCCL.jl, though it might be worth opening an issue there and asking if the feature is welcome.

[avik-pal]

Closing since this is going to directly happen in NCCL.jl