OpenMP parallelism for a small loop

awslabs · Mar 4, 2024 · 1ec7cb0 · 1ec7cb0
1 parent 5cb5a79
commit 1ec7cb0
Showing 1 changed file with 12 additions and 6 deletions.
diff --git a/palace/fem/bilinearform.cpp b/palace/fem/bilinearform.cpp
@@ -260,14 +260,20 @@ std::unique_ptr<ceed::Operator> DiscreteLinearOperator::PartialAssemble() const
   Vector test_multiplicity(test_fespace.GetVSize());
   test_multiplicity = 0.0;
   auto *h_mult = test_multiplicity.HostReadWrite();
-  mfem::Array<int> dofs;
-  for (int i = 0; i < test_fespace.GetMesh().GetNE(); i++)
+  PalacePragmaOmp(parallel)
   {
-    test_fespace.Get().GetElementVDofs(i, dofs);
-    for (int j = 0; j < dofs.Size(); j++)
+    mfem::Array<int> dofs;
+    mfem::DofTransformation dof_trans;
+    PalacePragmaOmp(for schedule(static))
+    for (int i = 0; i < test_fespace.GetMesh().GetNE(); i++)
     {
-      const int k = dofs[j];
-      h_mult[(k >= 0) ? k : -1 - k] += 1.0;
+      test_fespace.Get().GetElementVDofs(i, dofs, dof_trans);
+      for (int j = 0; j < dofs.Size(); j++)
+      {
+        const int k = dofs[j];
+        PalacePragmaOmp(atomic update)
+        h_mult[(k >= 0) ? k : -1 - k] += 1.0;
+      }
     }
   }
   test_multiplicity.UseDevice(true);