fix missing std::array and std::ignore

examples/CMakeLists.txt

@@ -59,6 +59,7 @@ if (MFEM_USE_MPI)
     dfem_minimal_example.cpp
     dfem_test_diffusion_2d.cpp
     dfem_test_diffusion_3d.cpp
+    dfem_test_diffusion_3d_refactor.cpp
     dfem_test_ordering.cpp
     dfem_test_vector_diffusion.cpp
     dfem_test_elasticity.cpp
@@ -138,6 +139,7 @@ target_link_libraries(test_dfem ClangEnzymeFlags)
 target_link_libraries(dfem_laghos ClangEnzymeFlags)
 target_link_libraries(dfem_minimal_example ClangEnzymeFlags)
 target_link_libraries(dfem_test_diffusion_3d ClangEnzymeFlags)
+target_link_libraries(dfem_test_diffusion_3d_refactor ClangEnzymeFlags)
 target_link_libraries(dfem_test_nonlinear_diffusion_3d ClangEnzymeFlags)
 target_link_libraries(dfem_test_nonlinear_elasticity_3d ClangEnzymeFlags)
 

examples/dfem/cxxtimer.hpp

@@ -0,0 +1,184 @@
+/*
+
+MIT License
+
+Copyright (c) 2017 André L. Maravilha
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+*/
+
+#ifndef CXX_TIMER_HPP
+#define CXX_TIMER_HPP
+
+#include <chrono>
+
+
+namespace cxxtimer {
+
+/**
+ * This class works as a stopwatch.
+ */
+class Timer {
+
+public:
+
+    /**
+     * Constructor.
+     *
+     * @param   start
+     *          If true, the timer is started just after construction.
+     *          Otherwise, it will not be automatically started.
+     */
+    Timer(bool start = false);
+
+    /**
+     * Copy constructor.
+     *
+     * @param   other
+     *          The object to be copied.
+     */
+    Timer(const Timer& other) = default;
+
+    /**
+     * Transfer constructor.
+     *
+     * @param   other
+     *          The object to be transferred.
+     */
+    Timer(Timer&& other) = default;
+
+    /**
+     * Destructor.
+     */
+    virtual ~Timer() = default;
+
+    /**
+     * Assignment operator by copy.
+     *
+     * @param   other
+     *          The object to be copied.
+     *
+     * @return  A reference to this object.
+     */
+    Timer& operator=(const Timer& other) = default;
+
+    /**
+     * Assignment operator by transfer.
+     *
+     * @param   other
+     *          The object to be transferred.
+     *
+     * @return  A reference to this object.
+     */
+    Timer& operator=(Timer&& other) = default;
+
+    /**
+     * Start/resume the timer.
+     */
+    void start();
+
+    /**
+     * Stop/pause the timer.
+     */
+    void stop();
+
+    /**
+     * Reset the timer.
+     */
+    void reset();
+
+    /**
+     * Return the elapsed time.
+     *
+     * @param   duration_t
+     *          The duration type used to return the time elapsed. If not
+     *          specified, it returns the time as represented by
+     *          std::chrono::milliseconds.
+     *
+     * @return  The elapsed time.
+     */
+    template <class duration_t = std::chrono::milliseconds>
+    typename duration_t::rep count() const;
+
+private:
+
+    bool started_;
+    bool paused_;
+    std::chrono::steady_clock::time_point reference_;
+    std::chrono::duration<long double> accumulated_;
+};
+
+}
+
+
+inline cxxtimer::Timer::Timer(bool start) :
+        started_(false), paused_(false),
+        reference_(std::chrono::steady_clock::now()),
+        accumulated_(std::chrono::duration<long double>(0)) {
+    if (start) {
+        this->start();
+    }
+}
+
+inline void cxxtimer::Timer::start() {
+    if (!started_) {
+        started_ = true;
+        paused_ = false;
+        accumulated_ = std::chrono::duration<long double>(0);
+        reference_ = std::chrono::steady_clock::now();
+    } else if (paused_) {
+        reference_ = std::chrono::steady_clock::now();
+        paused_ = false;
+    }
+}
+
+inline void cxxtimer::Timer::stop() {
+    if (started_ && !paused_) {
+        std::chrono::steady_clock::time_point now = std::chrono::steady_clock::now();
+        accumulated_ = accumulated_ + std::chrono::duration_cast< std::chrono::duration<long double> >(now - reference_);
+        paused_ = true;
+    }
+}
+
+inline void cxxtimer::Timer::reset() {
+    if (started_) {
+        started_ = false;
+        paused_ = false;
+        reference_ = std::chrono::steady_clock::now();
+        accumulated_ = std::chrono::duration<long double>(0);
+    }
+}
+
+template <class duration_t>
+typename duration_t::rep cxxtimer::Timer::count() const {
+    if (started_) {
+        if (paused_) {
+            return std::chrono::duration_cast<duration_t>(accumulated_).count();
+        } else {
+            return std::chrono::duration_cast<duration_t>(
+                    accumulated_ + (std::chrono::steady_clock::now() - reference_)).count();
+        }
+    } else {
+        return duration_t(0).count();
+    }
+}
+
+
+#endif

examples/dfem/dfem.hpp

@@ -1,3 +1,4 @@
 #pragma once
 
 #include "dfem_differentiable_operator.hpp"
+#include "dfem_element_operator.hpp"

examples/dfem/dfem_differentiable_operator.hpp

@@ -20,6 +20,7 @@
 #include "dfem_util.hpp"
 #include "dfem_interpolate.hpp"
 #include "dfem_qfunction.hpp"
+#include "dfem_qfunction_dual.hpp"
 #include "dfem_integrate.hpp"
 
 namespace mfem
@@ -32,7 +33,8 @@ template <
    size_t num_solutions,
    size_t num_parameters,
    size_t num_fields = num_solutions + num_parameters,
-   size_t num_kernels = mfem::tuple_size<kernels_tuple>::value
+   size_t num_kernels = mfem::tuple_size<kernels_tuple>::value,
+   typename autodiff_t = AutoDiff::NativeDualNumber
    >
 class DifferentiableOperator : public Operator
 {
@@ -239,11 +241,10 @@ public:
                           std::array<FieldDescriptor, num_parameters> p,
                           kernels_tuple ks,
                           ParMesh &m,
-                          const IntegrationRule &integration_rule) :
+                          autodiff_t ad = AutoDiff::NativeDualNumber{}) :
       kernels(ks),
       mesh(m),
       dim(mesh.Dimension()),
