Comment and move things around for clarity

Derivatives become nans if solver upper bound depends on state.
Need to debug this with Bill. But derivatives are now being computed,
and appear correct.

fem/CMakeLists.txt

@@ -213,6 +213,7 @@ 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

@@ -0,0 +1,232 @@
+// 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<NB*NB>(NE, NB, NB, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall_2D(NE, NB, NB, [=] MFEM_HOST_DEVICE (int e)
    {
       // Perform change of basis if needed
       if (CHANGE_BASIS)

fem/fem.hpp

@@ -53,6 +53,7 @@
 #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,16 +1004,13 @@ 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, "");
-   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 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);
    auto y = Reshape(y_.ReadWrite(), D1D, D1D, D1D, NE);
    MFEM_VERIFY(D1D <= Q1D, "THREAD_DIRECT requires D1D <= Q1D");
-
-   mfem::forall_3D<T_Q1D*T_Q1D*T_Q1D>(NE,
-                                      Q1D, Q1D, Q1D,
-                                      [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall_3D(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, "");
-   const auto b = b_.Read();
-   const auto d = d_.Read();
-   const auto x = x_.Read();
+   auto b = b_.Read();
+   auto d = d_.Read();
+   auto x = x_.Read();
    auto y = y_.ReadWrite();
-   mfem::forall_2D<T_Q1D*T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE (int e)
+   mfem::forall_2D(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, "");
-   const auto B = Reshape(basis.Read(), Q1D, D1D);
-   const auto D = Reshape(padata.Read(), Q1D, NE);
+   auto B = Reshape(basis.Read(), Q1D, D1D);
+   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<T_Q1D*T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
+   mfem::forall_2D(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<T_Q1D*T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
+   mfem::forall_2D(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<T_Q1D*T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
+   mfem::forall_2D(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<T_Q1D*T_Q1D>(NE, Q1D, Q1D, [=] MFEM_HOST_DEVICE(int e)
+   mfem::forall_2D(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,7 +44,6 @@
 #endif
 
 #if !defined(MFEM_USE_CUDA_OR_HIP)
-constexpr bool mfem_use_gpu = false;
 #define MFEM_DEVICE
 #define MFEM_HOST
 #define MFEM_LAMBDA
@@ -53,7 +52,6 @@ constexpr bool mfem_use_gpu = false;
 #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,11 +20,9 @@
 
 #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,12 +295,11 @@ using hip_threads_z =
 #endif
 
 #if defined(MFEM_USE_RAJA) && defined(RAJA_ENABLE_CUDA) && defined(__CUDACC__)
-template <typename DBODY>
+template <const int BLOCKS = MFEM_CUDA_BLOCKS, typename DBODY>
 void RajaCuWrap1D(const int N, DBODY &&d_body)
 {
    //true denotes asynchronous kernel
-   RAJA::forall<RAJA::cuda_exec<MFEM_CUDA_BLOCKS,true>>(RAJA::RangeSegment(0,N),
-                                                        d_body);
+   RAJA::forall<RAJA::cuda_exec<BLOCKS,true>>(RAJA::RangeSegment(0,N),d_body);
 }
 
 template <typename DBODY>
@@ -363,18 +362,18 @@ struct RajaCuWrap;
 template <>
 struct RajaCuWrap<1>
 {
-   template <typename DBODY>
+   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
-      RajaCuWrap1D(N, d_body);
+      RajaCuWrap1D<BLCK>(N, d_body);
    }
 };
 
 template <>
 struct RajaCuWrap<2>
 {
-   template <typename DBODY>
+   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -385,7 +384,7 @@ struct RajaCuWrap<2>
 template <>
 struct RajaCuWrap<3>
 {
-   template <typename DBODY>
+   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -396,12 +395,11 @@ struct RajaCuWrap<3>
 #endif
 
 #if defined(MFEM_USE_RAJA) && defined(RAJA_ENABLE_HIP) && defined(__HIP__)
-template <typename DBODY>
+template <const int BLOCKS = MFEM_HIP_BLOCKS, typename DBODY>
 void RajaHipWrap1D(const int N, DBODY &&d_body)
 {
    //true denotes asynchronous kernel
-   RAJA::forall<RAJA::hip_exec<MFEM_HIP_BLOCKS,true>>(RAJA::RangeSegment(0,N),
-                                                      d_body);
+   RAJA::forall<RAJA::hip_exec<BLOCKS,true>>(RAJA::RangeSegment(0,N),d_body);
 }
 
 template <typename DBODY>
@@ -464,18 +462,18 @@ struct RajaHipWrap;
 template <>
 struct RajaHipWrap<1>
 {
-   template <typename DBODY>
+   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
-      RajaHipWrap1D(N, d_body);
+      RajaHipWrap1D<BLCK>(N, d_body);
    }
 };
 
 template <>
 struct RajaHipWrap<2>
 {
-   template <typename DBODY>
+   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -486,7 +484,7 @@ struct RajaHipWrap<2>
 template <>
 struct RajaHipWrap<3>
 {
-   template <typename DBODY>
+   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -586,31 +584,12 @@ 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)
 {
@@ -625,8 +604,6 @@ 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);
@@ -634,19 +611,6 @@ 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)
@@ -658,35 +622,24 @@ void CuWrap3D(const int N, DBODY &&d_body,
    MFEM_GPU_CHECK(cudaGetLastError());
 }
 
-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)
-{
-   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>
+struct CuWrap;
 
-template <int Dim, int MAX_THREADS_PER_BLOCK> struct CuWrap;
-
-template <int MAX_THREADS_PER_BLOCK>
-struct CuWrap<1, MAX_THREADS_PER_BLOCK>
+template <>
+struct CuWrap<1>
 {
-   template <typename DBODY>
+   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
-      CuWrap1D<MFEM_CUDA_BLOCKS>(N, d_body);
+      CuWrap1D<BLCK>(N, d_body);
    }
 };
 
 template <>
-struct CuWrap<2, 0>
+struct CuWrap<2>
 {
-   template <typename DBODY>
+   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -694,22 +647,10 @@ struct CuWrap<2, 0>
    }
 };
 
-template <int MAX_THREADS_PER_BLOCK>
-struct CuWrap<2, 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)
-   {
-      static_assert(MAX_THREADS_PER_BLOCK > 0);
-      CuWrap2DLaunchBounds<MAX_THREADS_PER_BLOCK>(N, d_body, X, Y, Z);
-   }
-};
-
 template <>
-struct CuWrap<3, 0>
+struct CuWrap<3>
 {
-   template <typename DBODY>
+   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -717,17 +658,6 @@ struct CuWrap<3, 0>
    }
 };
 
-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__)
 
 
@@ -750,31 +680,13 @@ 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 <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>
+template <const int BLCK = MFEM_HIP_BLOCKS, typename DBODY>
 void HipWrap1D(const int N, DBODY &&d_body)
 {
    if (N==0) { return; }
@@ -788,27 +700,12 @@ 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)
@@ -820,36 +717,24 @@ void HipWrap3D(const int N, DBODY &&d_body,
    MFEM_GPU_CHECK(hipGetLastError());
 }
 
-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)
-{
-   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>
+struct HipWrap;
 
-template <int Dim, int MAX_THREADS_PER_BLOCK> struct HipWrap;
-
-template <int MAX_THREADS_PER_BLOCK>
-struct HipWrap<1, MAX_THREADS_PER_BLOCK>
+template <>
+struct HipWrap<1>
 {
-   template <typename DBODY>
+   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
-      HipWrap1D<MFEM_HIP_BLOCKS>(N, d_body);
+      HipWrap1D<BLCK>(N, d_body);
    }
 };
 
 template <>
