More use of workspace vectors

awslabs · Feb 19, 2024 · 6a95817 · 6a95817
1 parent dc8fbe8
commit 6a95817
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 36 deletions.
diff --git a/palace/drivers/eigensolver.cpp b/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();

diff --git a/palace/linalg/divfree.cpp b/palace/linalg/divfree.cpp
@@ -91,10 +91,9 @@ DivFreeSolver<VecType>::DivFreeSolver(
     const int mg_smooth_order =
         std::max(h1_fespaces.GetFinestFESpace().GetMaxElementOrder(), 2);
     pc = std::make_unique<GeometricMultigridSolver<OperType>>(
-        h1_fespaces.GetFinestFESpace().GetComm(),
-
-        std::move(amg), h1_fespaces.GetProlongationOperators(), nullptr, 1, 1,
-        mg_smooth_order, 1.0, 0.0, true);
+        h1_fespaces.GetFinestFESpace().GetComm(), std::move(amg),
+        h1_fespaces.GetProlongationOperators(), nullptr, 1, 1, mg_smooth_order, 1.0, 0.0,
+        true);
   }
   else
   {

diff --git a/palace/linalg/operator.cpp b/palace/linalg/operator.cpp
@@ -636,10 +636,9 @@ double SpectralNorm(MPI_Comm comm, const ComplexOperator &A, bool herm, double t
   int it = 0;
   double res = 0.0;
   double l = 0.0, l0 = 0.0;
-  ComplexVector u(A.Height()), v(A.Height());
-  u.UseDevice(true);
-  v.UseDevice(true);
-  SetRandom(comm, u);
+  auto u = workspace::NewVector<ComplexVector>(A.Height());
+  auto v = workspace::NewVector<ComplexVector>(A.Height());
+  SetRandom<ComplexVector>(comm, u);
   Normalize(comm, u);
   while (it < max_it)
   {

diff --git a/palace/linalg/slepc.cpp b/palace/linalg/slepc.cpp
@@ -137,50 +137,56 @@ inline VecType PetscVecType()
   return VECSTANDARD;
 }
 
-struct MatShellContext
-{
-  const ComplexOperator &A;
-  ComplexVector &x, &y;
-};
-
 PetscErrorCode __mat_apply_shell(Mat A, Vec x, Vec y)
 {
   PetscFunctionBeginUser;
-  MatShellContext *ctx;
+  ComplexOperator *ctx;
   PetscCall(MatShellGetContext(A, (void **)&ctx));
   MFEM_VERIFY(ctx, "Invalid PETSc shell matrix context for SLEPc!");
 
-  PetscCall(FromPetscVec(x, ctx->x));
-  ctx->A.Mult(ctx->x, ctx->y);
-  PetscCall(ToPetscVec(ctx->y, y));
+  PetscInt m, n;
+  PalacePetscCall(MatGetLocalSize(A, &m, &n));
+  auto x1 = workspace::NewVector<ComplexVector>(n);
+  auto y1 = workspace::NewVector<ComplexVector>(m);
+  PetscCall(FromPetscVec(x, x1));
+  ctx->Mult(x1, y1);
+  PetscCall(ToPetscVec(y1, y));
 
   PetscFunctionReturn(PETSC_SUCCESS);
 }
 
 PetscErrorCode __mat_apply_transpose_shell(Mat A, Vec x, Vec y)
 {
   PetscFunctionBeginUser;
-  MatShellContext *ctx;
+  ComplexOperator *ctx;
   PetscCall(MatShellGetContext(A, (void **)&ctx));
   MFEM_VERIFY(ctx, "Invalid PETSc shell matrix context for SLEPc!");
 
-  PetscCall(FromPetscVec(x, ctx->x));
-  ctx->A.MultTranspose(ctx->x, ctx->y);
-  PetscCall(ToPetscVec(ctx->y, y));
+  PetscInt m, n;
+  PalacePetscCall(MatGetLocalSize(A, &m, &n));
+  auto x1 = workspace::NewVector<ComplexVector>(m);
+  auto y1 = workspace::NewVector<ComplexVector>(n);
+  PetscCall(FromPetscVec(x, x1));
+  ctx->MultTranspose(x1, y1);
+  PetscCall(ToPetscVec(y1, y));
 
   PetscFunctionReturn(PETSC_SUCCESS);
 }
 
 PetscErrorCode __mat_apply_hermitian_transpose_shell(Mat A, Vec x, Vec y)
 {
   PetscFunctionBeginUser;
-  MatShellContext *ctx;
+  ComplexOperator *ctx;
   PetscCall(MatShellGetContext(A, (void **)&ctx));
   MFEM_VERIFY(ctx, "Invalid PETSc shell matrix context for SLEPc!");
 
-  PetscCall(FromPetscVec(x, ctx->x));
-  ctx->A.MultHermitianTranspose(ctx->x, ctx->y);
-  PetscCall(ToPetscVec(ctx->y, y));
+  PetscInt m, n;
+  PalacePetscCall(MatGetLocalSize(A, &m, &n));
+  auto x1 = workspace::NewVector<ComplexVector>(m);
+  auto y1 = workspace::NewVector<ComplexVector>(n);
+  PetscCall(FromPetscVec(x, x1));
+  ctx->MultHermitianTranspose(x1, y1);
+  PetscCall(ToPetscVec(y1, y));
 
   PetscFunctionReturn(PETSC_SUCCESS);
 };
@@ -241,14 +247,9 @@ PetscReal GetMaxSingularValue(MPI_Comm comm, const ComplexOperator &A, bool herm
   // or SVD solvers, namely MATOP_MULT and MATOP_MULT_HERMITIAN_TRANSPOSE (if the matrix
   // is not Hermitian).
   MFEM_VERIFY(A.Height() == A.Width(), "Spectral norm requires a square matrix!");
-  const PetscInt n = A.Height();
-  ComplexVector x(n), y(n);
-  x.UseDevice(true);
-  y.UseDevice(true);
-  MatShellContext ctx = {A, x, y};
   Mat A0;
-  PalacePetscCall(
-      MatCreateShell(comm, n, n, PETSC_DECIDE, PETSC_DECIDE, (void *)&ctx, &A0));
+  const PetscInt n = A.Height();
+  PalacePetscCall(MatCreateShell(comm, n, n, PETSC_DECIDE, PETSC_DECIDE, (void *)&A, &A0));
   PalacePetscCall(MatShellSetOperation(A0, MATOP_MULT, (void (*)(void))__mat_apply_shell));
   PalacePetscCall(MatShellSetVecType(A0, PetscVecType()));
   if (herm)