Use mfem::Workspace for temporary Vector allocations

cmake/ExternalHYPRE.cmake

@@ -141,7 +141,7 @@ if(PALACE_WITH_CUDA)
     "--with-cuda-home=${CUDA_DIR}"
     "--enable-curand"
     "--enable-cusparse"
-    "--enable-device-memory-pool"
+    # "--enable-device-memory-pool"
   )
   if(NOT "${CMAKE_CUDA_ARCHITECTURES}" STREQUAL "")
     list(GET CMAKE_CUDA_ARCHITECTURES 0 HYPRE_CUDA_ARCH)

cmake/ExternalMFEM.cmake

@@ -364,6 +364,7 @@ set(MFEM_PATCH_FILES
   "${CMAKE_SOURCE_DIR}/extern/patch/mfem/patch_par_tet_mesh_fix.diff"
   "${CMAKE_SOURCE_DIR}/extern/patch/mfem/patch_mesh_prism_vtu_fix.diff"
   "${CMAKE_SOURCE_DIR}/extern/patch/mfem/patch_workspace_vectors.diff"
+  "${CMAKE_SOURCE_DIR}/extern/patch/mfem/patch_workspace_vectors_2.diff"
   "${CMAKE_SOURCE_DIR}/extern/patch/mfem/patch_hypre_runtime_compute_policy.diff"
 )
 

extern/patch/mfem/patch_workspace_vectors.diff

@@ -1068,10 +1068,10 @@ index 000000000..b5c0ebae8
 +   WorkspaceVector(WorkspaceVector &&other);
 +
 +   /// No copy constructor.
-+   WorkspaceVector(WorkspaceVector &other) = delete;
++   WorkspaceVector(const WorkspaceVector &other) = delete;
 +
 +   /// Copy assignment: copy contents of vector, not metadata.
