Allreduce implementation, various smaller extra changes

csrc/all_gather.cpp

@@ -18,7 +18,7 @@
 //  - comm: MPI communicator (default MPI_COMM_WORLD).
 void recursiveDoublingAllGatherGPU(void* output, 
                                   const void* input, 
-                                  int total_elems, 
+                                  int64_t total_elems, 
                                   void* recv_buf,  // Same as output size
                                   MPI_Comm comm) {
 
@@ -27,7 +27,7 @@ void recursiveDoublingAllGatherGPU(void* output,
     MPI_Comm_size(comm, &size);
 
     assert(total_elems % size == 0 && "Input tensor size must be divisible by number of processes");
-    int block_size = total_elems / size;
+    int64_t block_size = total_elems / size;
 
     auto stream = at::cuda::getCurrentCUDAStream();
 
@@ -71,3 +71,50 @@ void recursiveDoublingAllGatherGPU(void* output,
 
     CUDA_CHECK(cudaEventDestroy(stream_sync_event));
 }
+
+void ringAllGatherGPU(void* output,
+                      const void* input,
+                      int64_t total_elems,
+                      MPI_Comm comm) {
+    int rank, size;
+    MPI_Comm_rank(comm, &rank);
+    MPI_Comm_size(comm, &size);
+
+    assert(total_elems % size == 0 && "Input tensor size must be divisible by number of processes");
+    int64_t block_size = total_elems / size;
+    // printf("[Rank %d] block_size = %d\n", rank, block_size);
+
+    auto stream = at::cuda::getCurrentCUDAStream();
+    cudaEvent_t stream_sync_event;
+
+    // Copy local input into its designated block in the output buffer.
+    CUDA_CHECK(cudaMemcpyAsync(static_cast<char*>(output) + rank * block_size, 
+                             input, 
+                             block_size, 
+                             cudaMemcpyDeviceToDevice, 
+                             stream));
+
+    CUDA_CHECK(cudaEventCreateWithFlags(&stream_sync_event, cudaEventDisableTiming));
+
+    // P-1 rounds each sending N/P data (where P is num processes, N is total data size)
+    for (int step = 0; step < size - 1; step++) {
+        // Compute block indices
+        int send_idx = (rank - step + size) % size;
+        int recv_idx = (rank - step - 1 + size) % size;
+        int send_peer = (rank + 1) % size;
+        int recv_peer = (rank - 1 + size) % size;
+
+        // Record an event on the cuda stream.
+        CUDA_CHECK(cudaEventRecord(stream_sync_event, stream));
+        // Wait for the copy to complete.
+        CUDA_CHECK(cudaEventSynchronize(stream_sync_event));
+
+        // Send the block to the right neighbor and receive from the left neighbor.
+        MPI_Sendrecv(static_cast<char*>(output) + send_idx * block_size, block_size, MPI_BYTE, send_peer, 0,
+                     static_cast<char*>(output) + recv_idx * block_size, block_size, MPI_BYTE, recv_peer, 0,
+                     comm, MPI_STATUS_IGNORE);
+    }
+
+    // destroy event
+    CUDA_CHECK(cudaEventDestroy(stream_sync_event));
+}

csrc/all_gather.h

@@ -11,8 +11,13 @@
 
 void recursiveDoublingAllGatherGPU(void* output, 
                                   const void* input, 
-                                  int total_elems, 
+                                  int64_t total_elems, 
                                   void* recv_buf,  
                                   MPI_Comm comm = MPI_COMM_WORLD);
 
+void ringAllGatherGPU(void* output,
+                      const void* input,
+                      int64_t total_elems,
+                      MPI_Comm comm);
+
 #endif // ALL_GATHER_H

csrc/all_reduce.cpp

@@ -0,0 +1,50 @@
+// Copyright 2025 Parallel Software and Systems Group, University of Maryland.
+// See the top-level LICENSE file for details.
+//
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+
+#include <cassert>
+#include <cmath>
+
+#include "all_gather.h"
+#include "all_reduce.h"
+#include "common.h"
+#include "reduce_scatter.h"
+
+
+// Performs an all-reduce on GPU tensors via recursive-halving reduce-scatter followed by recursive-doubling all-gather.
+//  - output: CUDA device pointer where the final gathered tensor will be stored.
+//  - input: CUDA device pointer to the local block of size block_size.
+//  - total_elems: total number of elements in output (P * block_size).
+//  - buf: main working buffer for the algorithm
+//  - recv_buf: buffer to receive data before it is processed
+//  - comm: MPI communicator (default MPI_COMM_WORLD).
+void recursiveHalvingDoublingAllReduceGPU(float* output,
+                                const float* input,
+                                int64_t total_elems,
+                                float* buf,                 // Same as input size
+                                float* recv_buf,            // Same as input size
+                                float* intermediate_buf,    // Input size / world size
+                                MPI_Comm comm) {
+    recursiveHalvingReduceScatterGPU(intermediate_buf, input, total_elems, buf, recv_buf, comm);
+
+    // allgather uses void* so multiply total_elems by size of float dtype
+    recursiveDoublingAllGatherGPU(output, intermediate_buf, total_elems*sizeof(float), recv_buf, comm);
+}
+
+// Performs an all-reduce on GPU tensors via ring reduce-scatter followed by ring all-gather.
+void ringAllReduceGPU(float* output,
+                    const float* input,
+                    int64_t total_elems,
+                    float* intermediate_buf,        // Input size / world size
+                    float* d_buf,                   // Input size
+                    float* d_send,                  // Input size / world size
+                    float* d_tmp,                   // Input size / world size
+                    MPI_Comm comm) {
+    ringReduceScatterGPU(intermediate_buf, input, total_elems, d_buf, d_send, d_tmp);
+
+    // allgather uses void* so multiply total_elems by size of float dtype
+    ringAllGatherGPU(output, intermediate_buf, total_elems*sizeof(float), comm);
+}
+

csrc/all_reduce.h

@@ -0,0 +1,29 @@
+// Copyright 2025 Parallel Software and Systems Group, University of Maryland.
+// See the top-level LICENSE file for details.
+//
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+
+#ifndef ALL_REDUCE_H
+#define ALL_REDUCE_H
+
+#include <mpi.h>
+
+void recursiveHalvingDoublingAllReduceGPU(float* output, 
+                                  const float* input, 
+                                  int64_t total_elems, 
+                                  float* buf,  
+                                  float* recv_buf,  
+                                  float* intermediate_buf,  
+                                  MPI_Comm comm = MPI_COMM_WORLD);
+
+void ringAllReduceGPU(float* output, 
+                    const float* input, 
+                    int64_t total_elems, 
+                    float* intermediate_buf,        // Input size / world size
+                    float* d_buf,                   // Input size
+                    float* d_send,                  // Input size / world size
+                    float* d_tmp,                   // Input size / world size
+                    MPI_Comm comm = MPI_COMM_WORLD);
+
+#endif // ALL_REDUCE_H

csrc/common.cu

@@ -8,15 +8,15 @@
 
 
 // Kernel for vector addition.
-__global__ void vectorAddKernel(float* a, const float* b, int n) {
+__global__ void vectorAddKernel(float* a, const float* b, int64_t n) {
     int idx = blockIdx.x * blockDim.x + threadIdx.x;
     if (idx < n) {
         a[idx] += b[idx];
     }
 }
 
 // Function to launch the kernel.
-void vectorAdd(float* a, const float* b, int n, cudaStream_t stream) {
+void vectorAdd(float* a, const float* b, int64_t n, cudaStream_t stream) {
     int threads = 256;
     int blocks = (n + threads - 1) / threads;
 

csrc/common.h

@@ -21,9 +21,9 @@
 } while(0)
 
 // Kernel for vector addition.
-__global__ void vectorAddKernel(float* a, const float* b, int n);
+__global__ void vectorAddKernel(float* a, const float* b, int64_t n);
 
 // Function to launch the kernel.
-void vectorAdd(float* a, const float* b, int n, cudaStream_t stream);
+void vectorAdd(float* a, const float* b, int64_t n, cudaStream_t stream);
 
 #endif // COMMON_H

csrc/pccl_mpi_extension.cpp

@@ -10,6 +10,7 @@
 #include <mpi4py/mpi4py.h>
 #include "reduce_scatter.h"
 #include "all_gather.h"
+#include "all_reduce.h"
 
 
 namespace py = pybind11;
@@ -142,13 +143,104 @@ void all_gather_mpi(const torch::Tensor& output_tensor,
             tmp_wrkspace_tensor_1.data_ptr(),
             //tmp_wrkspace_tensor_2.data_ptr(),
             comm);
+    } else if (algorithm == "ring") {
+        // always use torch tensors. do NOT use malloc.
+        // malloc's have high overheads and will slow your communication down
+        // torch mallocs memory in advance and manages it internally.
+        // therefore these calls are low overheads
+        // auto tmp_wrkspace_tensor_1 = torch::empty_like(output_tensor);
+        // No workspace tensors needed for ring all-gather algorithm
+        ringAllGatherGPU(output_ptr, 
+            input_ptr, 
+            total_elems * dtype_size,
+            comm);
     } else {
     TORCH_CHECK(false, "Unknown algorithm specified for all_gather_mpi: ", algorithm);
     }
 }
 
+void all_reduce_mpi(const torch::Tensor& output_tensor, 
+    const torch::Tensor& input_tensor, 
+    py::object py_comm,
+    const std::string& algorithm = "recursive")
+{
+    TORCH_CHECK(output_tensor.is_contiguous(), "output tensor must be contiguous.");
+    TORCH_CHECK(input_tensor.is_contiguous(), "input tensor must be contiguous.");
+
+    // Ensure 1D tensors.
+    TORCH_CHECK(output_tensor.dim() == 1, "output tensor must be 1D");
+    TORCH_CHECK(input_tensor.dim() == 1, "input tensor must be 1D");
+
+    // Ensure input and output dtypes are the same
+    TORCH_CHECK(input_tensor.dtype() == output_tensor.dtype(),
+                "Input and output tensors must have the same dtype.");
+
+    // Get MPI rank/size.
+    int rank, size;
+    // Get reference to base communicator
+    MPI_Comm comm = ((PyMPIIntracommObject*)(py_comm.ptr()))->__pyx_base.ob_mpi;
+
+    MPI_Comm_rank(comm, &rank);
+    MPI_Comm_size(comm, &size);
+
+    // Input tensor has one block
+    int64_t block_size = input_tensor.numel();
+    int64_t total_elems = block_size;
+    // Ensure output tensor is same size as input tensor.
+    TORCH_CHECK(output_tensor.numel() == block_size,
+    "Output tensor must have same size as input tensor");
+
+    // Ensure input tensor divisible by world size.
+    TORCH_CHECK(block_size % size == 0, 
+        "Input tensor size must be divisible by world_size for recursive halving algorithm");
+
+    // Get raw device pointers (assumes tensors reside on GPU).
+    float* output_ptr = output_tensor.data_ptr<float>();
+    const float* input_ptr  = input_tensor.data_ptr<float>();
+
+    // Call the corresponding GPU reduce-scatter algorithm.
+    if (algorithm == "recursive") {
+        // always use torch tensors. do NOT use malloc.
+        // malloc's have high overheads and will slow your communication down
+        // torch mallocs memory in advance and manages it internally.
+        // therefore these calls are low overheads
+        auto tmp_wrkspace_tensor_1 = torch::empty_like(input_tensor);
+        auto tmp_wrkspace_tensor_2 = torch::empty_like(input_tensor);
+        auto tmp_wrkspace_tensor_4 = torch::empty({block_size / size}, input_tensor.options());
+
+        recursiveHalvingDoublingAllReduceGPU(output_ptr, 
+            input_ptr, 
+            total_elems,
+            tmp_wrkspace_tensor_1.data_ptr<float>(),
+            tmp_wrkspace_tensor_2.data_ptr<float>(),
+            tmp_wrkspace_tensor_4.data_ptr<float>(),
+            comm);
+    } else if (algorithm == "ring") {
+        // always use torch tensors. do NOT use malloc.
+        // malloc's have high overheads and will slow your communication down
+        // torch mallocs memory in advance and manages it internally.
+        // therefore these calls are low overheads
+        auto tmp_wrkspace_tensor_1 = torch::empty({block_size / size}, input_tensor.options());
+        auto tmp_wrkspace_tensor_2 = torch::empty_like(input_tensor);
+        auto tmp_wrkspace_tensor_3 = torch::empty({block_size / size}, input_tensor.options());
+        auto tmp_wrkspace_tensor_4 = torch::empty({block_size / size}, input_tensor.options());
+
+        ringAllReduceGPU(output_ptr,
+            input_ptr,
+            total_elems,
+            tmp_wrkspace_tensor_1.data_ptr<float>(),
+            tmp_wrkspace_tensor_2.data_ptr<float>(),
+            tmp_wrkspace_tensor_3.data_ptr<float>(),
+            tmp_wrkspace_tensor_4.data_ptr<float>(),
+            comm);
+    } else {
+    TORCH_CHECK(false, "Unknown algorithm specified for all_reduce_mpi: ", algorithm);
+    }
+}
+
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("reduce_scatter_mpi", reduce_scatter_mpi);
     m.def("all_gather_mpi", all_gather_mpi);
+    m.def("all_reduce_mpi", all_reduce_mpi);
 }

csrc/reduce_scatter.cpp

@@ -19,7 +19,7 @@
 // The reduction operation is elementwise addition.
 void recursiveHalvingReduceScatterGPU(float* output, 
     const float* input, 
-    int total_elems,
+    int64_t total_elems,
     float* buf, // same as size of input 
     float* recv_buf, // same as size of input 
     MPI_Comm comm) {
@@ -29,7 +29,7 @@ void recursiveHalvingReduceScatterGPU(float* output,
     MPI_Comm_size(comm, &size);
 
     assert(total_elems % size == 0 && "Input tensor size must be divisible by number of processes");
-    int block_size = total_elems / size;
+    int64_t block_size = total_elems / size;
     auto stream = at::cuda::getCurrentCUDAStream();
 
     // copy the input into buf
@@ -55,7 +55,7 @@ void recursiveHalvingReduceScatterGPU(float* output,
         // The current buffer holds 'current_blocks' contiguous blocks.
         int half = current_blocks / 2;
         // Number of elements to send/receive in this round.
-        int count = half * block_size;
+        int64_t count = half * block_size;
 
         if ((rank % group_size) < (group_size / 2)) {
             // Lower half: keep the lower half and send the upper half.
@@ -102,7 +102,7 @@ void recursiveHalvingReduceScatterGPU(float* output,
 
 void ringReduceScatterGPU(float* output, 
     const float* input, 
-    int total_elems, 
+    int64_t total_elems, 
     float* d_buf, // same as size of input
     float* d_send, // same as size of output
     float* d_tmp, // same as size of output
@@ -113,7 +113,7 @@ void ringReduceScatterGPU(float* output,
     MPI_Comm_size(comm, &size);
 
     assert(total_elems % size == 0 && "Input tensor size must be divisible by number of processes");
-    int block_size = total_elems / size;
+    int64_t block_size = total_elems / size;
 
     auto stream = at::cuda::getCurrentCUDAStream();
     cudaEvent_t stream_sync_event;

csrc/reduce_scatter.h

@@ -11,14 +11,14 @@
 
 void recursiveHalvingReduceScatterGPU(float* output, 
     const float* input, 
-    int total_elems,
+    int64_t total_elems,
     float* buf, 
     float* recv_buf, 
     MPI_Comm comm = MPI_COMM_WORLD);
 
 void ringReduceScatterGPU(float* output, 
     const float* input, 
-    int total_elems, 
+    int64_t total_elems, 
     float* d_buf, 
     float* d_send, 
     float* d_tmp,