-      integration_rule(integration_rule),
       solutions(s),
       parameters(p)
    {
@@ -287,7 +288,6 @@ public:
    kernels_tuple kernels;
    ParMesh &mesh;
    const int dim;
-   const IntegrationRule &integration_rule;
 
    std::array<FieldDescriptor, num_solutions> solutions;
    std::array<FieldDescriptor, num_parameters> parameters;
@@ -323,7 +323,8 @@ template <
    size_t num_solutions,
    size_t num_parameters,
    size_t num_fields,
-   size_t num_kernels
+   size_t num_kernels,
+   typename autodiff_t
    >
 template <
    typename kernel_t
@@ -332,7 +333,8 @@ void DifferentiableOperator<kernels_tuple,
      num_solutions,
      num_parameters,
      num_fields,
-     num_kernels>::Action::create_action_callback(
+     num_kernels,
+     autodiff_t>::Action::create_action_callback(
         kernel_t kernel,
         mult_func_t &func)
 {
@@ -354,7 +356,7 @@ void DifferentiableOperator<kernels_tuple,
 
    const int num_elements = GetNumEntities<Entity::Element>(op.mesh);
    const int num_entities = GetNumEntities<entity_t>(op.mesh);
-   const int num_qp = op.integration_rule.GetNPoints();
+   const int num_qp = kernel.integration_rule.GetNPoints();
 
    // All solutions T-vector sizes make up the width of the operator, since
    // they are explicitly provided in Mult() for example.
@@ -377,7 +379,7 @@ void DifferentiableOperator<kernels_tuple,
    std::vector<const DofToQuad*> dtq;
    for (const auto &field : op.fields)
    {
-      dtq.emplace_back(GetDofToQuad<entity_t>(field, op.integration_rule,
+      dtq.emplace_back(GetDofToQuad<entity_t>(field, kernel.integration_rule,
                                               doftoquad_mode));
    }
    const int q1d = (int)floor(pow(num_qp, 1.0/op.mesh.Dimension()) + 0.5);
@@ -404,8 +406,7 @@ void DifferentiableOperator<kernels_tuple,
                             mfem::get<hardcoded_output_idx>(output_fops).vdim /
                             num_entities;
 
-   auto ir_weights = Reshape(this->op.integration_rule.GetWeights().Read(),
-                             num_qp);
+   auto ir_weights = Reshape(kernel.integration_rule.GetWeights().Read(), num_qp);
 
    auto input_size_on_qp = get_input_size_on_qp(kernel.inputs,
                                                 std::make_index_sequence<kernel.num_kinputs> {});
@@ -558,7 +559,8 @@ template <
    size_t num_solutions,
    size_t num_parameters,
    size_t num_fields,
-   size_t num_kernels
+   size_t num_kernels,
+   typename autodiff_t
    >
 template <
    size_t derivative_idx
@@ -570,8 +572,9 @@ void DifferentiableOperator<kernels_tuple,
      num_solutions,
      num_parameters,
      num_fields,
-     num_kernels>::Derivative<derivative_idx>::create_callback(kernel_t kernel,
-                                                               mult_func_t &func)
+     num_kernels,
+     autodiff_t>::Derivative<derivative_idx>::create_callback(kernel_t kernel,
+                                                              mult_func_t &func)
 {
    using entity_t = typename kernel_t::entity_t;
 
@@ -591,13 +594,13 @@ void DifferentiableOperator<kernels_tuple,
 
    const int num_elements = GetNumEntities<Entity::Element>(op.mesh);
    const int num_entities = GetNumEntities<entity_t>(op.mesh);
-   const int num_qp = op.integration_rule.GetNPoints();
+   const int num_qp = kernel.integration_rule.GetNPoints();
 
    // assume only a single element type for now
    std::vector<const DofToQuad*> dtq;
    for (const auto &field : op.fields)
    {
-      dtq.emplace_back(GetDofToQuad<entity_t>(field, op.integration_rule,
+      dtq.emplace_back(GetDofToQuad<entity_t>(field, kernel.integration_rule,
                                               doftoquad_mode));
    }
    const int q1d = dtq[0]->nqpt;
@@ -624,8 +627,7 @@ void DifferentiableOperator<kernels_tuple,
                             mfem::get<hardcoded_output_idx>(output_fops).vdim /
                             num_entities;
 
-   auto ir_weights = Reshape(this->op.integration_rule.GetWeights().Read(),
-                             num_qp);
+   auto ir_weights = Reshape(kernel.integration_rule.GetWeights().Read(), num_qp);
 
    auto input_size_on_qp = get_input_size_on_qp(kernel.inputs,
                                                 std::make_index_sequence<kernel.num_kinputs> {});
@@ -750,20 +752,35 @@ void DifferentiableOperator<kernels_tuple,
                MFEM_FOREACH_THREAD(qz, z, q1d)
                {
                   const int q = qx + q1d * (qy + q1d * qz);
-
-                  auto kernel_args = decay_tuple<typename kernel_t::kf_param_ts> {};
-                  auto kernel_shadow_args = decay_tuple<typename kernel_t::kf_param_ts> {};
-
                   auto r = Reshape(&residual_shmem(0, q), da_size_on_qp);
-                  apply_kernel_fwddiff_enzyme(
-                     r,
-                     kernel.func,
-                     kernel_args,
-                     input_shmem,
-                     kernel_shadow_args,
-                     shadow_shmem,
-                     q);
-                  // printf(">>>>> WARNING: AD DISABLED\n");
+                  auto kernel_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+
+                  if constexpr (std::is_same_v<autodiff_t, AutoDiff::EnzymeForward>)
+                  {
+                     auto kernel_shadow_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+                     apply_kernel_fwddiff_enzyme(
+                        r,
+                        kernel.func,
+                        kernel_args,
+                        kernel_shadow_args,
+                        input_shmem,
+                        shadow_shmem,
+                        q);
+                  }
+                  else if constexpr (std::is_same_v<autodiff_t, AutoDiff::NativeDualNumber>)
+                  {
+                     apply_kernel_native_dual(
+                        r,
+                        kernel.func,
+                        kernel_args,
+                        input_shmem,
+                        shadow_shmem,
+                        q);
+                  }
+                  else
+                  {
+                     static_assert(always_false<autodiff_t>, "unknown autodiff type");
+                  }
                }
             }
          }
@@ -800,7 +817,4 @@ void DifferentiableOperator<kernels_tuple,
    }
 }
 
-// #include "dfem_assemble_vector.icc"
-// #include "dfem_assemble_hypreparmatrix.icc"
-
-}
+} // namespace mfem

examples/dfem/dfem_element_operator.hpp

@@ -0,0 +1,79 @@
+#include "dfem_util.hpp"
+
+namespace mfem
+{
+
+template <typename func_t, typename input_t, typename output_t, typename dependency_map_t>
+struct ElementOperator;
+
+template <typename func_t, typename... input_ts, typename... output_ts, typename dependency_map_t>
+struct ElementOperator<func_t, mfem::tuple<input_ts...>, mfem::tuple<output_ts...>, dependency_map_t>
+{
+   using entity_t = Entity::Element;
+
+   func_t qfunc;
+
+   mfem::tuple<input_ts...> inputs;
+   mfem::tuple<output_ts...> outputs;
+
+   dependency_map_t dependency_map;
+
+   using qf_param_ts = typename create_function_signature<
+                       decltype(&func_t::operator())>::type::parameter_ts;
+   using qf_output_t = typename create_function_signature<
+                       decltype(&func_t::operator())>::type::return_t;
+
+   static constexpr size_t num_inputs =
+      mfem::tuple_size<decltype(inputs)>::value;
+   static constexpr size_t num_outputs =
+      mfem::tuple_size<decltype(outputs)>::value;
+
+   ElementOperator(func_t qfunc,
+                   mfem::tuple<input_ts...> inputs,
+                   mfem::tuple<output_ts...> outputs)
+      : qfunc(qfunc), inputs(inputs), outputs(outputs),
+        dependency_map(make_dependency_map(inputs))
+   {
+      // Consistency checks
+      if constexpr (num_outputs > 1)
+      {
+         static_assert(always_false<func_t>,
+                       "more than one output per kernel is not supported right now");
+      }
+
+      constexpr size_t num_qfinputs = mfem::tuple_size<qf_param_ts>::value;
+      static_assert(num_qfinputs == num_inputs,
+                    "kernel function inputs and descriptor inputs have to match");
+
+      constexpr size_t num_qf_outputs = mfem::tuple_size<qf_output_t>::value;
+      static_assert(num_qf_outputs == num_qf_outputs,
+                    "kernel function outputs and descriptor outputs have to match");
+   }
+};
+
+template <typename func_t, typename... input_ts, typename... output_ts>
+ElementOperator(func_t, mfem::tuple<input_ts...>, mfem::tuple<output_ts...>)
+-> ElementOperator<func_t, mfem::tuple<input_ts...>, mfem::tuple<output_ts...>,
+decltype(make_dependency_map(std::declval<mfem::tuple<input_ts...>>()))>;
+
+// template <typename func_t, typename input_t, typename output_t>
+// struct BoundaryElementOperator : public
+//    ElementOperator<func_t, input_t, output_t>
+// {
+// public:
+//    using entity_t = Entity::BoundaryElement;
+//    BoundaryElementOperator(func_t func, input_t inputs, output_t outputs)
+//       : ElementOperator<func_t, input_t, output_t>(func, inputs, outputs) {}
+// };
+
+// template <typename func_t, typename input_t, typename output_t>
+// struct FaceOperator : public
+//    ElementOperator<func_t, input_t, output_t>
+// {
+// public:
+//    using entity_t = Entity::Face;
+//    FaceOperator(func_t func, input_t inputs, output_t outputs)
+//       : ElementOperator<func_t, input_t, output_t>(func, inputs, outputs) {}
+// };
+
+} // namespace mfem

examples/dfem/dfem_fieldoperator.hpp

@@ -2,14 +2,17 @@
 
 #include <string>
 
+namespace mfem
+{
+
+template <int FIELD_ID = -1>
 class FieldOperator
 {
 public:
-   FieldOperator(std::string field_label = "", int size_on_qp = 0) :
-      field_label(field_label),
+   constexpr FieldOperator(int size_on_qp = 0) :
       size_on_qp(size_on_qp) {};
 
-   std::string field_label;
+   static constexpr int GetFieldId() { return FIELD_ID; }
 
    int size_on_qp = -1;
 
@@ -18,101 +21,207 @@ public:
    int vdim = -1;
 };
 
-class None : public FieldOperator
+template <int FIELD_ID = -1>
+class None : public FieldOperator<FIELD_ID>
 {
 public:
-   None(std::string field_label) :
-      FieldOperator(field_label) {}
+   constexpr None() : FieldOperator<FIELD_ID>() {}
 };
 
-class Weight : public FieldOperator
+template< typename T >
+struct is_none_fop
+{
+   static const bool value = false;
+};
+
+template <int FIELD_ID>
+struct is_none_fop<None<FIELD_ID>>
+{
+   static const bool value = true;
+};
+
+template <typename T>
+struct DisableAD
+{
+   T& operator()() const { return fop; }
+   T fop;
+};
+
+class Weight : public FieldOperator<-1>
 {
 public:
-   Weight() : FieldOperator("quadrature_weights") {};
+   constexpr Weight() : FieldOperator<-1>() {};
 };
 
-class Value : public FieldOperator
+template< typename T >
+struct is_weight_fop
+{
+   static const bool value = false;
+};
+
+template <>
+struct is_weight_fop<Weight>
+{
+   static const bool value = true;
+};
+
+template <int FIELD_ID = -1>
+class Value : public FieldOperator<FIELD_ID>
 {
 public:
-   Value(std::string field_label) : FieldOperator(field_label) {};
+   constexpr Value() : FieldOperator<FIELD_ID>() {};
 };
 
-class Gradient : public FieldOperator
+template< typename T >
+struct is_value_fop
+{
+   static const bool value = false;
+};
+
+template <int FIELD_ID>
+struct is_value_fop<Value<FIELD_ID>>
+{
+   static const bool value = true;
+};
+
+template <typename T>
+struct is_value_fop<DisableAD<T>>
+{
+   static const bool value = is_value_fop<T>::value;
+};
+
+template <int FIELD_ID = -1>
+class Gradient : public FieldOperator<FIELD_ID>
 {
 public:
-   Gradient(std::string field_label) : FieldOperator(field_label) {};
+   constexpr Gradient() : FieldOperator<FIELD_ID>() {};
 };
 
-class Curl : public FieldOperator
+template< typename T >
+struct is_gradient_fop
 {
-public:
-   Curl(std::string field_label) : FieldOperator(field_label) {};
+   static const bool value = false;
 };
 
-class Div : public FieldOperator
+template <int FIELD_ID>
+struct is_gradient_fop<Gradient<FIELD_ID>>
 {
-public:
-   Div(std::string field_label) : FieldOperator(field_label) {};
+   static const bool value = true;
 };
 
-class FaceValueLeft : public FieldOperator
-{
-public:
-   FaceValueLeft(std::string field_label) : FieldOperator(field_label) {};
-};
+// class FieldOperator
+// {
+// public:
+//    FieldOperator(std::string field_label = "", int size_on_qp = 0) :
+//       field_label(field_label),
+//       size_on_qp(size_on_qp) {};
 
-class FaceValueRight : public FieldOperator
-{
-public:
-   FaceValueRight(std::string field_label) : FieldOperator(field_label) {};
-};
+//    std::string field_label;
 
-class FaceNormal : public FieldOperator
-{
-public:
-   FaceNormal(std::string field_label) : FieldOperator(field_label) {};
-};
+//    int size_on_qp = -1;
 
-class One : public FieldOperator
-{
-public:
-   One(std::string field_label) : FieldOperator(field_label) {};
-};
+//    int dim = -1;
 
-namespace BareFieldOperator
-{
+//    int vdim = -1;
+// };
 
-struct Base
-{
-   Base(FieldOperator &o)
-   {
-      size_on_qp = o.size_on_qp;
-      dim = o.dim;
-      vdim = o.vdim;
-   };
-   int size_on_qp = -1;
-   int dim = -1;
-   int vdim = -1;
-};
+// class None : public FieldOperator
+// {
+// public:
+//    None(std::string field_label) :
+//       FieldOperator(field_label) {}
+// };
 
-struct None : Base
-{
-   None(FieldOperator &o) : Base(o) {}
-};
+// class Weight : public FieldOperator
+// {
+// public:
+//    Weight() : FieldOperator("quadrature_weights") {};
+// };
 
-struct Weight : Base
-{
-   Weight(FieldOperator &o) : Base(o) {}
-};
+// class Value : public FieldOperator
+// {
+// public:
+//    Value(std::string field_label) : FieldOperator(field_label) {};
+// };
 
-struct Value : Base
-{
-   Value(FieldOperator &o) : Base(o) {}
-};
+// class Gradient : public FieldOperator
+// {
+// public:
+//    Gradient(std::string field_label) : FieldOperator(field_label) {};
+// };
 
-struct Gradient : Base
-{
-   Gradient(FieldOperator &o) : Base(o) {}
-};
+// class Curl : public FieldOperator
+// {
+// public:
+//    Curl(std::string field_label) : FieldOperator(field_label) {};
+// };
 
-}
+// class Div : public FieldOperator
+// {
+// public:
+//    Div(std::string field_label) : FieldOperator(field_label) {};
+// };
+
+// class FaceValueLeft : public FieldOperator
+// {
+// public:
+//    FaceValueLeft(std::string field_label) : FieldOperator(field_label) {};
+// };
+
+// class FaceValueRight : public FieldOperator
+// {
+// public:
+//    FaceValueRight(std::string field_label) : FieldOperator(field_label) {};
+// };
+
+// class FaceNormal : public FieldOperator
+// {
+// public:
+//    FaceNormal(std::string field_label) : FieldOperator(field_label) {};
+// };
+
+// class One : public FieldOperator
+// {
+// public:
+//    One(std::string field_label) : FieldOperator(field_label) {};
+// };
+
+// namespace BareFieldOperator
+// {
+
+// struct Base
+// {
+//    Base(FieldOperator &o)
+//    {
+//       size_on_qp = o.size_on_qp;
+//       dim = o.dim;
+//       vdim = o.vdim;
+//    };
+//    int size_on_qp = -1;
+//    int dim = -1;
+//    int vdim = -1;
+// };
+
+// struct None : Base
+// {
+//    None(FieldOperator &o) : Base(o) {}
+// };
+
+// struct Weight : Base
+// {
+//    Weight(FieldOperator &o) : Base(o) {}
+// };
+
+// struct Value : Base
+// {
+//    Value(FieldOperator &o) : Base(o) {}
+// };
+
+// struct Gradient : Base
+// {
+//    Gradient(FieldOperator &o) : Base(o) {}
+// };
+
+// }
+
+} // namespace mfem

examples/dfem/dfem_integrate.hpp

@@ -17,7 +17,7 @@ void map_quadrature_data_to_fields_impl(DeviceTensor<2, double> &y,
    auto G = dtq.G;
    // assuming the quadrature point residual has to "play nice with
    // the test function"
-   if constexpr (std::is_same_v<std::decay_t<output_t>, BareFieldOperator::Value>)
+   if constexpr (std::is_same_v<std::decay_t<output_t>, Value<>>)
    {
       const auto [num_qp, cdim, num_dof] = B.GetShape();
       const int vdim = output.vdim > 0 ? output.vdim : cdim ;
@@ -35,7 +35,7 @@ void map_quadrature_data_to_fields_impl(DeviceTensor<2, double> &y,
       }
    }
    else if constexpr (
-      std::is_same_v<std::decay_t<output_t>, BareFieldOperator::Gradient>)
+      std::is_same_v<std::decay_t<output_t>, Gradient<>>)
    {
       const auto [num_qp, dim, num_dof] = G.GetShape();
       const int vdim = output.vdim;
@@ -67,7 +67,7 @@ void map_quadrature_data_to_fields_impl(DeviceTensor<2, double> &y,
    //    }
    // }
    else if constexpr (
-      std::is_same_v<std::decay_t<output_t>, BareFieldOperator::None>)
+      std::is_same_v<std::decay_t<output_t>, None<>>)
    {
       const auto [vdim, dim, num_qp] = G.GetShape();
       auto cc = Reshape(&f(0, 0, 0), num_qp * vdim);
@@ -95,7 +95,7 @@ void map_quadrature_data_to_fields_tensor_impl(DeviceTensor<2, double> &y,
    auto B = dtq.B;
    auto G = dtq.G;
 
-   if constexpr (std::is_same_v<std::decay_t<output_t>, BareFieldOperator::Value>)
+   if constexpr (is_value_fop<std::decay_t<output_t>>::value)
    {
       const auto [q1d, unused, d1d] = B.GetShape();
       const int vdim = output.vdim;
@@ -162,8 +162,7 @@ void map_quadrature_data_to_fields_tensor_impl(DeviceTensor<2, double> &y,
          MFEM_SYNC_THREAD;
       }
    }
-   else if constexpr (
-      std::is_same_v<std::decay_t<output_t>, BareFieldOperator::Gradient>)
+   else if constexpr (is_gradient_fop<std::decay_t<output_t>>::value)
    {
       const auto [q1d, unused, d1d] = G.GetShape();
       const int vdim = output.vdim;
@@ -242,8 +241,7 @@ void map_quadrature_data_to_fields_tensor_impl(DeviceTensor<2, double> &y,
          MFEM_SYNC_THREAD;
       }
    }
-   else if constexpr (
-      std::is_same_v<std::decay_t<output_t>, BareFieldOperator::None>)
+   else if constexpr (is_none_fop<std::decay_t<output_t>>::value)
    {
       const auto [q1d, unused, d1d] = B.GetShape();
       auto fqp = Reshape(&f(0, 0, 0), output.size_on_qp, q1d, q1d, q1d);

examples/dfem/dfem_interpolate.hpp

@@ -1,7 +1,6 @@
 #pragma once
 
 #include "dfem_util.hpp"
-#include <type_traits>
 
 namespace mfem
 {
@@ -19,8 +18,7 @@ void map_field_to_quadrature_data_tensor_product(
    auto B = dtq.B;
    auto G = dtq.G;
 
-   if constexpr (
-      std::is_same_v<std::decay_t<field_operator_t>, BareFieldOperator::Value>)
+   if constexpr (is_value_fop<std::decay_t<field_operator_t>>::value)
    {
       auto [q1d, unused, d1d] = B.GetShape();
       const int vdim = input.vdim;
@@ -84,7 +82,7 @@ void map_field_to_quadrature_data_tensor_product(
       }
    }
    else if constexpr (
-      std::is_same_v<std::decay_t<field_operator_t>, BareFieldOperator::Gradient>)
+      is_gradient_fop<std::decay_t<field_operator_t>>::value)
    {
       const auto [q1d, unused, d1d] = B.GetShape();
       const int vdim = input.vdim;
@@ -166,7 +164,7 @@ void map_field_to_quadrature_data_tensor_product(
    }
    // TODO: Create separate function for clarity
    else if constexpr (
-      std::is_same_v<std::decay_t<field_operator_t>, BareFieldOperator::Weight>)
+      std::is_same_v<std::decay_t<field_operator_t>, Weight>)
    {
       const int num_qp = integration_weights.GetShape()[0];
       // TODO: eeek
@@ -185,8 +183,7 @@ void map_field_to_quadrature_data_tensor_product(
       }
       MFEM_SYNC_THREAD;
    }
-   else if constexpr (
-      std::is_same_v<std::decay_t<field_operator_t>, BareFieldOperator::None>)
+   else if constexpr (is_none_fop<std::decay_t<field_operator_t>>::value)
    {
       const int q1d = B.GetShape()[0];
       auto field = Reshape(&field_e[0], input.size_on_qp, q1d * q1d * q1d);
@@ -211,7 +208,7 @@ void map_field_to_quadrature_data(
 {
    auto B = dtq.B;
    auto G = dtq.G;
-   if constexpr (std::is_same_v<field_operator_t, BareFieldOperator::Value>)
+   if constexpr (is_value_fop<field_operator_t>::value)
    {
       auto [num_qp, dim, num_dof] = B.GetShape();
       const int vdim = input.vdim;
@@ -230,8 +227,7 @@ void map_field_to_quadrature_data(
          }
       }
    }
-   else if constexpr (
-      std::is_same_v<field_operator_t, BareFieldOperator::Gradient>)
+   else if constexpr (is_gradient_fop<field_operator_t>::value)
    {
       const auto [num_qp, dim, num_dof] = G.GetShape();
       const int vdim = input.vdim;
@@ -268,7 +264,7 @@ void map_field_to_quadrature_data(
    //    }
    // }
    // TODO: Create separate function for clarity
-   else if constexpr (std::is_same_v<field_operator_t, BareFieldOperator::Weight>)
+   else if constexpr (std::is_same_v<field_operator_t, Weight>)
    {
       const int num_qp = integration_weights.GetShape()[0];
       auto f = Reshape(&field_qp[0], num_qp);
@@ -277,7 +273,7 @@ void map_field_to_quadrature_data(
          f(qp) = integration_weights(qp);
       }
    }
-   else if constexpr (std::is_same_v<field_operator_t, BareFieldOperator::None>)
+   else if constexpr (is_none_fop<field_operator_t>::value)
    {
       auto [num_qp, unused, num_dof] = B.GetShape();
       const int size_on_qp = input.size_on_qp;
@@ -296,33 +292,39 @@ void map_field_to_quadrature_data(
    }
 }
 
-template <typename T = NonTensorProduct, size_t num_kinputs, typename field_operator_ts, std::size_t... i>
+template <typename T = NonTensorProduct, typename field_operator_ts, size_t num_inputs, size_t num_fields>
 MFEM_HOST_DEVICE inline
 void map_fields_to_quadrature_data(
-   std::array<DeviceTensor<2>, num_kinputs> &fields_qp,
-   const std::array<DeviceTensor<1>, num_kinputs> &fields_e,
-   const std::array<DofToQuadMap, num_kinputs> &dtqmaps,
+   std::array<DeviceTensor<2>, num_inputs> &fields_qp,
+   const std::array<DeviceTensor<1>, num_fields> &fields_e,
+   const std::array<DofToQuadMap, num_inputs> &dtqmaps,
+   const std::array<int, num_inputs> &input_to_field,
    const field_operator_ts &fops,
    const DeviceTensor<1, const double> &integration_weights,
-   const std::array<DeviceTensor<1>, 6> &scratch_mem,
-   std::index_sequence<i...>)
+   const std::array<DeviceTensor<1>, 6> &scratch_mem)
 {
-   if constexpr (std::is_same_v<T, TensorProduct>)
+   for_constexpr<num_inputs>([&](auto i)
    {
-
-      (map_field_to_quadrature_data_tensor_product(fields_qp[i],
-                                                   dtqmaps[i], fields_e[i],
-                                                   mfem::get<i>(fops), integration_weights,
-                                                   scratch_mem),
-       ...);
-   }
-   else
-   {
-      (map_field_to_quadrature_data(fields_qp[i],
-                                    dtqmaps[i], fields_e[i],
-                                    mfem::get<i>(fops), integration_weights),
-       ...);
-   }
+      if constexpr (std::is_same_v<T, TensorProduct>)
+      {
+         map_field_to_quadrature_data_tensor_product(
+            fields_qp[i],
+            dtqmaps[i],
+            fields_e[input_to_field[i]],
+            mfem::get<i>(fops),
+            integration_weights,
+            scratch_mem);
+      }
+      else
+      {
+         map_field_to_quadrature_data(
+            fields_qp[i],
+            dtqmaps[i],
+            fields_e[i],
+            mfem::get<i>(fops),
+            integration_weights);
+      }
+   });
 }
 
 template <typename T, typename field_operator_t>
@@ -375,26 +377,27 @@ void map_fields_to_quadrature_data_conditional(
     ...);
 }
 
-template <typename T = NonTensorProduct, size_t num_kinputs, typename field_operator_ts, std::size_t... i>
+template <typename T = NonTensorProduct, size_t num_inputs, typename field_operator_ts>
 MFEM_HOST_DEVICE
 void map_direction_to_quadrature_data_conditional(
-   std::array<DeviceTensor<2>, num_kinputs> &directions_qp,
+   std::array<DeviceTensor<2>, num_inputs> &directions_qp,
    const DeviceTensor<1> &direction_e,
-   const std::array<DofToQuadMap, num_kinputs> &dtqmaps,
+   const std::array<DofToQuadMap, num_inputs> &dtqmaps,
    field_operator_ts fops,
    const DeviceTensor<1, const double> &integration_weights,
    const std::array<DeviceTensor<1>, 6> &scratch_mem,
-   const std::array<bool, num_kinputs> &conditions,
-   std::index_sequence<i...>)
+   const std::array<bool, num_inputs> &conditions)
 {
-   (map_field_to_quadrature_data_conditional<T>(directions_qp[i],
-                                                direction_e,
-                                                dtqmaps[i],
-                                                mfem::get<i>(fops),
-                                                integration_weights,
-                                                scratch_mem,
-                                                conditions[i]),
-    ...);
+   for_constexpr<num_inputs>([&](auto i)
+   {
+      map_field_to_quadrature_data_conditional<T>(directions_qp[i],
+                                                  direction_e,
+                                                  dtqmaps[i],
+                                                  mfem::get<i>(fops),
+                                                  integration_weights,
+                                                  scratch_mem,
+                                                  conditions[i]);
+   });
 }
 
 }

examples/dfem/dfem_qfunction.hpp

@@ -1,21 +1,36 @@
 #pragma once
 #include "dfem_util.hpp"
 
+#ifdef MFEM_USE_ENZYME
+#include <enzyme/utils>
+#include <enzyme/enzyme>
+#endif
+
 namespace mfem
 {
 
-MFEM_HOST_DEVICE inline
+template <typename T0, typename T1>
+MFEM_HOST_DEVICE
+void process_kf_arg(const T0 &, T1 &)
+{
+   static_assert(always_false<T0, T1>,
+                 "process_kf_arg not implemented for arg type");
+}
+
+template <typename T>
+MFEM_HOST_DEVICE
 void process_kf_arg(
-   const DeviceTensor<1> &u,
-   double &arg)
+   const DeviceTensor<1, T> &u,
+   T &arg)
 {
    arg = u(0);
 }
 
-MFEM_HOST_DEVICE inline
+template <typename T>
+MFEM_HOST_DEVICE
 void process_kf_arg(
-   const DeviceTensor<1> &u,
-   internal::tensor<double> &arg)
+   const DeviceTensor<1, T> &u,
+   internal::tensor<T> &arg)
 {
    arg(0) = u(0);
 }
@@ -32,11 +47,11 @@ void process_kf_arg(
    }
 }
 
-template <int n, int m>
+template <typename T, int n, int m>
 MFEM_HOST_DEVICE
 void process_kf_arg(
    const DeviceTensor<1> &u,
-   internal::tensor<double, n, m> &arg)
+   internal::tensor<T, n, m> &arg)
 {
    for (int i = 0; i < m; i++)
    {
@@ -45,39 +60,23 @@ void process_kf_arg(
          arg(j, i) = u((i * m) + j);
       }
    }
-   // assuming col major layout. translating to row major.
-   // i + N_i*j
-   // arg(0, 0) = u(0);
-   // arg(0, 1) = u(0 + 2 * 1);
-   // arg(1, 0) = u(1 + 2 * 0);
-   // arg(1, 1) = u(1 + 2 * 1);
 }
 
 template <typename arg_type>
 MFEM_HOST_DEVICE
 void process_kf_arg(const DeviceTensor<2> &u, arg_type &arg, int qp)
 {
-   // out << "qp: " << qp << "\n";
-   // for (int i = 0; i < u.GetShape()[0] * u.GetShape()[1]; i++)
-   // {
-   //    out << (&u(0, 0))[i] << " ";
-   // }
-   // out << "\n";
-
    const auto u_qp = Reshape(&u(0, qp), u.GetShape()[0]);
-   // for (int i = 0; i < u_qp.GetShape()[0]; i++)
-   // {
-   //    out << (&u_qp(0))[i] << " ";
-   // }
-   // out << "\n";
-
    process_kf_arg(u_qp, arg);
 }
 
 template <size_t num_fields, typename kf_args, std::size_t... i>
 MFEM_HOST_DEVICE
-void process_kf_args(const std::array<DeviceTensor<2>, num_fields> &u,
-                     kf_args &args, int qp, std::index_sequence<i...>)
+void process_kf_args(
+   const std::array<DeviceTensor<2>, num_fields> &u,
+   kf_args &args,
+   const int &qp,
+   std::index_sequence<i...>)
 {
    (process_kf_arg(u[i], mfem::get<i>(args), qp), ...);
 }
@@ -125,7 +124,6 @@ void process_kf_result(
    DeviceTensor<1, T> &r,
    const internal::tensor<T, n, m> &x)
 {
-   // out << "x: " << x << "\n";
    for (size_t i = 0; i < n; i++)
    {
       for (size_t j = 0; j < m; j++)
@@ -133,25 +131,24 @@ void process_kf_result(
          r(i + n * j) = x(i, j);
       }
    }
-
-   // out << "r: ";
-   // for (int i = 0; i < r.GetShape()[0]; i++)
-   // {
-   //    out << r(i) << " ";
-   // }
-   // out << "\n\n";
 }
 
-template <typename T> inline
-void process_kf_arg(const DeviceTensor<1> &u, const DeviceTensor<1> &v,
-                    double &arg)
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   const DeviceTensor<1> &v,
+   double &arg)
 {
    arg = u(0);
 }
 
-template <int n, int m> inline
-void process_kf_arg(const DeviceTensor<1> &u, const DeviceTensor<1> &v,
-                    internal::tensor<double, n, m> &arg)
+template <int n, int m>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   const DeviceTensor<1> &v,
+   internal::tensor<double, n, m> &arg)
 {
    for (int i = 0; i < m; i++)
    {
@@ -162,23 +159,6 @@ void process_kf_arg(const DeviceTensor<1> &u, const DeviceTensor<1> &v,
    }
 }
 
-template <typename arg_type> inline
-void process_kf_arg(const DeviceTensor<2> &u, const DeviceTensor<2> &v,
-                    arg_type &arg, int qp)
-{
-   const auto u_qp = Reshape(&u(0, qp), u.GetShape()[0]);
-   const auto v_qp = Reshape(&v(0, qp), v.GetShape()[0]);
-   process_kf_arg(u_qp, v_qp, arg);
-}
-
-template <size_t num_fields, typename kf_args, std::size_t... i> inline
-void process_kf_args(std::array<DeviceTensor<2>, num_fields> &u,
-                     std::array<DeviceTensor<2>, num_fields> &v,
-                     kf_args &args, int qp, std::index_sequence<i...>)
-{
-   (process_kf_arg(u[i], v[i], mfem::get<i>(args), qp), ...);
-}
-
 template <typename kernel_func_t, typename kernel_args_ts, size_t num_args>
 MFEM_HOST_DEVICE inline
 void apply_kernel(
@@ -194,16 +174,18 @@ void apply_kernel(
    process_kf_result(f_qp, mfem::get<0>(mfem::apply(kf, args)));
 }
 
+#ifdef MFEM_USE_ENZYME
 // Version for active function arguments only
 //
 // This is an Enzyme regression and can be removed in later versions.
 template <typename kernel_t, typename arg_ts, std::size_t... Is,
           typename inactive_arg_ts>
-inline auto fwddiff_apply_enzyme_indexed(kernel_t kernel, arg_ts &&args,
-                                         arg_ts &&shadow_args,
-                                         std::index_sequence<Is...>,
-                                         inactive_arg_ts &&inactive_args,
-                                         std::index_sequence<>)
+MFEM_HOST_DEVICE inline
+auto fwddiff_apply_enzyme_indexed(kernel_t kernel, arg_ts &&args,
+                                  arg_ts &&shadow_args,
+                                  std::index_sequence<Is...>,
+                                  inactive_arg_ts &&inactive_args,
+                                  std::index_sequence<>)
 {
    using kf_return_t = typename create_function_signature<
                        decltype(&kernel_t::operator())>::type::return_t;
@@ -215,11 +197,12 @@ inline auto fwddiff_apply_enzyme_indexed(kernel_t kernel, arg_ts &&args,
 // Interleave function arguments for enzyme
 template <typename kernel_t, typename arg_ts, std::size_t... Is,
           typename inactive_arg_ts, std::size_t... Js>
-inline auto fwddiff_apply_enzyme_indexed(kernel_t kernel, arg_ts &&args,
-                                         arg_ts &&shadow_args,
-                                         std::index_sequence<Is...>,
-                                         inactive_arg_ts &&inactive_args,
-                                         std::index_sequence<Js...>)
+MFEM_HOST_DEVICE inline
+auto fwddiff_apply_enzyme_indexed(kernel_t kernel, arg_ts &&args,
+                                  arg_ts &&shadow_args,
+                                  std::index_sequence<Is...>,
+                                  inactive_arg_ts &&inactive_args,
+                                  std::index_sequence<Js...>)
 {
    using kf_return_t = typename create_function_signature<
                        decltype(&kernel_t::operator())>::type::return_t;
@@ -230,9 +213,10 @@ inline auto fwddiff_apply_enzyme_indexed(kernel_t kernel, arg_ts &&args,
 }
 
 template <typename kernel_t, typename arg_ts, typename inactive_arg_ts>
-inline auto fwddiff_apply_enzyme(kernel_t kernel, arg_ts &&args,
-                                 arg_ts &&shadow_args,
-                                 inactive_arg_ts &&inactive_args)
+MFEM_HOST_DEVICE inline
+auto fwddiff_apply_enzyme(kernel_t kernel, arg_ts &&args,
+                          arg_ts &&shadow_args,
+                          inactive_arg_ts &&inactive_args)
 {
    auto arg_indices = std::make_index_sequence<
                       mfem::tuple_size<std::remove_reference_t<arg_ts>>::value> {};
@@ -250,8 +234,8 @@ void apply_kernel_fwddiff_enzyme(
    DeviceTensor<1, double> &f_qp,
    const kf_t &kf,
    kernel_arg_ts &args,
-   const std::array<DeviceTensor<2>, num_args> &u,
    kernel_arg_ts &shadow_args,
+   const std::array<DeviceTensor<2>, num_args> &u,
    const std::array<DeviceTensor<2>, num_args> &v,
    int qp_idx)
 {
@@ -264,5 +248,6 @@ void apply_kernel_fwddiff_enzyme(
    process_kf_result(f_qp,
                      mfem::get<0>(fwddiff_apply_enzyme(kf, args, shadow_args, mfem::tuple<> {})));
 }
+#endif // MFEM_USE_ENZYME
 
-}
+} // namespace mfem

examples/dfem/dfem_qfunction_dual.hpp

@@ -0,0 +1,187 @@
+#pragma once
+#include "dfem_util.hpp"
+#include "dfem_qfunction.hpp"
+
+namespace mfem
+{
+
+MFEM_HOST_DEVICE
+template <typename T0, typename T1, typename T2>
+void process_kf_arg(const T0 &, const T1 &, T2 &)
+{
+   static_assert(always_false<T0, T1, T2>,
+                 "process_kf_arg not implemented for arg type");
+}
+
+template <typename T>
+MFEM_HOST_DEVICE
+void process_kf_arg(
+   const DeviceTensor<1, T> &u,
+   const DeviceTensor<1, T> &v,
+   T &arg)
+{
+   arg = u(0);
+}
+
+template <typename T, int n, int m>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   internal::tensor<internal::dual<T, T>, n, m> &arg)
+{
+   for (int i = 0; i < m; i++)
+   {
+      for (int j = 0; j < n; j++)
+      {
+         arg(j, i).value = u((i * m) + j);
+      }
+   }
+}
+
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   internal::dual<T, T> &arg)
+{
+   arg.value = u(0);
+}
+
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   const DeviceTensor<1> &v,
+   internal::dual<T, T> &arg)
+{
+   arg.value = u(0);
+   arg.gradient = v(0);
+}
+
+template <typename T, int n>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   const DeviceTensor<1> &v,
+   internal::tensor<internal::dual<T, T>, n> &arg)
+{
+   for (int i = 0; i < n; i++)
+   {
+      arg(i).value = u(i);
+      arg(i).gradient = v(i);
+   }
+}
+
+template <typename T, int n, int m>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   const DeviceTensor<1> &v,
+   internal::tensor<internal::dual<T, T>, n, m> &arg)
+{
+   for (int i = 0; i < m; i++)
+   {
+      for (int j = 0; j < n; j++)
+      {
+         arg(j, i).value = u((i * m) + j);
+         arg(j, i).gradient = v((i * m) + j);
+      }
+   }
+}
+
+template <typename T, int n>
+MFEM_HOST_DEVICE inline
+void process_kf_result(
+   DeviceTensor<1, T> &r,
+   const internal::tensor<internal::dual<T, T>, n> &x)
+{
+   for (size_t i = 0; i < n; i++)
+   {
+      r(i) = x(i).value;
+   }
+}
+
+template <typename T, int n, int m>
+MFEM_HOST_DEVICE inline
+void process_kf_result(
+   DeviceTensor<1, T> &r,
+   const internal::tensor<internal::dual<T, T>, n, m> &x)
+{
+   for (size_t i = 0; i < n; i++)
+   {
+      for (size_t j = 0; j < m; j++)
+      {
+         r(i + n * j) = x(i, j).value;
+      }
+   }
+}
+
+template <typename arg_type>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<2> &u,
+   const DeviceTensor<2> &v,
+   arg_type &arg,
+   const int &qp)
+{
+   const auto u_qp = Reshape(&u(0, qp), u.GetShape()[0]);
+   const auto v_qp = Reshape(&v(0, qp), v.GetShape()[0]);
+   process_kf_arg(u_qp, v_qp, arg);
+}
+
+template <size_t num_args, typename kf_args, std::size_t... Is>
+MFEM_HOST_DEVICE inline
+void process_kf_args(
+   const std::array<DeviceTensor<2>, num_args> &u,
+   const std::array<DeviceTensor<2>, num_args> &v,
+   kf_args &args,
+   const int &qp,
+   std::index_sequence<Is...>)
+{
+   (process_kf_arg(u[Is], v[Is], mfem::get<Is>(args), qp), ...);
+}
+
+template <typename T, int n, int m>
+MFEM_HOST_DEVICE inline
+void process_derivative_from_native_dual(
+   DeviceTensor<1, T> &r,
+   const internal::tensor<internal::dual<T, T>, n, m> &x)
+{
+   for (size_t i = 0; i < n; i++)
+   {
+      for (size_t j = 0; j < m; j++)
+      {
+         r(i + n * j) = x(i, j).gradient;
+      }
+   }
+}
+
+template <typename T, int n>
+MFEM_HOST_DEVICE inline
+void process_derivative_from_native_dual(
+   DeviceTensor<1, T> &r,
+   const internal::tensor<internal::dual<T, T>, n> &x)
+{
+   for (size_t i = 0; i < n; i++)
+   {
+      r(i) = x(i).gradient;
+   }
+}
+
+template <typename kf_t, typename kernel_arg_ts, size_t num_args>
+MFEM_HOST_DEVICE inline
+void apply_kernel_native_dual(
+   DeviceTensor<1, double> &f_qp,
+   const kf_t &kf,
+   kernel_arg_ts &args,
+   const std::array<DeviceTensor<2>, num_args> &u,
+   const std::array<DeviceTensor<2>, num_args> &v,
+   const int &qp_idx)
+{
+   process_kf_args(u, v, args, qp_idx,
+                   std::make_index_sequence<mfem::tuple_size<kernel_arg_ts>::value> {});
+   auto r = mfem::get<0>(mfem::apply(kf, args));
+   process_derivative_from_native_dual(f_qp, r);
+}
+
+} // namespace mfem

