From 2dfdabcbcff6a2cec2502c5fae41f434e2db7bf4 Mon Sep 17 00:00:00 2001 From: Julian Andrej Date: Fri, 8 May 2026 14:39:54 -0700 Subject: [PATCH] cleanup --- .../local_qf/derivative_apply_transpose.hpp | 99 +++--- .../backends/local_qf/derivative_assemble.hpp | 22 +- .../local_qf/derivative_assemble_diagonal.hpp | 332 ++++++++++-------- .../backends/local_qf/derivative_setup.hpp | 6 +- fem/dfem/doperator.cpp | 8 +- fem/dfem/doperator.hpp | 26 +- 6 files changed, 252 insertions(+), 241 deletions(-) diff --git a/fem/dfem/backends/local_qf/derivative_apply_transpose.hpp b/fem/dfem/backends/local_qf/derivative_apply_transpose.hpp index 3c65bd53e2..fa9d754722 100644 --- a/fem/dfem/backends/local_qf/derivative_apply_transpose.hpp +++ b/fem/dfem/backends/local_qf/derivative_apply_transpose.hpp @@ -14,10 +14,6 @@ namespace mfem::future namespace LocalQFImpl { -// Applies J^T using the cached Jacobian from DerivativeSetup. -// Inputs/outputs play swapped roles compared to DerivativeApply: -// direction_l : concatenated OUTPUT L-space direction (test space) -// ye : element-space result in INPUT space (trial/derivative field) template< int derivative_id, typename qfunc_t, @@ -154,10 +150,10 @@ struct DerivativeApplyTranspose // [dir_at_qp region] [result_at_qp region] [scratch] const int output_size_on_qp = std::accumulate(out_qp_size.begin(), out_qp_size.end(), 0); - dir_at_qp_offset = 0; + dir_at_qp_offset = 0; result_at_qp_offset = output_size_on_qp * num_qp; - scratch_offset = result_at_qp_offset - + trial_vdim * total_trial_op_dim * num_qp; + scratch_offset = + result_at_qp_offset + trial_vdim * total_trial_op_dim * num_qp; // Conservative scratch: 6 * q1d^3 (matches get_shmem_info). const int scratch_per_buf = (q1d > 0) ? q1d * q1d * q1d : 1; @@ -182,13 +178,11 @@ struct DerivativeApplyTranspose "DerivativeApplyTranspose: derivative field not found in infds"); // Number of DOFs per element for the derivative input field. - { - const auto &fd = ctx.infds[deriv_infd_idx]; - auto R = get_restriction(fd, dof_ordering); - MFEM_ASSERT(R != nullptr, - "LocalQFBackend: missing element restriction for deriv input"); - deriv_in_elem_sz = num_entities ? (R->Height() / num_entities) : 0; - } + const auto &fd = ctx.infds[deriv_infd_idx]; + auto R = get_restriction(fd, dof_ordering); + MFEM_ASSERT(R != nullptr, + "LocalQFBackend: missing element restriction for deriv input"); + deriv_in_elem_sz = num_entities ? (R->Height() / num_entities) : 0; } void operator()( @@ -198,30 +192,26 @@ struct DerivativeApplyTranspose { if (ctx.attr.Size() == 0) { return; } - // --- Restrict OUTPUT direction from L-space to element space --- - // direction_l is the concatenation of L-space vectors for each outfd. + int l_offset = 0; + int e_offset = 0; + for_constexpr([&](auto o) { - int l_offset = 0; - int e_offset = 0; - for_constexpr([&](auto o) - { - const size_t outfd_idx = output_to_outfd[o]; - const auto &fd = ctx.outfds[outfd_idx]; - const int l_size = GetVSize(fd); + const size_t outfd_idx = output_to_outfd[o]; + const auto &fd = ctx.outfds[outfd_idx]; + const int l_size = GetVSize(fd); - Vector dir_o_l(const_cast(direction_l->GetData()) + l_offset, - l_size); + Vector dir_o_l(const_cast(direction_l->GetData()) + l_offset, + l_size); - // Temporary element-space slice for output o. - const int e_size = out_elem_dof_size[o] * num_entities; - Vector dir_o_e(dir_out_e.GetData() + e_offset, e_size); + // Temporary element-space slice for output o. + const int e_size = out_elem_dof_size[o] * num_entities; + Vector dir_o_e(dir_out_e.GetData() + e_offset, e_size); - restriction(fd, dir_o_l, dir_o_e, dof_ordering); + restriction(fd, dir_o_l, dir_o_e, dof_ordering); - l_offset += l_size; - e_offset += e_size; - }); - } + l_offset += l_size; + e_offset += e_size; + }); // Wrap output direction element data as [elem_dof, entity]. // Offsets into dir_out_e for each output field. @@ -307,8 +297,6 @@ struct DerivativeApplyTranspose scratch_buf_size); } - // --- Step 1: Map OUTPUT direction element DOFs → QPs --- - // Concatenate per-output dir_at_qp slices. int qp_offset = 0; for_constexpr([&](auto o) { @@ -344,9 +332,9 @@ struct DerivativeApplyTranspose }); MFEM_SYNC_THREAD; - // --- Step 2: Contract J^T --- - // Forward: result(i,k,q) = sum_{j,m} J(i,k,j,m,q) * dir(j,m,q) - // Transpose: result(j,m,q) = sum_{i,k} J(i,k,j,m,q) * dir(i,k,q) + // contract J^T + // forward: result(i,k,q) = sum_{j,m} J(i,k,j,m,q) * dir(j,m,q) + // transpose: result(j,m,q) = sum_{i,k} J(i,k,j,m,q) * dir(i,k,q) MFEM_FOREACH_THREAD(q, x, num_qp_local) { for (int j = 0; j < trial_vdim_local; j++) @@ -390,7 +378,6 @@ struct DerivativeApplyTranspose } MFEM_SYNC_THREAD; - // --- Step 3: Map result QPs → INPUT derivative element DOFs --- auto ye_deriv = DeviceTensor<2, real_t>( ye_deriv_ptr, deriv_in_elem_sz_local, num_entities_local); auto result_dof = Reshape(&ye_deriv(0, e), deriv_in_elem_sz_local / trial_vdim_local, @@ -416,26 +403,26 @@ struct DerivativeApplyTranspose } IntegratorContext ctx; - qfunc_t qfunc; - inputs_t inputs; + qfunc_t qfunc; + inputs_t inputs; outputs_t outputs; const Vector &qp_cache; std::array output_to_outfd; - std::array input_to_field; + std::array input_to_field; std::array output_to_field; - int dimension = 0; + int dimension = 0; int num_entities = 0; - int num_qp = 0; - int q1d = 0; + int num_qp = 0; + int q1d = 0; bool use_sum_factorization = false; ElementDofOrdering dof_ordering = ElementDofOrdering::LEXICOGRAPHIC; ThreadBlocks thread_blocks; - std::vector dtqs; - std::array input_dtq_maps; + std::vector dtqs; + std::array input_dtq_maps; std::array output_dtq_maps; std::array out_qp_size; @@ -447,17 +434,17 @@ struct DerivativeApplyTranspose std::vector input_size_on_qp; std::array input_is_dependent; - int trial_vdim = 1; - int total_trial_op_dim = 0; - size_t deriv_input_idx = SIZE_MAX; - size_t deriv_infd_idx = SIZE_MAX; - int deriv_in_elem_sz = 0; + int trial_vdim = 1; + int total_trial_op_dim = 0; + size_t deriv_input_idx = SIZE_MAX; + size_t deriv_infd_idx = SIZE_MAX; + int deriv_in_elem_sz = 0; - // Shared memory layout offsets (in reals). - int dir_at_qp_offset = 0; + // Shared memory layout offsets (in real_t). + int dir_at_qp_offset = 0; int result_at_qp_offset = 0; - int scratch_offset = 0; - int shmem_per_elem = 0; + int scratch_offset = 0; + int shmem_per_elem = 0; mutable Vector shmem_cache; mutable Vector dir_out_e; diff --git a/fem/dfem/backends/local_qf/derivative_assemble.hpp b/fem/dfem/backends/local_qf/derivative_assemble.hpp index 616d0e00ac..39f2852e9a 100644 --- a/fem/dfem/backends/local_qf/derivative_assemble.hpp +++ b/fem/dfem/backends/local_qf/derivative_assemble.hpp @@ -25,7 +25,7 @@ template< struct DerivativeAssemble { static constexpr auto inout_tuple = - merge_mfem_tuples_as_empty_std_tuple(inputs_t {}, outputs_t {}); + merge_mfem_tuples_as_empty_std_tuple(inputs_t {}, outputs_t {}); static constexpr auto filtered_inout_tuple = filter_fields(inout_tuple); static constexpr size_t nfields = count_unique_field_ids(filtered_inout_tuple); @@ -108,9 +108,9 @@ struct DerivativeAssemble // Get test and trial spaces test_fes = std::get_if( - &ctx.unionfds[test_field_idx].data); + &ctx.unionfds[test_field_idx].data); trial_fes = std::get_if( - &ctx.unionfds[trial_field_idx].data); + &ctx.unionfds[trial_field_idx].data); MFEM_ASSERT(test_fes != nullptr && *test_fes != nullptr, "LocalQFBackend: test space is not a ParFiniteElementSpace"); @@ -144,16 +144,19 @@ struct DerivativeAssemble // Accumulate trial_op_dim for all dependent inputs if (input_is_dependent[i]) { - total_trial_op_dim += get(this->inputs).size_on_qp / get(this->inputs).vdim; + total_trial_op_dim += get(this->inputs).size_on_qp / get + (this->inputs).vdim; } }); MFEM_ASSERT(trial_vdim > 0, "LocalQFBackend: could not determine trial vdim"); - MFEM_ASSERT(total_trial_op_dim > 0, "LocalQFBackend: no dependent inputs found"); + MFEM_ASSERT(total_trial_op_dim > 0, + "LocalQFBackend: no dependent inputs found"); num_trial_dof = (*trial_fes)->GetFE(0)->GetDof(); num_trial_dof_1d = (dimension > 0) ? - static_cast(std::floor(std::pow(num_trial_dof, 1.0 / dimension) + 0.5)) : 0; + static_cast(std::floor(std::pow(num_trial_dof, + 1.0 / dimension) + 0.5)) : 0; // Setup DofToQuad maps dtqs.reserve(ctx.unionfds.size()); @@ -203,7 +206,8 @@ struct DerivativeAssemble { if (input_is_dependent[i]) { - inputs_trial_op_dim[i] = get(this->inputs).size_on_qp / get(this->inputs).vdim; + inputs_trial_op_dim[i] = get(this->inputs).size_on_qp / get + (this->inputs).vdim; } else { @@ -212,9 +216,7 @@ struct DerivativeAssemble }); } - void operator()( - std::vector &fields_e, - SparseMatrix *&A) const + void operator()(SparseMatrix *&A) const { if (ctx.attr.Size() == 0) { return; } diff --git a/fem/dfem/backends/local_qf/derivative_assemble_diagonal.hpp b/fem/dfem/backends/local_qf/derivative_assemble_diagonal.hpp index 36f94b042a..bd6931ea89 100644 --- a/fem/dfem/backends/local_qf/derivative_assemble_diagonal.hpp +++ b/fem/dfem/backends/local_qf/derivative_assemble_diagonal.hpp @@ -1,6 +1,7 @@ #pragma once #include "../../integrator_ctx.hpp" +#include "../../../kernel_dispatch.hpp" #include @@ -25,67 +26,20 @@ struct DerivativeAssembleDiagonal static constexpr size_t nfields = count_unique_field_ids(filtered_inout_tuple); template - MFEM_HOST_DEVICE static inline real_t EvalFactor2D( + MFEM_HOST_DEVICE static inline real_t EvalFactor1D( const DofToQuadMap &dtq, - const int dof, - const int td1d, const int k, - const int qx, - const int qy) + const int axis, + const int q, + const int d) { - const auto B = dtq.B; - const auto G = dtq.G; - - // This is the reverse mapping from the 1D index dof - // to the 2D (ix, iy) indices on the dofs. - const int ix = dof % td1d; - const int iy = dof / td1d; - if constexpr (is_value_fop::value) { - return (k == 0) ? B(qx, 0, ix) * B(qy, 0, iy) : 0.0; + return (k == 0) ? dtq.B(q, 0, d) : 0.0; } else if constexpr (is_gradient_fop::value) { - if (k == 0) { return G(qx, 0, ix) * B(qy, 0, iy); } - if (k == 1) { return B(qx, 0, ix) * G(qy, 0, iy); } - return 0.0; - } - else - { - return 0.0; - } - } - - template - MFEM_HOST_DEVICE static inline real_t EvalFactor3D( - const DofToQuadMap &dtq, - const int dof, - const int td1d, - const int k, - const int qx, - const int qy, - const int qz) - { - const auto B = dtq.B; - const auto G = dtq.G; - - // This is the reverse mapping from the 1D index dof - // to the 3D (ix, iy, iz) indices on the dofs. - const int ix = dof % td1d; - const int iy = (dof / td1d) % td1d; - const int iz = dof / (td1d * td1d); - - if constexpr (is_value_fop::value) - { - return (k == 0) ? B(qx, 0, ix) * B(qy, 0, iy) * B(qz, 0, iz) : 0.0; - } - else if constexpr (is_gradient_fop::value) - { - if (k == 0) { return G(qx, 0, ix) * B(qy, 0, iy) * B(qz, 0, iz); } - if (k == 1) { return B(qx, 0, ix) * G(qy, 0, iy) * B(qz, 0, iz); } - if (k == 2) { return B(qx, 0, ix) * B(qy, 0, iy) * G(qz, 0, iz); } - return 0.0; + return (k == axis) ? dtq.G(q, 0, d) : dtq.B(q, 0, d); } else { @@ -228,36 +182,11 @@ struct DerivativeAssembleDiagonal if (!is_square) { return; } - if (use_sum_factorization) - { - const auto &dm = - dynamic_cast(*(*test_fes)->GetFE(0)).GetDofMap(); - dofmap_h = dm; - } - - elem_vdofs.SetSize(num_test_dof * test_vdim * num_entities); - auto elem_vdofs_h = Reshape(elem_vdofs.HostWrite(), - num_test_dof * test_vdim, num_entities); - for (int e = 0; e < num_entities; e++) - { - Array vdofs; - (*test_fes)->GetElementVDofs(e, vdofs); - - for (int vd = 0; vd < test_vdim; vd++) - { - for (int i = 0; i < num_test_dof; i++) - { - const int native_i = (dofmap_h.Size() > 0) ? dofmap_h[i] : i; - elem_vdofs_h(i + vd * num_test_dof, e) = - vdofs[native_i + vd * num_test_dof]; - } - } - } + Ye_mem.SetSize(num_test_dof * test_vdim * num_entities); + Ye_mem.UseDevice(true); } - void operator()( - std::vector &/*fields_e*/, - Vector &diag_l) const + void operator()(Vector &diag_e) const { if (!is_square) { return; } if (ctx.attr.Size() == 0) { return; } @@ -268,6 +197,10 @@ struct DerivativeAssembleDiagonal "for tensor-product 2D/3D elements only"); return; } + MFEM_VERIFY(num_test_dof_1d == num_trial_dof_1d, + "DerivativeAssembleDiagonal requires matching tensor dofs"); + MFEM_VERIFY(num_test_dof_1d <= DeviceDofQuadLimits::Get().MAX_D1D, ""); + MFEM_VERIFY(q1d <= DeviceDofQuadLimits::Get().MAX_Q1D, ""); const bool has_attr = ctx.attr.Size() > 0; const auto d_attr = ctx.attr.Read(); @@ -277,8 +210,7 @@ struct DerivativeAssembleDiagonal qp_cache.Read(), residual_size_on_qp, num_qp, num_entities); const int num_dofs_per_elem = num_test_dof * test_vdim; - auto diag = diag_l.ReadWrite(); - auto elem_vdofs_d = Reshape(elem_vdofs.Read(), num_dofs_per_elem, num_entities); + auto Ye = Reshape(Ye_mem.ReadWrite(), num_dofs_per_elem, num_entities); const auto output_dtq_map_local = output_dtq_maps[0]; const auto input_dtq_maps_local = input_dtq_maps; @@ -301,6 +233,8 @@ struct DerivativeAssembleDiagonal mfem::forall(num_entities, [=] MFEM_HOST_DEVICE (int e) mutable { + constexpr int MD1 = DofQuadLimits::MAX_D1D; + constexpr int MQ1 = DofQuadLimits::MAX_Q1D; if (has_attr && !d_attr[d_elem_attr[e] - 1]) { return; } auto qpdc = Reshape(&cache_tensor(0, 0, e), @@ -308,83 +242,178 @@ struct DerivativeAssembleDiagonal trial_vdim_local, total_trial_op_dim_local, num_qp_local); - for (int vd = 0; vd < test_vdim_local; vd++) + if (dimension_local == 2) { - for (int dof = 0; dof < num_test_dof_local; dof++) + real_t QD[MQ1][MD1]; + for (int vd = 0; vd < test_vdim_local; vd++) { - real_t val = 0.0; - for (int q = 0; q < num_qp_local; q++) + auto Y = Reshape(&Ye(vd * num_test_dof_local, e), + num_test_dof_1d_local, num_test_dof_1d_local); + for (int dy = 0; dy < num_test_dof_1d_local; dy++) { - // This is the reverse mapping from the 1D index q - // to the 3D (qx, qy, qz) indices on the quadrature points. - const int qx = q % q1d_local; - const int qy = (dimension_local == 2) - ? q / q1d_local - : (q / q1d_local) % q1d_local; - const int qz = q / (q1d_local * q1d_local); - - for (int k = 0; k < test_op_dim_local; k++) + for (int dx = 0; dx < num_test_dof_1d_local; dx++) { - real_t psi = 0.0; - if (dimension_local == 2) - { - psi = EvalFactor2D(output_dtq_map_local, - dof, num_test_dof_1d_local, - k, qx, qy); - } - else - { - psi = EvalFactor3D(output_dtq_map_local, - dof, num_test_dof_1d_local, - k, qx, qy, qz); - } - - if (psi == 0.0) { continue; } - - real_t trial_contraction = 0.0; - int m_offset = 0; - [[maybe_unused]] const auto &inputs_ref = inputs_local; - for_constexpr([&](auto s) - { - using fop_t = std::decay_t(inputs_ref))>; - - const int trial_op_dim = itod[static_cast(s)]; - if (trial_op_dim == 0) { return; } - - const auto &dtq = input_dtq_maps_local[s]; - for (int m = 0; m < trial_op_dim; m++) - { - real_t phi = 0.0; - if (dimension_local == 2) - { - phi = EvalFactor2D(dtq, dof, - num_trial_dof_1d_local, - m, qx, qy); - } - else - { - phi = EvalFactor3D(dtq, dof, - num_trial_dof_1d_local, - m, qx, qy, qz); - } - trial_contraction += qpdc(vd, k, vd, m_offset + m, q) * phi; - } - - m_offset += trial_op_dim; - }); - - val += psi * trial_contraction; + Y(dx,dy) = 0.0; } } + for (int k = 0; k < test_op_dim_local; k++) + { + int m_offset = 0; + [[maybe_unused]] const auto &inputs_ref = inputs_local; + for_constexpr([&](auto s) + { + using fop_t = std::decay_t(inputs_ref))>; + const int trial_op_dim = itod[static_cast(s)]; + if (trial_op_dim == 0) { return; } - const int local = dof + vd * num_test_dof_local; - const int gdof = elem_vdofs_d(local, e); - const int abs_gdof = (gdof >= 0) ? gdof : (-1 - gdof); - const real_t sign = (gdof >= 0) ? 1.0 : -1.0; - AtomicAdd(diag[abs_gdof], sign * val); + for (int m = 0; m < trial_op_dim; m++) + { + for (int qx = 0; qx < q1d_local; qx++) + { + for (int dy = 0; dy < num_test_dof_1d_local; dy++) + { + QD[qx][dy] = 0.0; + for (int qy = 0; qy < q1d_local; qy++) + { + const int q = qx + qy * q1d_local; + const real_t Ly = + EvalFactor1D(output_dtq_map_local, + k, 1, qy, dy); + const real_t Ry = + EvalFactor1D(input_dtq_maps_local[s], + m, 1, qy, dy); + QD[qx][dy] += Ly * Ry * + qpdc(vd, k, vd, m_offset + m, q); + } + } + } + for (int dy = 0; dy < num_test_dof_1d_local; dy++) + { + for (int dx = 0; dx < num_test_dof_1d_local; dx++) + { + real_t accum = 0.0; + for (int qx = 0; qx < q1d_local; qx++) + { + const real_t Lx = + EvalFactor1D(output_dtq_map_local, + k, 0, qx, dx); + const real_t Rx = + EvalFactor1D(input_dtq_maps_local[s], + m, 0, qx, dx); + accum += Lx * Rx * QD[qx][dy]; + } + Y(dx,dy) += accum; + } + } + } + m_offset += trial_op_dim; + }); + } + } + } + else + { + real_t QQD[MQ1][MQ1][MD1]; + real_t QDD[MQ1][MD1][MD1]; + for (int vd = 0; vd < test_vdim_local; vd++) + { + auto Y = Reshape(&Ye(vd * num_test_dof_local, e), + num_test_dof_1d_local, num_test_dof_1d_local, + num_test_dof_1d_local); + for (int dz = 0; dz < num_test_dof_1d_local; dz++) + { + for (int dy = 0; dy < num_test_dof_1d_local; dy++) + { + for (int dx = 0; dx < num_test_dof_1d_local; dx++) + { + Y(dx,dy,dz) = 0.0; + } + } + } + for (int k = 0; k < test_op_dim_local; k++) + { + int m_offset = 0; + [[maybe_unused]] const auto &inputs_ref = inputs_local; + for_constexpr([&](auto s) + { + using fop_t = std::decay_t(inputs_ref))>; + const int trial_op_dim = itod[static_cast(s)]; + if (trial_op_dim == 0) { return; } + + for (int m = 0; m < trial_op_dim; m++) + { + for (int qx = 0; qx < q1d_local; qx++) + { + for (int qy = 0; qy < q1d_local; qy++) + { + for (int dz = 0; dz < num_test_dof_1d_local; dz++) + { + QQD[qx][qy][dz] = 0.0; + for (int qz = 0; qz < q1d_local; qz++) + { + const int q = qx + (qy + qz * q1d_local) * q1d_local; + const real_t Lz = + EvalFactor1D(output_dtq_map_local, + k, 2, qz, dz); + const real_t Rz = + EvalFactor1D(input_dtq_maps_local[s], + m, 2, qz, dz); + QQD[qx][qy][dz] += Lz * Rz * + qpdc(vd, k, vd, m_offset + m, q); + } + } + } + } + for (int qx = 0; qx < q1d_local; qx++) + { + for (int dz = 0; dz < num_test_dof_1d_local; dz++) + { + for (int dy = 0; dy < num_test_dof_1d_local; dy++) + { + QDD[qx][dy][dz] = 0.0; + for (int qy = 0; qy < q1d_local; qy++) + { + const real_t Ly = + EvalFactor1D(output_dtq_map_local, + k, 1, qy, dy); + const real_t Ry = + EvalFactor1D(input_dtq_maps_local[s], + m, 1, qy, dy); + QDD[qx][dy][dz] += Ly * Ry * QQD[qx][qy][dz]; + } + } + } + } + for (int dz = 0; dz < num_test_dof_1d_local; dz++) + { + for (int dy = 0; dy < num_test_dof_1d_local; dy++) + { + for (int dx = 0; dx < num_test_dof_1d_local; dx++) + { + real_t accum = 0.0; + for (int qx = 0; qx < q1d_local; qx++) + { + const real_t Lx = + EvalFactor1D(output_dtq_map_local, + k, 0, qx, dx); + const real_t Rx = + EvalFactor1D(input_dtq_maps_local[s], + m, 0, qx, dx); + accum += Lx * Rx * QDD[qx][dy][dz]; + } + Y(dx,dy,dz) += accum; + } + } + } + } + m_offset += trial_op_dim; + }); + } } } }); + + diag_e += Ye_mem; } IntegratorContext ctx; @@ -428,8 +457,7 @@ struct DerivativeAssembleDiagonal std::array input_is_dependent; std::array inputs_trial_op_dim {}; - Array dofmap_h; - Array elem_vdofs; + mutable Vector Ye_mem; int residual_size_on_qp = 0; }; diff --git a/fem/dfem/backends/local_qf/derivative_setup.hpp b/fem/dfem/backends/local_qf/derivative_setup.hpp index d18b3ce598..c85add77cd 100644 --- a/fem/dfem/backends/local_qf/derivative_setup.hpp +++ b/fem/dfem/backends/local_qf/derivative_setup.hpp @@ -396,9 +396,7 @@ struct DerivativeSetup std::make_integer_sequence {}); } - void operator()( - const std::vector &xe, - const Vector &direction_l) const + void operator()(const std::vector &xe) const { if (ctx.attr.Size() == 0) { return; } @@ -407,8 +405,6 @@ struct DerivativeSetup static_assert(tuple_size::value == ninputs + noutputs, "qfunc parameter count must match inputs+outputs"); - // Don't need direction_l for full Jacobian computation (unused) - std::array, nfields> wrapped_fields_e; for (size_t uf = 0; uf < nfields; uf++) { diff --git a/fem/dfem/doperator.cpp b/fem/dfem/doperator.cpp index 5bdade9f67..25978c032b 100644 --- a/fem/dfem/doperator.cpp +++ b/fem/dfem/doperator.cpp @@ -122,17 +122,13 @@ std::shared_ptr DifferentiableOperator::GetDerivative( // Prolong to L-space prolongation(infds, bx, infields_l); - // Get the direction field (the field corresponding to derivative_id) - Vector direction_l(GetVSize(infds[dfidx])); - prolongation(infds[dfidx], bx.GetBlock(dfidx), direction_l); - // Restrict to element space restriction(infds, infields_l, infields_e); // Run all setup callbacks to populate the cache for (const auto &setup_callback : derivative_setup_callbacks[derivative_id]) { - setup_callback(infields_e, direction_l); + setup_callback(infields_e); } return std::make_shared( @@ -223,7 +219,7 @@ std::shared_ptr DifferentiableOperator::GetDerivative( // Run all setup callbacks to populate the cache for (const auto &setup_callback : derivative_setup_callbacks[derivative_id]) { - setup_callback(infields_e, direction_l); + setup_callback(infields_e); } return std::make_shared( diff --git a/fem/dfem/doperator.hpp b/fem/dfem/doperator.hpp index 7fd0683fe6..6fba74f525 100644 --- a/fem/dfem/doperator.hpp +++ b/fem/dfem/doperator.hpp @@ -30,7 +30,7 @@ using action_t = /// @brief Type alias for a function that computes the cache for the action of a derivative using derivative_setup_t = - std::function &, const Vector &)>; + std::function &)>; /// @brief Type alias for a function that computes the action of a derivative using derivative_action_t = @@ -39,17 +39,16 @@ using derivative_action_t = /// @brief Type alias for a function that assembles the SparseMatrix of a /// derivative operator using assemble_derivative_sparsematrix_callback_t = - std::function &, SparseMatrix *&)>; + std::function; /// @brief Type alias for a function that assembles the HypreParMatrix of a /// derivative operator using assemble_derivative_hypreparmatrix_callback_t = - std::function &, HypreParMatrix *&)>; + std::function; /// @brief Type alias for a function that assembles the diagonal of a derivative -/// operator into an L-vector -using assemble_diagonal_callback_t = - std::function &, Vector &)>; +/// operator into an E-vector +using assemble_diagonal_callback_t = std::function; /// @brief Type alias for a function that applies the appropriate restriction to /// the solution and parameters @@ -251,7 +250,7 @@ public: for (const auto &f : assemble_derivative_sparsematrix_callbacks) { - f(fields_e, A); + f(A); } } @@ -266,7 +265,7 @@ public: for (const auto &f : assemble_derivative_hypreparmatrix_callbacks) { - f(fields_e, A); + f(A); } } @@ -277,21 +276,24 @@ public: { MFEM_ASSERT(!assemble_diagonal_callbacks.empty(), "derivative can't assemble diagonal"); + MFEM_ASSERT(outfds.size() == 1, + "AssembleDiagonal currently requires a single output field"); const auto *test_pf = std::get_if(&outfds[0].data); MFEM_ASSERT(test_pf && *test_pf, "AssembleDiagonal: test field must be a ParFiniteElementSpace"); - Vector diag_l((*test_pf)->GetVSize()); - diag_l = 0.0; + prepare_residual(outfds, daction_e); + for (auto *v : daction_e) { *v = 0.0; } for (const auto &f : assemble_diagonal_callbacks) { - f(fields_e, diag_l); + f(*daction_e[0]); } - (*test_pf)->GetProlongationMatrix()->MultTranspose(diag_l, diag); + restriction_transpose(outfds, daction_e, daction_l); + prolongation_transpose(outfds[0], *daction_l[0], diag); } private: