feat(parallel): splits atomicadd for cplx types

labmec · May 28, 2024 · 7688a2b · 7688a2b
1 parent e795710
commit 7688a2b
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 13 deletions.
diff --git a/Util/TPZParallelUtils.cpp b/Util/TPZParallelUtils.cpp
@@ -13,8 +13,11 @@ namespace pzutils{
     }
 };
 
-template int pzutils::AtomicAdd<int>(int&, int);
-template int64_t pzutils::AtomicAdd<int64_t>(int64_t&, int64_t);
-template float pzutils::AtomicAdd<float>(float&, float);
-template double pzutils::AtomicAdd<double>(double&, double);
-template long double pzutils::AtomicAdd<long double>(long double&, long double);
+template void pzutils::AtomicAdd<int>(int&, int);
+template void pzutils::AtomicAdd<int64_t>(int64_t&, int64_t);
+template void pzutils::AtomicAdd<float>(float&, float);
+template void pzutils::AtomicAdd<double>(double&, double);
+template void pzutils::AtomicAdd<long double>(long double&, long double);
+template void pzutils::AtomicAdd<std::complex<float>>(std::complex<float>&, std::complex<float>);
+template void pzutils::AtomicAdd<std::complex<double>>(std::complex<double>&, std::complex<double>);
+template void pzutils::AtomicAdd<std::complex<long double>>(std::complex<long double>&, std::complex<long double>);
diff --git a/Util/TPZParallelUtils.h b/Util/TPZParallelUtils.h
@@ -15,6 +15,7 @@
 #include <vector>
 #include <thread>
 #include <atomic>
+#include <complex>
 
 namespace pzutils{
 //!Sets number of mkl threads for current threads and return previous value
@@ -75,17 +76,27 @@ as in
 #pragma omp atomic
 a += b
 */
+
 template<typename T>
-inline T AtomicAdd( T & sum, T val )
+inline void AtomicAdd( T & sum, T val )
 {
-  std::atomic<T> & asum = AtomicCast(sum);
-  T current = asum.load();
-  T desired{0};
-  do{
-    desired = current+val;
+  if constexpr (std::is_same_v<T,std::complex<float>> ||
+                std::is_same_v<T,std::complex<double>> ||
+                std::is_same_v<T,std::complex<long double>>){
+    const auto real = val.real();
+    AtomicAdd (reinterpret_cast<typename T::value_type(&)[2]>(sum)[0], real);
+    const auto imag = val.imag();
+    AtomicAdd (reinterpret_cast<typename T::value_type(&)[2]>(sum)[1], imag);
+  }else{
+    std::atomic<T> & asum = AtomicCast(sum);
+    T current = asum.load();
+    T desired{0};
+    do{
+      desired = current+val;
+    }
+    while (!asum.compare_exchange_weak(current, desired));
   }
-  while (!asum.compare_exchange_weak(current, desired));
-  return current;
 }
+
 }//namespace
 #endif