From 1ec7cb04ce978f6547361185681ecb5208621b3c Mon Sep 17 00:00:00 2001
From: Sebastian Grimberg <sjg@amazon.com>
Date: Fri, 1 Mar 2024 11:14:44 -0800
Subject: [PATCH] OpenMP parallelism for a small loop

---
 palace/fem/bilinearform.cpp | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/palace/fem/bilinearform.cpp b/palace/fem/bilinearform.cpp
index 1661cf9c9..b97f33d93 100644
--- 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);