-struct HipWrap<2, 0>
+struct HipWrap<2>
 {
-   template <typename DBODY>
+   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -857,21 +742,10 @@ struct HipWrap<2, 0>
    }
 };
 
-template <int MAX_THREADS_PER_BLOCK>
-struct HipWrap<2, 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)
-   {
-      HipWrap2DLaunchBounds<MAX_THREADS_PER_BLOCK>(N, d_body, X, Y, Z);
-   }
-};
-
 template <>
-struct HipWrap<3, 0>
+struct HipWrap<3>
 {
-   template <typename DBODY>
+   template <const int BLCK = MFEM_CUDA_BLOCKS, typename DBODY>
    static void run(const int N, DBODY &&d_body,
                    const int X, const int Y, const int Z, const int G)
    {
@@ -879,24 +753,11 @@ struct HipWrap<3, 0>
    }
 };
 
-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__)
 
 
-///////////////////////////////////////////////////////////////////////////////
-/// Forall host & device kernel dispatch
-template <int DIM, int MAX_THREADS_PER_BLOCK = 0,
-          typename d_lambda, typename h_lambda>
+/// The forall kernel body wrapper
+template <const int DIM, 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,
@@ -929,7 +790,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, MAX_THREADS_PER_BLOCK>::run(N, d_body, X, Y, Z, G);
+      return CuWrap<DIM>::run(N, d_body, X, Y, Z, G);
    }
 #endif
 
@@ -937,7 +798,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, MAX_THREADS_PER_BLOCK>::run(N, d_body, X, Y, Z, G);
+      return HipWrap<DIM>::run(N, d_body, X, Y, Z, G);
    }
 #endif
 
