From 5283c1bc06ea081dbc8c6814a1488f31caadfb4d Mon Sep 17 00:00:00 2001
From: lucasdelimanogueira <lucasln96@hotmail.com>
Date: Sat, 25 May 2024 08:52:20 -0300
Subject: [PATCH] allreduce mean

---
 norch/csrc/distributed.cpp       | 13 +++++++++++++
 norch/distributed/distributed.py |  7 +++++++
 2 files changed, 20 insertions(+)

diff --git a/norch/csrc/distributed.cpp b/norch/csrc/distributed.cpp
index 6b25365..cf84b85 100644
--- a/norch/csrc/distributed.cpp
+++ b/norch/csrc/distributed.cpp
@@ -56,6 +56,19 @@ void allreduce_sum_tensor(Tensor* tensor) {
     cudaStreamDestroy(stream);
 }
 
+void allreduce_mean_tensor(Tensor* tensor) {
+    cudaStream_t stream;
+    cudaStreamCreate(&stream);
+
+    // Perform NCCL AllReduce operation to calculate the mean of all tensors across all processes
+    NCCL_CHECK(ncclAllReduce(tensor->data, tensor->data, tensor->size, ncclFloat, ncclSum, nccl_comm, stream));
+
+    tensor_div_scalar(tensor, tensor->data);
+    
+    cudaStreamSynchronize(stream);
+    cudaStreamDestroy(stream);
+}
+
 void end_process_group() {
     MPI_CHECK(MPI_Finalize());
     NCCL_CHECK(ncclCommDestroy(nccl_comm));
diff --git a/norch/distributed/distributed.py b/norch/distributed/distributed.py
index 9c797f8..f84a7fa 100644
--- a/norch/distributed/distributed.py
+++ b/norch/distributed/distributed.py
@@ -22,3 +22,10 @@ def allreduce_sum_tensor(tensor):
     Tensor._C.allreduce_sum_tensor.restype = None
     
     Tensor._C.allreduce_sum_tensor(tensor.tensor)
+
+def allreduce_mean_tensor(tensor):
+
+    Tensor._C.allreduce_mean_tensor.argtypes = [ctypes.POINTER(CTensor)]
+    Tensor._C.allreduce_mean_tensor.restype = None
+    
+    Tensor._C.allreduce_mean_tensor(tensor.tensor)