Pass custom_data through assemblers to kernel functions.

Author: sclaus2

cpp/dolfinx/fem/Form.h

@@ -55,6 +55,8 @@ struct integral_data
   /// @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 @@ struct integral_data
                                            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,
+                void* custom_data = nullptr)
       : 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 @@ struct integral_data
   /// @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.
+  void* custom_data = nullptr;
 };
 
 /// @brief A representation of finite element variational forms.
@@ -391,6 +399,39 @@ class Form
     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).
+  ///
+  /// @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 nullptr if not set.
+  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 Set the custom data pointer for an integral.
+  ///
+  /// @param[in] type Integral type.
+  /// @param[in] id Integral subdomain ID.
+  /// @param[in] kernel_idx Index of the kernel.
+  /// @param[in] data Custom data pointer to set.
+  void set_custom_data(IntegralType type, int id, int kernel_idx, void* data)
+  {
+    auto it = _integrals.find({type, id, kernel_idx});
+    if (it == _integrals.end())
+      throw std::runtime_error("Requested integral not found.");
+    it->second.custom_data = data;
+  }
+
   /// @brief Get types of integrals in the form.
   /// @return Integrals types.
   std::set<IntegralType> integral_types() const

cpp/dolfinx/fem/assemble_matrix_impl.h

@@ -60,6 +60,7 @@ using mdspan2_t = md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>>;
 /// 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>
 void assemble_cells_matrix(
     la::MatSet<T> auto mat_set, mdspan2_t x_dofmap,
@@ -74,7 +75,7 @@ void assemble_cells_matrix(
     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, void* custom_data = nullptr)
 {
   if (cells.empty())
     return;
@@ -109,7 +110,7 @@ void assemble_cells_matrix(
     // Tabulate tensor
     std::ranges::fill(Ae, 0);
     kernel(Ae.data(), &coeffs(c, 0), constants.data(), cdofs.data(), nullptr,
-           nullptr, nullptr);
+           nullptr, custom_data);
 
     // Compute A = P_0 \tilde{A} P_1^T (dof transformation)
     P0(Ae, cell_info0, cell0, ndim1);  // B = P0 \tilde{A}
@@ -198,6 +199,7 @@ void assemble_cells_matrix(
 /// 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>
 void assemble_entities(
     la::MatSet<T> auto mat_set, mdspan2_t x_dofmap,
@@ -221,7 +223,8 @@ void assemble_entities(
     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,
+    void* custom_data = nullptr)
 {
   if (entities.empty())
     return;
@@ -259,7 +262,7 @@ void assemble_entities(
     // Tabulate tensor
     std::ranges::fill(Ae, 0);
     kernel(Ae.data(), &coeffs(f, 0), constants.data(), cdofs.data(),
-           &local_entity, &perm, nullptr);
+           &local_entity, &perm, custom_data);
     P0(Ae, cell_info0, cell0, ndim1);
     P1T(Ae, cell_info1, cell1, ndim0);
 
@@ -363,7 +366,8 @@ void assemble_interior_facets(
         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,
+    void* custom_data = nullptr)
 {
   if (facets.empty())
     return;
@@ -440,7 +444,7 @@ void assemble_interior_facets(
                           : std::array{perms(cells[0], local_facet[0]),
                                        perms(cells[1], local_facet[1])};
     kernel(Ae.data(), &coeffs(f, 0, 0), constants.data(), cdofs.data(),
-           local_facet.data(), perm.data(), nullptr);
+           local_facet.data(), perm.data(), custom_data);
 
     // Local element layout is a 2x2 block matrix with structure
     //
@@ -605,12 +609,13 @@ void assemble_matrix(
       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});
+      void* custom_data = a.custom_data(IntegralType::cell, i, cell_type_idx);
       assert(cells.size() * cstride == coeffs.size());
       impl::assemble_cells_matrix(
           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;
@@ -646,6 +651,7 @@ void assemble_matrix(
       assert(fn);
       auto& [coeffs, cstride]
           = coefficients.at({IntegralType::interior_facet, i});
+      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);
@@ -661,7 +667,7 @@ void assemble_matrix(
            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,
@@ -688,6 +694,7 @@ void assemble_matrix(
         auto fn = a.kernel(itg_type, i, 0);
         assert(fn);
         auto& [coeffs, cstride] = coefficients.at({itg_type, i});
+        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);
@@ -700,7 +707,7 @@ void assemble_matrix(
             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);
       }
     }
   }

cpp/dolfinx/fem/assemble_scalar_impl.h

@@ -30,7 +30,8 @@ T assemble_cells(mdspan2_t x_dofmap,
                  std::span<const std::int32_t> cells, FEkernel<T> auto fn,
                  std::span<const T> constants,
                  md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
-                 std::span<scalar_value_t<T>> cdofs_b)
+                 std::span<scalar_value_t<T>> cdofs_b,
+                 void* custom_data = nullptr)
 {
   T value(0);
   if (cells.empty())
@@ -49,7 +50,7 @@ T assemble_cells(mdspan2_t x_dofmap,
       std::copy_n(&x(x_dofs[i], 0), 3, std::next(cdofs_b.begin(), 3 * i));
 
     fn(&value, &coeffs(index, 0), constants.data(), cdofs_b.data(), nullptr,
-       nullptr, nullptr);
+       nullptr, custom_data);
   }
 
   return value;
@@ -77,7 +78,7 @@ T assemble_entities(
     FEkernel<T> auto fn, std::span<const T> constants,
     md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
     md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms,
-    std::span<scalar_value_t<T>> cdofs_b)
+    std::span<scalar_value_t<T>> cdofs_b, void* custom_data = nullptr)
 {
   T value(0);
   if (entities.empty())
@@ -99,7 +100,7 @@ T assemble_entities(
     // Permutations
     std::uint8_t perm = perms.empty() ? 0 : perms(cell, local_entity);
     fn(&value, &coeffs(f, 0), constants.data(), cdofs_b.data(), &local_entity,
-       &perm, nullptr);
+       &perm, custom_data);
   }
 
   return value;
@@ -120,7 +121,7 @@ T assemble_interior_facets(
                                     md::dynamic_extent>>
         coeffs,
     md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms,
-    std::span<scalar_value_t<T>> cdofs_b)
+    std::span<scalar_value_t<T>> cdofs_b, void* custom_data = nullptr)
 {
   T value(0);
   if (facets.empty())
@@ -150,7 +151,7 @@ T assemble_interior_facets(
                           : std::array{perms(cells[0], local_facet[0]),
                                        perms(cells[1], local_facet[1])};
     fn(&value, &coeffs(f, 0, 0), constants.data(), cdofs_b.data(),
-       local_facet.data(), perm.data(), nullptr);
+       local_facet.data(), perm.data(), custom_data);
   }
 
   return value;
@@ -178,11 +179,12 @@ T assemble_scalar(
     auto fn = M.kernel(IntegralType::cell, i, 0);
     assert(fn);
     auto& [coeffs, cstride] = coefficients.at({IntegralType::cell, i});
+    void* custom_data = M.custom_data(IntegralType::cell, i, 0);
     std::span<const std::int32_t> cells = M.domain(IntegralType::cell, i, 0);
     assert(cells.size() * cstride == coeffs.size());
     value += impl::assemble_cells(
         x_dofmap, x, cells, fn, constants,
-        md::mdspan(coeffs.data(), cells.size(), cstride), cdofs_b);
+        md::mdspan(coeffs.data(), cells.size(), cstride), cdofs_b, custom_data);
   }
 
   mesh::CellType cell_type = mesh->topology()->cell_type();
