diff --git a/.github/workflows/fenicsx-refs.env b/.github/workflows/fenicsx-refs.env index 22c7e44a7e7..693a736666f 100644 --- a/.github/workflows/fenicsx-refs.env +++ b/.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 diff --git a/cpp/dolfinx/fem/Expression.h b/cpp/dolfinx/fem/Expression.h index eb524f0a234..3132598f638 100644 --- a/cpp/dolfinx/fem/Expression.h +++ b/cpp/dolfinx/fem/Expression.h @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -66,6 +67,8 @@ class Expression /// 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>>& coefficients, const std::vector>>& @@ -78,11 +81,12 @@ class Expression fn, const std::vector& value_shape, std::shared_ptr> argument_space - = nullptr) + = nullptr, + std::optional custom_data = std::nullopt) : _argument_space(argument_space), _coefficients(coefficients), _constants(constants), _fn(fn), _value_shape(value_shape), - _x_ref(std::vector(X.begin(), X.end()), Xshape) - + _x_ref(std::vector(X.begin(), X.end()), Xshape), + _custom_data(custom_data) { for (auto& c : _coefficients) { @@ -152,6 +156,13 @@ class Expression 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 custom_data() const { return _custom_data; } + /// @brief Value size of the Expression result. int value_size() const { @@ -189,5 +200,8 @@ class Expression // Evaluation points on reference cell std::pair, std::array> _x_ref; + + // Custom user data pointer passed to the kernel function + std::optional _custom_data = std::nullopt; }; } // namespace dolfinx::fem diff --git a/cpp/dolfinx/fem/Form.h b/cpp/dolfinx/fem/Form.h index 0d3e5b012ae..a73f52fb9cc 100644 --- a/cpp/dolfinx/fem/Form.h +++ b/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 requires std::is_convertible_v< std::remove_cvref_t, @@ -64,9 +66,10 @@ struct integral_data std::vector> and std::is_convertible_v, std::vector> - integral_data(K&& kernel, V&& entities, W&& coeffs) + integral_data(K&& kernel, V&& entities, W&& coeffs, + std::optional custom_data = std::nullopt) : kernel(std::forward(kernel)), entities(std::forward(entities)), - coeffs(std::forward(coeffs)) + coeffs(std::forward(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 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 custom_data = std::nullopt; }; /// @brief A representation of finite element variational forms. @@ -391,6 +399,30 @@ 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). + /// + /// @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 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 integral_types() const diff --git a/cpp/dolfinx/fem/assemble_expression_impl.h b/cpp/dolfinx/fem/assemble_expression_impl.h index 618c37af314..da4192a8477 100644 --- a/cpp/dolfinx/fem/assemble_expression_impl.h +++ b/cpp/dolfinx/fem/assemble_expression_impl.h @@ -11,12 +11,14 @@ #include "traits.h" #include "utils.h" #include +#include #include #include #include #include #include #include +#include #include namespace dolfinx::fem::impl @@ -58,6 +60,8 @@ namespace dolfinx::fem::impl /// @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 void tabulate_expression( std::span values, fem::FEkernel auto fn, @@ -68,7 +72,8 @@ void tabulate_expression( md::mdspan> coeffs, std::span constants, fem::MDSpan2 auto entities, std::span cell_info, - fem::DofTransformKernel auto P0) + fem::DofTransformKernel auto P0, + std::optional custom_data = std::nullopt) { static_assert(entities.rank() == 1 or entities.rank() == 2); @@ -91,8 +96,10 @@ void tabulate_expression( 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(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 @@ void tabulate_expression( std::copy_n(std::next(x.begin(), 3 * x_dofs[i]), 3, std::next(coord_dofs.begin(), 3 * i)); } + std::array entity_local_index{ + entities(e, 1), static_cast(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 @@ void tabulate_expression( /// 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 void tabulate_expression( std::span values, fem::FEkernel auto fn, @@ -157,7 +169,8 @@ void tabulate_expression( fem::MDSpan2 auto entities, std::optional< std::pair>, std::size_t>> - element) + element, + std::optional custom_data = std::nullopt) { std::function, std::span, std::int32_t, int)> @@ -187,6 +200,6 @@ void tabulate_expression( tabulate_expression(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 diff --git a/cpp/dolfinx/fem/assemble_matrix_impl.h b/cpp/dolfinx/fem/assemble_matrix_impl.h index ca887d16d86..5be00942575 100644 --- a/cpp/dolfinx/fem/assemble_matrix_impl.h +++ b/cpp/dolfinx/fem/assemble_matrix_impl.h @@ -12,12 +12,14 @@ #include "traits.h" #include "utils.h" #include +#include #include #include #include #include #include #include +#include #include #include #include @@ -62,6 +64,7 @@ using mdspan2_t = md::mdspan>; /// 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 void assemble_cells_matrix( la::MatSet auto mat_set, mdspan2_t x_dofmap, @@ -76,7 +79,8 @@ void assemble_cells_matrix( std::span bc1, FEkernel auto kernel, md::mdspan> coeffs, std::span constants, std::span cell_info0, - std::span cell_info1) + std::span cell_info1, + std::optional custom_data = std::nullopt) { if (cells.empty()) return; @@ -133,8 +137,9 @@ void assemble_cells_matrix( // 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(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 @@ 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 void assemble_entities( la::MatSet auto mat_set, mdspan2_t x_dofmap, @@ -249,7 +255,8 @@ void assemble_entities( md::mdspan> coeffs, std::span constants, std::span cell_info0, std::span cell_info1, - md::mdspan> perms) + md::mdspan> perms, + std::optional custom_data = std::nullopt) { if (entities.empty()) return; @@ -309,8 +316,10 @@ void assemble_entities( // Tabulate tensor std::ranges::fill(Ae, 0); + std::array entity_local_index{ + local_entity, static_cast(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 @@ void assemble_entities( /// 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 void assemble_interior_facets( la::MatSet auto mat_set, mdspan2_t x_dofmap, @@ -418,7 +429,8 @@ void assemble_interior_facets( coeffs, std::span constants, std::span cell_info0, std::span cell_info1, - md::mdspan> perms) + md::mdspan> perms, + std::optional custom_data = std::nullopt) { if (facets.empty()) return; @@ -514,8 +526,10 @@ void assemble_interior_facets( ? std::array{0, 0} : std::array{perms(cells[0], local_facet[0]), perms(cells[1], local_facet[1])}; + std::array entity_local_index{ + local_facet[0], local_facet[1], static_cast(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 @@ 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}); + std::optional 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> facet_perms; @@ -727,6 +743,8 @@ void assemble_matrix( assert(fn); auto& [coeffs, cstride] = coefficients.at({IntegralType::interior_facet, i}); + std::optional 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 @@ 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, @@ -769,6 +787,7 @@ void assemble_matrix( auto fn = a.kernel(itg_type, i, 0); assert(fn); auto& [coeffs, cstride] = coefficients.at({itg_type, i}); + std::optional 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 @@ 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); } } } diff --git a/cpp/dolfinx/fem/assemble_scalar_impl.h b/cpp/dolfinx/fem/assemble_scalar_impl.h index 37166b46ea6..c609dde7041 100644 --- a/cpp/dolfinx/fem/assemble_scalar_impl.h +++ b/cpp/dolfinx/fem/assemble_scalar_impl.h @@ -11,12 +11,14 @@ #include "FunctionSpace.h" #include "utils.h" #include +#include #include #include #include #include #include #include +#include #include namespace dolfinx::fem::impl @@ -30,7 +32,8 @@ T assemble_cells(mdspan2_t x_dofmap, std::span cells, FEkernel auto fn, std::span constants, md::mdspan> coeffs, - std::span> cdofs_b) + std::span> cdofs_b, + std::optional custom_data = std::nullopt) { T value(0); if (cells.empty()) @@ -48,8 +51,9 @@ T assemble_cells(mdspan2_t x_dofmap, for (std::size_t i = 0; i < x_dofs.size(); ++i) 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); + std::int32_t entity_local_index = static_cast(index); + fn(&value, &coeffs(index, 0), constants.data(), cdofs_b.data(), + &entity_local_index, nullptr, custom_data.value_or(nullptr)); } return value; @@ -77,7 +81,8 @@ T assemble_entities( FEkernel auto fn, std::span constants, md::mdspan> coeffs, md::mdspan> perms, - std::span> cdofs_b) + std::span> cdofs_b, + std::optional custom_data = std::nullopt) { T value(0); if (entities.empty()) @@ -98,8 +103,10 @@ 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); + std::array entity_local_index{ + local_entity, static_cast(f)}; + fn(&value, &coeffs(f, 0), constants.data(), cdofs_b.data(), + entity_local_index.data(), &perm, custom_data.value_or(nullptr)); } return value; @@ -120,7 +127,8 @@ T assemble_interior_facets( md::dynamic_extent>> coeffs, md::mdspan> perms, - std::span> cdofs_b) + std::span> cdofs_b, + std::optional custom_data = std::nullopt) { T value(0); if (facets.empty()) @@ -149,8 +157,10 @@ T assemble_interior_facets( ? std::array{0, 0} : std::array{perms(cells[0], local_facet[0]), perms(cells[1], local_facet[1])}; + std::array entity_local_index{ + local_facet[0], local_facet[1], static_cast(f)}; fn(&value, &coeffs(f, 0, 0), constants.data(), cdofs_b.data(), - local_facet.data(), perm.data(), nullptr); + entity_local_index.data(), perm.data(), custom_data.value_or(nullptr)); } return value; @@ -178,11 +188,12 @@ T assemble_scalar( auto fn = M.kernel(IntegralType::cell, i, 0); assert(fn); auto& [coeffs, cstride] = coefficients.at({IntegralType::cell, i}); + std::optional custom_data = M.custom_data(IntegralType::cell, i, 0); std::span 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 +215,8 @@ T assemble_scalar( assert(fn); auto& [coeffs, cstride] = coefficients.at({IntegralType::interior_facet, i}); + std::optional 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 +233,7 @@ T assemble_scalar( md::mdspan>( 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 +249,7 @@ T assemble_scalar( auto fn = M.kernel(itg_type, i, 0); assert(fn); auto& [coeffs, cstride] = coefficients.at({itg_type, i}); + std::optional custom_data = M.custom_data(itg_type, i, 0); std::span entities = M.domain(itg_type, i, 0); @@ -248,7 +262,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); } } diff --git a/cpp/dolfinx/fem/assemble_vector_impl.h b/cpp/dolfinx/fem/assemble_vector_impl.h index b6f971b2c95..841a75f4cae 100644 --- a/cpp/dolfinx/fem/assemble_vector_impl.h +++ b/cpp/dolfinx/fem/assemble_vector_impl.h @@ -13,6 +13,7 @@ #include "traits.h" #include "utils.h" #include +#include #include #include #include @@ -60,6 +61,8 @@ using mdspan2_t = md::mdspan>; /// coefficient for cell `i`. /// @param[in] cell_info0 Cell permutation information for the test /// function mesh. +/// @param[in] custom_data Optional pointer to user-supplied data passed to +/// the kernel at runtime. template ::value_type> requires std::is_same_v::value_type, T> @@ -73,7 +76,8 @@ void assemble_cells( std::tuple> dofmap, FEkernel auto kernel, std::span constants, md::mdspan> coeffs, - std::span cell_info0) + std::span cell_info0, + std::optional custom_data = std::nullopt) { if (cells.empty()) return; @@ -99,8 +103,9 @@ void assemble_cells( // Tabulate vector for cell std::ranges::fill(be, 0); + std::int32_t entity_local_index = static_cast(index); kernel(be.data(), &coeffs(index, 0), constants.data(), cdofs.data(), - nullptr, nullptr, nullptr); + &entity_local_index, nullptr, custom_data.value_or(nullptr)); P0(be, cell_info0, c0, 1); // Scatter cell vector to 'global' vector array @@ -153,6 +158,8 @@ void assemble_cells( /// function mesh. /// @param[in] perms Entity permutation integer. Empty if entity /// permutations are not required. +/// @param[in] custom_data Optional pointer to user-supplied data passed to +/// the kernel at runtime. template ::value_type> requires std::is_same_v::value_type, T> @@ -171,7 +178,8 @@ void assemble_entities( FEkernel auto kernel, std::span constants, md::mdspan> coeffs, std::span cell_info0, - md::mdspan> perms) + md::mdspan> perms, + std::optional custom_data = std::nullopt) { if (entities.empty()) return; @@ -202,8 +210,10 @@ void assemble_entities( // Tabulate element vector std::ranges::fill(be, 0); + std::array entity_local_index{ + local_entity, static_cast(f)}; kernel(be.data(), &coeffs(f, 0), constants.data(), cdofs.data(), - &local_entity, &perm, nullptr); + entity_local_index.data(), &perm, custom_data.value_or(nullptr)); P0(be, cell_info0, cell0, 1); // Add element vector to global vector @@ -248,6 +258,8 @@ void assemble_entities( /// 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 ::value_type> requires std::is_same_v::value_type, T> @@ -268,7 +280,8 @@ void assemble_interior_facets( md::dynamic_extent>> coeffs, std::span cell_info0, - md::mdspan> perms) + md::mdspan> perms, + std::optional custom_data = std::nullopt) { using X = scalar_value_t; @@ -319,8 +332,10 @@ void assemble_interior_facets( ? std::array{0, 0} : std::array{perms(cells[0], local_facet[0]), perms(cells[1], local_facet[1])}; + std::array entity_local_index{ + local_facet[0], local_facet[1], static_cast(f)}; kernel(be.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)); if (cells0[0] >= 0) P0(be, cell_info0, cells0[0], 1); @@ -610,24 +625,29 @@ void assemble_vector( std::span cells = L.domain(IntegralType::cell, i, cell_type_idx); std::span cells0 = L.domain_arg(IntegralType::cell, 0, i, cell_type_idx); auto& [coeffs, cstride] = coefficients.at({IntegralType::cell, i}); + void* custom_data = L.custom_data(IntegralType::cell, i, cell_type_idx) + .value_or(nullptr); assert(cells.size() * cstride == coeffs.size()); if (bs == 1) { impl::assemble_cells<1>( P0, b, x_dofmap, x, cells, {dofs, bs, cells0}, fn, constants, - md::mdspan(coeffs.data(), cells.size(), cstride), cell_info0); + md::mdspan(coeffs.data(), cells.size(), cstride), cell_info0, + custom_data); } else if (bs == 3) { impl::assemble_cells<3>( P0, b, x_dofmap, x, cells, {dofs, bs, cells0}, fn, constants, - md::mdspan(coeffs.data(), cells.size(), cstride), cell_info0); + md::mdspan(coeffs.data(), cells.size(), cstride), cell_info0, + custom_data); } else { - impl::assemble_cells( - P0, b, x_dofmap, x, cells, {dofs, bs, cells0}, fn, constants, - md::mdspan(coeffs.data(), cells.size(), cstride), cell_info0); + impl::assemble_cells(P0, b, x_dofmap, x, cells, {dofs, bs, cells0}, fn, + constants, + md::mdspan(coeffs.data(), cells.size(), cstride), + cell_info0, custom_data); } } @@ -661,6 +681,8 @@ void assemble_vector( assert(fn); auto& [coeffs, cstride] = coefficients.at({IntegralType::interior_facet, i}); + void* custom_data + = L.custom_data(IntegralType::interior_facet, i, 0).value_or(nullptr); std::span facets = L.domain(IntegralType::interior_facet, i, 0); std::span facets1 = L.domain_arg(IntegralType::interior_facet, 0, i, 0); assert((facets.size() / 4) * 2 * cstride == coeffs.size()); @@ -673,7 +695,7 @@ void assemble_vector( mdspanx22_t(facets1.data(), facets1.size() / 4, 2, 2)}, fn, constants, mdspanx2x_t(coeffs.data(), facets.size() / 4, 2, cstride), - cell_info0, facet_perms); + cell_info0, facet_perms, custom_data); } else if (bs == 3) { @@ -684,7 +706,7 @@ void assemble_vector( mdspanx22_t(facets1.data(), facets1.size() / 4, 2, 2)}, fn, constants, mdspanx2x_t(coeffs.data(), facets.size() / 4, 2, cstride), - cell_info0, facet_perms); + cell_info0, facet_perms, custom_data); } else { @@ -695,7 +717,7 @@ void assemble_vector( mdspanx22_t(facets1.data(), facets1.size() / 4, 2, 2)}, fn, constants, mdspanx2x_t(coeffs.data(), facets.size() / 4, 2, cstride), - cell_info0, facet_perms); + cell_info0, facet_perms, custom_data); } } @@ -712,6 +734,7 @@ void assemble_vector( auto fn = L.kernel(itg_type, i, 0); assert(fn); auto& [coeffs, cstride] = coefficients.at({itg_type, i}); + void* custom_data = L.custom_data(itg_type, i, 0).value_or(nullptr); std::span e = L.domain(itg_type, i, 0); mdspanx2_t entities(e.data(), e.size() / 2, 2); std::span e1 = L.domain_arg(itg_type, 0, i, 0); @@ -722,7 +745,7 @@ void assemble_vector( impl::assemble_entities<1>( P0, b, x_dofmap, x, entities, {dofs, bs, entities1}, fn, constants, md::mdspan(coeffs.data(), entities.extent(0), cstride), - cell_info0, perms); + cell_info0, perms, custom_data); } else if (bs == 3) { @@ -730,7 +753,7 @@ void assemble_vector( P0, b, x_dofmap, x, entities, {dofs, bs, entities1}, fn, constants, md::mdspan(coeffs.data(), entities.size() / 2, cstride), - cell_info0, perms); + cell_info0, perms, custom_data); } else { @@ -738,7 +761,7 @@ void assemble_vector( P0, b, x_dofmap, x, entities, {dofs, bs, entities1}, fn, constants, md::mdspan(coeffs.data(), entities.size() / 2, cstride), - cell_info0, perms); + cell_info0, perms, custom_data); } } } diff --git a/cpp/dolfinx/fem/assembler.h b/cpp/dolfinx/fem/assembler.h index 9ba3b6c886f..786f9911ac4 100644 --- a/cpp/dolfinx/fem/assembler.h +++ b/cpp/dolfinx/fem/assembler.h @@ -72,7 +72,8 @@ void tabulate_expression( { auto [X, Xshape] = e.X(); impl::tabulate_expression(values, e.kernel(), Xshape, e.value_size(), coeffs, - constants, mesh, entities, element); + constants, mesh, entities, element, + e.custom_data()); } /// @brief Evaluate an Expression on cells or facets. diff --git a/cpp/dolfinx/fem/utils.h b/cpp/dolfinx/fem/utils.h index ab5361c85fb..7d152263194 100644 --- a/cpp/dolfinx/fem/utils.h +++ b/cpp/dolfinx/fem/utils.h @@ -953,7 +953,8 @@ Expression create_expression( const ufcx_expression& e, const std::vector>>& coefficients, const std::vector>>& constants, - std::shared_ptr> argument_space = nullptr) + std::shared_ptr> argument_space = nullptr, + std::optional custom_data = std::nullopt) { if (!coefficients.empty()) { @@ -1008,7 +1009,7 @@ Expression create_expression( assert(tabulate_tensor); return Expression(coefficients, constants, std::span(X), Xshape, - tabulate_tensor, value_shape, argument_space); + tabulate_tensor, value_shape, argument_space, custom_data); } /// @brief Create Expression from UFC input (with named coefficients and @@ -1019,7 +1020,8 @@ Expression create_expression( const std::map>>& coefficients, const std::map>>& constants, - std::shared_ptr> argument_space = nullptr) + std::shared_ptr> argument_space = nullptr, + std::optional custom_data = std::nullopt) { // Place coefficients in appropriate order std::vector>> coeff_map; @@ -1057,7 +1059,8 @@ Expression create_expression( } } - return create_expression(e, coeff_map, const_map, argument_space); + return create_expression(e, coeff_map, const_map, argument_space, + custom_data); } } // namespace dolfinx::fem diff --git a/python/demo/demo_static-condensation.py b/python/demo/demo_static-condensation.py index 3b594061bb4..c26fc906a1c 100644 --- a/python/demo/demo_static-condensation.py +++ b/python/demo/demo_static-condensation.py @@ -243,7 +243,7 @@ def tabulate_A(A_, w_, c_, coords_, entity_local_index, permutation=ffi.NULL, cu formtype = form_cpp_class(dtype) # type: ignore cells = np.arange(msh.topology.index_map(msh.topology.dim).size_local) -integrals = {IntegralType.cell: [(0, tabulate_A.address, cells, np.array([], dtype=np.int8))]} +integrals = {IntegralType.cell: [(0, tabulate_A.address, cells, np.array([], dtype=np.int8), None)]} a_cond = Form( formtype([U._cpp_object, U._cpp_object], integrals, [], [], False, [], mesh=msh._cpp_object) ) diff --git a/python/dolfinx/wrappers/dolfinx_wrappers/fem.h b/python/dolfinx/wrappers/dolfinx_wrappers/fem.h index d208ff30b1b..aeb36ec91f1 100644 --- a/python/dolfinx/wrappers/dolfinx_wrappers/fem.h +++ b/python/dolfinx/wrappers/dolfinx_wrappers/fem.h @@ -684,19 +684,33 @@ void declare_objects(nb::module_& m, std::string type) std::uintptr_t fn_addr, const std::vector& value_shape, std::shared_ptr> - argument_space) + argument_space, + std::optional custom_data) { auto tabulate_expression_ptr = (void (*)(T*, const T*, const T*, const typename geom_type::value_type*, const int*, const std::uint8_t*, void*))fn_addr; + std::optional data = std::nullopt; + if (custom_data) + data = reinterpret_cast(*custom_data); new (ex) dolfinx::fem::Expression( coefficients, constants, std::span(X.data(), X.size()), {X.shape(0), X.shape(1)}, tabulate_expression_ptr, value_shape, - argument_space); + argument_space, data); }, nb::arg("coefficients"), nb::arg("constants"), nb::arg("X"), - nb::arg("fn"), nb::arg("value_shape"), nb::arg("argument_space")) + nb::arg("fn"), nb::arg("value_shape"), nb::arg("argument_space"), + nb::arg("custom_data").none() = nb::none()) + .def("custom_data", + [](const dolfinx::fem::Expression& self) + -> std::optional + { + std::optional data = self.custom_data(); + if (!data) + return std::nullopt; + return reinterpret_cast(*data); + }) .def("X", [](const dolfinx::fem::Expression& self) { @@ -725,15 +739,20 @@ void declare_objects(nb::module_& m, std::string type) coefficients, const std::vector>>& constants, - std::shared_ptr> argument_space) + std::shared_ptr> argument_space, + std::optional custom_data) { const ufcx_expression* p = reinterpret_cast(expression); - return dolfinx::fem::create_expression(*p, coefficients, - constants, argument_space); + std::optional data = std::nullopt; + if (custom_data) + data = reinterpret_cast(*custom_data); + return dolfinx::fem::create_expression( + *p, coefficients, constants, argument_space, data); }, nb::arg("expression"), nb::arg("coefficients"), nb::arg("constants"), nb::arg("argument_space").none(), + nb::arg("custom_data").none() = nb::none(), "Create Expression from a pointer to ufc_form."); } @@ -756,8 +775,8 @@ void declare_form(nb::module_& m, std::string type) std::vector, nb::c_contig>, - nb::ndarray, - nb::c_contig>>>>& integrals, + nb::ndarray, nb::c_contig>, + std::optional>>>& integrals, const std::vector>>& coefficients, const std::vector< @@ -773,17 +792,20 @@ void declare_form(nb::module_& m, std::string type) // Loop over kernel for each entity type for (auto& [type, kernels] : integrals) { - for (auto& [id, ptr, e, c] : kernels) + for (auto& [id, ptr, e, c, custom_data] : kernels) { auto kn_ptr = (void (*)(T*, const T*, const T*, const typename geom_type::value_type*, const int*, const std::uint8_t*, void*))ptr; + std::optional data = std::nullopt; + if (custom_data) + data = reinterpret_cast(*custom_data); _integrals.insert( {{type, id, 0}, {kn_ptr, std::vector(e.data(), e.data() + e.size()), - std::vector(c.data(), c.data() + c.size())}}); + std::vector(c.data(), c.data() + c.size()), data}}); } } @@ -847,6 +869,18 @@ void declare_form(nb::module_& m, std::string type) .def_prop_ro("function_spaces", &dolfinx::fem::Form::function_spaces) .def("num_integrals", &dolfinx::fem::Form::num_integrals) + .def( + "custom_data", + [](const dolfinx::fem::Form& self, + dolfinx::fem::IntegralType type, int idx, + int kernel_idx) -> std::optional + { + std::optional data = self.custom_data(type, idx, kernel_idx); + if (!data) + return std::nullopt; + return reinterpret_cast(*data); + }, + nb::arg("type"), nb::arg("idx"), nb::arg("kernel_idx")) .def_prop_ro("integral_types", &dolfinx::fem::Form::integral_types) .def_prop_ro("needs_facet_permutations", &dolfinx::fem::Form::needs_facet_permutations) diff --git a/python/test/unit/fem/test_custom_data.py b/python/test/unit/fem/test_custom_data.py new file mode 100644 index 00000000000..09789669dd8 --- /dev/null +++ b/python/test/unit/fem/test_custom_data.py @@ -0,0 +1,613 @@ +"""Unit tests for custom_data functionality in assembly.""" + +# Copyright (C) 2025 Susanne Claus +# +# This file is part of DOLFINx (https://www.fenicsproject.org) +# +# SPDX-License-Identifier: LGPL-3.0-or-later + +from mpi4py import MPI + +import numpy as np +import pytest + +import ffcx.codegeneration.utils as codegen_utils +from dolfinx import la +from dolfinx.fem import ( + Form, + IntegralType, + assemble_matrix, + assemble_scalar, + assemble_vector, + compute_integration_domains, + form_cpp_class, + functionspace, +) +from dolfinx.mesh import create_unit_square, exterior_facet_indices + +numba = pytest.importorskip("numba") +ufcx_signature = codegen_utils.numba_ufcx_kernel_signature + + +def voidptr_to_dtype_ptr(dtype): + """Return a Numba void pointer caster for the scalar dtype.""" + if np.dtype(dtype) == np.float32: + numba_dtype = numba.types.float32 + elif np.dtype(dtype) == np.float64: + numba_dtype = numba.types.float64 + else: + raise NotImplementedError(f"Unsupported dtype: {dtype}") + return codegen_utils._create_voidptr_to_dtype_ptr_caster(numba_dtype) + + +def tabulate_rank1_with_custom_data(dtype, xdtype): + """Kernel that reads a scaling factor from custom_data. + + Note: custom_data must be set to a valid pointer before assembly. + """ + voidptr_to_scalar_ptr = voidptr_to_dtype_ptr(dtype) + + @numba.cfunc(ufcx_signature(dtype, xdtype), nopython=True) + def tabulate(b_, w_, c_, coords_, local_index, orientation, custom_data): + b = numba.carray(b_, (3), dtype=dtype) + coordinate_dofs = numba.carray(coords_, (3, 3), dtype=xdtype) + + # Cast void* to the scalar pointer type and read the scale value + typed_ptr = voidptr_to_scalar_ptr(custom_data) + scale = typed_ptr[0] + + x0, y0 = coordinate_dofs[0, :2] + x1, y1 = coordinate_dofs[1, :2] + x2, y2 = coordinate_dofs[2, :2] + + # 2x Element area Ae + Ae = abs((x0 - x1) * (y2 - y1) - (y0 - y1) * (x2 - x1)) + b[:] = scale * Ae / 6.0 + + return tabulate + + +def tabulate_rank2_with_custom_data(dtype, xdtype): + """Kernel that reads a scaling factor from custom_data for matrix assembly. + + Note: custom_data must be set to a valid pointer before assembly. + """ + voidptr_to_scalar_ptr = voidptr_to_dtype_ptr(dtype) + + @numba.cfunc(ufcx_signature(dtype, xdtype), nopython=True) + def tabulate(A_, w_, c_, coords_, entity_local_index, cell_orientation, custom_data): + A = numba.carray(A_, (3, 3), dtype=dtype) + coordinate_dofs = numba.carray(coords_, (3, 3), dtype=xdtype) + + # Cast void* to the scalar pointer type and read the scale value + typed_ptr = voidptr_to_scalar_ptr(custom_data) + scale = typed_ptr[0] + + x0, y0 = coordinate_dofs[0, :2] + x1, y1 = coordinate_dofs[1, :2] + x2, y2 = coordinate_dofs[2, :2] + + # 2x Element area Ae + Ae = abs((x0 - x1) * (y2 - y1) - (y0 - y1) * (x2 - x1)) + B = np.array([y1 - y2, y2 - y0, y0 - y1, x2 - x1, x0 - x2, x1 - x0], dtype=dtype).reshape( + 2, 3 + ) + A[:, :] = scale * np.dot(B.T, B) / (2 * Ae) + + return tabulate + + +def tabulate_scalar_with_loop_index(dtype, xdtype, integral_type): + """Kernel that accumulates a custom-data value selected by loop index.""" + voidptr_to_scalar_ptr = voidptr_to_dtype_ptr(dtype) + + @numba.cfunc(ufcx_signature(dtype, xdtype), nopython=True) + def tabulate(A_, w_, c_, coords_, entity_local_index, orientation, custom_data): + A = numba.carray(A_, (1,), dtype=dtype) + data = voidptr_to_scalar_ptr(custom_data) + + if integral_type == 0: + entity = numba.carray(entity_local_index, (1,), dtype=np.int32) + A[0] += data[entity[0]] + elif integral_type == 1: + entity = numba.carray(entity_local_index, (2,), dtype=np.int32) + A[0] += data[entity[1]] + 10 * entity[0] + else: + entity = numba.carray(entity_local_index, (3,), dtype=np.int32) + A[0] += data[entity[2]] + 10 * entity[0] + 100 * entity[1] + + return tabulate + + +def tabulate_vector_with_loop_index(dtype, xdtype, integral_type): + """Kernel that writes a custom-data value selected by loop index.""" + voidptr_to_scalar_ptr = voidptr_to_dtype_ptr(dtype) + + @numba.cfunc(ufcx_signature(dtype, xdtype), nopython=True) + def tabulate(b_, w_, c_, coords_, entity_local_index, orientation, custom_data): + b = numba.carray(b_, (3,), dtype=dtype) + data = voidptr_to_scalar_ptr(custom_data) + b[:] = 0 + + if integral_type == 0: + entity = numba.carray(entity_local_index, (1,), dtype=np.int32) + b[0] = data[entity[0]] + elif integral_type == 1: + entity = numba.carray(entity_local_index, (2,), dtype=np.int32) + b[0] = data[entity[1]] + 10 * entity[0] + else: + entity = numba.carray(entity_local_index, (3,), dtype=np.int32) + b[0] = data[entity[2]] + 10 * entity[0] + 100 * entity[1] + + return tabulate + + +def tabulate_matrix_with_loop_index(dtype, xdtype, integral_type): + """Kernel that writes a custom-data value selected by loop index.""" + voidptr_to_scalar_ptr = voidptr_to_dtype_ptr(dtype) + + @numba.cfunc(ufcx_signature(dtype, xdtype), nopython=True) + def tabulate(A_, w_, c_, coords_, entity_local_index, orientation, custom_data): + A = numba.carray(A_, (3, 3), dtype=dtype) + data = voidptr_to_scalar_ptr(custom_data) + A[:, :] = 0 + + if integral_type == 0: + entity = numba.carray(entity_local_index, (1,), dtype=np.int32) + A[0, 0] = data[entity[0]] + elif integral_type == 1: + entity = numba.carray(entity_local_index, (2,), dtype=np.int32) + A[0, 0] = data[entity[1]] + 10 * entity[0] + else: + entity = numba.carray(entity_local_index, (3,), dtype=np.int32) + A[0, 0] = data[entity[2]] + 10 * entity[0] + 100 * entity[1] + + return tabulate + + +def loop_index_entities(mesh, integral_type): + """Return local integration entities for a loop-index test.""" + tdim = mesh.topology.dim + if integral_type == IntegralType.cell: + num_cells = mesh.topology.index_map(tdim).size_local + return np.arange(num_cells, dtype=np.int32) + + mesh.topology.create_connectivity(tdim - 1, tdim) + exterior_facets = exterior_facet_indices(mesh.topology) + if integral_type == IntegralType.exterior_facet: + return compute_integration_domains(integral_type, mesh.topology, exterior_facets) + + num_facets = mesh.topology.index_map(tdim - 1).size_local + facets = np.arange(num_facets, dtype=np.int32) + interior_facets = np.setdiff1d(facets, exterior_facets).astype(np.int32) + return compute_integration_domains(integral_type, mesh.topology, interior_facets) + + +def expected_loop_index_sum(mesh, integral_type, entities, data): + """Return the expected global sum for the loop-index test kernels.""" + if integral_type == IntegralType.cell: + local_value = np.sum(data) + elif integral_type == IntegralType.exterior_facet: + entity = entities.reshape(-1, 2) + local_value = np.sum(data) + 10 * np.sum(entity[:, 1]) + else: + entity = entities.reshape(-1, 4) + local_value = np.sum(data) + 10 * np.sum(entity[:, 1]) + 100 * np.sum(entity[:, 3]) + + return mesh.comm.allreduce(local_value, op=MPI.SUM) + + +def assemble_loop_index_sum(mesh, V, dtype, integral_type, kernel, entities, data, rank): + """Assemble a custom kernel and return the global sum of assembled entries.""" + active_coeffs = np.array([], dtype=np.int8) + integrals = {integral_type: [(0, kernel.address, entities, active_coeffs, data.ctypes.data)]} + formtype = form_cpp_class(dtype) + + if rank == 0: + F = Form(formtype([], integrals, [], [], False, [], mesh=mesh._cpp_object)) + local_value = assemble_scalar(F) + elif rank == 1: + F = Form(formtype([V._cpp_object], integrals, [], [], False, [], mesh=mesh._cpp_object)) + b = assemble_vector(F) + b.scatter_reverse(la.InsertMode.add) + num_owned = V.dofmap.index_map.size_local * V.dofmap.index_map_bs + local_value = np.sum(b.array[:num_owned]) + else: + F = Form( + formtype( + [V._cpp_object, V._cpp_object], + integrals, + [], + [], + False, + [], + mesh=mesh._cpp_object, + ) + ) + A = assemble_matrix(F) + A.scatter_reverse() + local_value = A.to_scipy().sum() + + return mesh.comm.allreduce(local_value, op=MPI.SUM) + + +@pytest.mark.parametrize("dtype", [np.float64, np.float32]) +def test_custom_data_vector_assembly(dtype): + """Test that custom_data is correctly passed to kernels during vector assembly.""" + xdtype = np.real(dtype(0)).dtype + k1 = tabulate_rank1_with_custom_data(dtype, xdtype) + + mesh = create_unit_square(MPI.COMM_WORLD, 13, 13, dtype=xdtype) + V = functionspace(mesh, ("Lagrange", 1)) + + tdim = mesh.topology.dim + num_cells = mesh.topology.index_map(tdim).size_local + mesh.topology.index_map(tdim).num_ghosts + cells = np.arange(num_cells, dtype=np.int32) + active_coeffs = np.array([], dtype=np.int8) + + # Create custom_data with scale=1.0 + scale_value = np.array([1.0], dtype=dtype) + scale_ptr = scale_value.ctypes.data + + # Pass custom_data at form creation time via the 5th element of the integrals tuple + integrals = {IntegralType.cell: [(0, k1.address, cells, active_coeffs, scale_ptr)]} + formtype = form_cpp_class(dtype) + L = Form(formtype([V._cpp_object], integrals, [], [], False, [], mesh=mesh._cpp_object)) + + # Assemble with scale=1.0 + b1 = assemble_vector(L) + b1.scatter_reverse(la.InsertMode.add) + norm1 = la.norm(b1) + + # Verify we can read back the custom_data pointer + assert L._cpp_object.custom_data(IntegralType.cell, 0, 0) == scale_ptr + + # Update custom_data to scale=2.0 (by modifying the underlying array) + scale_value[0] = 2.0 + b2 = assemble_vector(L) + b2.scatter_reverse(la.InsertMode.add) + norm2 = la.norm(b2) + + # The norm with scale=2 should be 2x the norm with scale=1 + assert np.isclose(norm2, 2.0 * norm1) + + # Test with scale=3.0 + scale_value[0] = 3.0 + b3 = assemble_vector(L) + b3.scatter_reverse(la.InsertMode.add) + norm3 = la.norm(b3) + + assert np.isclose(norm3, 3.0 * norm1) + + +@pytest.mark.parametrize("dtype", [np.float64, np.float32]) +def test_custom_data_matrix_assembly(dtype): + """Test that custom_data is correctly passed to kernels during matrix assembly.""" + xdtype = np.real(dtype(0)).dtype + k2 = tabulate_rank2_with_custom_data(dtype, xdtype) + + mesh = create_unit_square(MPI.COMM_WORLD, 13, 13, dtype=xdtype) + V = functionspace(mesh, ("Lagrange", 1)) + + cells = np.arange(mesh.topology.index_map(mesh.topology.dim).size_local, dtype=np.int32) + active_coeffs = np.array([], dtype=np.int8) + + # Create custom_data with scale=1.0 + scale_value = np.array([1.0], dtype=dtype) + + # Pass custom_data at form creation time via the 5th element + integrals = { + IntegralType.cell: [(0, k2.address, cells, active_coeffs, scale_value.ctypes.data)] + } + formtype = form_cpp_class(dtype) + a = Form( + formtype( + [V._cpp_object, V._cpp_object], + integrals, + [], + [], + False, + [], + mesh=mesh._cpp_object, + ) + ) + + # Assemble with scale=1.0 + A1 = assemble_matrix(a) + A1.scatter_reverse() + norm1 = np.sqrt(A1.squared_norm()) + + # Update custom_data to scale=2.0 (by modifying the underlying array) + scale_value[0] = 2.0 + A2 = assemble_matrix(a) + A2.scatter_reverse() + norm2 = np.sqrt(A2.squared_norm()) + + # The norm with scale=2 should be 2x the norm with scale=1 + assert np.isclose(norm2, 2.0 * norm1) + + +@pytest.mark.parametrize("dtype", [np.float64, np.float32]) +@pytest.mark.parametrize( + "integral_type, integral_code", + [ + (IntegralType.cell, 0), + (IntegralType.exterior_facet, 1), + (IntegralType.interior_facet, 2), + ], +) +@pytest.mark.parametrize("rank", [0, 1, 2]) +def test_custom_data_loop_index_all_assemblers(dtype, integral_type, integral_code, rank): + """Test that entity_local_index carries loop indices in all assemblers. + + The test kernels use the loop index to select per-entity custom data. For + facet kernels they also check that the existing local facet/entity entries + remain available before the appended loop index. + """ + xdtype = np.real(dtype(0)).dtype + mesh = create_unit_square(MPI.COMM_WORLD, 4, 3, dtype=xdtype) + V = functionspace(mesh, ("Lagrange", 1)) + entities = loop_index_entities(mesh, integral_type) + num_entities = entities.size if integral_type == IntegralType.cell else entities.size // 2 + if integral_type == IntegralType.interior_facet: + num_entities = entities.size // 4 + + data = np.arange(1, num_entities + 1, dtype=dtype) + expected = expected_loop_index_sum(mesh, integral_type, entities, data) + + if rank == 0: + kernel = tabulate_scalar_with_loop_index(dtype, xdtype, integral_code) + elif rank == 1: + kernel = tabulate_vector_with_loop_index(dtype, xdtype, integral_code) + else: + kernel = tabulate_matrix_with_loop_index(dtype, xdtype, integral_code) + + actual = assemble_loop_index_sum(mesh, V, dtype, integral_type, kernel, entities, data, rank) + assert np.isclose(actual, expected) + + +@pytest.mark.parametrize("dtype", [np.float64, np.float32]) +def test_custom_data_default_nullptr(dtype): + """Test that custom_data defaults to nullptr (None) when not provided.""" + xdtype = np.real(dtype(0)).dtype + + # Define a simple kernel that doesn't use custom_data + @numba.cfunc(ufcx_signature(dtype, xdtype), nopython=True) + def tabulate_simple(b_, w_, c_, coords_, local_index, orientation, custom_data): + b = numba.carray(b_, (3), dtype=dtype) + coordinate_dofs = numba.carray(coords_, (3, 3), dtype=xdtype) + + x0, y0 = coordinate_dofs[0, :2] + x1, y1 = coordinate_dofs[1, :2] + x2, y2 = coordinate_dofs[2, :2] + + Ae = abs((x0 - x1) * (y2 - y1) - (y0 - y1) * (x2 - x1)) + b[:] = Ae / 6.0 + + mesh = create_unit_square(MPI.COMM_WORLD, 5, 5, dtype=xdtype) + V = functionspace(mesh, ("Lagrange", 1)) + + tdim = mesh.topology.dim + num_cells = mesh.topology.index_map(tdim).size_local + mesh.topology.index_map(tdim).num_ghosts + cells = np.arange(num_cells, dtype=np.int32) + active_coeffs = np.array([], dtype=np.int8) + + # Pass None as custom_data (5th element) to get std::nullopt + integrals = {IntegralType.cell: [(0, tabulate_simple.address, cells, active_coeffs, None)]} + formtype = form_cpp_class(dtype) + L = Form(formtype([V._cpp_object], integrals, [], [], False, [], mesh=mesh._cpp_object)) + + # custom_data should be None (std::nullopt) when passed as None + assert L._cpp_object.custom_data(IntegralType.cell, 0, 0) is None + + +@pytest.mark.parametrize("dtype", [np.float64, np.float32]) +def test_custom_data_struct(dtype): + """Test passing a struct with multiple values through custom_data.""" + xdtype = np.real(dtype(0)).dtype + + # Define a kernel that reads two values from custom_data + voidptr_to_scalar_ptr = voidptr_to_dtype_ptr(dtype) + + @numba.cfunc(ufcx_signature(dtype, xdtype), nopython=True) + def tabulate_with_struct(b_, w_, c_, coords_, local_index, orientation, custom_data): + b = numba.carray(b_, (3), dtype=dtype) + coordinate_dofs = numba.carray(coords_, (3, 3), dtype=xdtype) + + # Cast void* to the scalar pointer type and read two values: [scale, offset] + typed_ptr = voidptr_to_scalar_ptr(custom_data) + scale = typed_ptr[0] + offset = typed_ptr[1] + + x0, y0 = coordinate_dofs[0, :2] + x1, y1 = coordinate_dofs[1, :2] + x2, y2 = coordinate_dofs[2, :2] + + Ae = abs((x0 - x1) * (y2 - y1) - (y0 - y1) * (x2 - x1)) + b[:] = scale * Ae / 6.0 + offset + + mesh = create_unit_square(MPI.COMM_WORLD, 5, 5, dtype=xdtype) + V = functionspace(mesh, ("Lagrange", 1)) + + tdim = mesh.topology.dim + # Only iterate over local cells (not ghosts) to avoid double-counting + # contributions after scatter_reverse + num_local_cells = mesh.topology.index_map(tdim).size_local + cells = np.arange(num_local_cells, dtype=np.int32) + active_coeffs = np.array([], dtype=np.int8) + + # Test 1: scale=1.0, offset=0.0 (baseline) + struct_data = np.array([1.0, 0.0], dtype=dtype) + + # Pass custom_data at form creation time + integrals = { + IntegralType.cell: [ + (0, tabulate_with_struct.address, cells, active_coeffs, struct_data.ctypes.data) + ] + } + formtype = form_cpp_class(dtype) + L = Form(formtype([V._cpp_object], integrals, [], [], False, [], mesh=mesh._cpp_object)) + + b_baseline = assemble_vector(L) + b_baseline.scatter_reverse(la.InsertMode.add) + norm_baseline = la.norm(b_baseline) + + # Test 2: scale=2.0, offset=0.0 - should double the norm + struct_data[0] = 2.0 + struct_data[1] = 0.0 + b_scaled = assemble_vector(L) + b_scaled.scatter_reverse(la.InsertMode.add) + norm_scaled = la.norm(b_scaled) + assert np.isclose(norm_scaled, 2.0 * norm_baseline) + + # Test 3: scale=0.0, offset=1.0 - pure offset contribution + struct_data[0] = 0.0 + struct_data[1] = 1.0 + b_offset = assemble_vector(L) + b_offset.scatter_reverse(la.InsertMode.add) + # With offset=1.0, each DOF gets contribution from each cell it touches + # Interior nodes touch 6 cells, edge nodes touch 3-4, corner nodes touch 1-2 + # The sum of all contributions equals 3 * num_local_cells (3 DOFs per cell, offset=1.0 each) + # Sum only local DOFs and gather across processes + local_sum = np.sum(b_offset.array[: V.dofmap.index_map.size_local * V.dofmap.index_map_bs]) + total_contribution = mesh.comm.allreduce(local_sum, op=MPI.SUM) + total_cells = mesh.comm.allreduce(num_local_cells, op=MPI.SUM) + assert np.isclose(total_contribution, 3.0 * total_cells) + + +@pytest.mark.parametrize("dtype", [np.float64, np.float32]) +def test_custom_data_multiple_parameters(dtype): + """Test custom_data with multiple parameters demonstrating complex data passing. + + This test shows how to pass multiple values through custom_data, simulating + the use case of passing runtime-computed parameters like material properties + or integration parameters. + """ + xdtype = np.real(dtype(0)).dtype + + # Kernel that uses three parameters: coefficient, exponent base, and additive term + # Computes: coeff * base^area + additive + voidptr_to_scalar_ptr = voidptr_to_dtype_ptr(dtype) + + @numba.cfunc(ufcx_signature(dtype, xdtype), nopython=True) + def tabulate_with_params(b_, w_, c_, coords_, local_index, orientation, custom_data): + b = numba.carray(b_, (3), dtype=dtype) + coordinate_dofs = numba.carray(coords_, (3, 3), dtype=xdtype) + + # Read three parameters from custom_data + typed_ptr = voidptr_to_scalar_ptr(custom_data) + coeff = typed_ptr[0] + power = typed_ptr[1] + additive = typed_ptr[2] + + x0, y0 = coordinate_dofs[0, :2] + x1, y1 = coordinate_dofs[1, :2] + x2, y2 = coordinate_dofs[2, :2] + + Ae = abs((x0 - x1) * (y2 - y1) - (y0 - y1) * (x2 - x1)) + # Use power as a simple multiplier (avoiding actual power function complexity) + val = coeff * power * Ae / 6.0 + additive + b[:] = val + + mesh = create_unit_square(MPI.COMM_WORLD, 4, 4, dtype=xdtype) + V = functionspace(mesh, ("Lagrange", 1)) + + tdim = mesh.topology.dim + num_cells = mesh.topology.index_map(tdim).size_local + mesh.topology.index_map(tdim).num_ghosts + cells = np.arange(num_cells, dtype=np.int32) + active_coeffs = np.array([], dtype=np.int8) + + # Test with specific parameters: coeff=2, power=3, additive=0 + params = np.array([2.0, 3.0, 0.0], dtype=dtype) + + # Pass custom_data at form creation time + integrals = { + IntegralType.cell: [ + (0, tabulate_with_params.address, cells, active_coeffs, params.ctypes.data) + ] + } + formtype = form_cpp_class(dtype) + L = Form(formtype([V._cpp_object], integrals, [], [], False, [], mesh=mesh._cpp_object)) + + b1 = assemble_vector(L) + b1.scatter_reverse(la.InsertMode.add) + norm1 = la.norm(b1) + + # Change parameters: coeff=1, power=6, additive=0 (should give same result: 1*6 = 2*3) + params[0] = 1.0 + params[1] = 6.0 + b2 = assemble_vector(L) + b2.scatter_reverse(la.InsertMode.add) + norm2 = la.norm(b2) + + assert np.isclose(norm1, norm2) + + +@pytest.mark.parametrize("dtype", [np.float64, np.float32]) +def test_custom_data_global_parameter_update(dtype): + """Test updating custom_data between assemblies for parameter studies. + + This demonstrates a key use case: running multiple assemblies with + different parameter values without recreating the form. The custom_data + points to a parameter that can be modified between assembly calls. + """ + xdtype = np.real(dtype(0)).dtype + + # Kernel that reads a diffusion coefficient from custom_data + voidptr_to_scalar_ptr = voidptr_to_dtype_ptr(dtype) + + @numba.cfunc(ufcx_signature(dtype, xdtype), nopython=True) + def tabulate_diffusion(b_, w_, c_, coords_, local_index, orientation, custom_data): + b = numba.carray(b_, (3), dtype=dtype) + coordinate_dofs = numba.carray(coords_, (3, 3), dtype=xdtype) + + # Read diffusion coefficient from custom_data + typed_ptr = voidptr_to_scalar_ptr(custom_data) + kappa = typed_ptr[0] # Diffusion coefficient + + x0, y0 = coordinate_dofs[0, :2] + x1, y1 = coordinate_dofs[1, :2] + x2, y2 = coordinate_dofs[2, :2] + + # 2x Element area Ae + Ae = abs((x0 - x1) * (y2 - y1) - (y0 - y1) * (x2 - x1)) + # Simple load vector scaled by diffusion coefficient + b[:] = kappa * Ae / 6.0 + + mesh = create_unit_square(MPI.COMM_WORLD, 8, 8, dtype=xdtype) + V = functionspace(mesh, ("Lagrange", 1)) + + tdim = mesh.topology.dim + num_cells = mesh.topology.index_map(tdim).size_local + mesh.topology.index_map(tdim).num_ghosts + cells = np.arange(num_cells, dtype=np.int32) + active_coeffs = np.array([], dtype=np.int8) + + # Parameter array - this can be updated between assemblies + kappa = np.array([1.0], dtype=dtype) + + integrals = { + IntegralType.cell: [ + (0, tabulate_diffusion.address, cells, active_coeffs, kappa.ctypes.data) + ] + } + formtype = form_cpp_class(dtype) + L = Form(formtype([V._cpp_object], integrals, [], [], False, [], mesh=mesh._cpp_object)) + + # Store results for different kappa values + results = [] + kappa_values = [0.1, 1.0, 10.0, 100.0] + + for k in kappa_values: + kappa[0] = k + b = assemble_vector(L) + b.scatter_reverse(la.InsertMode.add) + results.append(la.norm(b)) + + # Verify linear scaling: norm should scale linearly with kappa + # norm(kappa=k) / norm(kappa=1) should equal k + base_norm = results[1] # kappa=1.0 + rtol = np.sqrt(np.finfo(dtype).eps) + for i, k in enumerate(kappa_values): + expected_ratio = k / 1.0 + actual_ratio = results[i] / base_norm + assert np.isclose(actual_ratio, expected_ratio, rtol=rtol) diff --git a/python/test/unit/fem/test_custom_jit_kernels.py b/python/test/unit/fem/test_custom_jit_kernels.py index 242e8674695..3f798f66c1c 100644 --- a/python/test/unit/fem/test_custom_jit_kernels.py +++ b/python/test/unit/fem/test_custom_jit_kernels.py @@ -93,9 +93,9 @@ def test_numba_assembly(dtype): active_coeffs = np.array([], dtype=np.int8) integrals = { IntegralType.cell: [ - (0, k2.address, cells, active_coeffs), - (1, k2.address, np.arange(0), active_coeffs), - (2, k2.address, np.arange(0), active_coeffs), + (0, k2.address, cells, active_coeffs, None), + (1, k2.address, np.arange(0), active_coeffs, None), + (2, k2.address, np.arange(0), active_coeffs, None), ] } formtype = form_cpp_class(dtype) @@ -104,7 +104,7 @@ def test_numba_assembly(dtype): [V._cpp_object, V._cpp_object], integrals, [], [], False, [], mesh=mesh._cpp_object ) ) - integrals = {IntegralType.cell: [(0, k1.address, cells, active_coeffs)]} + integrals = {IntegralType.cell: [(0, k1.address, cells, active_coeffs, None)]} L = Form(formtype([V._cpp_object], integrals, [], [], False, [], mesh=mesh._cpp_object)) A = dolfinx.fem.assemble_matrix(a) @@ -135,7 +135,9 @@ def test_coefficient(dtype): num_cells = mesh.topology.index_map(tdim).size_local + mesh.topology.index_map(tdim).num_ghosts active_coeffs = np.array([0], dtype=np.int8) integrals = { - IntegralType.cell: [(0, k1.address, np.arange(num_cells, dtype=np.int32), active_coeffs)] + IntegralType.cell: [ + (0, k1.address, np.arange(num_cells, dtype=np.int32), active_coeffs, None) + ] } formtype = form_cpp_class(dtype) L = Form( @@ -275,7 +277,7 @@ def test_cffi_assembly(): ptrA = ffi.cast("intptr_t", ffi.addressof(lib, "tabulate_tensor_poissonA")) active_coeffs = np.array([], dtype=np.int8) - integrals = {IntegralType.cell: [(0, ptrA, cells, active_coeffs)]} + integrals = {IntegralType.cell: [(0, ptrA, cells, active_coeffs, None)]} a = Form( _cpp.fem.Form_float64( [V._cpp_object, V._cpp_object], integrals, [], [], False, [], mesh=mesh._cpp_object @@ -283,7 +285,7 @@ def test_cffi_assembly(): ) ptrL = ffi.cast("intptr_t", ffi.addressof(lib, "tabulate_tensor_poissonL")) - integrals = {IntegralType.cell: [(0, ptrL, cells, active_coeffs)]} + integrals = {IntegralType.cell: [(0, ptrL, cells, active_coeffs, None)]} L = Form( _cpp.fem.Form_float64([V._cpp_object], integrals, [], [], False, [], mesh=mesh._cpp_object) )