Merge pull request #1158 from ProfFan/fix/tbb

Fix the race!
borglab · Apr 7, 2022 · a0a8010 · a0a8010
2 parents 1232552 + af1dedb
commit a0a8010
Showing 1 changed file with 26 additions and 2 deletions.
diff --git a/gtsam/inference/ClusterTree-inst.h b/gtsam/inference/ClusterTree-inst.h
@@ -15,6 +15,10 @@
 #include <gtsam/base/timing.h>
 #include <gtsam/base/treeTraversal-inst.h>
 
+#ifdef GTSAM_USE_TBB
+#include <mutex>
+#endif
+
 namespace gtsam {
 
 /* ************************************************************************* */
@@ -120,12 +124,25 @@ struct EliminationData {
   size_t myIndexInParent;
   FastVector<sharedFactor> childFactors;
   boost::shared_ptr<BTNode> bayesTreeNode;
+#ifdef GTSAM_USE_TBB
+  boost::shared_ptr<std::mutex> writeLock;
+#endif
 
   EliminationData(EliminationData* _parentData, size_t nChildren) :
-      parentData(_parentData), bayesTreeNode(boost::make_shared<BTNode>()) {
+      parentData(_parentData), bayesTreeNode(boost::make_shared<BTNode>())
+#ifdef GTSAM_USE_TBB
+      , writeLock(boost::make_shared<std::mutex>())
+#endif
+    {
     if (parentData) {
+#ifdef GTSAM_USE_TBB
+      parentData->writeLock->lock();
+#endif
       myIndexInParent = parentData->childFactors.size();
       parentData->childFactors.push_back(sharedFactor());
+#ifdef GTSAM_USE_TBB
+      parentData->writeLock->unlock();
+#endif
     } else {
       myIndexInParent = 0;
     }
@@ -196,8 +213,15 @@ struct EliminationData {
         nodesIndex_.insert(std::make_pair(j, myData.bayesTreeNode));
 
       // Store remaining factor in parent's gathered factors
-      if (!eliminationResult.second->empty())
+      if (!eliminationResult.second->empty()) {
+#ifdef GTSAM_USE_TBB
+        myData.parentData->writeLock->lock();
+#endif
         myData.parentData->childFactors[myData.myIndexInParent] = eliminationResult.second;
+#ifdef GTSAM_USE_TBB
+        myData.parentData->writeLock->unlock();
+#endif
+      }
     }
   };
 };