style

Add toggle for SuperLU_DIST device offload

fem/CMakeLists.txt

@@ -213,7 +213,6 @@ set(HDRS
   dfem/qfunction_apply.hpp
   dfem/qfunction_transform.hpp
   dfem/tuple.hpp
-  dfem/univarsolvers.hpp
   dfem/util.hpp
   eltrans.hpp
   estimators.hpp

fem/dfem/univarsolvers.hpp

@@ -1,232 +0,0 @@
-// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
-// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
-// LICENSE and NOTICE for details. LLNL-CODE-806117.
-//
-// This file is part of the MFEM library. For more information and source code
-// availability visit https://mfem.org.
-//
-// MFEM is free software; you can redistribute it and/or modify it under the
-// terms of the BSD-3 license. We welcome feedback and contributions, see file
-// CONTRIBUTING.md for details.
-
-/**
- * @file univarsolvers.hpp
- *
- * @brief Solvers of functions of a single variable suitable for use in ∂FEM q-functions.
- */
-
-#ifndef MFEM_UNIVARSOLVERS
-#define MFEM_UNIVARSOLVERS
-
-#include <cmath>
-#include <limits>
-
-#include "../config/config.hpp"
-#include "general/enzyme.hpp"
-#include "general/error.hpp"
-
-#ifdef MFEM_USE_ENZYME
-
-namespace mfem {
-
-namespace future {
-/// Representation of bound constraints
-struct Bounds {
-  real_t lower, upper;
-};
-
-
-/// Settings for univariate solver
-struct SolverSettings {
-  real_t residual_abs_tol = 1e-10; ///< Tolerance for convergence check on absolute value of residual
-  real_t residual_rel_tol = 0.0;   ///< Tolerance for convergence check on absolute value of current residual relative to absolute value of residual at initial guess
-  Bounds bounds; ///< Bounds on root
-};
-} // namespace future
-
-namespace internal {
-/// @cond DO_NOT_DOCUMENT
-
-using future::SolverSettings;
-
-template <auto f, typename T>
-__attribute__((noinline))
-MFEM_HOST_DEVICE void SolveNewtonBisection_impl(const real_t* x0_ptr, const T* p_ptr, const SolverSettings* settings_ptr, real_t* x_ptr)
-{
-    // It would be better to have the max iterations in the settings instead of
-    // hard-coded.
-    constexpr int max_iters = 50;
-
-    const real_t& x0 = *x0_ptr;
-    const T& p = *p_ptr;
-    const SolverSettings& settings = *settings_ptr;
-    const real_t& left_bracket = settings.bounds.lower;
-    const real_t& right_bracket = settings.bounds.upper;
-    real_t& x = *x_ptr;
-    using std::abs;
-
-    auto fprime = [&p](real_t x) {
-        real_t x_dot = 1.0;
-        return __enzyme_fwddiff<real_t>((void*)+f, enzyme_dup, x, x_dot, enzyme_const, p);
-    };
-
-    real_t fl = f(left_bracket, p);
-    real_t fh = f(right_bracket, p);
-
-    // handle corner cases where one of the brackets is the root
-    if (abs(fl) < settings.residual_abs_tol)
-    {
-        x = left_bracket;
-        return;
-    }
-    else if (abs(fh) < settings.residual_abs_tol)
-    {
-        x = right_bracket;
-        return;
-    }
-
-    if (fl * fh > 0)
-    {
-        MFEM_WARNING("Root is not bracketed, solver may diverge.");
-    }
-
-    // clamp initial guess within root brackets
-    x = x0 > right_bracket? right_bracket : x0 < left_bracket? left_bracket : x0;
-    std::cout << "why is this not printing?" << std::endl;
-    std::cout << "x0 = " << x0 << " initial x set to " << x << std::endl;
-
-    // Orient search so that f(xl) < 0
-    real_t xl = left_bracket;
-    real_t xh = right_bracket;
-    if (fl > 0.0) {
-        xl = right_bracket;
-        xh = left_bracket;
-        real_t tmp = fl;
-        fl = fh;
-        fh = tmp;
-    }
-
-    real_t dx_old = abs(right_bracket - left_bracket);
-    real_t dx = dx_old;
-    x = x0;
-    real_t r = f(x, p);
-    real_t dr_dx = fprime(x);
-    real_t r_old = r;
-    for (int i = 0; i < max_iters; i++) {
-        if ((((x - xh) * dr_dx - r)*((x - xl)*dr_dx - r) >= 0.0) ||  // Newton out of range
-            (std::abs(2.0*r) > std::abs(dx_old*dr_dx))) {            // Newton decreasing bracket slower than bisection
-            // Take bisection step
-            dx_old = dx;
-            dx = 0.5*(xh - xl);
-            real_t x_old = x;
-            x = xl + dx;
-            if (x == x_old) return;
-        } else {
-            // Take Newton step
-            dx_old = dx;
-            dx = -r/dr_dx;
-            real_t x_old = x;
-            x += dx;
-            if (x == x_old) return;
-        }
-
-        // update residual and jacobian
-        r = f(x, p);
-        dr_dx = fprime(x);
-
-        // Check convergence
-        if (abs(r) < settings.residual_rel_tol*r_old || abs(r) < settings.residual_abs_tol) {
-            return;
-        }
-
-        // Update bracket
-        if (r < 0.0) {
-            xl = x;
-            fl = r;
-        } else {
-            xh = x;
-            fh = r;
-        }
-    }
-    MFEM_ABORT("Univariate solve did not converge.");
-}
-
-template <auto f, typename T>
-void SolveNewtonBisection_impl_fwddiff(const real_t* x0, const real_t* /* unused shadow */,
-                                       const T* p, const T* dp,
-                                       const SolverSettings* settings, const SolverSettings* /* unused shadow */,
-                                       real_t* x, real_t* dx)
-{
-  SolveNewtonBisection_impl<f>(x0, p, settings, x);
-  real_t dfdx = __enzyme_fwddiff<real_t>((void*)+f, enzyme_dup, *x, 1.0, enzyme_const, *p);
-  real_t dfdp = __enzyme_fwddiff<real_t>((void*)+f, enzyme_const, *x, enzyme_dup, *p, *dp);
-  *dx = -dfdp/dfdx;
-}
-
-
-template<auto f, typename T>
-void SolveNewtonBisection_impl_aug(const real_t* x0, real_t* x0_bar,
-                                   const T* p, T* p_bar,
-                                   const SolverSettings* settings, SolverSettings* settings_bar,
-                                   real_t* x, real_t* x_bar)
-{
-    SolveNewtonBisection_impl<f>(x0, p, settings, x);
-    std::cout << "augmented forward, x = " << *x << std::endl;
-}
-
-// Change the residual function to return-by-reference so that there is a
-// slot to provide the downstream cotangent (ie the shadow for y)
-// in the reverse mode call.
-template<auto f, typename T>
-void wrapper(real_t x, T& p, real_t& y)
-{
-    y = f(x, p);
-}
-
-template<auto f, typename T>
-void SolveNewtonBisection_impl_rev(const real_t* x0, real_t* x0_bar,
-                                   const T* p, T* p_bar,
-                                   const SolverSettings* settings, SolverSettings* settings_bar,
-                                   real_t* x, real_t* x_bar)
-{
-    real_t drdx = __enzyme_fwddiff<real_t>((void*)+f, enzyme_dup, *x, 1.0, enzyme_const, *p);
-    std::cout << "reverse pass" << std::endl;
-    std::cout << "x = " << *x << std::endl;
-    std::cout << "x from tape = " << *x << std::endl;
-    std::cout << "drdx = " << drdx << std::endl;
-    real_t lambda = -(*x_bar / drdx);
-    std::cout << "lambda = " << lambda << std::endl;
-    real_t r;
-    __enzyme_autodiff<void>((void*)wrapper<f, T>, enzyme_const, *x, enzyme_dup, p, p_bar, enzyme_dupnoneed, &r, &lambda);
-
-    std::cout << "p_bar = " << *p_bar << std::endl;
-
-    // TODO: Make enzyme treat these as enzyme_const
-    // The solution has no sensitivity to these parameters.
-    *x0_bar = 0.0;
-    *settings_bar = SolverSettings{};
-    std::cout << "settings_bar.bounds.upper = " << settings_bar->bounds.upper << std::endl;
-}
-
-/// @endcond
-} // namespace internal
-
-
-namespace future {
-
-/**
- * @brief Find the root of a univariate funtion
- */
-template<auto f, typename T>
-MFEM_HOST_DEVICE __attribute__((always_inline)) real_t SolveNewtonBisection(real_t x0, T p, SolverSettings settings) {
-  real_t x;
-  internal::SolveNewtonBisection_impl<f>(&x0, &p, &settings, &x);
-  return x;
-}
-
-
-} // namespace future
-} // namespace mfem
-
-#endif // MFEM_USE_ENZYME
-#endif // MFEM_UNIVARSOLVERS

fem/dgmassinv_kernels.hpp

@@ -387,7 +387,7 @@ void DGMassInverse::DGMassCGIteration(const Vector &b_, Vector &u_) const
 
    static constexpr int NB = Q1D ? Q1D : 1; // block size
 
-   mfem::forall_2D(NE, NB, NB, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall_2D<NB*NB>(NE, NB, NB, [=] MFEM_HOST_DEVICE (int e)
    {
       // Perform change of basis if needed
       if (CHANGE_BASIS)

fem/fem.hpp

@@ -53,7 +53,6 @@
 #include "particleset.hpp"
 
 #include "dfem/doperator.hpp"
-#include "dfem/univarsolvers.hpp"
 
 #ifdef MFEM_USE_MPI
 #include "pfespace.hpp"

fem/integ/bilininteg_diffusion_kernels.hpp

@@ -1004,13 +1004,16 @@ inline void SmemPADiffusionApply3D(const int NE,
    const int max_d1d = T_D1D ? T_D1D : DeviceDofQuadLimits::Get().MAX_D1D;
    MFEM_VERIFY(D1D <= max_d1d, "");
    MFEM_VERIFY(Q1D <= max_q1d, "");
-   auto b = Reshape(b_.Read(), Q1D, D1D);
-   auto g = Reshape(g_.Read(), Q1D, D1D);
-   auto d = Reshape(d_.Read(), Q1D, Q1D, Q1D, symmetric ? 6 : 9, NE);
-   auto x = Reshape(x_.Read(), D1D, D1D, D1D, NE);
+   const auto b = Reshape(b_.Read(), Q1D, D1D);
+   const auto g = Reshape(g_.Read(), Q1D, D1D);
+   const auto d = Reshape(d_.Read(), Q1D, Q1D, Q1D, symmetric ? 6 : 9, NE);
+   const auto x = Reshape(x_.Read(), D1D, D1D, D1D, NE);
    auto y = Reshape(y_.ReadWrite(), D1D, D1D, D1D, NE);
    MFEM_VERIFY(D1D <= Q1D, "THREAD_DIRECT requires D1D <= Q1D");
-   mfem::forall_3D(NE, Q1D, Q1D, Q1D, [=] MFEM_HOST_DEVICE (int e)
+
+   mfem::forall_3D<T_Q1D*T_Q1D*T_Q1D>(NE,
+                                      Q1D, Q1D, Q1D,
+                                      [=] MFEM_HOST_DEVICE (int e)
    {
       const int D1D = T_D1D ? T_D1D : d1d;
       const int Q1D = T_Q1D ? T_Q1D : q1d;

fem/integ/bilininteg_mass_kernels.hpp

@@ -1133,11 +1133,11 @@ inline void SmemPAMassApply3D(const int NE,
    const int max_d1d = T_D1D ? T_D1D : DeviceDofQuadLimits::Get().MAX_D1D;
    MFEM_VERIFY(D1D <= max_d1d, "");
    MFEM_VERIFY(Q1D <= max_q1d, "");
-   auto b = b_.Read();
-   auto d = d_.Read();
-   auto x = x_.Read();
+   const auto b = b_.Read();
+   const auto d = d_.Read();
+   const auto x = x_.Read();
    auto y = y_.ReadWrite();
-   mfem::forall_2D(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall_2D<T_Q1D*T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE (int e)
    {
       internal::SmemPAMassApply3D_Element<T_D1D,T_Q1D>(e, NE, b, d, x, y, d1d, q1d);
    });
@@ -1156,8 +1156,8 @@ inline void EAMassAssemble1D(const int NE,
    const int Q1D = T_Q1D ? T_Q1D : q1d;
    MFEM_VERIFY(D1D <= DeviceDofQuadLimits::Get().MAX_D1D, "");
    MFEM_VERIFY(Q1D <= DeviceDofQuadLimits::Get().MAX_Q1D, "");
-   auto B = Reshape(basis.Read(), Q1D, D1D);
-   auto D = Reshape(padata.Read(), Q1D, NE);
+   const auto B = Reshape(basis.Read(), Q1D, D1D);
+   const auto D = Reshape(padata.Read(), Q1D, NE);
    auto M = Reshape(add ? eadata.ReadWrite() : eadata.Write(), D1D, D1D, NE);
    mfem::forall_2D(NE, D1D, D1D, [=] MFEM_HOST_DEVICE (int e)
    {

fem/integ/bilininteg_vecdiffusion_pa.hpp

@@ -54,7 +54,7 @@ void SmemPAVectorDiffusionApply2D(const int NE,
    const auto XE = Reshape(x.Read(), D1D, D1D, SDIM, NE);
    auto YE = Reshape(y.ReadWrite(), D1D, D1D, SDIM, NE);
 
-   mfem::forall_2D(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
+   mfem::forall_2D<T_Q1D*T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
    {
       constexpr int MD1 = T_D1D > 0 ? SetMaxOf(T_D1D) : DofQuadLimits::MAX_T1D;
       constexpr int MQ1 = T_Q1D > 0 ? SetMaxOf(T_Q1D) : DofQuadLimits::MAX_T1D;
@@ -120,7 +120,7 @@ void SmemPAVectorDiffusionApply3D(const int NE,
    const auto XE = Reshape(x.Read(), D1D, D1D, D1D, SDIM, NE);
    auto YE = Reshape(y.ReadWrite(), D1D, D1D, D1D, SDIM, NE);
 
-   mfem::forall_2D(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
+   mfem::forall_2D<T_Q1D*T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
    {
       constexpr int MD1 = T_D1D > 0 ? SetMaxOf(T_D1D) : DofQuadLimits::MAX_T1D;
       constexpr int MQ1 = T_Q1D > 0 ? SetMaxOf(T_Q1D) : DofQuadLimits::MAX_T1D;

fem/integ/bilininteg_vecmass_pa.hpp

@@ -51,7 +51,7 @@ void SmemPAVectorMassApply2D(const int NE,
    const auto X = Reshape(x.Read(), D1D, D1D, VDIM, NE);
    auto Y = Reshape(y.ReadWrite(), D1D, D1D, VDIM, NE);
 
-   mfem::forall_2D(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
+   mfem::forall_2D<T_Q1D*T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
    {
       constexpr int MD1 = T_D1D > 0 ? SetMaxOf(T_D1D) : DofQuadLimits::MAX_T1D;
       constexpr int MQ1 = T_Q1D > 0 ? SetMaxOf(T_Q1D) : DofQuadLimits::MAX_T1D;
@@ -119,7 +119,7 @@ void SmemPAVectorMassApply3D(const int NE,
    const auto X = Reshape(x.Read(), D1D, D1D, D1D, VDIM, NE);
    auto Y = Reshape(y.ReadWrite(), D1D, D1D, D1D, VDIM, NE);
 
-   mfem::forall_2D(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
+   mfem::forall_2D<T_Q1D*T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
    {
       constexpr int MD1 = T_D1D > 0 ? SetMaxOf(T_D1D) : DofQuadLimits::MAX_T1D;
       constexpr int MQ1 = T_Q1D > 0 ? SetMaxOf(T_Q1D) : DofQuadLimits::MAX_T1D;

general/backends.hpp

@@ -44,6 +44,7 @@
 #endif
 
 #if !defined(MFEM_USE_CUDA_OR_HIP)
+constexpr bool mfem_use_gpu = false;
 #define MFEM_DEVICE
 #define MFEM_HOST
 #define MFEM_LAMBDA
@@ -52,6 +53,7 @@
 #define MFEM_DEVICE_SYNC
 // MFEM_STREAM_SYNC is used for UVM and MPI GPU-Aware kernels
 #define MFEM_STREAM_SYNC
+#define MFEM_LAUNCH_BOUNDS(...)
 #endif
 
 #if !((defined(MFEM_USE_CUDA) && defined(__CUDA_ARCH__)) || \

general/cuda.hpp

@@ -20,9 +20,11 @@
 
 #if defined(MFEM_USE_CUDA) && defined(__CUDACC__)
 #define MFEM_USE_CUDA_OR_HIP
+constexpr bool mfem_use_gpu = true;
 #define MFEM_DEVICE __device__
 #define MFEM_HOST __host__
 #define MFEM_LAMBDA __host__
+#define MFEM_LAUNCH_BOUNDS __launch_bounds__
 // #define MFEM_HOST_DEVICE __host__ __device__ // defined in config/config.hpp
 #define MFEM_DEVICE_SYNC MFEM_GPU_CHECK(cudaDeviceSynchronize())
 #define MFEM_STREAM_SYNC MFEM_GPU_CHECK(cudaStreamSynchronize(0))

general/forall.hpp

@@ -295,11 +295,12 @@ using hip_threads_z =
 #endif
 
 #if defined(MFEM_USE_RAJA) && defined(RAJA_ENABLE_CUDA) && defined(__CUDACC__)
-template <const int BLOCKS = MFEM_CUDA_BLOCKS, typename DBODY>
+template <typename DBODY>
 void RajaCuWrap1D(const int N, DBODY &&d_body)
 {
    //true denotes asynchronous kernel
-   RAJA::forall<RAJA::cuda_exec<BLOCKS,true>>(RAJA::RangeSegment(0,N),d_body);
+   RAJA::forall<RAJA::cuda_exec<MFEM_CUDA_BLOCKS,true>>(RAJA::RangeSegment(0,N),
+                                                        d_body);
 }
 
 template <typename DBODY>
@@ -362,18 +363,18 @@ struct RajaCuWrap;
 template <>
 struct RajaCuWrap<1>
 {
-   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
+   template <typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
-      RajaCuWrap1D<BLCK>(N, d_body);
+      RajaCuWrap1D(N, d_body);
    }
 };
 
 template <>
 struct RajaCuWrap<2>
 {
-   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
+   template <typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -384,7 +385,7 @@ struct RajaCuWrap<2>
 template <>
 struct RajaCuWrap<3>
 {
-   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
+   template <typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -395,11 +396,12 @@ struct RajaCuWrap<3>
 #endif
 
 #if defined(MFEM_USE_RAJA) && defined(RAJA_ENABLE_HIP) && defined(__HIP__)
-template <const int BLOCKS = MFEM_HIP_BLOCKS, typename DBODY>
+template <typename DBODY>
 void RajaHipWrap1D(const int N, DBODY &&d_body)
 {
    //true denotes asynchronous kernel
-   RAJA::forall<RAJA::hip_exec<BLOCKS,true>>(RAJA::RangeSegment(0,N),d_body);
+   RAJA::forall<RAJA::hip_exec<MFEM_HIP_BLOCKS,true>>(RAJA::RangeSegment(0,N),
+                                                      d_body);
 }
 
 template <typename DBODY>
@@ -462,18 +464,18 @@ struct RajaHipWrap;
 template <>
 struct RajaHipWrap<1>
 {
-   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
+   template <typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
-      RajaHipWrap1D<BLCK>(N, d_body);
+      RajaHipWrap1D(N, d_body);
    }
 };
 
 template <>
 struct RajaHipWrap<2>
 {
-   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
+   template <typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -484,7 +486,7 @@ struct RajaHipWrap<2>
 template <>
 struct RajaHipWrap<3>
 {
-   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
+   template <typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -584,12 +586,31 @@ void CuKernel2D(const int N, BODY body)
    body(k);
 }
 
+// __launch_bounds__ second argument is omitted to get the default behavior
+template <int MAX_THREADS_PER_BLOCK, typename BODY>
+__global__
+MFEM_LAUNCH_BOUNDS(MAX_THREADS_PER_BLOCK)
+static void CuKernel2DLaunchBounds(const int N, BODY body)
+{
+   const int k = blockIdx.x*blockDim.z + threadIdx.z;
+   if (k >= N) { return; }
+   body(k);
+}
+
 template <typename BODY> __global__ static
 void CuKernel3D(const int N, BODY body)
 {
    for (int k = blockIdx.x; k < N; k += gridDim.x) { body(k); }
 }
 
+template <int MAX_THREADS_PER_BLOCK, typename BODY>
+__global__
+MFEM_LAUNCH_BOUNDS(MAX_THREADS_PER_BLOCK)
+static void CuKernel3DLaunchBounds(const int N, BODY body)
+{
+   for (int k = blockIdx.x; k < N; k += gridDim.x) { body(k); }
+}
+
 template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
 void CuWrap1D(const int N, DBODY &&d_body)
 {
@@ -604,6 +625,8 @@ void CuWrap2D(const int N, DBODY &&d_body,
               const int X, const int Y, const int BZ)
 {
    if (N==0) { return; }
+   // required for optimized GCC/NVCC builds to prevent runtime
+   // ODR/linkage violations of inlined templated kernel helpers
    MFEM_VERIFY(BZ>0, "");
    const int GRID = (N+BZ-1)/BZ;
    const dim3 BLCK(X,Y,BZ);
@@ -611,6 +634,19 @@ void CuWrap2D(const int N, DBODY &&d_body,
    MFEM_GPU_CHECK(cudaGetLastError());
 }
 
+template <int MAX_THREADS_PER_BLOCK, typename DBODY>
+void CuWrap2DLaunchBounds(const int N, DBODY &&d_body,
+                          const int X, const int Y, const int BZ)
+{
+   if (N==0) { return; }
+   MFEM_VERIFY(BZ>0, "");
+   const int GRID = (N+BZ-1)/BZ;
+   const dim3 BLCK(X,Y,BZ);
+   static_assert(MAX_THREADS_PER_BLOCK > 0);
+   CuKernel2DLaunchBounds<MAX_THREADS_PER_BLOCK><<<GRID,BLCK>>>(N, d_body);
+   MFEM_GPU_CHECK(cudaGetLastError());
+}
+
 template <typename DBODY>
 void CuWrap3D(const int N, DBODY &&d_body,
               const int X, const int Y, const int Z, const int G)
@@ -622,24 +658,35 @@ void CuWrap3D(const int N, DBODY &&d_body,
    MFEM_GPU_CHECK(cudaGetLastError());
 }
 
-template <int Dim>
-struct CuWrap;
-
-template <>
-struct CuWrap<1>
+template <int MAX_THREADS_PER_BLOCK, typename DBODY>
+void CuWrap3DLaunchBounds(const int N, DBODY &&d_body,
+                          const int X, const int Y, const int Z, const int G)
 {
-   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
+   if (N==0) { return; }
+   const int GRID = G == 0 ? N : G;
+   const dim3 BLCK(X,Y,Z);
+   static_assert(MAX_THREADS_PER_BLOCK > 0);
+   CuKernel3DLaunchBounds<MAX_THREADS_PER_BLOCK><<<GRID, BLCK>>>(N, d_body);
+   MFEM_GPU_CHECK(cudaGetLastError());
+}
+
+template <int Dim, int MAX_THREADS_PER_BLOCK> struct CuWrap;
+
+template <int MAX_THREADS_PER_BLOCK>
+struct CuWrap<1, MAX_THREADS_PER_BLOCK>
+{
+   template <typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
-      CuWrap1D<BLCK>(N, d_body);
+      CuWrap1D<MFEM_CUDA_BLOCKS>(N, d_body);
    }
 };
 
 template <>
-struct CuWrap<2>
+struct CuWrap<2, 0>
 {
-   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
+   template <typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -647,10 +694,22 @@ struct CuWrap<2>
    }
 };
 
-template <>
-struct CuWrap<3>
+template <int MAX_THREADS_PER_BLOCK>
+struct CuWrap<2, MAX_THREADS_PER_BLOCK>
 {
-   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
+   template <typename DBODY>
+   static void run(const int N, DBODY &&d_body,
+                   const int X, const int Y, const int Z, const int G)
+   {
+      static_assert(MAX_THREADS_PER_BLOCK > 0);
+      CuWrap2DLaunchBounds<MAX_THREADS_PER_BLOCK>(N, d_body, X, Y, Z);
+   }
+};
+
+template <>
+struct CuWrap<3, 0>
+{
+   template <typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -658,6 +717,17 @@ struct CuWrap<3>
    }
 };
 
+template <int MAX_THREADS_PER_BLOCK>
+struct CuWrap<3, MAX_THREADS_PER_BLOCK>
+{
+   template <typename DBODY>
+   static void run(const int N, DBODY &&d_body,
+                   const int X, const int Y, const int Z, const int G)
+   {
+      CuWrap3DLaunchBounds<MAX_THREADS_PER_BLOCK>(N, d_body, X, Y, Z, G);
+   }
+};
+
 #endif // defined(MFEM_USE_CUDA) && defined(__CUDACC__)
 
 
@@ -680,13 +750,31 @@ void HipKernel2D(const int N, BODY body)
    body(k);
 }
 
+template <int MAX_THREADS_PER_BLOCK, typename BODY>
+__global__
+MFEM_LAUNCH_BOUNDS(MAX_THREADS_PER_BLOCK)
+static void HipKernel2DLaunchBounds(const int N, BODY body)
+{
+   const int k = hipBlockIdx_x*hipBlockDim_z + hipThreadIdx_z;
+   if (k >= N) { return; }
+   body(k);
+}
+
 template <typename BODY> __global__ static
 void HipKernel3D(const int N, BODY body)
 {
    for (int k = hipBlockIdx_x; k < N; k += hipGridDim_x) { body(k); }
 }
 
-template <const int BLCK = MFEM_HIP_BLOCKS, typename DBODY>
+template <int MAX_THREADS_PER_BLOCK, typename BODY>
+__global__
+MFEM_LAUNCH_BOUNDS(MAX_THREADS_PER_BLOCK)
+static void HipKernel3DLaunchBounds(const int N, BODY body)
+{
+   for (int k = hipBlockIdx_x; k < N; k += hipGridDim_x) { body(k); }
+}
+
+template <int BLCK = MFEM_HIP_BLOCKS, typename DBODY>
 void HipWrap1D(const int N, DBODY &&d_body)
 {
    if (N==0) { return; }
@@ -700,12 +788,27 @@ void HipWrap2D(const int N, DBODY &&d_body,
                const int X, const int Y, const int BZ)
 {
    if (N==0) { return; }
+   MFEM_VERIFY(BZ>0, "");
    const int GRID = (N+BZ-1)/BZ;
    const dim3 BLCK(X,Y,BZ);
    hipLaunchKernelGGL(HipKernel2D,GRID,BLCK,0,nullptr,N,d_body);
    MFEM_GPU_CHECK(hipGetLastError());
 }
 
+template <int MAX_THREADS_PER_BLOCK, typename DBODY>
+void HipWrap2DLaunchBounds(const int N, DBODY &&d_body,
+                           const int X, const int Y, const int BZ)
+{
+   if (N==0) { return; }
+   MFEM_VERIFY(BZ>0, "");
+   const int GRID = (N+BZ-1)/BZ;
+   const dim3 BLCK(X,Y,BZ);
+   static_assert(MAX_THREADS_PER_BLOCK > 0);
+   HipKernel2DLaunchBounds<MAX_THREADS_PER_BLOCK><<<dim3(GRID), dim3(BLCK), 0, 0>>>
+   (N, d_body);
+   MFEM_GPU_CHECK(hipGetLastError());
+}
+
 template <typename DBODY>
 void HipWrap3D(const int N, DBODY &&d_body,
                const int X, const int Y, const int Z, const int G)
@@ -717,24 +820,36 @@ void HipWrap3D(const int N, DBODY &&d_body,
    MFEM_GPU_CHECK(hipGetLastError());
 }
 
-template <int Dim>
-struct HipWrap;
-
-template <>
-struct HipWrap<1>
+template <int MAX_THREADS_PER_BLOCK, typename DBODY>
+void HipWrap3DLaunchBounds(const int N, DBODY &&d_body,
+                           const int X, const int Y, const int Z, const int G)
 {
-   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
+   if (N==0) { return; }
+   const int GRID = G == 0 ? N : G;
+   const dim3 BLCK(X,Y,Z);
+   static_assert(MAX_THREADS_PER_BLOCK > 0);
+   HipKernel3DLaunchBounds<MAX_THREADS_PER_BLOCK><<<dim3(GRID), dim3(BLCK), 0, 0>>>
+   (N, d_body);
+   MFEM_GPU_CHECK(hipGetLastError());
+}
+
+template <int Dim, int MAX_THREADS_PER_BLOCK> struct HipWrap;
+
+template <int MAX_THREADS_PER_BLOCK>
+struct HipWrap<1, MAX_THREADS_PER_BLOCK>
+{
+   template <typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
-      HipWrap1D<BLCK>(N, d_body);
+      HipWrap1D<MFEM_HIP_BLOCKS>(N, d_body);
    }
 };
 
 template <>
-struct HipWrap<2>
+struct HipWrap<2, 0>
 {
-   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
+   template <typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -742,10 +857,21 @@ struct HipWrap<2>
    }
 };
 
-template <>
-struct HipWrap<3>
+template <int MAX_THREADS_PER_BLOCK>
+struct HipWrap<2, MAX_THREADS_PER_BLOCK>
 {
-   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
+   template <typename DBODY>
+   static void run(const int N, DBODY &&d_body,
+                   const int X, const int Y, const int Z, const int G)
+   {
+      HipWrap2DLaunchBounds<MAX_THREADS_PER_BLOCK>(N, d_body, X, Y, Z);
+   }
+};
+
+template <>
+struct HipWrap<3, 0>
+{
+   template <typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -753,11 +879,24 @@ struct HipWrap<3>
    }
 };
 
+template <int MAX_THREADS_PER_BLOCK>
+struct HipWrap<3, MAX_THREADS_PER_BLOCK>
+{
+   template <typename DBODY>
+   static void run(const int N, DBODY &&d_body,
+                   const int X, const int Y, const int Z, const int G)
+   {
+      HipWrap3DLaunchBounds<MAX_THREADS_PER_BLOCK>(N, d_body, X, Y, Z, G);
+   }
+};
+
 #endif // defined(MFEM_USE_HIP) && defined(__HIP__)
 
 
-/// The forall kernel body wrapper
-template <const int DIM, typename d_lambda, typename h_lambda>
+///////////////////////////////////////////////////////////////////////////////
+/// Forall host & device kernel dispatch
+template <int DIM, int MAX_THREADS_PER_BLOCK = 0,
+          typename d_lambda, typename h_lambda>
 inline void ForallWrap(const bool use_dev, const int N,
                        d_lambda &&d_body, h_lambda &&h_body,
                        const int X=0, const int Y=0, const int Z=0,
@@ -790,7 +929,7 @@ inline void ForallWrap(const bool use_dev, const int N,
    // If Backend::CUDA is allowed, use it
    if (Device::Allows(Backend::CUDA))
    {
-      return CuWrap<DIM>::run(N, d_body, X, Y, Z, G);
+      return CuWrap<DIM, MAX_THREADS_PER_BLOCK>::run(N, d_body, X, Y, Z, G);
    }
 #endif
 
@@ -798,7 +937,7 @@ inline void ForallWrap(const bool use_dev, const int N,
    // If Backend::HIP is allowed, use it
    if (Device::Allows(Backend::HIP))
    {
-      return HipWrap<DIM>::run(N, d_body, X, Y, Z, G);
+      return HipWrap<DIM, MAX_THREADS_PER_BLOCK>::run(N, d_body, X, Y, Z, G);
    }
 #endif
 
@@ -827,7 +966,9 @@ backend_cpu:
    for (int k = 0; k < N; k++) { h_body(k); }
 }
 
-template <const int DIM, typename lambda>
+///////////////////////////////////////////////////////////////////////////////
+/// Forall host & device kernel wrappers
+template <int DIM, typename lambda>
 inline void ForallWrap(const bool use_dev, const int N, lambda &&body,
                        const int X=0, const int Y=0, const int Z=0,
                        const int G=0)
@@ -835,6 +976,16 @@ inline void ForallWrap(const bool use_dev, const int N, lambda &&body,
    ForallWrap<DIM>(use_dev, N, body, body, X, Y, Z, G);
 }
 
+template <int DIM, int MAX_THREADS_PER_BLOCK, typename lambda>
+inline void ForallWrap(const bool use_dev, const int N, lambda &&body,
+                       const int X=0, const int Y=0, const int Z=0,
+                       const int G=0)
+{
+   ForallWrap<DIM, MAX_THREADS_PER_BLOCK>(use_dev, N, body, body, X, Y, Z, G);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// forall interfaces
 template<typename lambda>
 inline void forall(int N, lambda &&body) { ForallWrap<1>(true, N, body); }
 
@@ -843,7 +994,7 @@ inline void forall(int Nx, int Ny, lambda &&body)
 {
    if (Device::Allows(Backend::DEVICE_MASK))
    {
-      forall(Nx * Ny, [=] MFEM_HOST_DEVICE(int idx)
+      mfem::forall(Nx * Ny, [=] MFEM_HOST_DEVICE(int idx)
       {
          int j = idx / Nx;
          int i = idx % Nx;
@@ -879,7 +1030,7 @@ inline void forall(int Nx, int Ny, int Nz, lambda &&body)
 {
    if (Device::Allows(Backend::DEVICE_MASK))
    {
-      forall(Nx * Ny * Nz, [=] MFEM_HOST_DEVICE(int idx)
+      mfem::forall(Nx * Ny * Nz, [=] MFEM_HOST_DEVICE(int idx)
       {
          int i = idx % Nx;
          int j = idx / Nx;
@@ -927,6 +1078,12 @@ inline void forall_2D(int N, int X, int Y, lambda &&body)
    ForallWrap<2>(true, N, body, X, Y, 1);
 }
 
+template<int MAX_THREADS_PER_BLOCK, typename lambda>
+inline void forall_2D(int N, int X, int Y, lambda &&body)
+{
+   ForallWrap<2, MAX_THREADS_PER_BLOCK>(true, N, body, X, Y, 1);
+}
+
 template<typename lambda>
 inline void forall_2D_batch(int N, int X, int Y, int BZ, lambda &&body)
 {
@@ -939,6 +1096,12 @@ inline void forall_3D(int N, int X, int Y, int Z, lambda &&body)
    ForallWrap<3>(true, N, body, X, Y, Z, 0);
 }
 
+template<int MAX_THREADS_PER_BLOCK, typename lambda>
+inline void forall_3D(int N, int X, int Y, int Z, lambda &&body)
+{
+   ForallWrap<3, MAX_THREADS_PER_BLOCK>(true, N, body, X, Y, Z, 0);
+}
+
 template<typename lambda>
 inline void forall_3D_grid(int N, int X, int Y, int Z, int G, lambda &&body)
 {

general/hip.hpp

@@ -20,9 +20,11 @@
 
 #if defined(MFEM_USE_HIP) && defined(__HIP__)
 #define MFEM_USE_CUDA_OR_HIP
+constexpr bool mfem_use_gpu = true;
 #define MFEM_DEVICE __device__
 #define MFEM_HOST __host__
 #define MFEM_LAMBDA __host__ __device__
+#define MFEM_LAUNCH_BOUNDS __launch_bounds__
 // #define MFEM_HOST_DEVICE __host__ __device__ // defined in config/config.hpp
 #define MFEM_DEVICE_SYNC MFEM_GPU_CHECK(hipDeviceSynchronize())
 #define MFEM_STREAM_SYNC MFEM_GPU_CHECK(hipStreamSynchronize(0))

linalg/CMakeLists.txt

@@ -55,6 +55,7 @@ list(APPEND HDRS
   dinvariants.hpp
   dtensor.hpp
   dual.hpp
+  eigensolver.hpp
   filteredsolver.hpp
   handle.hpp
   invariants.hpp

linalg/eigensolver.hpp

@@ -0,0 +1,203 @@
+// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
+// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
+// LICENSE and NOTICE for details. LLNL-CODE-806117.
+//
+// This file is part of the MFEM library. For more information and source code
+// availability visit https://mfem.org.
+//
+// MFEM is free software; you can redistribute it and/or modify it under the
+// terms of the BSD-3 license. We welcome feedback and contributions, see file
+// CONTRIBUTING.md for details.
+
+/**
+ * @file eigensolver.hpp
+ *
+ * @brief This file contains a common interface for all eigensolver classes
+ */
+
+#ifndef MFEM_EIGENSOLVER
+#define MFEM_EIGENSOLVER
+
+#ifdef MFEM_HYPRE
+#include "hypre.hpp"
+#endif
+
+#ifdef MFEM_SLEPC
+#include "slepc.hpp"
+#endif
+
+namespace mfem
+{
+
+enum class EigenSolverType
+{
+   HYPRE,
+   SLEPC,
+   INVALID_TYPE
+};
+
+/// Provides base class for MFEM Eigensolvers
+class EigenSolverBase
+{
+public:
+   EigenSolverBase() {}
+
+   /// Destructor
+   virtual ~EigenSolverBase() = default;
+
+   /// Solves the eigenvalue problem
+   virtual void Solve() = 0;
+
+   /// Set the required number of modes
+   virtual void SetNumModes(int num_Modes)
+   {
+      numModes=num_Modes;
+   }
+
+   /// @brief  Set the operator to the eigenvalue problem
+   /// @param A - operator
+   virtual void SetOperator(Operator& A) = 0;
+
+   /// @brief Sets operators for the generalized eigenvalue problem
+   /// @param A - operator
+   /// @param M - mass matrix
+   virtual void SetOperator(Operator& A, Operator& M)
+   {
+      MFEM_ABORT("Generalized eigensolver is not supported!");
+   }
+
+   /// Optional method - sets preconditioner for the
+   /// eigenvalue solver.
+   virtual void SetPreconditioner(Solver& precond)
+   {
+      MFEM_ABORT("Preconditioner is not supported!");
+   }
+
+   /// Returns the converged eigenvalues
+   virtual void GetEigenvalues(Array<real_t>& eigen_vals) = 0;
+
+   /// Returns the vec_index eigenvector.
+   virtual void GetEigenvector(int vec_index, Vector& vector) = 0;
+
+   /// Returns the eigensolver type.
+   EigenSolverType GetSolverType() { return eigSolverType; }
+
+protected:
+   int numModes = 0;
+   EigenSolverType eigSolverType = EigenSolverType::INVALID_TYPE;
+};
+
+#ifdef MFEM_HYPRE
+class EigenSolverHypreLOBPCG : public EigenSolverBase
+{
+public:
+   EigenSolverHypreLOBPCG(MPI_Comm comm)
+   {
+      eigenSolver = std::make_unique<HypreLOBPCG>(comm);
+      eigSolverType = EigenSolverType::HYPRE;
+   }
+
+   ~EigenSolverHypreLOBPCG() {}
+
+   void Solve() override { eigenSolver->Solve(); }
+   void SetNumModes(int num_Modes) override
+   {
+      eigenSolver->SetNumModes(num_Modes);
+      numModes = num_Modes;
+   }
+
+   void SetOperator(Operator& A) override { eigenSolver->SetOperator(A); }
+
+   void SetOperator(Operator& A, Operator& M) override
+   {
+      eigenSolver->SetOperator(A);
+      eigenSolver->SetMassMatrix(M);
+   }
+
+   void SetPreconditioner(Solver& precond) override { eigenSolver->SetPreconditioner(precond); }
+   void GetEigenvalues(Array<real_t>& eigen_vals) override { eigenSolver->GetEigenvalues(eigen_vals); }
+   void GetEigenvector(int vec_index, Vector& vector) override
+   {
+      const HypreParVector& eigenvec = eigenSolver->GetEigenvector(vec_index);
+      vector = eigenvec;
+   }
+
+   void SetTol(real_t tol) { eigenSolver->SetTol(tol); }
+   void SetRelTol(real_t rel_tol) { eigenSolver->SetRelTol(rel_tol); }
+   void SetMaxIter(int max_iter) { eigenSolver->SetMaxIter(max_iter); }
+   void SetPrintLevel(int logging) { eigenSolver->SetPrintLevel(logging); }
+   void SetRandomSeed(int seed) { eigenSolver->SetRandomSeed(seed); }
+   void SetPrecondUsageMode(int usage_mode) { eigenSolver->SetPrecondUsageMode(usage_mode); }
+
+private:
+   std::unique_ptr<HypreLOBPCG> eigenSolver = nullptr;
+};
+#endif
+
+#ifdef MFEM_SLEPC
+class EigenSolverSlepc : public EigenSolverBase
+{
+public:
+   EigenSolverSlepc(MPI_Comm comm)
+   {
+      eigSolverType = EigenSolverType::SLEPC;
+      eigenSolver = std::make_unique<SlepcEigenSolver>(comm);
+
+      eigenSolver->SetWhichEigenpairs(SlepcEigenSolver::TARGET_REAL);
+      eigenSolver->SetTarget(0.0);
+      eigenSolver->SetSpectralTransformation(SlepcEigenSolver::SHIFT_INVERT);
+   }
+
+   ~EigenSolverSlepc() {}
+
+   void Solve() override { eigenSolver->Solve(); }
+   void SetNumModes(int num_Modes) override
+   {
+      eigenSolver->SetNumModes(num_Modes);
+      numModes = num_Modes;
+   }
+   /// @brief  Set the operator to the slepc eigenvalue problem. This method deep copies data to create a PetscParMatrix
+   /// @param A - operator, must be of type HypreParMatrix.
+   void SetOperator(Operator& A) override
+   {
+      petscMatA = std::make_unique<PetscParMatrix>
+                  (dynamic_cast<HypreParMatrix*>(&A));
+      eigenSolver->SetOperator(*petscMatA);
+   }
+   /// @brief  Set the operators to the slepc eigenvalue problem. This method deep copies data to create a PetscParMatrix
+   /// @param A - operator, must be of type HypreParMatrix.
+   /// @param M - operator, must be of type HypreParMatrix.
+   void SetOperator(Operator& A, Operator& M) override
+   {
+      petscMatA = std::make_unique<PetscParMatrix>
+                  (dynamic_cast<const HypreParMatrix*>(&A));
+      petscMatM = std::make_unique<PetscParMatrix>
+                  (dynamic_cast<const HypreParMatrix*>(&M));
+
+      eigenSolver->SetOperators(*petscMatA, *petscMatM);
+   }
+   void SetPreconditioner([[maybe_unused]] Solver& precond) override {}
+   void GetEigenvalues(Array<real_t>& eigen_vals) override
+   {
+      eigen_vals.SetSize(numModes);
+      for (int ik = 0; ik < numModes; ik++)
+      {
+         eigenSolver->GetEigenvalue(static_cast<unsigned int>(ik), eigen_vals[ik]);
+      }
+   }
+   void GetEigenvector( int vec_index, Vector& vector) override
+   { eigenSolver->GetEigenvector(vec_index, vector); }
+
+   void SetTol(real_t tol) { eigenSolver->SetTol(tol); }
+   void SetMaxIter(int max_iter) { eigenSolver->SetMaxIter(max_iter); }
+
+private:
+   std::unique_ptr<SlepcEigenSolver> eigenSolver = nullptr;
+   std::unique_ptr<PetscParMatrix> petscMatA = nullptr;
+   std::unique_ptr<PetscParMatrix> petscMatM = nullptr;
+};
+#endif
+
+}  // namespace mfem
+
+#endif

linalg/superlu.cpp

@@ -363,14 +363,19 @@ void SuperLUSolver::Init(MPI_Comm comm)
    // Set default options:
    //    options.Fact = DOFACT;
    //    options.Equil = YES;
+   //    options.ParSymbFact = NO;
    //    options.ColPerm = METIS_AT_PLUS_A;
    //    options.RowPerm = LargeDiag_MC64;
    //    options.ReplaceTinyPivot = NO;
-   //    options.Trans = NOTRANS;
    //    options.IterRefine = SLU_DOUBLE;
+   //    options.Trans = NOTRANS;
    //    options.SolveInitialized = NO;
    //    options.RefineInitialized = NO;
    //    options.PrintStat = YES;
+   //    options.lookahead_etree = NO;
+   //    options.num_lookaheads = 10;
+   //    options.superlu_acc_offload = 1;
+   //    options.SymPattern = NO;
    superlu_dist_options_t *options = (superlu_dist_options_t *)optionsPtr_;
    set_default_options_dist(options);
 #if SUPERLU_DIST_MAJOR_VERSION > 7 || \
@@ -472,6 +477,12 @@ void SuperLUSolver::SetFact(superlu::Fact fact)
    options->Fact = opt;
 }
 
+void SuperLUSolver::SetDeviceOffload(bool offload)
+{
+   superlu_dist_options_t *options = (superlu_dist_options_t *)optionsPtr_;
+   options->superlu_acc_offload = offload;
+}
+
 void SuperLUSolver::SetOperator(const Operator &op)
 {
    // Verify that we have a compatible operator

linalg/superlu.hpp

@@ -250,7 +250,8 @@ public:
        work (default false) */
    void SetSymmetricPattern(bool sym);
 
-   /** @brief Specify whether to perform parallel symbolic factorization.
+   /** @brief Specify whether to perform parallel symbolic factorization
+       (default false)
        @note If true SuperLU will use superlu::PARMETIS for the Column
        Permutation regardless of the setting */
    void SetParSymbFact(bool par);
@@ -263,6 +264,10 @@ public:
        superlu::FACTORED*/
    void SetFact(superlu::Fact fact);
 
+   /** @brief Specify whether to offload numerical factorization onto the device
+       (default true if SuperLU_DIST has been compiled with GPU support) */
+   void SetDeviceOffload(bool offload);
+
    // Processor grid for SuperLU_DIST.
    const int nprow_, npcol_, npdep_;
 

tests/benchmarks/CMakeLists.txt

@@ -31,11 +31,6 @@ function(add_benchmark name)
        set_property(SOURCE ${${NAME}_BENCH_SRCS} PROPERTY LANGUAGE CUDA)
     endif(MFEM_USE_CUDA)
 
-    if (MFEM_USE_HIP)
-        set_property(SOURCE ${${NAME}_BENCH_SRCS} PROPERTY LANGUAGE 
-                     HIP_SOURCE_PROPERTY_FORMAT TRUE)
-    endif(MFEM_USE_HIP)
-
     add_executable(bench_${name} ${${NAME}_BENCH_SRCS})
     target_link_libraries(bench_${name} mfem pthread)
     add_dependencies(${MFEM_ALL_BENCHMARKS_TARGET_NAME} bench_${name})

tests/benchmarks/bench_ceed.cpp

@@ -8,23 +8,89 @@
 // MFEM is free software; you can redistribute it and/or modify it under the
 // terms of the BSD-3 license. We welcome feedback and contributions, see file
 // CONTRIBUTING.md for details.
+//
+//
+//  This benchmark contains the implementation of the CEED's bake-off problems:
+//  high-order kernels/benchmarks designed to test and compare the performance
+//  of high-order codes.
+//
+//  See: https://ceed.exascaleproject.org/bps
 
-#include "bench.hpp"
+#include "bench.hpp" // IWYU pragma: keep
 
 #ifdef MFEM_USE_BENCHMARK
 
-/*
-  This benchmark contains the implementation of the CEED's bake-off problems:
-  high-order kernels/benchmarks designed to test and compare the performance
-  of high-order codes.
+#include <cassert>
+#include <string>
 
-  See: ceed.exascaleproject.org/bps and github.com/CEED/benchmarks
-*/
-template <int VDIM, bool GLL>
+#include "fem/qinterp/det.hpp" // IWYU pragma: keep
+#include "fem/qinterp/grad.hpp" // IWYU pragma: keep
+#include "fem/integ/lininteg_domain_kernels.hpp" // IWYU pragma: keep
+#include "fem/integ/bilininteg_vecdiffusion_pa.hpp" // IWYU pragma: keep
+
+// Custom benchmark arguments generator
+static void CustomArguments(bmi::Benchmark *b) noexcept
+{
+   constexpr int MAX_NDOFS = 16 * 1024 * (mfem_use_gpu ? 1024 : 8);
+
+   const auto orders = { 7, 6, 5, 4, 3, 2, 1 };
+
+   constexpr auto ndofs = [](int n) constexpr noexcept -> int
+   {
+      return (n + 1) * (n + 1) * (n + 1);
+   };
+
+   constexpr auto inc = [](int n) constexpr noexcept -> int
+   {
+      return n < 160 ?  4 : n < 240 ?  8 : n < 320 ? 16 : 32;
+   };
+
+   for (auto p : orders)
+   {
+      for (int n = 16; ndofs(n) <= MAX_NDOFS; n += inc(n))
+      {
+         b->Args({p, n});
+      }
+   }
+}
+
+// Register kernel specializations used in the benchmarks
+static void AddKernelSpecializations()
+{
+   using DET = QuadratureInterpolator::DetKernels;
+   DET::Specialization<3, 3, 2, 2>::Add();
+   DET::Specialization<3, 3, 2, 3>::Add();
+   DET::Specialization<3, 3, 2, 5>::Add();
+   DET::Specialization<3, 3, 2, 6>::Add();
+   DET::Specialization<3, 3, 5, 5>::Add();
+   // Others might exceed memory limits
+
+   using GRAD = QuadratureInterpolator::GradKernels;
+   GRAD::Specialization<3, QVectorLayout::byNODES, false, 3, 2, 2>::Add();
+   GRAD::Specialization<3, QVectorLayout::byNODES, false, 3, 2, 7>::Add();
+   GRAD::Specialization<3, QVectorLayout::byNODES, false, 3, 2, 8>::Add();
+   GRAD::Specialization<3, QVectorLayout::byNODES, false, 3, 2, 9>::Add();
+
+   using LIN = DomainLFIntegrator::AssembleKernels;
+   LIN::Specialization<3, 7, 7>::Add();
+   LIN::Specialization<3, 6, 6>::Add();
+   LIN::Specialization<3, 8, 8>::Add();
+
+   using VDIFF = VectorDiffusionIntegrator::ApplyPAKernels;
+   VDIFF::Specialization<3, 3, 3, 3>::Add();
+   VDIFF::Specialization<3, 3, 4, 4>::Add();
+   VDIFF::Specialization<3, 3, 5, 5>::Add();
+   VDIFF::Specialization<3, 3, 6, 6>::Add();
+   VDIFF::Specialization<3, 3, 7, 7>::Add();
+   VDIFF::Specialization<3, 3, 8, 8>::Add();
+}
+
+// Bake-off base class
+template <int BFI, int VDIM, bool GLL>
 struct BakeOff
 {
-   static constexpr int DIM = 3;
-   const int N, p, q;
+   inline static constexpr int DIM = 3;
+   const int p, c, q, n, nx, ny, nz;
    Mesh mesh;
    H1_FECollection fec;
    FiniteElementSpace fes;
@@ -38,12 +104,15 @@ struct BakeOff
    GridFunction x, y;
    BilinearForm a;
    double mdofs{};
+   BilinearFormIntegrator *bfi;
 
-   BakeOff(int p):
-      N(Device::IsEnabled() ? 32 : 4),
-      p(p),
-      q(2 * p + (GLL ? -1 : 3)),
-      mesh(Mesh::MakeCartesian3D(N, N, N, Element::HEXAHEDRON)),
+   BakeOff(int p, int side):
+      p(p), c(side), q(2 * p + (GLL ? -1 : 3)),
+      n((assert(c >= p), c / p)),
+      nx(n + (p * (n + 1) * p * n * p * n < c * c * c ? 1 : 0)),
+      ny(n + (p * (n + 1) * p * (n + 1) * p * n < c * c * c ? 1 : 0)),
+      nz(n),
+      mesh(Mesh::MakeCartesian3D(nx, ny, nz, Element::HEXAHEDRON)),
       fec(p, DIM, BasisType::GaussLobatto),
       fes(&mesh, &fec, VDIM, VDIM == 3 ? Ordering::byVDIM : Ordering::byNODES),
       geom_type(mesh.GetTypicalElementGeometry()),
@@ -58,22 +127,41 @@ struct BakeOff
       a(&fes)
    {
       x = 0.0;
+      if constexpr (BFI == 1)
+      {
+         bfi = new MassIntegrator(one, ir);
+      }
+      else if constexpr (BFI == 2)
+      {
+         bfi = new VectorMassIntegrator(one, ir);
+      }
+      else if constexpr (BFI == 3 || BFI == 5)
+      {
+         bfi = new DiffusionIntegrator(one, ir);
+      }
+      else if constexpr (BFI == 4 || BFI == 6)
+      {
+         bfi = new VectorDiffusionIntegrator(one, ir);
+      }
+      else
+      {
+         static_assert(BFI >= 1 && BFI <= 6, "Invalid BilinearFormIntegrator");
+      }
+      a.AddDomainIntegrator(bfi);
    }
 
    virtual void benchmark() = 0;
 
-   double SumMdofs() const { return mdofs; }
+   [[nodiscard]] double SumMdofs() const noexcept { return mdofs; }
 
-   double MDofs() const { return 1e-6 * dofs; }
+   [[nodiscard]] double MDofs() const noexcept { return 1e-6 * dofs; }
 };
 
-/// Bake-off Problems (BPs)
-template <typename BFI, int VDIM, bool GLL>
-struct Problem : public BakeOff<VDIM, GLL>
+// Bake-off Problems (BPs)
+template <int BFI, int VDIM, bool GLL>
+struct BP : public BakeOff<BFI, VDIM, GLL>
 {
-   const double rtol = 1e-12;
-   const int max_it = 32;
-   const int print_lvl = -1;
+   const int max_it = 32, print_lvl = -1;
 
    Array<int> ess_tdof_list;
    Array<int> ess_bdr;
@@ -82,44 +170,56 @@ struct Problem : public BakeOff<VDIM, GLL>
    Vector B, X;
    CGSolver cg;
 
-   using BakeOff<VDIM, GLL>::a;
-   using BakeOff<VDIM, GLL>::ir;
-   using BakeOff<VDIM, GLL>::one;
-   using BakeOff<VDIM, GLL>::mesh;
-   using BakeOff<VDIM, GLL>::fes;
-   using BakeOff<VDIM, GLL>::x;
-   using BakeOff<VDIM, GLL>::y;
-   using BakeOff<VDIM, GLL>::mdofs;
+   using base = BakeOff<BFI, VDIM, GLL>;
+   using base::a;
+   using base::ir;
+   using base::one;
+   using base::mesh;
+   using base::fes;
+   using base::x;
+   using base::y;
+   using base::mdofs;
+   using base::unit_vec;
+   using base::bfi;
 
-   Problem(int order):
-      BakeOff<VDIM, GLL>(order),
+   BP(int p, int side) noexcept: base(p, side),
       ess_bdr(mesh.bdr_attributes.Max()),
       b(&fes)
    {
       ess_bdr = 1;
       fes.GetEssentialTrueDofs(ess_bdr, ess_tdof_list);
-      if (VDIM == 1)
+
+      if constexpr (VDIM == 1)
       {
-         b.AddDomainIntegrator(new DomainLFIntegrator(this->one));
+         b.AddDomainIntegrator(new DomainLFIntegrator(one));
       }
       else
       {
-         b.AddDomainIntegrator(new VectorDomainLFIntegrator(this->unit_vec));
+         b.AddDomainIntegrator(new VectorDomainLFIntegrator(unit_vec));
       }
       b.UseFastAssembly(true);
       b.Assemble();
 
       a.SetAssemblyLevel(AssemblyLevel::PARTIAL);
-      a.AddDomainIntegrator(new BFI(one, ir));
       a.Assemble();
       a.FormLinearSystem(ess_tdof_list, x, b, A, X, B);
 
-      cg.SetRelTol(rtol);
       cg.SetOperator(*A);
+      cg.SetAbsTol(0.0);
+      cg.iterative_mode = false;
+      {
+         cg.SetPrintLevel(-1);
+         cg.SetMaxIter(1000);
+         cg.SetRelTol(1e-8);
+         cg.Mult(B, X);
+         MFEM_VERIFY(cg.GetConverged(), "CG solver did not converge!");
+      }
+      cg.SetRelTol(0.0);
       cg.SetMaxIter(max_it);
       cg.SetPrintLevel(print_lvl);
-      cg.iterative_mode = false;
-      MFEM_DEVICE_SYNC;
+
+      benchmark();
+      mdofs = 0.0;
    }
 
    void benchmark() override
@@ -130,104 +230,115 @@ struct Problem : public BakeOff<VDIM, GLL>
    }
 };
 
-/// Bake-off Problems (BPs)
-#define BakeOff_Problem(i, Kernel, VDIM, p_eq_q)                     \
-   static void BP##i(bm::State &state)                               \
-   {                                                                 \
-      Problem<Kernel##Integrator, VDIM, p_eq_q> ker(state.range(0)); \
-      while (state.KeepRunning()) { ker.benchmark(); }               \
-      state.counters["MDof/s"] =                                     \
-         bm::Counter(ker.SumMdofs(), bm::Counter::kIsRate);          \
-   }                                                                 \
-   BENCHMARK(BP##i)->DenseRange(1, 6)->Unit(bm::kMillisecond);
-
-/// BP1: scalar PCG with mass matrix, q=p+2
-BakeOff_Problem(1, Mass, 1, false)
-
-/// BP2: vector PCG with mass matrix, q=p+2
-BakeOff_Problem(2, VectorMass, 3, false)
-
-/// BP3: scalar PCG with stiffness matrix, q=p+2
-BakeOff_Problem(3, Diffusion, 1, false)
-
-/// BP4: vector PCG with stiffness matrix, q=p+2
-BakeOff_Problem(4, VectorDiffusion, 3, false)
-
-/// BP5: scalar PCG with stiffness matrix, q=p+1
-BakeOff_Problem(5, Diffusion, 1, true)
-
-/// BP6: vector PCG with stiffness matrix, q=p+1
-BakeOff_Problem(6, VectorDiffusion, 3, true)
-
-/// Bake-off Kernels (BKs)
-template <typename BFI, int VDIM, bool GLL>
-struct Kernel : public BakeOff<VDIM, GLL>
+// Bake-off Kernels (BKs)
+template <int BFI, int VDIM, bool GLL>
+struct BK : public BakeOff<BFI, VDIM, GLL>
 {
-   using BakeOff<VDIM, GLL>::a;
-   using BakeOff<VDIM, GLL>::ir;
-   using BakeOff<VDIM, GLL>::one;
-   using BakeOff<VDIM, GLL>::fes;
-   using BakeOff<VDIM, GLL>::x;
-   using BakeOff<VDIM, GLL>::y;
-   using BakeOff<VDIM, GLL>::mdofs;
+   Vector xe, ye;
 
-   Kernel(int order): BakeOff<VDIM, GLL>(order)
+   using base = BakeOff<BFI, VDIM, GLL>;
+   using base::ir;
+   using base::one;
+   using base::bfi;
+   using base::fes;
+   using base::mdofs;
+
+   BK(int order, int side) noexcept: base(order, side)
    {
-      x.Randomize(1);
-      a.SetAssemblyLevel(AssemblyLevel::PARTIAL);
-      a.AddDomainIntegrator(new BFI(one, ir));
-      a.Assemble();
-      a.Mult(x, y);
-      MFEM_DEVICE_SYNC;
+      bfi->AssemblePA(fes);
+
+      const Table &el2dof = fes.GetElementToDofTable();
+      const int e_size = el2dof.Size_of_connections()*fes.GetVDim();
+      const auto R = fes.GetElementRestriction(ElementDofOrdering::LEXICOGRAPHIC);
+      MFEM_VERIFY(e_size == R->Height(), "Input/Output E-vector size mismatch!");
+
+      xe.SetSize(R->Height());
+      ye.SetSize(R->Height());
+      xe.UseDevice(true);
+      ye.UseDevice(true);
+
+      xe.Randomize(1);
+      xe.Read();
+      ye = 0.0;
+
+      benchmark();
+      mdofs = 0.0;
    }
 
    void benchmark() override
    {
-      a.Mult(x, y);
+      bfi->AddMultPA(xe, ye);
       MFEM_DEVICE_SYNC;
       mdofs += this->MDofs();
    }
 };
 
-/// Generic CEED BKi
-#define BakeOff_Kernel(i, KER, VDIM, GLL)                      \
-   static void BK##i(bm::State &state)                         \
-   {                                                           \
-      Kernel<KER##Integrator, VDIM, GLL> ker(state.range(0));  \
-      while (state.KeepRunning()) { ker.benchmark(); }         \
-      state.counters["MDof/s"] =                               \
-         bm::Counter(ker.SumMdofs(), bm::Counter::kIsRate);    \
-   }                                                           \
-   BENCHMARK(BK##i)->DenseRange(1, 6)->Unit(bm::kMillisecond);
+// Benchmarks
+template <typename T>
+static void Benchmark(bm::State& state) noexcept
+{
+   T run(state.range(0), state.range(1));
+   while (state.KeepRunning()) { run.benchmark(); }
+   state.counters["Dofs"] = bm::Counter(run.dofs);
+   state.counters["MDof/s"] = bm::Counter(run.SumMdofs(), bm::Counter::kIsRate);
+   state.counters["Order"] = bm::Counter(state.range(0));
+}
 
-/// BK1: scalar E-vector-to-E-vector evaluation of mass matrix, q=p+2
-BakeOff_Kernel(1, Mass, 1, false)
+#define REGISTER(PK, BFI, VDIM, GLL) \
+   BENCHMARK_TEMPLATE(Benchmark, PK<BFI, VDIM, GLL>) \
+   ->Name(#PK #BFI)->Apply(CustomArguments)->Unit(bm::kMillisecond)
 
-/// BK2: vector E-vector-to-E-vector evaluation of mass matrix, q=p+2
-BakeOff_Kernel(2, VectorMass, 3, false)
+// BP1: scalar PCG with mass matrix, q=p+2
+REGISTER(BP, 1, 1, false);
 
-/// BK3: scalar E-vector-to-E-vector evaluation of stiffness matrix, q=p+2
-BakeOff_Kernel(3, Diffusion, 1, false)
+// BP2: vector PCG with mass matrix, q=p+2
+REGISTER(BP, 2, 3, false);
 
-/// BK4: vector E-vector-to-E-vector evaluation of stiffness matrix, q=p+2
-BakeOff_Kernel(4, VectorDiffusion, 3, false)
+// BP3: scalar PCG with stiffness matrix, q=p+2
+REGISTER(BP, 3, 1, false);
 
-/// BK5: scalar E-vector-to-E-vector evaluation of stiffness matrix, q=p+1
-BakeOff_Kernel(5, Diffusion, 1, true)
+// BP4: vector PCG with stiffness matrix, q=p+2
+REGISTER(BP, 4, 3, false);
 
-/// BK6: vector E-vector-to-E-vector evaluation of stiffness matrix, q=p+1
-BakeOff_Kernel(6, VectorDiffusion, 3, true)
+// BP5: scalar PCG with stiffness matrix, q=p+1
+REGISTER(BP, 5, 1, true);
+
+// BP6: vector PCG with stiffness matrix, q=p+1
+REGISTER(BP, 6, 3, true);
+
+// BK1: scalar E-vector-to-E-vector evaluation of mass matrix, q=p+2
+REGISTER(BK, 1, 1, false);
+
+// BK2: vector E-vector-to-E-vector evaluation of mass matrix, q=p+2
+REGISTER(BK, 2, 3, false);
+
+// BK3: scalar E-vector-to-E-vector evaluation of stiffness matrix, q=p+2
+REGISTER(BK, 3, 1, false);
+
+// BK4: vector E-vector-to-E-vector evaluation of stiffness matrix, q=p+2
+REGISTER(BK, 4, 3, false);
+
+// BK5: scalar E-vector-to-E-vector evaluation of stiffness matrix, q=p+1
+REGISTER(BK, 5, 1, true);
+
+// BK6: vector E-vector-to-E-vector evaluation of stiffness matrix, q=p+1
+REGISTER(BK, 6, 3, true);
 
 /**
- * @brief main entry point
- * --benchmark_filter=BK1/6
- * --benchmark_context=device=cpu
+ * @brief CEED Bake-off Problems main entry point
+ * Command line options:
+ *    --benchmark_context=device=gpu
+ *    --benchmark_filter=BP1
+ *    --benchmark_out_format=csv
+ *    --benchmark_out=bp1.csv
  */
 int main(int argc, char *argv[])
 {
    bm::ConsoleReporter CR;
    bm::Initialize(&argc, argv);
 
+   AddKernelSpecializations();
+
    // Device setup, cpu by default
    std::string device_config = "cpu";
    auto global_context = bmi::GetGlobalContext();
@@ -240,12 +351,16 @@ int main(int argc, char *argv[])
          device_config = device->second;
       }
    }
+
    Device device(device_config.c_str());
    device.Print();
 
-   if (bm::ReportUnrecognizedArguments(argc, argv)) { return 1; }
+   if (bm::ReportUnrecognizedArguments(argc, argv)) { return EXIT_FAILURE; }
+
    bm::RunSpecifiedBenchmarks(&CR);
-   return 0;
+   bm::Shutdown();
+
+   return EXIT_SUCCESS;
 }
 
 #endif // MFEM_USE_BENCHMARK

tests/unit/CMakeLists.txt

@@ -72,7 +72,6 @@ set(UNIT_TESTS_SRCS
   linalg/test_operator.cpp
   linalg/test_particlevector.cpp
   linalg/test_sparsesmoothers.cpp
-  linalg/test_univarsolver.cpp
   linalg/test_vector.cpp
   mesh/mesh_test_utils.cpp
   mesh/test_exodus_reader.cpp

tests/unit/linalg/test_univarsolver.cpp

@@ -1,358 +0,0 @@
-// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
-// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
-// LICENSE and NOTICE for details. LLNL-CODE-806117.
-//
-// This file is part of the MFEM library. For more information and source code
-// availability visit https://mfem.org.
-//
-// MFEM is free software; you can redistribute it and/or modify it under the
-// terms of the BSD-3 license. We welcome feedback and contributions, see file
-// CONTRIBUTING.md for details.
-
-#include <algorithm>
-#include <cmath>
-
-#include "mfem.hpp"
-#include "unit_tests.hpp"
-
-using mfem::real_t;
-using namespace mfem::future;
-
-MFEM_HOST_DEVICE inline real_t FlowResistance(real_t eqps, real_t sigma_y, real_t n, real_t ep_0)
-{
-    return sigma_y*(1.0 + std::pow((eqps)/ep_0, n));
-}
-
-using J2PlasticityParameters = tuple<real_t, real_t, real_t, real_t, real_t, real_t>;
-
-// Residual function that is solved in the plasticity model.
-// Made a free function to facilitate Enzyme differentiation.
-real_t J2PlasticityResidual(real_t delta_eqps, J2PlasticityParameters p)
-{
-    auto [eqps, q, G, sigma_y, n, ep_0] = p;
-    return q - 3.0*G*delta_eqps - FlowResistance(eqps + delta_eqps, sigma_y, n, ep_0);
-}
-
-struct J2Plasticity {
-  static constexpr int dim = 3;  ///< spatial dimension
-  static constexpr int N_INTERNAL_STATES = 10;
-  static constexpr real_t tol = 1e-10;  ///< relative tolerance on residual mag to judge convergence of return map
-
-  real_t E;        ///< Young's modulus
-  real_t nu;       ///< Poisson's ratio
-  real_t sigma_y;  ///< Yield strength
-  real_t n;        ///< Hardening index
-  real_t ep_0;     ///< Reference plastic strain
-  
-  /// @brief variables required to characterize the hysteresis response
-  struct InternalState {
-    tensor<real_t, dim, dim> plastic_strain;
-    real_t accumulated_plastic_strain;
-  };
-
-  /// Internal state variables in a flattened array for storing in a global field
-  using PackedInternalState = mfem::future::tensor<real_t, N_INTERNAL_STATES>;
-
-  MFEM_HOST_DEVICE static inline InternalState unpack_internal_state(
-      const mfem::future::tensor<real_t, N_INTERNAL_STATES>& packed_state)
-  {
-    auto plastic_strain =
-        mfem::future::make_tensor<dim, dim>([&packed_state](int i, int j) { return packed_state[dim * i + j]; });
-    real_t accumulated_plastic_strain = packed_state[N_INTERNAL_STATES - 1];
-    return {plastic_strain, accumulated_plastic_strain};
-  }
-
-  MFEM_HOST_DEVICE static inline PackedInternalState pack_internal_state(
-      const mfem::future::tensor<real_t, dim, dim>& plastic_strain, real_t accumulated_plastic_strain)
-  {
-    PackedInternalState packed_state{};
-    for (int i = 0, ij = 0; i < dim; i++) {
-      for (int j = 0; j < dim; j++, ij++) {
-        packed_state[ij] = plastic_strain[i][j];
-      }
-    }
-    packed_state[N_INTERNAL_STATES - 1] = accumulated_plastic_strain;
-    return packed_state;
-  }
-
-  // Compute the new stress and the internal state variables
-  MFEM_HOST_DEVICE inline tuple<tensor<real_t, dim, dim>, PackedInternalState>
-  update(tensor<real_t, dim, dim> dudxi,
-         PackedInternalState internal_state,
-         tensor<real_t, dim, dim> J,
-         real_t w) const
-  {
-    auto invJ = inv(J);
-    const auto dudX = dudxi * invJ;
-    auto I = IdentityMatrix<dim>();
-    const real_t K = E / (3.0 * (1.0 - 2.0 * nu));
-    const real_t G = 0.5 * E / (1.0 + nu);
-
-    auto [plastic_strain, accumulated_plastic_strain] = unpack_internal_state(internal_state);
-    
-    auto el_strain = sym(dudX) - plastic_strain;
-    auto p = K * tr(el_strain);
-    auto s = 2.0 * G * dev(el_strain);
-    auto q = std::sqrt(1.5) * norm(s);
-    real_t denom = q > 0.0? q : 1.0;
-    auto Np = 1.5 * s / denom;
-
-    if (q > FlowResistance(accumulated_plastic_strain, sigma_y, n, ep_0)) {
-        real_t lb = 0.0;
-        real_t ub = (q - FlowResistance(accumulated_plastic_strain, sigma_y, n, ep_0))/(3*G);
-        SolverSettings settings{.residual_abs_tol = 1e-10*sigma_y, .residual_rel_tol = 1e-10, 
-                                .bounds{.lower = lb, .upper = ub}};
-        real_t delta_eqps = SolveNewtonBisection<J2PlasticityResidual>(
-            0.5*(lb + ub), make_tuple(accumulated_plastic_strain, q, G, sigma_y, n, ep_0), settings);
-        accumulated_plastic_strain += delta_eqps;
-        plastic_strain += delta_eqps * Np;
-        s -= 2.0 * G * delta_eqps * Np;
-    }
-    auto Q_new = pack_internal_state(plastic_strain, accumulated_plastic_strain);
-    auto stress = s + p * I;
-    const real_t dV = det(J)*w;
-    // Question: if I use make_tuple as in this comment, I get a segfault in
-    // derivatives of this function. Is this expected?
-    // return make_tuple(stress*transpose(invJ)*dV, Q_new);
-    return {stress*transpose(invJ)*dV, Q_new};
-  }
-  
-  MFEM_HOST_DEVICE inline tensor<real_t, dim, dim>
-  stress(tensor<real_t, dim, dim> dudxi,
-         PackedInternalState internal_state,
-         tensor<real_t, dim, dim> J,
-         real_t w) const
-  {
-    auto [stress, internal_state_new] = update(dudxi, internal_state, J, w);
-    return stress;
-  }
-
-  MFEM_HOST_DEVICE inline PackedInternalState
-  internal_state_new(tensor<real_t, dim, dim> dudxi,
-                     PackedInternalState internal_state,
-                     tensor<real_t, dim, dim> J,
-                     real_t w) const
-  {
-    auto [stress, internal_state_new] = update(dudxi, internal_state, J, w);
-    return internal_state_new;
-  }
-};
-
-
-// Register the custom derivatives for the solver.
-// This needs to be done for every residual function that the solver is applied on,
-// since the SolveNewtonBisection_impl is a function template, and we need a real
-// function with an address to specify the custom derivative.
-
-// Forward mode
-__attribute__((used))
-void *  __enzyme_register_derivative_newton_bisection_on_j2[2] = {
-  (void*) mfem::internal::SolveNewtonBisection_impl<J2PlasticityResidual, J2PlasticityParameters>,
-  (void*) mfem::internal::SolveNewtonBisection_impl_fwddiff<J2PlasticityResidual, J2PlasticityParameters>
-};
-
- // Reverse mode
-__attribute__((used))
-void* __enzyme_register_gradient_SolveNewtonBisectionJ2[3] = {
-  (void*)mfem::internal::SolveNewtonBisection_impl<J2PlasticityResidual, J2PlasticityParameters>,
-  (void*)mfem::internal::SolveNewtonBisection_impl_aug<J2PlasticityResidual, J2PlasticityParameters>,
-  (void*)mfem::internal::SolveNewtonBisection_impl_rev<J2PlasticityResidual, J2PlasticityParameters>
-};
-
-// create a free function for Enzyme to differentiate in the tests
-tensor<real_t, 3, 3> ComputeStress(
-    J2Plasticity* material, tensor<real_t, 3, 3> dudxi, 
-    J2Plasticity::PackedInternalState Q, tensor<real_t, 3, 3> J, real_t w)
-{
-    return material->stress(dudxi, Q, J, w);
-}
-
-template <int dim>
-real_t elementwise_norm(tensor<real_t, dim, dim> A) {
-  real_t maxval = 0;
-  for (int i = 0; i < dim; i++) {
-    for (int j = 0; j < dim; j++) {
-      maxval = std::max(std::abs(A[i][j]), maxval);
-    }
-  }
-  return maxval;
-}
-
-void ComputeStressRef(const J2Plasticity* material, const tensor<real_t, 3, 3>& dudxi,
-                      const J2Plasticity::PackedInternalState& Q,
-                      const tensor<real_t, 3, 3>& J, real_t w,
-                      tensor<real_t, 3, 3>& sigma)
-{
-  sigma = material->stress(dudxi, Q, J, w);
-}
-
-
-
-TEST_CASE("Univariate function solver in a qfunction", "[univar]")
-{
-    J2Plasticity material{.E = 70.0e3, .nu = 0.34, .sigma_y = 240.0, .n = 0.15, .ep_0 = 1e-3};
-    tensor<real_t, 3, 3> H{{{0.947667  , 0.9785799 , 0.33229148},
-                            {0.46866846, 0.5698887 , 0.16550303},
-                            {0.3101946 , 0.68948054, 0.74676657}}};
-    J2Plasticity::PackedInternalState Q{};
-    const tensor<real_t, 3, 3> J = IdentityMatrix<3>();
-    const real_t w = 1.0;
-
-    SECTION("Correctness")
-    {
-      // Checks that stress after update is on the yield surface.
-      auto [stress, Q_new] = material.update(H, Q, IdentityMatrix<3>(), 1.0);
-      real_t mises = std::sqrt(1.5)*norm(dev(stress));
-      real_t eqps = Q_new[9];
-      // This test only makes sense if the displacement gradient is big enough to 
-      // cuase yielding.
-      REQUIRE(eqps > 1e-9);
-      real_t Y = FlowResistance(eqps, material.sigma_y, material.n, material.ep_0);
-      CHECK(mises == MFEM_Approx(Y, 0.0, 1e-8));
-    }
-
-    SECTION("JVP")
-    {
-      // Compare forward mode derivative to finite difference approximation
-
-      tensor<real_t, 3, 3> H_dot{{{1.0, 0.0 , 0.0},
-                                  {0.0, 0.0 , 0.0},
-                                  {0.0, 0.0 , 0.0}}};
-
-      // Enzyme directional derivative (uses custom derivative of solver)
-      auto sigma_dot = __enzyme_fwddiff<tensor<real_t, 3, 3>>((void*)ComputeStress, 
-                                                              enzyme_const, &material,
-                                                              enzyme_dup, H, H_dot,
-                                                              enzyme_const, Q,
-                                                              enzyme_const, J,
-                                                              enzyme_const, w);
-      // sigma_dot = ∂sigma / ∂H[0, 0]
-      REQUIRE(sigma_dot[0][0] > 0.0);
-
-      // Finite difference derivative approximation
-      constexpr int dim = 3;
-      real_t eps = 1e-5;
-      tensor<real_t, 3, 3> sigma = ComputeStress(&material, H, Q, J, w);
-      tensor<real_t, 3, 3> sigma_p = ComputeStress(&material, H + eps*H_dot, Q, J, w);
-      tensor<real_t, 3, 3> sigma_dot_h = (1.0/eps)*(sigma_p - sigma);
-
-      tensor<real_t, 3, 3> rel_error = sigma_dot - sigma_dot_h;
-      for (int i = 0; i < dim; i++) {
-        for (int j = 0; j < dim; j++) {
-          real_t denom = sigma[i][j] != 0? sigma[i][j] : 1.0;
-          rel_error[i][j] /= denom;
-        }
-      }
-
-      CHECK(elementwise_norm(rel_error) < 1e-5);
-    }
-
-    SECTION("VJP")
-    {
-      // compare reverse mode derivative to finite differences
-
-      tensor<real_t, 3, 3> sigma;
-      ComputeStressRef(&material, H, Q, J, w, sigma);
-      double epsilon = 1e-6;
-      tensor<real_t, 3, 3> dH{{{1.0, 0.0, 0.0}, {0.0, 0.0, 0.0}, {0.0, 0.0, 0.0}}};
-      auto H_p = H + epsilon*dH;
-      tensor<real_t, 3, 3> sigma_p;
-      ComputeStressRef(&material, H_p, Q, J, w, sigma_p);
-      auto sigma_dot_h = (sigma_p - sigma)/epsilon;
-      // Note: sigma_dot_h[i,j] = ∂sigma[i,j]/∂H[0,0]
-
-      tensor<real_t, 3, 3> sigma_bar{{{1.0, 0.0, 0.0},
-                                      {0.0, 0.0, 0.0},
-                                      {0.0, 0.0, 0.0}}};
-
-      J2Plasticity material_bar;
-      tensor<real_t, 3, 3> H_bar{};
-      J2Plasticity::PackedInternalState Q_bar{};
-      tensor<real_t, 3, 3> J_bar{};
-      __enzyme_autodiff<void>(
-        (void*)ComputeStressRef,  enzyme_const, &material, enzyme_dup, &H, &H_bar,
-        enzyme_dup, &Q, &Q_bar, enzyme_dup, &J, &J_bar, enzyme_const, w,
-        enzyme_dup, &sigma, &sigma_bar);
-
-      // H_bar[ij] = ∂sigma[0,0]/∂H[i,j]
-      // For this model, we expect the major symmetries in the tangent operator.
-      // Hence H_bar \approx sigma_dot_h
-
-      const double abs_tol = 1e-12;
-      const double rel_tol = 5e-6;
-      for (int i = 0; i < 3; i++) {
-        for (int j = 0; j < 3; j++) {
-          CHECK(H_bar[i][j] == MFEM_Approx(sigma_dot_h[i][j], abs_tol, rel_tol));
-        }
-      }
-    }
-}
-
-
-real_t nthroot_res(real_t x, tuple<real_t, real_t> p)
-{
-    auto [index, radicand] = p;
-    return std::pow(x, index) - radicand;
-}
-
-__attribute__((used))
-void* __enzyme_register_gradient_solver[3] = {
-  (void*)mfem::internal::SolveNewtonBisection_impl<nthroot_res, tuple<real_t, real_t>>,
-  (void*)mfem::internal::SolveNewtonBisection_impl_aug<nthroot_res, tuple<real_t, real_t>>,
-  (void*)mfem::internal::SolveNewtonBisection_impl_rev<nthroot_res, tuple<real_t, real_t>>
-};
-
-real_t mysqrt(real_t x)
-{
-  real_t x0 = x;
-  real_t index = 2.0;
-  real_t ub = std::max(1.0, x);
-  SolverSettings settings{.bounds = {.lower = 0, .upper = ub}};
-  return SolveNewtonBisection<nthroot_res>(x0, make_tuple(index, x), settings);
-}
-
-TEST_CASE("Univariate solver reverse mode", "[univar]")
-{
-  real_t x = 2.0;
-  real_t y = mysqrt(x);
-  std::cout << "x = " << x << " sqrt(x) = " << y << std::endl;
-  CHECK(y == MFEM_Approx(M_SQRT2, 0.0, 1e-8));
-
-  std::cout << "Computing derivative" << std::endl;
-  real_t dydx = __enzyme_autodiff<real_t>((void*)mysqrt, enzyme_out, x);
-  CHECK(dydx == MFEM_Approx(0.5/std::sqrt(2.0)));
-}
-
-TEST_CASE("Univariate function solver robustness", "[univar]")
-{
-    SECTION("Simple case")
-    {
-        auto Nthroot = [](real_t x, real_t n) {
-            real_t x0 = std::max(x, 1.0);
-            SolverSettings settings{.bounds{0.0, x0}};
-            return SolveNewtonBisection<nthroot_res>(x0, make_tuple(n, x), settings);
-        };
-        real_t x = 8.0;
-        real_t y = Nthroot(x, 3.0);
-        REQUIRE(y == MFEM_Approx(2.0));
-    }
-
-    SECTION("Stiff problem")
-    {
-        auto f = [](real_t x, real_t p) { return std::pow(x, p) - 1.0; };
-        real_t x0 = 0.1;
-        real_t p = 50;
-        SolverSettings settings{.bounds{.lower = 0.0, .upper = 5.1}};
-        real_t x = SolveNewtonBisection<+f>(x0, p, settings);
-        REQUIRE(x == MFEM_Approx(1.0));
-    }
-
-    SECTION("Works where Newton diverges")
-    {
-        auto f = [](double x, int) { return std::atan(x); };
-        real_t x0 = 1.5;
-        SolverSettings settings{.bounds{.lower = 0.0, .upper = 2.0}};
-        real_t x = SolveNewtonBisection<+f>(x0, int{}, settings);
-        CHECK(std::abs(x) == MFEM_Approx(0.0));
-    }
-}