Pass entity index to integration kernel for cell-dependent custom data

.github/workflows/fenicsx-refs.env

@@ -3,4 +3,4 @@ basix_ref=main
 ufl_repository=FEniCS/ufl
 ufl_ref=main
 ffcx_repository=FEniCS/ffcx
-ffcx_ref=main
+ffcx_ref=sclaus2/add_voidptr_cast_intrinsic

cpp/dolfinx/fem/Expression.h

@@ -15,6 +15,7 @@
 #include <dolfinx/mesh/Mesh.h>
 #include <functional>
 #include <memory>
+#include <optional>
 #include <span>
 #include <utility>
 #include <vector>
@@ -66,7 +67,9 @@
   /// computed a 1-form expression, e.g. can be used to create a matrix
   /// that when applied to a degree-of-freedom vector gives the
   /// expression values at the evaluation points.
+  /// @param[in] custom_data Optional custom user data pointer passed to
+  /// the kernel function.
   Expression(const std::vector<std::shared_ptr<
                  const Function<scalar_type, geometry_type>>>& coefficients,
              const std::vector<std::shared_ptr<const Constant<scalar_type>>>&
                  constants,
@@ -78,11 +81,12 @@
                  fn,
              const std::vector<std::size_t>& value_shape,
              std::shared_ptr<const FunctionSpace<geometry_type>> argument_space
-             = nullptr)
+             = nullptr,
+             std::optional<void*> custom_data = std::nullopt)
       : _argument_space(argument_space), _coefficients(coefficients),
         _constants(constants), _fn(fn), _value_shape(value_shape),
-        _x_ref(std::vector<geometry_type>(X.begin(), X.end()), Xshape)
-
+        _x_ref(std::vector<geometry_type>(X.begin(), X.end()), Xshape),
+        _custom_data(custom_data)
   {
     for (auto& c : _coefficients)
     {
@@ -152,6 +156,13 @@
     return _fn;
   }
 
+  /// @brief Get the custom data pointer for the Expression.
+  ///
+  /// The custom data pointer is passed to the kernel function during
+  /// expression tabulation.
+  /// @return Custom data pointer, or std::nullopt if not set.
+  std::optional<void*> custom_data() const { return _custom_data; }
+
   /// @brief Value size of the Expression result.
   int value_size() const
   {
@@ -189,5 +200,8 @@
 
   // Evaluation points on reference cell
   std::pair<std::vector<geometry_type>, std::array<std::size_t, 2>> _x_ref;
+
+  // Custom user data pointer passed to the kernel function
+  std::optional<void*> _custom_data = std::nullopt;
 };
 } // namespace dolfinx::fem

cpp/dolfinx/fem/Form.h

@@ -55,6 +55,8 @@
   /// @param[in] entities Indices of entities to integrate over.
   /// @param[in] coeffs Indices of the coefficients that are present
   /// (active) in `kernel`.
+  /// @param[in] custom_data Optional custom user data pointer passed to
+  /// the kernel function.
   template <typename K, typename V, typename W>
     requires std::is_convertible_v<
                  std::remove_cvref_t<K>,
@@ -64,9 +66,10 @@
                                            std::vector<std::int32_t>>
                  and std::is_convertible_v<std::remove_cvref_t<W>,
                                            std::vector<int>>
-  integral_data(K&& kernel, V&& entities, W&& coeffs)
+  integral_data(K&& kernel, V&& entities, W&& coeffs,
+                std::optional<void*> custom_data = std::nullopt)
       : kernel(std::forward<K>(kernel)), entities(std::forward<V>(entities)),
-        coeffs(std::forward<W>(coeffs))
+        coeffs(std::forward<W>(coeffs)), custom_data(custom_data)
   {
   }
 
@@ -82,6 +85,11 @@
   /// @brief Indices of coefficients (from the form) that are in this
   /// integral.
   std::vector<int> coeffs;
+
+  /// @brief Custom user data pointer passed to the kernel function.
+  /// This can be used to pass runtime-computed data (e.g., per-cell
+  /// quadrature rules, material properties) to the kernel.
+  std::optional<void*> custom_data = std::nullopt;
 };
 
 /// @brief A representation of finite element variational forms.