examples/dfem/dfem_refactor.hpp

@@ -0,0 +1,530 @@
+#pragma once
+
+#include <mfem.hpp>
+#include <utility>
+#include "dfem_interpolate.hpp"
+#include "dfem_integrate.hpp"
+#include "dfem_qfunction.hpp"
+#include "dfem_qfunction_dual.hpp"
+#include "examples/dfem/dfem_util.hpp"
+
+namespace mfem
+{
+
+class DerivativeOperator : public Operator
+{
+   using derivative_action_t =
+      std::function<void(std::vector<Vector> &, const Vector &, Vector &)>;
+
+   using restriction_callback_t =
+      std::function<void(std::vector<Vector> &,
+                         const std::vector<Vector> &,
+                         std::vector<Vector> &)>;
+
+public:
+   DerivativeOperator(
+      const std::vector<derivative_action_t> &derivative_actions,
+      const FieldDescriptor &direction,
+      const std::vector<Vector *> &solutions_l,
+      const std::vector<Vector *> &parameters_l,
+      const std::vector<restriction_callback_t> &restriction_callbacks,
+      const std::function<void(Vector &, Vector &)> prolongation_transpose) :
+      derivative_actions(derivative_actions),
+      direction(direction),
+      restriction_callbacks(restriction_callbacks),
+      derivative_action_l(GetVSize(direction)),
+      prolongation_transpose(prolongation_transpose)
+   {
+      MFEM_ASSERT(derivative_actions.size() == restriction_callbacks.size(),
+                  "internal error");
+
+      derivative_action_l = 0.0;
+      this->solutions_l.resize(solutions_l.size());
+      this->parameters_l.resize(parameters_l.size());
+
+      for (int i = 0; i < solutions_l.size(); i++)
+      {
+         this->solutions_l[i] = *solutions_l[i];
+      }
+
+      for (int i = 0; i < parameters_l.size(); i++)
+      {
+         this->parameters_l[i] = *parameters_l[i];
+      }
+
+      fields_e.resize(solutions_l.size() + parameters_l.size());
+   }
+
+   void Mult(const Vector &x, Vector &y) const override
+   {
+      direction_t = x;
+      direction_t.SetSubVector(ess_tdof_list, 0.0);
+
+      prolongation(direction, direction_t, direction_l);
+      for (int i = 0; i < derivative_actions.size(); i++)
+      {
+         restriction_callbacks[i](solutions_l, parameters_l, fields_e);
+         derivative_actions[i](fields_e, direction_l, derivative_action_l);
+      }
+      prolongation_transpose(derivative_action_l, y);
+
+      y.SetSubVector(ess_tdof_list, 0.0);
+   };
+
+private:
+   std::vector<derivative_action_t> derivative_actions;
+
+   mutable std::vector<Vector> solutions_l;
+   std::vector<Vector> parameters_l;
+
+   FieldDescriptor direction;
+   mutable Vector direction_t;
+   mutable Vector direction_e;
+   mutable Vector direction_l;
+
+   mutable Vector derivative_action_e;
+   mutable Vector derivative_action_l;
+
+   mutable std::vector<Vector> fields_e;
+
+   Array<int> ess_tdof_list;
+   std::vector<restriction_callback_t> restriction_callbacks;
+   std::function<void(Vector &, Vector &)> prolongation_transpose;
+};
+
+class DifferentiableOperator : public Operator
+{
+   using action_t =
+      std::function<void(std::vector<Vector> &, const std::vector<Vector> &, Vector &)>;
+
+   using derivative_action_t =
+      std::function<void(std::vector<Vector> &, const Vector &, Vector &)>;
+
+   using restriction_callback_t =
+      std::function<void(std::vector<Vector> &,
+                         const std::vector<Vector> &,
+                         std::vector<Vector> &)>;
+
+public:
+   DifferentiableOperator(
+      const std::vector<FieldDescriptor> &solutions,
+      const std::vector<FieldDescriptor> &parameters,
+      const ParMesh &mesh);
+
+   void Mult(const Vector &x, Vector &y) const override
+   {
+      MFEM_ASSERT(!action_callbacks.empty(), "no integrators have been set");
+      prolongation(solutions, x, solutions_l);
+      for (auto &action : action_callbacks)
+      {
+         action(solutions_l, parameters_l, residual_l);
+      }
+      prolongation_transpose(residual_l, y);
+
+      y.SetSubVector(ess_tdof_list, 0.0);
+   }
+
+   template <
+      typename func_t,
+      typename... input_ts,
+      typename... output_ts,
+      typename derivative_indices_t>
+   void AddDomainIntegrator(
+      func_t qfunc,
+      mfem::tuple<input_ts...> inputs,
+      mfem::tuple<output_ts...> outputs,
+      const IntegrationRule &integration_rule,
+      const derivative_indices_t derivative_indices = {});
+
+   void SetParameters(std::vector<Vector *> p) const;
+
+   std::shared_ptr<DerivativeOperator> GetDerivative(
+      size_t derivative_idx,
+      std::vector<Vector *> solutions_l,
+      std::vector<Vector *> parameters_l)
+   {
+      MFEM_ASSERT(derivative_action_callbacks.find(derivative_idx) !=
+                  derivative_action_callbacks.end(),
+                  "no derivative action has been found for index " << derivative_idx);
+
+      return std::make_shared<DerivativeOperator>(
+                derivative_action_callbacks[derivative_idx],
+                fields[derivative_idx],
+                solutions_l,
+                parameters_l,
+                restriction_callbacks,
+                prolongation_transpose);
+   }
+
+private:
+   const ParMesh &mesh;
+
+   std::vector<action_t> action_callbacks;
+   std::map<size_t, std::vector<derivative_action_t>> derivative_action_callbacks;
+
+   std::vector<FieldDescriptor> solutions;
+   std::vector<FieldDescriptor> parameters;
+   // solutions and parameters
+   std::vector<FieldDescriptor> fields;
+
+   Array<int> ess_tdof_list;
+
+   mutable std::vector<Vector> solutions_l;
+   mutable std::vector<Vector> parameters_l;
+   mutable Vector residual_l;
+
+   mutable std::vector<Vector> fields_e;
+   mutable Vector residual_e;
+
+   std::function<void(Vector &, Vector &)> prolongation_transpose;
+   std::vector<restriction_callback_t> restriction_callbacks;
+};
+
+void DifferentiableOperator::SetParameters(std::vector<Vector *> p) const
+{
+   MFEM_ASSERT(parameters.size() == p.size(),
+               "number of parameters doesn't match descriptors");
+   for (int i = 0; i < parameters.size(); i++)
+   {
+      p[i]->Read();
+      parameters_l[i] = *p[i];
+   }
+}
+
+DifferentiableOperator::DifferentiableOperator(
+   const std::vector<FieldDescriptor> &solutions,
+   const std::vector<FieldDescriptor> &parameters,
+   const ParMesh &mesh) :
+   mesh(mesh),
+   solutions(solutions),
+   parameters(parameters)
+{
+   fields.resize(solutions.size() + parameters.size());
+   fields_e.resize(fields.size());
+   solutions_l.resize(solutions.size());
+   parameters_l.resize(parameters.size());
+
+   for (int i = 0; i < solutions.size(); i++)
+   {
+      fields[i] = solutions[i];
+   }
+
+   for (int i = 0; i < parameters.size(); i++)
+   {
+      fields[i + solutions.size()] = parameters[i];
+   }
+}
+
+template <
+   typename func_t,
+   typename... input_ts,
+   typename... output_ts,
+   typename derivative_indices_t = std::make_index_sequence<0>>
+void DifferentiableOperator::AddDomainIntegrator(
+   func_t qfunc,
+   mfem::tuple<input_ts...> inputs,
+   mfem::tuple<output_ts...> outputs,
+   const IntegrationRule &integration_rule,
+   const derivative_indices_t derivative_indices)
+{
+   using entity_t = Entity::Element;
+
+   static constexpr size_t num_inputs =
+      mfem::tuple_size<decltype(inputs)>::value;
+
+   static constexpr size_t num_outputs =
+      mfem::tuple_size<decltype(outputs)>::value;
+
+   using qf_param_ts = typename create_function_signature<
+                       decltype(&func_t::operator())>::type::parameter_ts;
+
+   using qf_output_t = typename create_function_signature<
+                       decltype(&func_t::operator())>::type::return_t;
+
+   // Consistency checks
+   if constexpr (num_outputs > 1)
+   {
+      static_assert(always_false<func_t>,
+                    "more than one output per kernel is not supported right now");
+   }
+
+   constexpr size_t num_qfinputs = mfem::tuple_size<qf_param_ts>::value;
+   static_assert(num_qfinputs == num_inputs,
+                 "kernel function inputs and descriptor inputs have to match");
+
+   constexpr size_t num_qf_outputs = mfem::tuple_size<qf_output_t>::value;
+   static_assert(num_qf_outputs == num_qf_outputs,
+                 "kernel function outputs and descriptor outputs have to match");
+
+   constexpr auto field_tuple = std::tuple_cat(std::tuple<input_ts...> {},
+                                               std::tuple<output_ts...> {});
+   constexpr auto filtered_field_tuple = filter_fields(field_tuple);
+   constexpr size_t num_fields = count_unique_field_ids(filtered_field_tuple);
+
+   constexpr auto dependency_map = make_dependency_map(mfem::tuple<input_ts...> {});
+
+   // Create the action callback
+   auto input_to_field = create_descriptors_to_fields_map<entity_t>(
+                            fields,
+                            inputs,
+                            std::make_index_sequence<num_inputs> {});
+
+   auto output_to_field = create_descriptors_to_fields_map<entity_t>(
+                             fields,
+                             outputs,
+                             std::make_index_sequence<num_outputs> {});
+
+   constexpr int hardcoded_output_idx = 0;
+   const int test_space_field_idx = output_to_field[hardcoded_output_idx];
+
+   ElementDofOrdering element_dof_ordering = ElementDofOrdering::LEXICOGRAPHIC;
+   DofToQuad::Mode doftoquad_mode = DofToQuad::Mode::TENSOR;
+
+   const Operator *R = get_restriction<entity_t>(fields[test_space_field_idx],
+                                                 element_dof_ordering);
+
+   // The explicit captures are necessary to avoid dependency on
+   // the specific instance of this class (this pointer).
+   auto restriction_callback =
+      [=, solutions = this->solutions, parameters = this->parameters]
+      (std::vector<Vector> &solutions_l,
+       const std::vector<Vector> &parameters_l,
+       std::vector<Vector> &fields_e)
+   {
+      restriction<entity_t>(solutions, solutions_l, fields_e,
+                            element_dof_ordering);
+      restriction<entity_t>(parameters, parameters_l, fields_e,
+                            element_dof_ordering,
+                            solutions.size());
+   };
+   restriction_callbacks.push_back(restriction_callback);
+
+   auto output_fop = mfem::get<hardcoded_output_idx>(outputs);
+
+   if constexpr (is_none_fop<decltype(output_fop)>::value)
+   {
+      prolongation_transpose = [&](Vector &r_local, Vector &y)
+      {
+         y = r_local;
+      };
+   }
+   // else if constexpr (std::is_same_v<decltype(output_fop), One>)
+   // {
+   //    prolongation_transpose = [&](Vector &r_local, Vector &y)
+   //    {
+   //       double local_sum = r_local.Sum();
+   //       MPI_Allreduce(&local_sum, y.GetData(), 1, MPI_DOUBLE, MPI_SUM,
+   //                     op.mesh.GetComm());
+   //       MFEM_ASSERT(y.Size() == 1, "output size doesn't match kernel description");
+   //    };
+   // }
+   else
+   {
+      auto P = get_prolongation(fields[test_space_field_idx]);
+      prolongation_transpose = [P](const Vector &r_local, Vector &y)
+      {
+         P->MultTranspose(r_local, y);
+      };
+   }
+
+   const int num_elements = GetNumEntities<Entity::Element>(mesh);
+   const int num_entities = GetNumEntities<entity_t>(mesh);
+   const int num_qp = integration_rule.GetNPoints();
+
+   size_t residual_lsize = GetVSize(fields[test_space_field_idx]);
+
+   // if constexpr (std::is_same_v<decltype(output_fop), One>)
+   // {
+   //    this->width = 1;
+   // }
+   // else
+   {
+      width = residual_lsize;
+   }
+
+   residual_l.SetSize(residual_lsize);
+
+   std::vector<const DofToQuad*> dtq;
+   for (const auto &field : fields)
+   {
+      dtq.emplace_back(GetDofToQuad<entity_t>(
+                          field,
+                          integration_rule,
+                          doftoquad_mode));
+   }
+   const int q1d = (int)floor(pow(num_qp, 1.0/mesh.Dimension()) + 0.5);
+
+   residual_e.SetSize(R->Height());
+
+   const int residual_size_on_qp =
+      GetSizeOnQP<entity_t>(mfem::get<hardcoded_output_idx>(outputs),
+                            fields[test_space_field_idx]);
+
+   auto input_dtq_maps =
+      create_dtq_maps<entity_t>(inputs, dtq, input_to_field);
+   auto output_dtq_maps =
+      create_dtq_maps<entity_t>(outputs, dtq, output_to_field);
+
+   const int test_vdim = mfem::get<hardcoded_output_idx>(outputs).vdim;
+   const int test_op_dim =
+      mfem::get<hardcoded_output_idx>(inputs).size_on_qp /
+      mfem::get<hardcoded_output_idx>(outputs).vdim;
+   const int num_test_dof = R->Height() /
+                            mfem::get<hardcoded_output_idx>(outputs).vdim /
+                            num_entities;
+
+   auto ir_weights = Reshape(integration_rule.GetWeights().Read(), num_qp);
+
+   auto input_size_on_qp =
+      get_input_size_on_qp(inputs, std::make_index_sequence<num_inputs> {});
+
+   auto action_shmem_info =
+      get_shmem_info<entity_t, num_fields, num_inputs, num_outputs>
+      (input_dtq_maps, output_dtq_maps, fields, num_entities, inputs, num_qp,
+       input_size_on_qp, residual_size_on_qp);
+
+   Vector shmem_cache(action_shmem_info.total_size);
+
+   // print_shared_memory_info(action_shmem_info);
+
+   action_callbacks.push_back(
+      [=](std::vector<Vector> &solutions_l,
+          const std::vector<Vector> &parameters_l,
+          Vector &residual_l) mutable
+   {
+      restriction_callback(solutions_l, parameters_l, fields_e);
+
+      residual_e = 0.0;
+      auto ye = Reshape(residual_e.ReadWrite(), test_vdim, num_test_dof, num_entities);
+
+      auto wrapped_fields_e = wrap_fields(fields_e,
+                                          action_shmem_info.field_sizes,
+                                          num_entities);
+
+      forall([=] MFEM_HOST_DEVICE (int e, void *shmem)
+      {
+         auto [input_dtq_shmem, output_dtq_shmem, fields_shmem, input_shmem,
+                                residual_shmem, scratch_shmem] =
+         unpack_shmem(shmem, action_shmem_info, input_dtq_maps, output_dtq_maps,
+                      wrapped_fields_e, num_qp, e);
+
+         map_fields_to_quadrature_data<TensorProduct>(
+            input_shmem, fields_shmem, input_dtq_shmem, input_to_field, inputs, ir_weights,
+            scratch_shmem);
+
+         call_qfunction<TensorProduct, qf_param_ts>(
+            qfunc, input_shmem, residual_shmem,
+            residual_size_on_qp, num_qp, q1d);
+
+         auto fhat = Reshape(&residual_shmem(0, 0), test_vdim, test_op_dim, num_qp);
+         auto y = Reshape(&ye(0, 0, e), num_test_dof, test_vdim);
+         map_quadrature_data_to_fields<TensorProduct>(y, fhat,
+                                                      mfem::get<0>(outputs),
+                                                      output_dtq_shmem[hardcoded_output_idx],
+                                                      scratch_shmem);
+      }, num_entities, q1d, q1d, q1d, action_shmem_info.total_size, shmem_cache.ReadWrite());
+
+      if constexpr (is_none_fop<decltype(output_fop)>::value)
+      {
+         residual_l = residual_e;
+      }
+      else
+      {
+         R->MultTranspose(residual_e, residual_l);
+      }
+   });
+
+   for_constexpr([&](auto derivative_idx)
+   {
+      // bool is_dependent = false;
+      // for_constexpr<num_inputs>([&](auto input_idx)
+      // {
+      //    constexpr auto input_is_dependent_on_field_idx =
+      //       std::get<derivative_idx>(std::get<input_idx>(dependency_map));
+
+      //    if constexpr (input_is_dependent_on_field_idx == 1)
+      //    {
+      //       is_dependent = true;
+      //    }
+      // });
+
+      // if (!is_dependent)
+      // {
+      //    derivative_action_callbacks[derivative_idx].push_back(
+      //       [=](const Vector &direction_l, Vector &y) mutable
+      //    {
+      //       y += 0.0;
+      //    });
+
+      //    return;
+      // }
+
+      auto direction = fields[derivative_idx];
+      size_t derivative_action_l_size = GetVSize(direction);
+
+      const int da_size_on_qp = GetSizeOnQP<entity_t>(
+                                   mfem::get<hardcoded_output_idx>(outputs),
+                                   fields[test_space_field_idx]);
+
+      auto shmem_info =
+         get_shmem_info<entity_t, num_fields, num_inputs, num_outputs>
+         (input_dtq_maps, output_dtq_maps, fields, num_entities, inputs, num_qp,
+          input_size_on_qp, residual_size_on_qp, derivative_idx);
+
+      Vector shmem_cache(shmem_info.total_size);
+
+      // print_shared_memory_info(shmem_info);
+
+      Vector direction_e;
+      Vector derivative_action_e(R->Height());
+      derivative_action_e = 0.0;
+
+      auto input_is_dependent = get_array_from_tuple(std::get<derivative_idx>
+                                                     (dependency_map));
+
+      derivative_action_callbacks[derivative_idx].push_back(
+         [=](std::vector<Vector> &fields_e, const Vector &direction_l,
+             Vector &derivative_action_l) mutable
+      {
+         restriction<entity_t>(direction, direction_l, direction_e, element_dof_ordering);
+         auto ye = Reshape(derivative_action_e.ReadWrite(), num_test_dof, test_vdim, num_entities);
+         auto wrapped_fields_e = wrap_fields(fields_e, shmem_info.field_sizes, num_entities);
+         auto wrapped_direction_e = Reshape(direction_e.ReadWrite(), shmem_info.direction_size, num_entities);
+         forall([=] MFEM_HOST_DEVICE (int e, double *shmem)
+         {
+            auto [input_dtq_shmem, output_dtq_shmem, fields_shmem, direction_shmem,
+                                   input_shmem, shadow_shmem, residual_shmem, scratch_shmem] =
+            unpack_shmem(shmem, shmem_info, input_dtq_maps,
+                         output_dtq_maps, wrapped_fields_e, wrapped_direction_e, num_qp, e);
+
+            map_fields_to_quadrature_data<TensorProduct>(
+               input_shmem, fields_shmem, input_dtq_shmem, input_to_field, inputs, ir_weights,
+               scratch_shmem);
+
+            zero_all(shadow_shmem);
+            map_direction_to_quadrature_data_conditional<TensorProduct>(
+               shadow_shmem, direction_shmem, input_dtq_shmem, inputs, ir_weights,
+               scratch_shmem, input_is_dependent);
+
+            call_qfunction_derivative_action<TensorProduct, qf_param_ts>(
+               qfunc, input_shmem, shadow_shmem, residual_shmem,
+               da_size_on_qp, num_qp, q1d);
+
+            auto fhat = Reshape(&residual_shmem(0, 0), test_vdim, test_op_dim, num_qp);
+            auto y = Reshape(&ye(0, 0, e), num_test_dof, test_vdim);
+            map_quadrature_data_to_fields<TensorProduct>(y, fhat,
+                                                         mfem::get<0>(outputs),
+                                                         output_dtq_shmem[hardcoded_output_idx],
+                                                         scratch_shmem);
+         }, num_entities, q1d, q1d, q1d, shmem_info.total_size, shmem_cache.ReadWrite());
+
+         R->MultTranspose(derivative_action_e, derivative_action_l);
+      });
+   }, derivative_indices);
+}
+
+
+} // namespace mfem
+
+// #include "dfem_refactor_action.hpp"
+// #include "dfem_refactor_derivatives.hpp"

