feat(strmatrixomportbb): better way to set up buffer mat

labmec · Jun 19, 2024 · b8815db · b8815db
1 parent 957e69c
commit b8815db
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 21 deletions.
diff --git a/StrMatrix/TPZStructMatrixOMPorTBB.cpp b/StrMatrix/TPZStructMatrixOMPorTBB.cpp
@@ -448,12 +448,16 @@ void TPZStructMatrixOMPorTBB<TVar>::AssemblingUsingTBBbutNotColoring(TPZBaseMatr
           [&](tbb::blocked_range<int64_t> r){
               TPZElementMatrixT<TVar> ek(cmesh,TPZElementMatrix::EK);
               TPZElementMatrixT<TVar> ef(cmesh,TPZElementMatrix::EF);
-              const auto bufsz = 400*400;
-              TVar*buf = new TVar[bufsz];
+              TVar* buf = nullptr;
+              if(fUserMatSize>0){
+                  buf = new TVar[fUserMatSize];
+              }
 
               {
-                  TPZFMatrix<TVar> auxmat(1,1,buf,bufsz);
-                  ek.SetUserAllocMat(&auxmat);
+                  if(buf){
+                      TPZFMatrix<TVar> auxmat(1,1,buf,fUserMatSize);
+                      ek.SetUserAllocMat(&auxmat);
+                  }
                   for (int64_t iel = r.begin(); iel < r.end(); iel++)
                   {
                       TPZCompEl *el = cmesh->Element(iel);
@@ -469,7 +473,7 @@ void TPZStructMatrixOMPorTBB<TVar>::AssemblingUsingTBBbutNotColoring(TPZBaseMatr
                   }
               }
               //matrix has been destroyed
-              delete [] buf;
+              if(buf){delete [] buf;}
         });
 #else
     DebugStop();
@@ -494,30 +498,36 @@ void TPZStructMatrixOMPorTBB<TVar>::AssemblingUsingOMPbutNotColoring(TPZBaseMatr
 #pragma omp parallel
     {
 
-      const auto bufsz = 400*400;
-      TVar*buf = new TVar[bufsz];
+      TVar* buf = nullptr;
+      if(fUserMatSize>0){
+          buf = new TVar[fUserMatSize];
+      }
       TPZElementMatrixT<TVar> ek(cmesh,TPZElementMatrix::EK);
       TPZElementMatrixT<TVar> ef(cmesh,TPZElementMatrix::EF);
       auto mklthreads = pzutils::SetNumThreadsLocalMKL(1);
       {
-	TPZFMatrix<TVar> auxmat(1,1,buf,bufsz);
-	ek.SetUserAllocMat(&auxmat);
+
+        if(buf){
+            TPZFMatrix<TVar> auxmat(1,1,buf,fUserMatSize);
+            ek.SetUserAllocMat(&auxmat);
+        }
 #pragma omp for schedule(dynamic,1)
         for (int64_t iel = 0; iel < nelem; iel++){
-	  {
-	    TPZCompEl *el = cmesh->Element(iel);
-	    if ((!el) ||
-		(nmatids != 0 &&
-		 !el->NeedsComputing(matids)))
-	      {
-		continue;
-	      }
-
-	    CalcStiffAndAssemble(mat,rhs,el,ek,ef);
-	  }
+          {
+            TPZCompEl *el = cmesh->Element(iel);
+            if ((!el) ||
+                (nmatids != 0 &&
+                 !el->NeedsComputing(matids)))
+              {
+                continue;
+              }
+
+            CalcStiffAndAssemble(mat,rhs,el,ek,ef);
+          }
         }
       }
-      delete [] buf;
+      //matrix has been destroyed
+      if(buf){delete [] buf;}
     }
 #else
     DebugStop();

diff --git a/StrMatrix/TPZStructMatrixOMPorTBB.h b/StrMatrix/TPZStructMatrixOMPorTBB.h
@@ -61,6 +61,10 @@ class TPZStructMatrixOMPorTBB : public virtual TPZStrMatParInterface{
     void Assemble(TPZBaseMatrix & mat, TPZBaseMatrix & rhs) override;
     void Assemble(TPZBaseMatrix & rhs) override;
 
+    /** @brief Sets buffer size to be preallocated for a TPZFMatrix provided
+        to TPZElementMatrixT<TVar> for the dependency matrix, useful
+    when dealing with large dependencies to avoid repeated dynamic allocation*/
+    void BufferSizeForUserMatrix(const int sz){fUserMatSize=sz;}
 public:
     int ClassId() const override;
     void Read(TPZStream &buf, void *context) override;
@@ -127,6 +131,9 @@ class TPZStructMatrixOMPorTBB : public virtual TPZStrMatParInterface{
     int fSomeoneIsSleeping;
 
     std::atomic<int64_t> fCurrentIndex;
+
+    /// user allocated mat for TPZElementMatrixT
+    int fUserMatSize{-1};
 };
 
 extern template class TPZStructMatrixOMPorTBB<STATE>;