@@ -391,6 +399,30 @@
     return it->second.kernel;
   }
 
+  /// @brief Get the custom data pointer for an integral.
+  ///
+  /// The custom data pointer is passed to the kernel function during
+  /// assembly. This can be used to pass runtime-computed data to
+  /// kernels (e.g., per-cell quadrature rules, material properties).
+  ///
+  /// @warning Void pointers are inherently unsafe and cannot be type or
+  /// bounds checked. Incorrect usage of this feature can and will lead
+  /// to undefined behaviour and crashes.
+  ///
+  /// @param[in] type Integral type.
+  /// @param[in] id Integral subdomain ID.
+  /// @param[in] kernel_idx Index of the kernel (we may have multiple
+  /// kernels for a given ID in mixed-topology meshes).
+  /// @return Custom data pointer for the integral, or std::nullopt if not set.
+  std::optional<void*> custom_data(IntegralType type, int id,
+                                   int kernel_idx) const
+  {
+    auto it = _integrals.find({type, id, kernel_idx});
+    if (it == _integrals.end())
+      throw std::runtime_error("Requested integral not found.");
+    return it->second.custom_data;
+  }
+
   /// @brief Get types of integrals in the form.
   /// @return Integrals types.
   std::set<IntegralType> integral_types() const

cpp/dolfinx/fem/assemble_expression_impl.h

@@ -11,12 +11,14 @@
 #include "traits.h"
 #include "utils.h"
 #include <algorithm>
+#include <array>
 #include <basix/mdspan.hpp>
 #include <dolfinx/common/IndexMap.h>
 #include <dolfinx/mesh/Geometry.h>
 #include <dolfinx/mesh/Mesh.h>
 #include <dolfinx/mesh/Topology.h>
 #include <memory>
+#include <optional>
 #include <vector>
 
 namespace dolfinx::fem::impl
@@ -58,6 +60,8 @@
 /// @param[in] P0 Degree-of-freedom transformation function. Applied when
 /// expressions includes an argument function that requires a
 /// transformation.
+/// @param[in] custom_data Optional pointer to user-supplied data passed
+/// to the kernel.
 template <dolfinx::scalar T, std::floating_point U>
 void tabulate_expression(
     std::span<T> values, fem::FEkernel<T> auto fn,
@@ -68,7 +72,8 @@
     md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
     std::span<const T> constants, fem::MDSpan2 auto entities,
     std::span<const std::uint32_t> cell_info,
-    fem::DofTransformKernel<T> auto P0)
+    fem::DofTransformKernel<T> auto P0,
+    std::optional<void*> custom_data = std::nullopt)
 {
   static_assert(entities.rank() == 1 or entities.rank() == 2);
 
@@ -91,8 +96,10 @@
         std::copy_n(std::next(x.begin(), 3 * x_dofs[i]), 3,
                     std::next(coord_dofs.begin(), 3 * i));
       }
+      std::int32_t entity_local_index = static_cast<std::int32_t>(e);
       fn(values_local.data(), &coeffs(e, 0), constants.data(),
-         coord_dofs.data(), nullptr, nullptr, nullptr);
+         coord_dofs.data(), &entity_local_index, nullptr,
+         custom_data.value_or(nullptr));
 
       P0(values_local, cell_info, entity, size0);
     }
@@ -105,8 +112,11 @@
         std::copy_n(std::next(x.begin(), 3 * x_dofs[i]), 3,
                     std::next(coord_dofs.begin(), 3 * i));
       }
+      std::array<std::int32_t, 2> entity_local_index{
+          entities(e, 1), static_cast<std::int32_t>(e)};
       fn(values_local.data(), &coeffs(e, 0), constants.data(),
-         coord_dofs.data(), &entities(e, 1), nullptr, nullptr);
+         coord_dofs.data(), entity_local_index.data(), nullptr,
+         custom_data.value_or(nullptr));
 
       P0(values_local, cell_info, entity, size0);
     }
@@ -148,6 +158,8 @@
 /// Used to computed a 1-form expression, e.g. can be used to create a
 /// matrix that when applied to a degree-of-freedom vector gives the
 /// expression values at the evaluation points.
