The same instance of SleefDFT can be called from multiple threads simultaneously.

Author: shibatch

src/dft/dft.cpp

@@ -482,8 +482,20 @@ class KShortest {
 
 template<typename real, typename real2, int MAXSHIFT, int MAXBUTWIDTH>
 void SleefDFTXX<real, real2, MAXSHIFT, MAXBUTWIDTH>::measurementRun(real *d, const real *s, const vector<Action> &path, uint64_t niter) {
-  const int tn = getThreadNum();
-  real *t[] = { x1[tn], x0[tn], d };
+  auto tid = this_thread::get_id();
+  real *t[] = { nullptr, nullptr, d };
+  {
+    unique_lock lock(mtx);
+    if (xn.count(tid) != 0) {
+      auto e = xn[tid];
+      t[0] = e.first;
+      t[1] = e.second;
+    } else {
+      t[0] = (real *)Sleef_malloc(sizeof(real) * 2 * (1L << log2len));
+      t[1] = (real *)Sleef_malloc(sizeof(real) * 2 * (1L << log2len));
+      xn[tid] = pair<real *, real *>{ t[0], t[1] };
+    }
+  }
 
   for(uint64_t i=0;i<niter;i++) {
     const real *lb = s;
@@ -1090,14 +1102,6 @@ SleefDFTXX<real, real2, MAXSHIFT, MAXBUTWIDTH>::SleefDFTXX(uint32_t n, const rea
   for(int level = log2len;level >= 1;level--) {
     perm[level] = (uint32_t *)Sleef_malloc(sizeof(uint32_t) * ((1 << log2len) + 8));
   }
-
-  x0 = (real **)malloc(sizeof(real *) * nThread);
-  x1 = (real **)malloc(sizeof(real *) * nThread);
-
-  for(int i=0;i<nThread;i++) {
-    x0[i] = (real *)Sleef_malloc(sizeof(real) * 2 * n);
-    x1[i] = (real *)Sleef_malloc(sizeof(real) * 2 * n);
-  }
 
   if ((mode & SLEEF_MODE_REAL) != 0) {
     rtCoef0 = (real *)Sleef_malloc(sizeof(real) * n);
@@ -1335,9 +1339,21 @@ void SleefDFTXX<real, real2, MAXSHIFT, MAXBUTWIDTH>::execute(const real *s0, rea
 
   //
 
-  const int tn = getThreadNum();
-  real *t[] = { x1[tn], x0[tn], d };
-  
+  auto tid = this_thread::get_id();
+  real *t[] = { nullptr, nullptr, d };
+  {
+    unique_lock lock(mtx);
+    if (xn.count(tid) != 0) {
+      auto e = xn[tid];
+      t[0] = e.first;
+      t[1] = e.second;
+    } else {
+      t[0] = (real *)Sleef_malloc(sizeof(real) * 2 * (1L << log2len));
+      t[1] = (real *)Sleef_malloc(sizeof(real) * 2 * (1L << log2len));
+      xn[tid] = pair<real *, real *>{ t[0], t[1] };
+    }
+  }
+
   const real *lb = s;
   int nb = 0;
 

src/dft/dftcommon.cpp

@@ -185,17 +185,10 @@ void SleefDFTXX<real, real2, MAXSHIFT, MAXBUTWIDTH>::freeTables() {
     tbl[N] = NULL;
   }
 
-  for(int i=0;i<nThread;i++) {
-    Sleef_free(x1[i]);
-    x1[i] = nullptr;
-    Sleef_free(x0[i]);
-    x0[i] = nullptr;
+  for(auto a : xn) {
+    Sleef_free(a.second.first);
+    Sleef_free(a.second.second);
   }
-
-  free(x1);
-  x1 = nullptr;
-  free(x0);
-  x0 = nullptr;
 }
 
 template<typename real, typename real2, int MAXSHIFT, int MAXBUTWIDTH>
@@ -632,12 +625,6 @@ namespace {
 
       waitUntilAllIdle(lock);
     }
-
-    int getThreadNum() {
-      auto id = this_thread::get_id();
-      if (thIdMap.count(id) == 0) return 0;
-      return thIdMap.at(id);
-    }
   };
 
 #ifdef SLEEF_ENABLE_PARALLELFOR
@@ -650,14 +637,6 @@ namespace sleef_internal {
 		   function<void(int64_t, int64_t, int64_t)> func_) {
 #ifdef SLEEF_ENABLE_PARALLELFOR
     parallelForManager.run(start_, end_, inc_, func_);
-#endif
-  }
-
-  int getThreadNum() {
-#ifndef SLEEF_ENABLE_PARALLELFOR
-      return omp_get_thread_num();
-#else
-      return parallelForManager.getThreadNum();
 #endif
   }
 }

src/dft/dftcommon.hpp

@@ -80,11 +80,13 @@ namespace sleef_internal {
 
     //
 
+    mutex mtx;
+
     real **tbl[MAXBUTWIDTH+1];
     real *rtCoef0, *rtCoef1;
     uint32_t **perm;
 
-    real **x0, **x1;
+    unordered_map<thread::id, pair<real *, real *>> xn;
 
     int isa = 0;
     int planMode = 0;
@@ -221,8 +223,6 @@ namespace sleef_internal {
   extern FILE *defaultVerboseFP;
 
   void parallelFor(int64_t start_, int64_t end_, int64_t inc_, std::function<void(int64_t, int64_t, int64_t)> func_);
-
-  int getThreadNum();
 }
 
 using namespace sleef_internal;