examples/dfem/dfem_refactor_action.hpp

@@ -0,0 +1,232 @@
+#pragma once
+
+#include "dfem_refactor.hpp"
+
+namespace mfem
+{
+
+template <typename element_operator_t, size_t num_fields>
+void DifferentiableOperator::instantiate_action(
+   element_operator_t element_operator, action_t &action)
+{
+   using entity_t = typename element_operator_t::entity_t;
+
+   auto kinput_to_field = create_descriptors_to_fields_map<entity_t>(
+                             fields,
+                             element_operator.inputs,
+                             std::make_index_sequence<element_operator.num_inputs> {});
+
+   auto koutput_to_field = create_descriptors_to_fields_map<entity_t>(
+                              fields,
+                              element_operator.outputs,
+                              std::make_index_sequence<element_operator.num_outputs> {});
+
+   constexpr int hardcoded_output_idx = 0;
+   const int test_space_field_idx = koutput_to_field[hardcoded_output_idx];
+
+   const Operator *R = get_restriction<entity_t>(fields[test_space_field_idx],
+                                                 element_dof_ordering);
+
+   auto output_fop = mfem::get<hardcoded_output_idx>(element_operator.outputs);
+
+   const int num_elements = GetNumEntities<Entity::Element>(mesh);
+   const int num_entities = GetNumEntities<entity_t>(mesh);
+   const int num_qp = integration_rule.GetNPoints();
+
+   this->width = GetTrueVSize(fields[test_space_field_idx]);
+   size_t residual_lsize = GetVSize(fields[test_space_field_idx]);
+
+   // if constexpr (std::is_same_v<decltype(output_fop), One>)
+   // {
+   //    this->width = 1;
+   // }
+   // else
+   {
+      this->width = residual_lsize;
+   }
+
+   residual_l.SetSize(residual_lsize);
+
+   // assume only a single element type for now
+   std::vector<const DofToQuad*> dtq;
+   for (const auto &field : fields)
+   {
+      dtq.emplace_back(GetDofToQuad<entity_t>(
+                          field,
+                          integration_rule,
+                          doftoquad_mode));
+   }
+   const int q1d = (int)floor(pow(num_qp, 1.0/mesh.Dimension()) + 0.5);
+
+   residual_e.SetSize(R->Height());
+
+   const int residual_size_on_qp = GetSizeOnQP<entity_t>(
+                                      mfem::get<hardcoded_output_idx>(element_operator.outputs),
+                                      fields[test_space_field_idx]);
+
+   auto input_dtq_maps = create_dtq_maps<entity_t>(element_operator.inputs, dtq,
+                                                   kinput_to_field);
+   auto output_dtq_maps = create_dtq_maps<entity_t>(element_operator.outputs, dtq,
+                                                    koutput_to_field);
+
+   // auto input_fops = create_bare_fops(element_operator.inputs);
+   // auto output_fops = create_bare_fops(element_operator.outputs);
+
+   const int test_vdim = mfem::get<hardcoded_output_idx>
+                         (element_operator.outputs).vdim;
+   const int test_op_dim =
+      mfem::get<hardcoded_output_idx>(element_operator.inputs).size_on_qp /
+      mfem::get<hardcoded_output_idx>(element_operator.outputs).vdim;
+   const int num_test_dof = R->Height() /
+                            mfem::get<hardcoded_output_idx>(element_operator.outputs).vdim /
+                            num_entities;
+
+   auto ir_weights = Reshape(integration_rule.GetWeights().Read(), num_qp);
+
+   auto input_size_on_qp = get_input_size_on_qp(
+                              element_operator.inputs,
+                              std::make_index_sequence<element_operator.num_inputs> {});
+
+   auto shmem_info =
+      get_shmem_info<entity_t, num_fields, element_operator.num_inputs, element_operator.num_outputs>
+      (input_dtq_maps,
+       output_dtq_maps,
+       fields,
+       num_entities,
+       element_operator.inputs,
+       num_qp,
+       input_size_on_qp,
+       residual_size_on_qp);
+
+   Vector shmem_cache(shmem_info.total_size);
+
+   print_shared_memory_info(shmem_info);
+
+   action = [=](const Vector &x, Vector &y) mutable
+   {
+      prolongation(solutions, x, solutions_l);
+
+      restriction<entity_t>(solutions, solutions_l, this->fields_e,
+                            element_dof_ordering);
+      restriction<entity_t>(parameters, parameters_l, this->fields_e,
+                            element_dof_ordering,
+                            solutions.size());
+
+      residual_e = 0.0;
+      auto ye = Reshape(residual_e.ReadWrite(), test_vdim, num_test_dof,
+                        num_entities);
+
+      auto wrapped_fields_e = wrap_fields(this->fields_e,
+                                          shmem_info.field_sizes,
+                                          num_entities);
+
+      forall([=] MFEM_HOST_DEVICE (int e, void *shmem)
+      {
+         // printf("\ne: %d\n", e);
+         // tic();
+         auto input_dtq_shmem = load_dtq_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::INPUT_DTQ],
+            shmem_info.input_dtq_sizes,
+            input_dtq_maps);
+
+         auto output_dtq_shmem = load_dtq_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::OUTPUT_DTQ],
+            shmem_info.output_dtq_sizes,
+            output_dtq_maps);
+
+         auto fields_shmem = load_field_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::FIELD],
+            shmem_info.field_sizes,
+            kinput_to_field,
+            element_operator.inputs,
+            wrapped_fields_e,
+            e,
+         std::make_index_sequence<element_operator.num_inputs> {});
+
+         // These functions don't copy, they simply create a `DeviceTensor` object
+         // that points to correct chunks of the shared memory pool.
+         auto input_shmem = load_input_mem(
+                               shmem,
+                               shmem_info.offsets[SharedMemory::Index::INPUT],
+                               shmem_info.input_sizes,
+                               num_qp);
+
+         auto residual_shmem = load_residual_mem(
+                                  shmem,
+                                  shmem_info.offsets[SharedMemory::Index::OUTPUT],
+                                  shmem_info.residual_size,
+                                  num_qp);
+
+         auto scratch_mem = load_scratch_mem(
+                               shmem,
+                               shmem_info.offsets[SharedMemory::Index::TEMP],
+                               shmem_info.temp_sizes);
+
+         MFEM_SYNC_THREAD;
+         // // printf("shmem load elapsed: %.1fus\n", toc() * 1e6);
+
+         // // tic();
+         map_fields_to_quadrature_data<TensorProduct>(
+            input_shmem, fields_shmem, input_dtq_shmem, element_operator.inputs, ir_weights,
+            scratch_mem,
+            std::make_index_sequence<element_operator.num_inputs> {});
+         // printf("interpolate elapsed: %.1fus\n", toc() * 1e6);
+
+         // // tic();
+         MFEM_FOREACH_THREAD(qx, x, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(qz, z, q1d)
+               {
+                  const int q = qx + q1d * (qy + q1d * qz);
+                  auto qf_args = decay_tuple<typename element_operator_t::qf_param_ts> {};
+                  auto r = Reshape(&residual_shmem(0, q), residual_size_on_qp);
+                  apply_kernel(r, element_operator.qfunc, qf_args, input_shmem, q);
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+         // // printf("qf elapsed: %.1fus\n", toc() * 1e6);
+
+         // // tic();
+         auto fhat = Reshape(&residual_shmem(0, 0), test_vdim, test_op_dim, num_qp);
+         auto y = Reshape(&ye(0, 0, e), num_test_dof, test_vdim);
+         map_quadrature_data_to_fields<TensorProduct>(y, fhat,
+                                                      mfem::get<0>(element_operator.outputs),
+                                                      output_dtq_shmem[hardcoded_output_idx],
+                                                      scratch_mem);
+         // printf("integrate elapsed: %.1fus\n", toc() * 1e6);
+
+      }, num_entities, q1d, q1d, q1d, shmem_info.total_size, shmem_cache.ReadWrite());
+
+      if constexpr (std::is_same_v<decltype(output_fop), None<>>)
+      {
+         residual_l = y;
+      }
+      else
+      {
+         R->MultTranspose(residual_e, residual_l);
+      }
+
+      if constexpr (std::is_same_v<decltype(output_fop), None<>>)
+      {
+         y = residual_l;
+      }
+      // else if constexpr (std::is_same_v<decltype(output_fop), One>)
+      // {
+      //    double local_sum = residual_l.Sum();
+      //    MPI_Allreduce(&local_sum, y.GetData(), 1, MPI_DOUBLE, MPI_SUM, mesh.GetComm());
+      //    MFEM_ASSERT(y.Size() == 1, "output size doesn't match kernel description");
+      // }
+      else
+      {
+         get_prolongation(fields[test_space_field_idx])->MultTranspose(residual_l, y);
+      }
+   };
+}
+
+}