+/// @param[in] custom_data Optional pointer to user-supplied data passed
+/// to the kernel.
 template <dolfinx::scalar T, std::floating_point U>
 void tabulate_expression(
     std::span<T> values, fem::FEkernel<T> auto fn,
@@ -157,7 +169,8 @@
     fem::MDSpan2 auto entities,
     std::optional<
         std::pair<std::reference_wrapper<const FiniteElement<U>>, std::size_t>>
-        element)
+        element,
+    std::optional<void*> custom_data = std::nullopt)
 {
   std::function<void(std::span<T>, std::span<const std::uint32_t>, std::int32_t,
                      int)>
@@ -187,6 +200,6 @@
   tabulate_expression<T, U>(values, fn, Xshape, value_size, num_argument_dofs,
                             mesh.geometry().dofmap(), mesh.geometry().x(),
                             coeffs, constants, entities, cell_info,
-                            post_dof_transform);
+                            post_dof_transform, custom_data);
 }
 } // namespace dolfinx::fem::impl

cpp/dolfinx/fem/assemble_matrix_impl.h

@@ -12,12 +12,14 @@
 #include "traits.h"
 #include "utils.h"
 #include <algorithm>
+#include <array>
 #include <dolfinx/la/utils.h>
 #include <dolfinx/mesh/Geometry.h>
 #include <dolfinx/mesh/Mesh.h>
 #include <dolfinx/mesh/Topology.h>
 #include <functional>
 #include <iterator>
+#include <optional>
 #include <span>
 #include <tuple>
 #include <vector>
@@ -62,6 +64,7 @@
 /// function mesh.
 /// @param cell_info1 Cell permutation information for the trial
 /// function mesh.
+/// @param custom_data Custom user data pointer passed to the kernel.
 template <dolfinx::scalar T, bool LiftingMode = false>
 void assemble_cells_matrix(
     la::MatSet<T> auto mat_set, mdspan2_t x_dofmap,
@@ -76,7 +79,8 @@
     std::span<const std::int8_t> bc1, FEkernel<T> auto kernel,
     md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
     std::span<const T> constants, std::span<const std::uint32_t> cell_info0,
-    std::span<const std::uint32_t> cell_info1)
+    std::span<const std::uint32_t> cell_info1,
+    std::optional<void*> custom_data = std::nullopt)
 {
   if (cells.empty())
     return;
@@ -133,8 +137,9 @@
 
     // Tabulate tensor
     std::ranges::fill(Ae, 0);
-    kernel(Ae.data(), &coeffs(c, 0), constants.data(), cdofs.data(), nullptr,
-           nullptr, nullptr);
+    std::int32_t entity_local_index = static_cast<std::int32_t>(c);
+    kernel(Ae.data(), &coeffs(c, 0), constants.data(), cdofs.data(),
+           &entity_local_index, nullptr, custom_data.value_or(nullptr));
 
     // Compute A = P_0 \tilde{A} P_1^T (dof transformation)
     P0(Ae, cell_info0, cell0, ndim1);  // B = P0 \tilde{A}
@@ -226,6 +231,7 @@
 /// function mesh.
 /// @param[in] perms Entity permutation integer. Empty if entity
 /// permutations are not required.
+/// @param custom_data Custom user data pointer passed to the kernel.
 template <dolfinx::scalar T, bool LiftingMode = false>
 void assemble_entities(
     la::MatSet<T> auto mat_set, mdspan2_t x_dofmap,
@@ -249,7 +255,8 @@
     md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
     std::span<const T> constants, std::span<const std::uint32_t> cell_info0,
     std::span<const std::uint32_t> cell_info1,
-    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms)
+    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms,
+    std::optional<void*> custom_data = std::nullopt)
 {
   if (entities.empty())
     return;
@@ -309,8 +316,10 @@
 
     // Tabulate tensor
     std::ranges::fill(Ae, 0);
+    std::array<std::int32_t, 2> entity_local_index{
+        local_entity, static_cast<std::int32_t>(f)};
     kernel(Ae.data(), &coeffs(f, 0), constants.data(), cdofs.data(),
-           &local_entity, &perm, nullptr);
+           entity_local_index.data(), &perm, custom_data.value_or(nullptr));
     P0(Ae, cell_info0, cell0, ndim1);
     P1T(Ae, cell_info1, cell1, ndim0);
 
