Performance improvements

setup.py

@@ -0,0 +1,32 @@
+"""
+from setuptools import setup, Extension
+from Cython.Build import cythonize
+import os
+
+CONDA_PREFIX = os.environ.get("CONDA_PREFIX", "")
+
+CUDA_INCLUDE = os.path.join(CONDA_PREFIX, "targets", "x86_64-linux", "include")
+
+
+
+ext = Extension(
+    name="cupyfix_backends.cuda.libs.cublas",
+    sources=[
+             "src/serinv/cupyfix_backends/cuda/libs/cublas.pyx"
+             ],
+    include_dirs=["cupyfix_backends/cuda/libs",
+                "cupyfix_backends/hip",
+                "cupyfix_backends/stub",
+                "cupyfix_backends/cuda",
+                "cupyfix_backends",
+                CUDA_INCLUDE
+                ],
+
+)
+
+setup(
+    name="cupyfix_backends",
+    ext_modules=cythonize([ext]),
+    packages=["src/serinv/cupyfix_backends.cuda.libs"],
+)
+"""

src/serinv/__init__.py

@@ -163,7 +163,7 @@ def _use_nccl(comm):
     return False
 
 
-def _get_nccl_parameters(arr, comm, op: str):
+def _get_nccl_parameters(arr, comm, rank, op: str):
     """Get the NCCL parameters for the given operation."""
     if np.iscomplexobj(arr):
         factor = 2
@@ -172,8 +172,8 @@ def _get_nccl_parameters(arr, comm, op: str):
 
     if backend_flags["nccl_avail"]:
         if op == "allgather":
-            count = (arr.size // comm.size) * factor
-            displacement = count * comm.rank * arr.dtype.itemsize
+            count = (arr.size // comm.size()) * factor
+            displacement = count * rank * (arr.dtype.itemsize // factor)
         elif op == "allreduce":
             count = arr.size * factor
             displacement = 0