examples/dfem/dfem_refactor_derivatives.hpp

@@ -0,0 +1,132 @@
+#pragma once
+
+#include "dfem_refactor.hpp"
+
+template<typename T, T... Ints>
+void print_sequence(std::integer_sequence<T, Ints...>)
+{
+   ((std::cout << Ints << " "), ...);
+   std::cout << std::endl;
+}
+
+namespace mfem
+{
+
+template <
+   typename element_operator_t,
+   size_t num_solutions,
+   size_t num_parameters,
+   size_t derivative_idx>
+DerivativeOperator::DerivativeOperator(
+   element_operator_t element_operator,
+   const std::array<FieldDescriptor, num_solutions> &solutions,
+   const std::array<FieldDescriptor, num_parameters> &parameters,
+   const std::vector<FieldDescriptor> &fields,
+   ParMesh &mesh,
+   const IntegrationRule &integration_rule,
+   const ElementDofOrdering &element_dof_ordering,
+   const DofToQuad::Mode &doftoquad_mode,
+   std::integral_constant<size_t, derivative_idx>)
+{
+   direction = fields[derivative_idx];
+
+   size_t derivative_action_l_size = 0;
+   for (auto &s : solutions)
+   {
+      derivative_action_l_size += GetVSize(s);
+      this->width += GetTrueVSize(s);
+   }
+   this->height = derivative_action_l_size;
+   derivative_action_l.SetSize(derivative_action_l_size);
+
+   constexpr size_t num_fields = num_solutions + num_parameters;
+   using entity_t = typename element_operator_t::entity_t;
+
+   auto kinput_to_field = create_descriptors_to_fields_map<entity_t>(
+                             fields,
+                             element_operator.inputs,
+                             std::make_index_sequence<element_operator.num_inputs> {});
+
+   auto koutput_to_field = create_descriptors_to_fields_map<entity_t>(
+                              fields,
+                              element_operator.outputs,
+                              std::make_index_sequence<element_operator.num_outputs> {});
+
+   constexpr int hardcoded_output_idx = 0;
+   const int test_space_field_idx = koutput_to_field[hardcoded_output_idx];
+
+   const Operator *R = get_restriction<entity_t>(fields[test_space_field_idx],
+                                                 element_dof_ordering);
+
+   auto output_fop = mfem::get<hardcoded_output_idx>(element_operator.outputs);
+
+   const int num_elements = GetNumEntities<Entity::Element>(mesh);
+   const int num_entities = GetNumEntities<entity_t>(mesh);
+   const int num_qp = integration_rule.GetNPoints();
+
+   // assume only a single element type for now
+   std::vector<const DofToQuad*> dtq;
+   for (const auto &field : fields)
+   {
+      dtq.emplace_back(GetDofToQuad<entity_t>(
+                          field,
+                          integration_rule,
+                          doftoquad_mode));
+   }
+   const int q1d = (int)floor(pow(num_qp, 1.0/mesh.Dimension()) + 0.5);
+
+   derivative_action_e.SetSize(R->Height());
+
+   const int da_size_on_qp = GetSizeOnQP<entity_t>(
+                                mfem::get<hardcoded_output_idx>(element_operator.outputs),
+                                fields[test_space_field_idx]);
+
+   auto input_dtq_maps = create_dtq_maps<entity_t>(element_operator.inputs, dtq,
+                                                   kinput_to_field);
+   auto output_dtq_maps = create_dtq_maps<entity_t>(element_operator.outputs, dtq,
+                                                    koutput_to_field);
+
+   const int test_vdim = mfem::get<hardcoded_output_idx>
+                         (element_operator.outputs).vdim;
+   const int test_op_dim =
+      mfem::get<hardcoded_output_idx>(element_operator.inputs).size_on_qp /
+      mfem::get<hardcoded_output_idx>(element_operator.outputs).vdim;
+   const int num_test_dof = R->Height() /
+                            mfem::get<hardcoded_output_idx>(element_operator.outputs).vdim /
+                            num_entities;
+
+   auto ir_weights = Reshape(integration_rule.GetWeights().Read(), num_qp);
+
+   auto input_size_on_qp = get_input_size_on_qp(
+                              element_operator.inputs,
+                              std::make_index_sequence<element_operator.num_inputs> {});
+
+   auto input_is_dependent = std::get<derivative_idx>
+                             (element_operator.dependency_map);
+
+   constexpr bool with_derivatives = true;
+   auto shmem_info =
+      get_shmem_info<entity_t, num_fields, element_operator.num_inputs, element_operator.num_outputs>
+      (input_dtq_maps,
+       output_dtq_maps,
+       fields,
+       num_entities,
+       element_operator.inputs,
+       num_qp,
+       input_size_on_qp,
+       da_size_on_qp,
+       derivative_idx);
+
+   Vector shmem_cache(shmem_info.total_size);
+
+   print_shared_memory_info(shmem_info);
+
+   action_callback = [=](const Vector &x, Vector &y) mutable
+   {
+      restriction<entity_t>(direction, direction_l, direction_e,
+                            element_dof_ordering);
+
+   };
+}
+
+} // namespace mfem