@@ -966,9 +827,7 @@ backend_cpu:
    for (int k = 0; k < N; k++) { h_body(k); }
 }
 
-///////////////////////////////////////////////////////////////////////////////
-/// Forall host & device kernel wrappers
-template <int DIM, typename lambda>
+template <const 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)
@@ -976,16 +835,6 @@ 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); }
 
@@ -994,7 +843,7 @@ inline void forall(int Nx, int Ny, lambda &&body)
 {
    if (Device::Allows(Backend::DEVICE_MASK))
    {
-      mfem::forall(Nx * Ny, [=] MFEM_HOST_DEVICE(int idx)
+      forall(Nx * Ny, [=] MFEM_HOST_DEVICE(int idx)
       {
          int j = idx / Nx;
          int i = idx % Nx;
@@ -1030,7 +879,7 @@ inline void forall(int Nx, int Ny, int Nz, lambda &&body)
 {
    if (Device::Allows(Backend::DEVICE_MASK))
    {
-      mfem::forall(Nx * Ny * Nz, [=] MFEM_HOST_DEVICE(int idx)
+      forall(Nx * Ny * Nz, [=] MFEM_HOST_DEVICE(int idx)
       {
          int i = idx % Nx;
          int j = idx / Nx;
@@ -1078,12 +927,6 @@ 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)
 {
@@ -1096,12 +939,6 @@ 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,11 +20,9 @@
 
 #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,7 +55,6 @@ list(APPEND HDRS
   dinvariants.hpp
   dtensor.hpp
   dual.hpp
-  eigensolver.hpp
   filteredsolver.hpp
   handle.hpp
   invariants.hpp

linalg/eigensolver.hpp

@@ -1,203 +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 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,19 +363,14 @@ 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.IterRefine = SLU_DOUBLE;
    //    options.Trans = NOTRANS;
+   //    options.IterRefine = SLU_DOUBLE;
    //    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 || \
@@ -477,12 +472,6 @@ 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,8 +250,7 @@ public:
        work (default false) */
    void SetSymmetricPattern(bool sym);
 
-   /** @brief Specify whether to perform parallel symbolic factorization
-       (default false)
+   /** @brief Specify whether to perform parallel symbolic factorization.
        @note If true SuperLU will use superlu::PARMETIS for the Column
        Permutation regardless of the setting */
    void SetParSymbFact(bool par);
@@ -264,10 +263,6 @@ 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,6 +31,11 @@ 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,89 +8,23 @@
 // 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" // IWYU pragma: keep
+#include "bench.hpp"
 
 #ifdef MFEM_USE_BENCHMARK
 
-#include <cassert>
-#include <string>
+/*
+  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 "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>
+  See: ceed.exascaleproject.org/bps and github.com/CEED/benchmarks
+*/
+template <int VDIM, bool GLL>
 struct BakeOff
 {
-   inline static constexpr int DIM = 3;
-   const int p, c, q, n, nx, ny, nz;
+   static constexpr int DIM = 3;
+   const int N, p, q;
    Mesh mesh;
    H1_FECollection fec;
    FiniteElementSpace fes;
@@ -104,15 +38,12 @@ struct BakeOff
    GridFunction x, y;
    BilinearForm a;
    double mdofs{};
-   BilinearFormIntegrator *bfi;
 
-   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)),
+   BakeOff(int p):
+      N(Device::IsEnabled() ? 32 : 4),
+      p(p),
+      q(2 * p + (GLL ? -1 : 3)),
+      mesh(Mesh::MakeCartesian3D(N, N, N, Element::HEXAHEDRON)),
       fec(p, DIM, BasisType::GaussLobatto),
       fes(&mesh, &fec, VDIM, VDIM == 3 ? Ordering::byVDIM : Ordering::byNODES),
       geom_type(mesh.GetTypicalElementGeometry()),
@@ -127,41 +58,22 @@ 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;
 
-   [[nodiscard]] double SumMdofs() const noexcept { return mdofs; }
+   double SumMdofs() const { return mdofs; }
 
-   [[nodiscard]] double MDofs() const noexcept { return 1e-6 * dofs; }
+   double MDofs() const { return 1e-6 * dofs; }
 };
 
-// Bake-off Problems (BPs)
-template <int BFI, int VDIM, bool GLL>
-struct BP : public BakeOff<BFI, VDIM, GLL>
+/// Bake-off Problems (BPs)
+template <typename BFI, int VDIM, bool GLL>
+struct Problem : public BakeOff<VDIM, GLL>
 {
-   const int max_it = 32, print_lvl = -1;
+   const double rtol = 1e-12;
+   const int max_it = 32;
+   const int print_lvl = -1;
 
    Array<int> ess_tdof_list;
    Array<int> ess_bdr;
@@ -170,56 +82,44 @@ struct BP : public BakeOff<BFI, VDIM, GLL>
    Vector B, X;
    CGSolver cg;
 
-   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;
+   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;
 
-   BP(int p, int side) noexcept: base(p, side),
+   Problem(int order):
+      BakeOff<VDIM, GLL>(order),
       ess_bdr(mesh.bdr_attributes.Max()),
       b(&fes)
    {
       ess_bdr = 1;
       fes.GetEssentialTrueDofs(ess_bdr, ess_tdof_list);
-
-      if constexpr (VDIM == 1)
+      if (VDIM == 1)
       {
-         b.AddDomainIntegrator(new DomainLFIntegrator(one));
+         b.AddDomainIntegrator(new DomainLFIntegrator(this->one));
       }
       else
       {
-         b.AddDomainIntegrator(new VectorDomainLFIntegrator(unit_vec));
+         b.AddDomainIntegrator(new VectorDomainLFIntegrator(this->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);
-
-      benchmark();
-      mdofs = 0.0;
+      cg.iterative_mode = false;
+      MFEM_DEVICE_SYNC;
    }
 
    void benchmark() override
@@ -230,115 +130,104 @@ struct BP : public BakeOff<BFI, VDIM, GLL>
    }
 };
 
-// Bake-off Kernels (BKs)
-template <int BFI, int VDIM, bool GLL>
-struct BK : public BakeOff<BFI, 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>
 {
-   Vector xe, ye;
+   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;
 
-   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)
+   Kernel(int order): BakeOff<VDIM, GLL>(order)
    {
-      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;
+      x.Randomize(1);
+      a.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+      a.AddDomainIntegrator(new BFI(one, ir));
+      a.Assemble();
+      a.Mult(x, y);
+      MFEM_DEVICE_SYNC;
    }
 
    void benchmark() override
    {
-      bfi->AddMultPA(xe, ye);
+      a.Mult(x, y);
       MFEM_DEVICE_SYNC;
       mdofs += this->MDofs();
    }
 };
 
-// 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));
-}
+/// 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);
 
-#define REGISTER(PK, BFI, VDIM, GLL) \
-   BENCHMARK_TEMPLATE(Benchmark, PK<BFI, VDIM, GLL>) \
-   ->Name(#PK #BFI)->Apply(CustomArguments)->Unit(bm::kMillisecond)
+/// BK1: scalar E-vector-to-E-vector evaluation of mass matrix, q=p+2
+BakeOff_Kernel(1, Mass, 1, false)
 
-// BP1: scalar PCG with mass matrix, q=p+2
-REGISTER(BP, 1, 1, false);
+/// BK2: vector E-vector-to-E-vector evaluation of mass matrix, q=p+2
+BakeOff_Kernel(2, VectorMass, 3, false)
 
-// BP2: vector PCG with mass matrix, q=p+2
-REGISTER(BP, 2, 3, false);
+/// BK3: scalar E-vector-to-E-vector evaluation of stiffness matrix, q=p+2
+BakeOff_Kernel(3, Diffusion, 1, false)
 
-// BP3: scalar PCG with stiffness matrix, q=p+2
-REGISTER(BP, 3, 1, false);
+/// BK4: vector E-vector-to-E-vector evaluation of stiffness matrix, q=p+2
+BakeOff_Kernel(4, VectorDiffusion, 3, false)
 
-// BP4: vector PCG with stiffness matrix, q=p+2
-REGISTER(BP, 4, 3, false);
+/// BK5: scalar E-vector-to-E-vector evaluation of stiffness matrix, q=p+1
+BakeOff_Kernel(5, Diffusion, 1, 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);
+/// BK6: vector E-vector-to-E-vector evaluation of stiffness matrix, q=p+1
+BakeOff_Kernel(6, VectorDiffusion, 3, true)
 
 /**
- * @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
+ * @brief main entry point
+ * --benchmark_filter=BK1/6
+ * --benchmark_context=device=cpu
  */
 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();
@@ -351,16 +240,12 @@ int main(int argc, char *argv[])
          device_config = device->second;
       }
    }
-
    Device device(device_config.c_str());
    device.Print();
 
-   if (bm::ReportUnrecognizedArguments(argc, argv)) { return EXIT_FAILURE; }
-
+   if (bm::ReportUnrecognizedArguments(argc, argv)) { return 1; }
    bm::RunSpecifiedBenchmarks(&CR);
-   bm::Shutdown();
-
-   return EXIT_SUCCESS;
+   return 0;
 }
 
 #endif // MFEM_USE_BENCHMARK

tests/unit/CMakeLists.txt

@@ -72,6 +72,7 @@ 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

@@ -0,0 +1,358 @@
+// 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));
+    }
+}