@@ -393,6 +402,8 @@
 /// function mesh.
 /// @param[in] perms Facet permutation integer. Empty if facet
 /// permutations are not required.
+/// @param[in] custom_data Optional pointer to user-supplied data passed
+/// to the kernel at runtime.
 template <dolfinx::scalar T, bool LiftingMode = false>
 void assemble_interior_facets(
     la::MatSet<T> auto mat_set, mdspan2_t x_dofmap,
@@ -418,7 +429,8 @@
         coeffs,
     std::span<const T> constants, std::span<const std::uint32_t> cell_info0,
     std::span<const std::uint32_t> cell_info1,
-    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms)
+    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms,
+    std::optional<void*> custom_data = std::nullopt)
 {
   if (facets.empty())
     return;
@@ -514,8 +526,10 @@
                           ? std::array<std::uint8_t, 2>{0, 0}
                           : std::array{perms(cells[0], local_facet[0]),
                                        perms(cells[1], local_facet[1])};
+    std::array<std::int32_t, 3> entity_local_index{
+        local_facet[0], local_facet[1], static_cast<std::int32_t>(f)};
     kernel(Ae.data(), &coeffs(f, 0, 0), constants.data(), cdofs.data(),
-           local_facet.data(), perm.data(), nullptr);
+           entity_local_index.data(), perm.data(), custom_data.value_or(nullptr));
 
     // Local element layout is a 2x2 block matrix with structure
     //
@@ -686,12 +700,14 @@
       std::span cells0 = a.domain_arg(IntegralType::cell, 0, i, cell_type_idx);
       std::span cells1 = a.domain_arg(IntegralType::cell, 1, i, cell_type_idx);
       auto& [coeffs, cstride] = coefficients.at({IntegralType::cell, i});
+      std::optional<void*> custom_data
+          = a.custom_data(IntegralType::cell, i, cell_type_idx);
       assert(cells.size() * cstride == coeffs.size());
       impl::assemble_cells_matrix<T, LiftingMode>(
           mat_set, x_dofmap, x, cells, {dofs0, bs0, cells0}, P0,
           {dofs1, bs1, cells1}, P1T, bc0, bc1, fn,
           md::mdspan(coeffs.data(), cells.size(), cstride), constants,
-          cell_info0, cell_info1);
+          cell_info0, cell_info1, custom_data);
     }
 
     md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> facet_perms;
@@ -727,6 +743,8 @@
       assert(fn);
       auto& [coeffs, cstride]
           = coefficients.at({IntegralType::interior_facet, i});
+      std::optional<void*> custom_data
+          = a.custom_data(IntegralType::interior_facet, i, 0);
 
       std::span facets = a.domain(IntegralType::interior_facet, i, 0);
       std::span facets0 = a.domain_arg(IntegralType::interior_facet, 0, i, 0);
@@ -742,7 +760,7 @@
            mdspanx22_t(facets1.data(), facets1.size() / 4, 2, 2)},
           P1T, bc0, bc1, fn,
           mdspanx2x_t(coeffs.data(), facets.size() / 4, 2, cstride), constants,
-          cell_info0, cell_info1, facet_perms);
+          cell_info0, cell_info1, facet_perms, custom_data);
     }
 
     for (auto itg_type : {fem::IntegralType::exterior_facet,
@@ -769,6 +787,7 @@
         auto fn = a.kernel(itg_type, i, 0);
         assert(fn);
         auto& [coeffs, cstride] = coefficients.at({itg_type, i});
+        std::optional<void*> custom_data = a.custom_data(itg_type, i, 0);
 
         std::span e = a.domain(itg_type, i, 0);
         mdspanx2_t entities(e.data(), e.size() / 2, 2);
@@ -781,7 +800,7 @@
             mat_set, x_dofmap, x, entities, {dofs0, bs0, entities0}, P0,
             {dofs1, bs1, entities1}, P1T, bc0, bc1, fn,
             md::mdspan(coeffs.data(), entities.extent(0), cstride), constants,
-            cell_info0, cell_info1, perms);
+            cell_info0, cell_info1, perms, custom_data);
       }
     }
   }