examples/dfem/dfem_test_macro.hpp

@@ -1,5 +1,5 @@
 #pragma once
-#include "dfem.hpp"
+#include "dfem_refactor.hpp"
 
 #define DFEM_TEST_MAIN(function)                                               \
    int main(int argc, char* argv[])                                            \

examples/dfem/old_derivative_action_callback.icc

@@ -0,0 +1,188 @@
+{
+   using entity_t = typename kernel_t::entity_t;
+
+   auto kinput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+   kernel.inputs, std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto koutput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+   kernel.outputs, std::make_index_sequence<kernel.num_koutputs> {});
+
+   constexpr int hardcoded_output_idx = 0;
+   const int test_space_field_idx = koutput_to_field[hardcoded_output_idx];
+
+   const Operator *R = get_restriction<entity_t>(op.fields[test_space_field_idx],
+                                                 element_dof_ordering);
+
+   auto output_fop = mfem::get<hardcoded_output_idx>(kernel.outputs);
+
+   const int num_elements = GetNumEntities<Entity::Element>(op.mesh);
+   const int num_entities = GetNumEntities<entity_t>(op.mesh);
+   const int num_qp = op.integration_rule.GetNPoints();
+
+   // assume only a single element type for now
+   std::vector<const DofToQuad*> dtq;
+   for (const auto &field : op.fields)
+   {
+      dtq.emplace_back(GetDofToQuad<entity_t>(field, op.integration_rule,
+                                              doftoquad_mode));
+   }
+   const int q1d = dtq[0]->nqpt;
+
+   derivative_action_e.SetSize(R->Height());
+
+   const int da_size_on_qp = GetSizeOnQP<entity_t>(
+      mfem::get<hardcoded_output_idx>(kernel.outputs),
+      op.fields[test_space_field_idx]);
+
+   auto input_dtq_maps = create_dtq_maps<entity_t>(kernel.inputs, dtq,
+                                                   kinput_to_field);
+   auto output_dtq_maps = create_dtq_maps<entity_t>(kernel.outputs, dtq,
+                                                    koutput_to_field);
+
+   auto input_fops = create_bare_fops(kernel.inputs);
+   auto output_fops = create_bare_fops(kernel.outputs);
+
+   const int test_vdim = mfem::get<hardcoded_output_idx>(output_fops).vdim;
+   const int test_op_dim =
+   mfem::get<hardcoded_output_idx>(output_fops).size_on_qp /
+   mfem::get<hardcoded_output_idx>(output_fops).vdim;
+   const int num_test_dof = R->Height() /
+   mfem::get<hardcoded_output_idx>(output_fops).vdim /
+   num_entities;
+
+   auto ir_weights = Reshape(this->op.integration_rule.GetWeights().Read(),
+                             num_qp);
+
+   auto input_size_on_qp = get_input_size_on_qp(kernel.inputs,
+   std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto shmem_info = get_shmem_info<entity_t>(input_dtq_maps,
+                                              output_dtq_maps,
+                                              op.fields,
+                                              num_entities,
+                                              kernel.inputs,
+                                              num_qp,
+                                              input_size_on_qp,
+                                              da_size_on_qp);
+
+   Vector shmem_cache(shmem_info.total_size);
+
+   func = [=](Vector &ye_mem) mutable
+   {
+      restriction<entity_t>(direction, direction_l, direction_e,
+                            op.element_dof_ordering, derivative_idx);
+
+      // Check which qf inputs are dependent on the dependent variable
+      std::array<bool, kernel.num_kinputs> kinput_is_dependent;
+      bool no_qfinput_is_dependent = true;
+      for (int i = 0; i < kinput_is_dependent.size(); i++)
+      {
+         if (kinput_to_field[i] == derivative_idx)
+         {
+            no_qfinput_is_dependent = false;
+            kinput_is_dependent[i] = true;
+            // out << "function input " << i << " is dependent on "
+            //     << op.fields[kinput_to_field[i]].field_label << "\n";
+         }
+         else
+         {
+            kinput_is_dependent[i] = false;
+         }
+      }
+
+      if (no_qfinput_is_dependent)
+      {
+         return;
+      }
+
+      // auto kernel_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+      // auto kernel_shadow_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+
+      // DeviceTensor<1, const double> integration_weights(
+      //    this->op.integration_rule.GetWeights().Read(), num_qp);
+
+      // Vector zero;
+      // GeometricFactorMaps geometric_factors
+      // {
+      //    DeviceTensor<3, const double>(zero.Read(), 0, 0, 0)
+      // };
+
+      // // Fields interpolated to the quadrature points in the order of
+      // // kernel function arguments
+      // auto input_qp = map_inputs_to_memory(input_qp_mem, num_qp,
+      //                                      kernel.inputs,
+      //                                      std::make_index_sequence<kernel.num_kinputs> {});
+
+      // auto directions_qp = map_inputs_to_memory(directions_qp_mem, num_qp,
+      //                                           kernel.inputs,
+      //                                           std::make_index_sequence<kernel.num_kinputs> {});
+
+      // constexpr int fixed_output_idx = 0;
+      // auto Bv = output_dtq_maps[fixed_output_idx];
+      // auto [num_test_qp, test_op_dim, num_test_dof] = Bv.GetShape();
+      // const int test_vdim = mfem::get<0>(kernel.outputs).vdim;
+      // DeviceTensor<3> ye = Reshape(ye_mem.ReadWrite(), num_test_dof, test_vdim, num_entities);
+
+      forall([=] MFEM_HOST_DEVICE (int e, double *shmem)
+      {
+         // map_fields_to_quadrature_data(
+         //    input_qp, e, this->fields_e,
+         //    kinput_to_field, input_dtq_maps,
+         //    integration_weights, geometric_factors, kernel.inputs,
+         //    std::make_index_sequence<kernel.num_kinputs> {});
+
+         // map_fields_to_quadrature_data_conditional(
+         //    directions_qp, e,
+         //    directions_e, kinput_to_field,
+         //    input_dtq_maps,
+         //    integration_weights,
+         //    geometric_factors,
+         //    kinput_is_dependent,
+         //    kernel.inputs,
+         //    std::make_index_sequence<kernel.num_kinputs> {});
+
+         // for (int qp = 0; qp < num_qp; qp++)
+         // {
+         //    auto f_qp = apply_kernel_fwddiff_enzyme(
+         //                   kernel.func,
+         //                   kernel_args,
+         //                   input_qp,
+         //                   kernel_shadow_args,
+         //                   directions_qp,
+         //                   qp);
+
+         //    auto r_qp = Reshape(&da_qp(0, qp, e), da_size_on_qp);
+         //    for (int i = 0; i < da_size_on_qp; i++)
+         //    {
+         //       r_qp(i) = f_qp(i);
+         //    }
+         // }
+
+         // DeviceTensor<3> fhat = Reshape(&da_qp(0, 0, e), test_vdim, test_op_dim, num_qp);
+         // DeviceTensor<2> y = Reshape(&ye(0, 0, e), num_test_dof, test_vdim);
+         // map_quadrature_data_to_fields(y, fhat,
+         //                               output_fop,
+         //                               output_dtq_maps[hardcoded_output_idx]);
+      }, num_entities, q1d, q1d, 1, shmem_info.total_size, shmem_cache.GetData());
+
+      R->MultTranspose(ye_mem, derivative_action_l);
+   };
+
+   if constexpr (std::is_same_v<decltype(output_fop), One>)
+   {
+      prolongation_transpose = [&](Vector &r_local, Vector &y)
+      {
+         double local_sum = r_local.Sum();
+         MPI_Allreduce(&local_sum, y.GetData(), 1, MPI_DOUBLE, MPI_SUM,
+                       op.mesh.GetComm());
+         MFEM_ASSERT(y.Size() == 1, "output size doesn't match kernel description");
+      };
+   }
+   else
+   {
+      auto P = get_prolongation(op.fields[test_space_field_idx]);
+      prolongation_transpose = [P](Vector &r_l, Vector &y)
+      {
+         P->MultTranspose(r_l, y);
+      };
+   }

examples/dfem_test_diffusion_3d.cpp

@@ -6,7 +6,7 @@
 
 using namespace mfem;
 using mfem::internal::tensor;
-
+using mfem::internal::dual;
 
 int test_diffusion_3d(
    std::string mesh_file, int refinements, int polynomial_order)
@@ -66,7 +66,7 @@ int test_diffusion_3d(
    {
       auto diffusion_mf_kernel =
          [] MFEM_HOST_DEVICE (
-            const tensor<double, dim>& dudxi,
+            const tensor<dual<real_t, real_t>, dim>& dudxi,
             const tensor<double, dim, dim>& J,
             const double& w)
       {
@@ -139,7 +139,7 @@ int test_diffusion_3d(
    {
       auto diffusion_apply_kernel =
          [] MFEM_HOST_DEVICE (
-            const tensor<double, dim>& dudxi,
+            const tensor<dual<real_t, real_t>, dim>& dudxi,
             const tensor<double, dim, dim>& qdata)
       {
          return mfem::tuple{dudxi * qdata};

examples/dfem_test_diffusion_3d_refactor.cpp

@@ -0,0 +1,309 @@
+#include "dfem/dfem_test_macro.hpp"
+#include "examples/dfem/dfem_fieldoperator.hpp"
+#include "examples/dfem/dfem_refactor.hpp"
+#include "fem/bilininteg.hpp"
+#include "general/tic_toc.hpp"
+#include <utility>
+
+using namespace mfem;
+using mfem::internal::tensor;
+using mfem::internal::dual;
+
+int test_diffusion_3d(
+   std::string mesh_file, int refinements, int polynomial_order)
+{
+   constexpr int num_samples = 100;
+   constexpr int dim = 3;
+   Mesh mesh_serial = Mesh(mesh_file);
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "incorrect mesh dimension");
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(polynomial_order);
+   mesh_serial.Clear();
+
+   out << "#el: " << mesh.GetNE() << "\n";
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction*>(mesh.GetNodes());
+   ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec);
+
+   out << "#dofs " << h1fes.GetTrueVSize() << "\n";
+
+   const IntegrationRule& ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(),
+                   h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(
+                      0)->GetDim() - 1);
+
+   printf("#ndof per el = %d\n", h1fes.GetFE(0)->GetDof());
+   printf("#nqp = %d\n", ir.GetNPoints());
+   printf("#q1d = %d\n", (int)floor(pow(ir.GetNPoints(), 1.0/dim) + 0.5));
+
+   ParametricSpace qdata_space(dim, dim * dim, ir.GetNPoints(),
+                               dim * dim * ir.GetNPoints() * mesh.GetNE());
+   ParametricFunction qdata(qdata_space);
+
+   ParGridFunction f1_g(&h1fes);
+   ParGridFunction rho_g(&h1fes);
+
+   auto f1 = [](const Vector& coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      const double z = coords(2);
+      return 2.345 + x + x*y + 1.25 * z*x;
+   };
+   FunctionCoefficient f1_c(f1);
+   f1_g.ProjectCoefficient(f1_c);
+
+   Vector x(f1_g), y(h1fes.GetTrueVSize());
+   {
+      std::shared_ptr<DerivativeOperator> dpotential;
+      {
+         auto diffusion_mf_kernel =
+            [] MFEM_HOST_DEVICE (
+               const tensor<real_t, dim>& dudxi,
+               const tensor<real_t, dim, dim>& J,
+               const real_t& w)
+         {
+            auto invJ = inv(J);
+            return mfem::tuple{dudxi * invJ * transpose(invJ) * det(J) * w};
+         };
+
+         constexpr int Potential = 0;
+         constexpr int Coordinates = 1;
+
+         auto input_operators = mfem::tuple{Gradient<Potential>{}, Gradient<Coordinates>{}, Weight{}};
+         auto output_operator = mfem::tuple{Gradient<Potential>{}};
+
+         auto solutions = std::vector{FieldDescriptor{Potential, &h1fes}};
+         auto parameters = std::vector{FieldDescriptor{Coordinates, &mesh_fes}};
+
+         DifferentiableOperator dop(solutions, parameters, mesh);
+         auto derivatives = std::integer_sequence<size_t, Potential> {};
+         dop.AddDomainIntegrator(
+            diffusion_mf_kernel, input_operators, output_operator, ir, derivatives);
+
+         dop.SetParameters({mesh_nodes});
+         StopWatch sw;
+         sw.Start();
+         for (int i = 0; i < num_samples; i++)
+         {
+            dop.Mult(x, y);
+         }
+         sw.Stop();
+         printf("dfem mf:       %fs\n", sw.RealTime() / num_samples);
+         y.HostRead();
+
+         dpotential = dop.GetDerivative(Potential, {&f1_g}, {mesh_nodes});
+      }
+      dpotential->Mult(x, y);
+   }
+
+   {
+      auto diffusion_setup_kernel =
+         [] MFEM_HOST_DEVICE (
+            const tensor<double, dim, dim>& J,
+            const double& w)
+      {
+         auto invJ = inv(J);
+         return mfem::tuple{invJ * transpose(invJ) * det(J) * w};
+      };
+
+      constexpr int Potential = 0;
+      constexpr int Coordinates = 1;
+      constexpr int QData = 2;
+
+      auto input_operators = mfem::tuple{Gradient<Coordinates>{}, Weight{}};
+      auto output_operator = mfem::tuple{None<QData>{}};
+
+      auto solutions = std::vector{FieldDescriptor{Potential, &h1fes}};
+      auto parameters = std::vector{FieldDescriptor{Coordinates, &mesh_fes},
+                                    FieldDescriptor{QData, &qdata_space}};
+
+      DifferentiableOperator dop(solutions, parameters, mesh);
+      dop.AddDomainIntegrator(
+         diffusion_setup_kernel, input_operators, output_operator, ir);
+
+      dop.SetParameters({mesh_nodes, &qdata});
+      StopWatch sw;
+      sw.Start();
+      for (int i = 0; i < num_samples; i++)
+      {
+         dop.Mult(x, qdata);
+      }
+      sw.Stop();
+      printf("dfem pa setup: %fs\n", sw.RealTime() / num_samples);
+      qdata.HostRead();
+   }
+
+   // printf("qdata: ");
+   // print_vector(qdata);
+
+   {
+      auto diffusion_apply_kernel =
+         [] MFEM_HOST_DEVICE (
+            const tensor<real_t, dim>& dudxi,
+            const tensor<double, dim, dim>& qdata)
+      {
+         return mfem::tuple{dudxi * qdata};
+      };
+
+      constexpr int Potential = 0;
+      constexpr int QData = 1;
+
+      auto input_operators = mfem::tuple{Gradient<Potential>{}, None<QData>{}};
+      auto output_operator = mfem::tuple{Gradient<Potential>{}};
+
+      auto solutions = std::vector{FieldDescriptor{Potential, &h1fes}};
+      auto parameters = std::vector{FieldDescriptor{QData, &qdata_space}};
+
+      DifferentiableOperator dop(solutions, parameters, mesh);
+      dop.AddDomainIntegrator(
+         diffusion_apply_kernel, input_operators, output_operator, ir);
+
+      dop.SetParameters({&qdata});
+      StopWatch sw;
+      sw.Start();
+      for (int i = 0; i < num_samples; i++)
+      {
+         dop.Mult(x, y);
+      }
+      sw.Stop();
+      printf("dfem pa apply: %fs\n", sw.RealTime() / num_samples);
+      y.HostRead();
+   }
+
+   // printf("y: ");
+   // print_vector(y);
+
+   Vector y2(h1fes.TrueVSize());
+   {
+      ParBilinearForm a(&h1fes);
+      auto diff_integ = new DiffusionIntegrator;
+      diff_integ->SetIntRule(&ir);
+      a.AddDomainIntegrator(diff_integ);
+      a.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+
+      OperatorPtr A;
+      StopWatch sw;
+      sw.Start();
+      a.Assemble();
+      a.Finalize();
+      Array<int> empty;
+      a.FormSystemMatrix(empty, A);
+      sw.Stop();
+      printf("mfem pa setup: %fs\n", sw.RealTime());
+
+      sw.Clear();
+      sw.Start();
+      y2 = 0.0;
+      for (int i = 0; i < num_samples; i++)
+      {
+         A->Mult(x, y2);
+      }
+      sw.Stop();
+      printf("mfem pa apply: %fs\n", sw.RealTime() / num_samples);
+      y2.HostRead();
+   }
+   // printf("y2: ");
+   // print_vector(y2);
+
+   Vector diff(y2);
+   diff -= y;
+   if (diff.Norml2() > 1e-15)
+   {
+      printf("y: ");
+      print_vector(y);
+      printf("y2: ");
+      print_vector(y2);
+      printf("diff: ");
+      print_vector(diff);
+      return 1;
+   }
+
+   // Test linearization here as well
+   // auto dFdu = dop.GetDerivativeWrt<0>({&f1_g}, {mesh_nodes});
+
+   // if (dFdu->Height() != h1fes.GetTrueVSize())
+   // {
+   //    out << "dFdu unexpected height of " << dFdu->Height() << "\n";
+   //    return 1;
+   // }
+
+   // dFdu->Mult(x, y);
+   // y.HostRead();
+   // a.Mult(x, y2);
+   // y2.HostRead();
+
+   // diff = y2;
+   // diff -= y;
+   // if (diff.Norml2() > 1e-10)
+   // {
+   //    print_vector(diff);
+   //    print_vector(y2);
+   //    print_vector(y);
+   //    return 1;
+   // }
+
+   // // fd jacobian test
+   // {
+   //    double eps = 1.0e-6;
+   //    Vector v(x), xpv(x), xmv(x), fxpv(x.Size()), fxmv(x.Size());
+   //    v *= eps;
+   //    xpv += v;
+   //    xmv -= v;
+   //    dop.Mult(xpv, fxpv);
+   //    dop.Mult(xmv, fxmv);
+   //    fxpv -= fxmv;
+   //    fxpv /= (2.0*eps);
+
+   //    fxpv -= y;
+   //    if (fxpv.Norml2() > eps)
+   //    {
+   //       out << "||dFdu_FD u^* - ex||_l2 = " << fxpv.Norml2() << "\n";
+   //       return 1;
+   //    }
+   // }
+
+   // f1_g.ProjectCoefficient(f1_c);
+   // rho_g.ProjectCoefficient(rho_c);
+   // auto dFdrho = dop.GetDerivativeWrt<1>({&f1_g}, {&rho_g, mesh_nodes});
+   // if (dFdrho->Height() != h1fes.GetTrueVSize())
+   // {
+   //    out << "dFdrho unexpected height of " << dFdrho->Height() << "\n";
+   //    return 1;
+   // }
+
+   // dFdrho->Mult(rho_g, y);
+
+   // // fd test
+   // {
+   //    double eps = 1.0e-6;
+   //    Vector v(rho_g), rhopv(rho_g), rhomv(rho_g), frhopv(x.Size()),
+   //    frhomv(x.Size()); v *= eps; rhopv += v; rhomv -= v;
+   //    dop.SetParameters({&rhopv, mesh_nodes});
+   //    dop.Mult(x, frhopv);
+   //    dop.SetParameters({&rhomv, mesh_nodes});
+   //    dop.Mult(x, frhomv);
+   //    frhopv -= frhomv;
+   //    frhopv /= (2.0*eps);
+
+   //    frhopv -= y;
+   //    if (frhopv.Norml2() > eps)
+   //    {
+   //       out << "||dFdu_FD u^* - ex||_l2 = " << frhopv.Norml2() << "\n";
+   //       return 1;
+   //    }
+   // }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_diffusion_3d);

examples/dfem_test_nonlinear_elasticity_3d.cpp

@@ -1,13 +1,10 @@
 #include "dfem/dfem.hpp"
 #include "dfem/dfem_test_macro.hpp"
-#include "fem/pfespace.hpp"
-#include "linalg/hypre.hpp"
-#include "linalg/operator.hpp"
-#include "linalg/solvers.hpp"
 #include <fstream>
 
 using namespace mfem;
 using mfem::internal::tensor;
+using mfem::internal::dual;
 
 class FDJacobian : public Operator
 {
@@ -181,19 +178,19 @@ int test_nonlinear_elasticity_3d(std::string mesh_file,
    ParGridFunction u(&h1fes);
 
    auto elasticity_kernel = [] MFEM_HOST_DEVICE
-                            (const tensor<real_t, dim, dim> &dudxi,
+                            (const tensor<dual<real_t, real_t>, dim, dim> &dudxi,
                              const tensor<real_t, dim, dim> &J,
-                             const double &w)
+                             const real_t &w)
    {
       // shear modulus
-      mfem::real_t D1 = 0.1e6;
+      real_t D1{0.1e6};
       // bulk modulus
-      mfem::real_t C1 = 1.0e6;
+      real_t C1{1.0e6};
       constexpr auto I = mfem::internal::IsotropicIdentity<dim>();
       auto invJ = inv(J);
       auto dudx = dudxi * invJ;
-      real_t F = det(I + dudx);
-      real_t p = -2.0 * D1 * F * (F - 1);
+      auto F = det(I + dudx);
+      auto p = -2.0 * D1 * F * (F - 1);
       auto devB = dev(dudx + transpose(dudx) + dot(dudx, transpose(dudx)));
       auto sigma = -(p / F) * I + 2.0 * (C1 / pow(F, 5.0 / 3.0)) * devB;
 
@@ -203,12 +200,14 @@ int test_nonlinear_elasticity_3d(std::string mesh_file,
    mfem::tuple argument_operators{Gradient{"displacement"}, Gradient{"coordinates"}, Weight{}};
    mfem::tuple output_operator{Gradient{"displacement"}};
 
-   ElementOperator op{elasticity_kernel, argument_operators, output_operator};
+   // B^T D(B0*dudxi, B1*J, B2*w)
+   ElementOperator op(elasticity_kernel, argument_operators, output_operator, ir);
 
    std::array solutions{FieldDescriptor{&h1fes, "displacement"}};
    std::array parameters{FieldDescriptor{&mesh_fes, "coordinates"}};
 
-   DifferentiableOperator dop{solutions, parameters, mfem::tuple{op}, mesh, ir};
+   DifferentiableOperator dop(solutions, parameters, mfem::tuple{op}, mesh,
+                              AutoDiff::NativeDualNumber{});
 
    ElasticityOperator elasticity(h1fes, dop, ess_tdof_list);
 
@@ -242,7 +241,7 @@ int test_nonlinear_elasticity_3d(std::string mesh_file,
    newton.SetOperator(elasticity);
    newton.SetRelTol(1e-6);
    newton.SetMaxIter(100);
-   newton.SetAdaptiveLinRtol();
+   // newton.SetAdaptiveLinRtol();
    newton.SetPrintLevel(IterativeSolver::PrintLevel().Iterations());
 
    elasticity.SetParameters(*mesh_nodes);

linalg/dtensor.hpp

@@ -12,78 +12,39 @@
 #ifndef MFEM_DTENSOR
 #define MFEM_DTENSOR
 
-#include "../general/backends.hpp"
+#include <numeric>
+#include <array>
+#include <tuple>
+
+#include "../config/config.hpp"
 
 namespace mfem
 {
 
-/// A Class to compute the real index from the multi-indices of a tensor
-template <int N, int Dim, typename T, typename... Args>
-class TensorInd
+/// ///////////////////////////////////////////////////////////////////////////
+template <int N, typename... Args>
+MFEM_HOST_DEVICE inline int ColMajor(const int (&dims)[N], Args... args)
 {
-public:
-   MFEM_HOST_DEVICE
-   static inline int result(const int* sizes, T first, Args... args)
-   {
-#if !(defined(MFEM_USE_CUDA) || defined(MFEM_USE_HIP))
-      MFEM_ASSERT(first<sizes[N-1],"Trying to access out of boundary.");
-#endif
-      return static_cast<int>(first + sizes[N - 1] * TensorInd < N + 1, Dim, Args... >
-                              ::result(sizes, args...));
-   }
-};
+   int offset = 0, i = 0, unused;
+   ((offset *= dims[N - ++i], offset += args, std::ignore = unused) = ...);
+   return offset;
+}
 
-// Terminal case
-template <int Dim, typename T, typename... Args>
-class TensorInd<Dim, Dim, T, Args...>
+/// ///////////////////////////////////////////////////////////////////////////
+template <int N, typename... Args>
+MFEM_HOST_DEVICE inline int RowMajor(const int (&dims)[N], Args... args)
 {
-public:
-   MFEM_HOST_DEVICE
-   static inline int result(const int* sizes, T first, Args... args)
-   {
-#if !(defined(MFEM_USE_CUDA) || defined(MFEM_USE_HIP))
-      MFEM_ASSERT(first<sizes[Dim-1],"Trying to access out of boundary.");
-#endif
-      return static_cast<int>(first);
-   }
-};
+   int offset = 0, i = 0;
+   return ((..., (offset *= dims[i++], offset += args, 0)), offset);
+}
 
-
-/// A class to initialize the size of a Tensor
-template <int N, int Dim, typename T, typename... Args>
-class Init
-{
-public:
-   MFEM_HOST_DEVICE
-   static inline int result(int* sizes, T first, Args... args)
-   {
-      sizes[N - 1] = first;
-      return first * Init < N + 1, Dim, Args... >::result(sizes, args...);
-   }
-};
-
-// Terminal case
-template <int Dim, typename T, typename... Args>
-class Init<Dim, Dim, T, Args...>
-{
-public:
-   MFEM_HOST_DEVICE
-   static inline int result(int* sizes, T first, Args... args)
-   {
-      sizes[Dim - 1] = first;
-      return first;
-   }
-};
-
-
-/// A basic generic Tensor class, appropriate for use on the GPU
-template<int Dim, typename Scalar = real_t>
+/// ///////////////////////////////////////////////////////////////////////////
+template <int N, typename T = real_t, bool Column = true>
 class DeviceTensor
 {
 protected:
-   int capacity;
-   Scalar *data;
-   int sizes[Dim];
+   int dims[N], size;
+   T *data;
 
 public:
    /// Default constructor
@@ -92,72 +53,236 @@ public:
    DeviceTensor() {}
 
    /// Constructor to initialize a tensor from the Scalar array data_
-   template <typename... Args> MFEM_HOST_DEVICE
-   DeviceTensor(Scalar* data_, Args... args)
-   {
-      static_assert(sizeof...(args) == Dim, "Wrong number of arguments");
-      // Initialize sizes, and compute the number of values
-      const long int nb = Init<1, Dim, Args...>::result(sizes, args...);
-      capacity = nb;
-      data = (capacity > 0) ? data_ : NULL;
-   }
+   template <typename... Args, typename = std::enable_if_t<(sizeof...(Args) == N)>>
+   MFEM_HOST_DEVICE DeviceTensor(T *data, Args... args):
+      dims{args...},
+      size(std::accumulate(dims, dims + N, 1, std::multiplies<int> {})),
+        data(size > 0 ? data : nullptr)
+   { }
 
    /// Copy constructor (default)
-   DeviceTensor(const DeviceTensor&) = default;
+   // DeviceTensor(const DeviceTensor &) = default;
 
    /// Copy assignment (default)
-   DeviceTensor& operator=(const DeviceTensor&) = default;
+   // DeviceTensor &operator=(const DeviceTensor &) = default;
 
    /// Conversion to `Scalar *`.
-   MFEM_HOST_DEVICE inline operator Scalar *() const { return data; }
+   MFEM_HOST_DEVICE inline operator T *() const { return data; }
+
+   /// Computes the offset of the tensor element at the given multi-indices
+   template <typename... Args, typename = std::enable_if_t<(sizeof...(Args) == N)>>
+   MFEM_HOST_DEVICE inline int Offset(Args... args) const
+   {
+      constexpr auto offset = Column ? ColMajor<N, Args...> : RowMajor<N, Args...>;
+      return offset(dims, args...);
+   }
 
    /// Const accessor for the data
-   template <typename... Args> MFEM_HOST_DEVICE inline
-   Scalar& operator()(Args... args) const
+   template <typename... Args, typename = std::enable_if_t<(sizeof...(Args) == N)>>
+   MFEM_HOST_DEVICE inline T &operator()(Args... args) const
    {
-      static_assert(sizeof...(args) == Dim, "Wrong number of arguments");
-      return data[ TensorInd<1, Dim, Args...>::result(sizes, args...) ];
+      return data[Offset(args...)];
    }
 
    /// Subscript operator where the tensor is viewed as a 1D array.
-   MFEM_HOST_DEVICE inline Scalar& operator[](int i) const
-   {
-      return data[i];
-   }
+   MFEM_HOST_DEVICE inline T &operator[](int i) const { return data[i]; }
 
-   MFEM_HOST_DEVICE inline std::array<int, Dim> GetShape() const
+   /// Returns the size of the tensor
+   MFEM_HOST_DEVICE inline int Size() const { return size; }
+
+   MFEM_HOST_DEVICE inline std::array<int, N> GetShape() const
    {
-      std::array<int, Dim> s;
-      for (int i = 0; i < Dim; i++)
+      std::array<int, N> s;
+      for (int i = 0; i < N; i++)
       {
-         s[i] = sizes[i];
+         s[i] = dims[i];
       }
       return s;
    }
 };
 
-
 /** @brief Wrap a pointer as a DeviceTensor with automatically deduced template
     parameters */
-template <typename T, typename... Dims> MFEM_HOST_DEVICE
-inline DeviceTensor<sizeof...(Dims),T> Reshape(T *ptr, Dims... dims)
+template <typename T, typename... Dims>
+MFEM_HOST_DEVICE inline DeviceTensor<sizeof...(Dims), T> Reshape(T *ptr,
+                                                                 Dims... dims)
 {
-   return DeviceTensor<sizeof...(Dims),T>(ptr, dims...);
+   return DeviceTensor<sizeof...(Dims), T>(ptr, dims...);
 }
 
+template <typename T, typename... Dims>
+MFEM_HOST_DEVICE inline DeviceTensor<sizeof...(Dims), T, false> RowReshape(
+   T *ptr,
+   Dims... dims)
+{
+   return DeviceTensor<sizeof...(Dims), T, false>(ptr, dims...);
+}
 
-typedef DeviceTensor<1,int> DeviceArray;
-typedef DeviceTensor<1,const int> ConstDeviceArray;
+using DeviceArray = DeviceTensor<1, int>;
+using ConstDeviceArray = DeviceTensor<1, const int>;
 
-typedef DeviceTensor<1,real_t> DeviceVector;
-typedef DeviceTensor<1,const real_t> ConstDeviceVector;
+using DeviceVector = DeviceTensor<1, real_t>;
+using ConstDeviceVector = DeviceTensor<1, const real_t>;
 
-typedef DeviceTensor<2,real_t> DeviceMatrix;
-typedef DeviceTensor<2,const real_t> ConstDeviceMatrix;
+using DeviceMatrix = DeviceTensor<2, real_t>;
+using ConstDeviceMatrix = DeviceTensor<2, const real_t>;
 
-typedef DeviceTensor<3,real_t> DeviceCube;
-typedef DeviceTensor<3,const real_t> ConstDeviceCube;
+using DeviceCube = DeviceTensor<3, real_t>;
+using ConstDeviceCube = DeviceTensor<3, const real_t>;
 
-} // mfem namespace
+} // namespace mfem
 
 #endif // MFEM_DTENSOR
+
+// #ifndef MFEM_DTENSOR
+// #define MFEM_DTENSOR
+
+// #include "../general/backends.hpp"
+
+// namespace mfem
+// {
+
+// /// A Class to compute the real index from the multi-indices of a tensor
+// template <int N, int Dim, typename T, typename... Args>
+// class TensorInd
+// {
+// public:
+//    MFEM_HOST_DEVICE
+//    static inline int result(const int* sizes, T first, Args... args)
+//    {
+// #if !(defined(MFEM_USE_CUDA) || defined(MFEM_USE_HIP))
+//       MFEM_ASSERT(first<sizes[N-1],"Trying to access out of boundary.");
+// #endif
+//       return static_cast<int>(first + sizes[N - 1] * TensorInd < N + 1, Dim, Args... >
+//                               ::result(sizes, args...));
+//    }
+// };
+
+// // Terminal case
+// template <int Dim, typename T, typename... Args>
+// class TensorInd<Dim, Dim, T, Args...>
+// {
+// public:
+//    MFEM_HOST_DEVICE
+//    static inline int result(const int* sizes, T first, Args... args)
+//    {
+// #if !(defined(MFEM_USE_CUDA) || defined(MFEM_USE_HIP))
+//       MFEM_ASSERT(first<sizes[Dim-1],"Trying to access out of boundary.");
+// #endif
+//       return static_cast<int>(first);
+//    }
+// };
+
+
+// /// A class to initialize the size of a Tensor
+// template <int N, int Dim, typename T, typename... Args>
+// class Init
+// {
+// public:
+//    MFEM_HOST_DEVICE
+//    static inline int result(int* sizes, T first, Args... args)
+//    {
+//       sizes[N - 1] = first;
+//       return first * Init < N + 1, Dim, Args... >::result(sizes, args...);
+//    }
+// };
+
+// // Terminal case
+// template <int Dim, typename T, typename... Args>
+// class Init<Dim, Dim, T, Args...>
+// {
+// public:
+//    MFEM_HOST_DEVICE
+//    static inline int result(int* sizes, T first, Args... args)
+//    {
+//       sizes[Dim - 1] = first;
+//       return first;
+//    }
+// };
+
+
+// /// A basic generic Tensor class, appropriate for use on the GPU
+// template<int Dim, typename Scalar = real_t>
+// class DeviceTensor
+// {
+// protected:
+//    int capacity;
+//    Scalar *data;
+//    int sizes[Dim];
+
+// public:
+//    /// Default constructor
+//    // DeviceTensor() = delete;
+//    MFEM_HOST_DEVICE
+//    DeviceTensor() {}
+
+//    /// Constructor to initialize a tensor from the Scalar array data_
+//    template <typename... Args> MFEM_HOST_DEVICE
+//    DeviceTensor(Scalar* data_, Args... args)
+//    {
+//       static_assert(sizeof...(args) == Dim, "Wrong number of arguments");
+//       // Initialize sizes, and compute the number of values
+//       const long int nb = Init<1, Dim, Args...>::result(sizes, args...);
+//       capacity = nb;
+//       data = (capacity > 0) ? data_ : NULL;
+//    }
+
+//    /// Copy constructor (default)
+//    DeviceTensor(const DeviceTensor&) = default;
+
+//    /// Copy assignment (default)
+//    DeviceTensor& operator=(const DeviceTensor&) = default;
+
+//    /// Conversion to `Scalar *`.
+//    MFEM_HOST_DEVICE inline operator Scalar *() const { return data; }
+
+//    /// Const accessor for the data
+//    template <typename... Args> MFEM_HOST_DEVICE inline
+//    Scalar& operator()(Args... args) const
+//    {
+//       static_assert(sizeof...(args) == Dim, "Wrong number of arguments");
+//       return data[ TensorInd<1, Dim, Args...>::result(sizes, args...) ];
+//    }
+
+//    /// Subscript operator where the tensor is viewed as a 1D array.
+//    MFEM_HOST_DEVICE inline Scalar& operator[](int i) const
+//    {
+//       return data[i];
+//    }
+
+//    MFEM_HOST_DEVICE inline std::array<int, Dim> GetShape() const
+//    {
+//       std::array<int, Dim> s;
+//       for (int i = 0; i < Dim; i++)
+//       {
+//          s[i] = sizes[i];
+//       }
+//       return s;
+//    }
+// };
+
+
+// /** @brief Wrap a pointer as a DeviceTensor with automatically deduced template
+//     parameters */
+// template <typename T, typename... Dims> MFEM_HOST_DEVICE
+// inline DeviceTensor<sizeof...(Dims),T> Reshape(T *ptr, Dims... dims)
+// {
+//    return DeviceTensor<sizeof...(Dims),T>(ptr, dims...);
+// }
+
+
+// typedef DeviceTensor<1,int> DeviceArray;
+// typedef DeviceTensor<1,const int> ConstDeviceArray;
+
+// typedef DeviceTensor<1,real_t> DeviceVector;
+// typedef DeviceTensor<1,const real_t> ConstDeviceVector;
+
+// typedef DeviceTensor<2,real_t> DeviceMatrix;
+// typedef DeviceTensor<2,const real_t> ConstDeviceMatrix;
+
+// typedef DeviceTensor<3,real_t> DeviceCube;
+// typedef DeviceTensor<3,const real_t> ConstDeviceCube;
+
+// } // mfem namespace
+
+// #endif // MFEM_DTENSOR