-+   WorkspaceVector& operator=(WorkspaceVector &other)
++   WorkspaceVector& operator=(const WorkspaceVector &other)
 +   {
 +      Vector::operator=(other);
 +      return *this;

extern/patch/mfem/patch_workspace_vectors_2.diff

@@ -0,0 +1,74 @@
+diff --git a/general/workspace.cpp b/general/workspace.cpp
+index ad8210a61..3ec077dfa 100644
+--- a/general/workspace.cpp
++++ b/general/workspace.cpp
+@@ -15,11 +15,11 @@ namespace mfem
+ {
+
+ WorkspaceVector::WorkspaceVector(
+-   internal::WorkspaceChunk &chunk_, int offset_, int n)
++   internal::WorkspaceChunk &chunk_, int offset_, int n, int padding)
+    : Vector(chunk_.GetData(), chunk_.GetOffset(), n),
+      chunk(chunk_),
+      offset(offset_),
+-     original_size(n)
++     original_size(n + padding)
+ {
+    UseDevice(true);
+ }
+@@ -48,8 +48,11 @@ WorkspaceChunk::WorkspaceChunk(int capacity)
+ WorkspaceVector WorkspaceChunk::NewVector(int n)
+ {
+    MFEM_ASSERT(HasCapacityFor(n), "Requested vector is too large.");
+-   WorkspaceVector vector(*this, offset, n);
+-   offset += n;
++   constexpr int alignment = 16;
++   const int s = (n * sizeof(double)) % alignment;
++   const int padding = (s > 0) ? (alignment - s) / sizeof(double) : 0;
++   WorkspaceVector vector(*this, offset, n, padding);
++   offset += n + padding;
+    vector_count += 1;
+    return vector;
+ }
+diff --git a/general/workspace.hpp b/general/workspace.hpp
+index b5c0ebae8..e41e71b11 100644
+--- a/general/workspace.hpp
++++ b/general/workspace.hpp
+@@ -50,7 +50,8 @@ class WorkspaceVector : public Vector
+    bool moved_from = false;
+
+    /// Private constructor, create with Workspace::NewVector() instead.
+-   WorkspaceVector(internal::WorkspaceChunk &chunk_, int offset_, int n);
++   WorkspaceVector(internal::WorkspaceChunk &chunk_, int offset_, int n,
++                   int padding = 0);
+
+ public:
+    /// @brief Move constructor. The moved-from WorkspaceVector has @a
+@@ -104,6 +105,9 @@ public:
+    /// Create a WorkspaceChunk with the given @a capacity.
+    WorkspaceChunk(int capacity);
+
++   /// Return the data offset.
++   int GetOffset() const { return offset; }
++
+    /// @brief Return the available capacity (i.e. the largest vector that will
+    /// fit in this chunk).
+    int GetAvailableCapacity() const { return data.Size() - offset; }
+@@ -115,15 +119,12 @@ public:
+    /// "original capacity" remains unchained.
+    int GetOriginalCapacity() const { return original_capacity; }
+
+-   /// Return the data offset.
+-   int GetOffset() const { return offset; }
++   /// Returns true if this chunk can fit a new vector of size @a n.
++   bool HasCapacityFor(int n) const { return n <= GetAvailableCapacity(); }
+
+    /// Sets whether the chunk is in the front of the list
+    void SetFront(bool front_) { front = front_; }
+
+-   /// Returns true if this chunk can fit a new vector of size @a n.
+-   bool HasCapacityFor(int n) const { return n <= GetAvailableCapacity(); }
+-
+    /// Returns true if this chunk is empty.
+    bool IsEmpty() const { return vector_count == 0; }
+

palace/drivers/eigensolver.cpp

@@ -19,6 +19,7 @@
 #include "utils/communication.hpp"
 #include "utils/iodata.hpp"
 #include "utils/timer.hpp"
+#include "utils/workspace.hpp"
 
 namespace palace
 {
@@ -177,7 +178,7 @@ EigenSolver::Solve(const std::vector<std::unique_ptr<Mesh>> &mesh) const
   // projected appropriately.
   if (iodata.solver.eigenmode.init_v0)
   {
-    ComplexVector v0;
+    auto v0 = workspace::NewVector<ComplexVector>(E.Size());
     if (iodata.solver.eigenmode.init_v0_const)
     {
       Mpi::Print(" Using constant starting vector\n");
@@ -196,8 +197,7 @@ EigenSolver::Solve(const std::vector<std::unique_ptr<Mesh>> &mesh) const
 
     // Debug
     // const auto &Grad = spaceop.GetGradMatrix();
-    // ComplexVector r0(Grad->Width());
-    // r0.UseDevice(true);
+    // auto r0 = workspace::NewVector<ComplexVector>(Grad.Width());
     // Grad.MultTranspose(v0.Real(), r0.Real());
     // Grad.MultTranspose(v0.Imag(), r0.Imag());
     // r0.Print();

palace/fem/fespace.hpp

@@ -31,9 +31,6 @@ class FiniteElementSpace
   mutable ceed::CeedObjectMap<CeedBasis> basis;
   mutable ceed::CeedObjectMap<CeedElemRestriction> restr, interp_restr, interp_range_restr;
 
-  // Temporary storage for operator applications.
-  mutable ComplexVector tx, lx, ly;
-
   bool HasUniqueInterpRestriction(const mfem::FiniteElement &fe) const
   {
     // For interpolation operators and tensor-product elements, we need native (not
@@ -63,9 +60,6 @@ class FiniteElementSpace
     : fespace(&mesh.Get(), std::forward<T>(args)...), mesh(mesh)
   {
     ResetCeedObjects();
-    tx.UseDevice(true);
-    lx.UseDevice(true);
-    ly.UseDevice(true);
   }
   virtual ~FiniteElementSpace() { ResetCeedObjects(); }
 
@@ -129,48 +123,6 @@ class FiniteElementSpace
                            mfem::Geometry::Type geom, const std::vector<int> &indices,
                            bool is_interp = false, bool is_interp_range = false);
 
-  template <typename VecType>
-  auto &GetTVector() const
-  {
-    tx.SetSize(GetTrueVSize());
-    if constexpr (std::is_same<VecType, ComplexVector>::value)
-    {
-      return tx;
-    }
-    else
-    {
-      return tx.Real();
-    }
-  }
-
-  template <typename VecType>
-  auto &GetLVector() const
-  {
-    lx.SetSize(GetVSize());
-    if constexpr (std::is_same<VecType, ComplexVector>::value)
-    {
-      return lx;
-    }
-    else
-    {
-      return lx.Real();
-    }
-  }
-
-  template <typename VecType>
-  auto &GetLVector2() const
-  {
-    ly.SetSize(GetVSize());
-    if constexpr (std::is_same<VecType, ComplexVector>::value)
-    {
-      return ly;
-    }
-    else
-    {
-      return ly.Real();
-    }
-  }
-
   // Get the associated MPI communicator.
   MPI_Comm GetComm() const { return fespace.GetComm(); }
 };

palace/fem/libceed/operator.cpp

@@ -10,6 +10,7 @@
 #include "fem/fespace.hpp"
 #include "linalg/hypre.hpp"
 #include "utils/omp.hpp"
+#include "utils/workspace.hpp"
 
 namespace palace::ceed
 {
@@ -37,7 +38,6 @@ Operator::Operator(int h, int w) : palace::Operator(h, w)
     u[id] = loc_u;
     v[id] = loc_v;
   }
-  temp.UseDevice(true);
 }
 
 Operator::~Operator()
@@ -174,15 +174,15 @@ void Operator::AddMult(const Vector &x, Vector &y, const double a) const
   MFEM_VERIFY(a == 1.0, "ceed::Operator::AddMult only supports coefficient = 1.0!");
   if (dof_multiplicity.Size() > 0)
   {
-    temp.SetSize(height);
-    temp = 0.0;
-    CeedAddMult(op, u, v, x, temp);
+    auto t = workspace::NewVector<Vector>(height);
+    t = 0.0;
+    CeedAddMult(op, u, v, x, t);
     {
       const auto *d_dof_multiplicity = dof_multiplicity.Read();
-      const auto *d_temp = temp.Read();
+      const auto *d_t = t.Read();
       auto *d_y = y.ReadWrite();
       mfem::forall(height, [=] MFEM_HOST_DEVICE(int i)
-                   { d_y[i] += d_dof_multiplicity[i] * d_temp[i]; });
+                   { d_y[i] += d_dof_multiplicity[i] * d_t[i]; });
     }
   }
   else
@@ -203,15 +203,15 @@ void Operator::AddMultTranspose(const Vector &x, Vector &y, const double a) cons
               "ceed::Operator::AddMultTranspose only supports coefficient = 1.0!");
   if (dof_multiplicity.Size() > 0)
   {
-    temp.SetSize(height);
+    auto t = workspace::NewVector<Vector>(height);
     {
       const auto *d_dof_multiplicity = dof_multiplicity.Read();
       const auto *d_x = x.Read();
-      auto *d_temp = temp.Write();
+      auto *d_t = t.Write();
       mfem::forall(height, [=] MFEM_HOST_DEVICE(int i)
-                   { d_temp[i] = d_dof_multiplicity[i] * d_x[i]; });
+                   { d_t[i] = d_dof_multiplicity[i] * d_x[i]; });
     }
-    CeedAddMult(op_t, v, u, temp, y);
+    CeedAddMult(op_t, v, u, t, y);
   }
   else
   {

palace/fem/libceed/operator.hpp

@@ -35,7 +35,6 @@ class Operator : public palace::Operator
   std::vector<CeedOperator> op, op_t;
   std::vector<CeedVector> u, v;
   Vector dof_multiplicity;
-  mutable Vector temp;
 
 public:
   Operator(int h, int w);