@@ -204,6 +206,7 @@ T assemble_scalar(
     assert(fn);
     auto& [coeffs, cstride]
         = coefficients.at({IntegralType::interior_facet, i});
+    void* custom_data = M.custom_data(IntegralType::interior_facet, i, 0);
     std::span facets = M.domain(IntegralType::interior_facet, i, 0);
 
     constexpr std::size_t num_adjacent_cells = 2;
@@ -220,7 +223,7 @@ T assemble_scalar(
         md::mdspan<const T, md::extents<std::size_t, md::dynamic_extent, 2,
                                         md::dynamic_extent>>(
             coeffs.data(), facets.size() / shape1, 2, cstride),
-        facet_perms, cdofs_b);
+        facet_perms, cdofs_b, custom_data);
   }
 
   for (auto itg_type : {fem::IntegralType::exterior_facet,
@@ -236,6 +239,7 @@ T assemble_scalar(
       auto fn = M.kernel(itg_type, i, 0);
       assert(fn);
       auto& [coeffs, cstride] = coefficients.at({itg_type, i});
+      void* custom_data = M.custom_data(itg_type, i, 0);
 
       std::span entities = M.domain(itg_type, i, 0);
 
@@ -248,7 +252,7 @@ T assemble_scalar(
               entities.data(), entities.size() / 2, 2),
           fn, constants,
           md::mdspan(coeffs.data(), entities.size() / 2, cstride), perms,
-          cdofs_b);
+          cdofs_b, custom_data);
     }
   }