diff --git a/source/MechanicalOperatorDevice.hh b/source/MechanicalOperatorDevice.hh new file mode 100644 index 00000000..ad09502c --- /dev/null +++ b/source/MechanicalOperatorDevice.hh @@ -0,0 +1,78 @@ +/* SPDX-FileCopyrightText: Copyright (c) 2026, the adamantine authors. + * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + */ + +#ifndef MECHANICAL_OPERATOR_DEVICE_HH +#define MECHANICAL_OPERATOR_DEVICE_HH + +#include +#include + +#include +#include +#include + +namespace adamantine +{ +template +class MechanicalOperatorDevice +{ +public: + MechanicalOperatorDevice( + MPI_Comm const &communicator, + MaterialProperty &material_properties); + + void reinit(dealii::DoFHandler const &dof_handler, + dealii::AffineConstraints const &affine_constraints); + + void vmult(dealii::LA::distributed::Vector &dst, + dealii::LA::distributed::Vector< + double, dealii::MemorySpace::Default> const &src) const; + + void vmult_add( + dealii::LA::distributed::Vector + &dst, + dealii::LA::distributed::Vector< + double, dealii::MemorySpace::Default> const &src) const; + + void initialize_dof_vector( + dealii::LA::distributed::Vector + &vector) const; + + dealii::Portable::MatrixFree const &get_matrix_free() const; + +private: + using kokkos_default = dealii::MemorySpace::Default::kokkos_space; + + MPI_Comm const &_communicator; + MaterialProperty &_material_properties; + + typename dealii::Portable::MatrixFree::AdditionalData + _matrix_free_data; + dealii::Portable::MatrixFree _matrix_free; + + Kokkos::View _lambda; + Kokkos::View _mu; + + std::map::cell_iterator, + std::vector> + _cell_it_to_mf_pos; + unsigned int _n_owned_cells = 0; +}; + +template +inline dealii::Portable::MatrixFree const & +MechanicalOperatorDevice::get_matrix_free() const +{ + return _matrix_free; +} + +} // namespace adamantine + +#include + +#endif diff --git a/source/MechanicalOperatorDevice.templates.hh b/source/MechanicalOperatorDevice.templates.hh new file mode 100644 index 00000000..c5c7bf8d --- /dev/null +++ b/source/MechanicalOperatorDevice.templates.hh @@ -0,0 +1,214 @@ +/* SPDX-FileCopyrightText: Copyright (c) 2026, the adamantine authors. + * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + */ + +#ifndef MECHANICAL_OPERATOR_DEVICE_TEMPLATES_HH +#define MECHANICAL_OPERATOR_DEVICE_TEMPLATES_HH + +#include +#include + +#include +#include +#include +#include +#include + +namespace adamantine +{ +// Small wrapper functor type expected by MatrixFree::cell_loop +template +class LocalMechanicalOperatorDevice +{ +public: + static const unsigned int n_q_points = + dealii::Utilities::pow(fe_degree + 1, dim); + + using kokkos_default = dealii::MemorySpace::Default::kokkos_space; + LocalMechanicalOperatorDevice(Kokkos::View lambda, + Kokkos::View mu) + : _lambda(lambda), _mu(mu) + { + } + + KOKKOS_FUNCTION void operator()( + typename dealii::Portable::MatrixFree::Data const *gpu_data, + const dealii::Portable::DeviceVector &src, + dealii::Portable::DeviceVector dst) const + { + dealii::Portable::FEEvaluation + fe_eval(gpu_data); + fe_eval.read_dof_values(src); + fe_eval.evaluate(dealii::EvaluationFlags::gradients); + + const int cell = fe_eval.get_current_cell_index(); + + gpu_data->for_each_quad_point( + [&](const int &q_point) + { + auto const grad_u = fe_eval.get_gradient(q_point); + auto const eps = (grad_u + dealii::transpose(grad_u)) * 0.5; + double const trace_eps = dealii::trace(eps); + + unsigned int const pos = gpu_data->local_q_point_id(cell, q_point); + + double const lambda = _lambda(pos); + double const mu = _mu(pos); + + dealii::Tensor<2, dim, double> stress; + for (unsigned int i = 0; i < dim; ++i) + for (unsigned int j = 0; j < dim; ++j) + { + stress[i][j] = 2.0 * mu * eps[i][j]; + if (i == j) + stress[i][j] += lambda * trace_eps; + } + + fe_eval.submit_gradient(stress, q_point); + }); + + fe_eval.integrate(dealii::EvaluationFlags::gradients); + fe_eval.distribute_local_to_global(dst); + } + +private: + Kokkos::View _lambda; + Kokkos::View _mu; +}; + +template +MechanicalOperatorDevice:: + MechanicalOperatorDevice( + MPI_Comm const &communicator, + MaterialProperty &material_properties) + : _communicator(communicator), _material_properties(material_properties) +{ + _matrix_free_data.mapping_update_flags = dealii::update_gradients; + //_matrix_free_data.tasks_parallel_scheme = + // dealii::MatrixFree::AdditionalData::partition_color; +} + +template +void MechanicalOperatorDevice:: + reinit(dealii::DoFHandler const &dof_handler, + dealii::AffineConstraints const &affine_constraints) +{ + dealii::IteratorFilters::ActiveFEIndexEqualTo filter(0, true); + _matrix_free.reinit(dealii::StaticMappingQ1::mapping, dof_handler, + affine_constraints, dealii::QGauss<1>(fe_degree + 1), + filter, _matrix_free_data); + + _n_owned_cells = + dynamic_cast const *>( + &dof_handler.get_triangulation()) + ->n_locally_owned_active_cells(); + + // Build mapping between cells and mf positions + _cell_it_to_mf_pos.clear(); + unsigned int constexpr n_dofs_1d = fe_degree + 1; + unsigned int constexpr n_q_points_per_cell = + dealii::Utilities::pow(n_dofs_1d, dim); + + auto graph = _matrix_free.get_colored_graph(); + unsigned int const n_colors = graph.size(); + for (unsigned int color = 0; color < n_colors; ++color) + { + auto gpu_data = _matrix_free.get_data(color); + unsigned int const n_cells = gpu_data.n_cells; + auto gpu_data_host = dealii::Portable::copy_mf_data_to_host( + gpu_data, _matrix_free_data.mapping_update_flags); + for (unsigned int cell_id = 0; cell_id < n_cells; ++cell_id) + { + auto cell = graph[color][cell_id]; + std::vector quad_pos(n_q_points_per_cell); + for (unsigned int i = 0; i < n_q_points_per_cell; ++i) + { + unsigned int const pos = + gpu_data_host.local_q_point_id(cell_id, n_q_points_per_cell, i); + quad_pos[i] = pos; + } + _cell_it_to_mf_pos[cell] = quad_pos; + } + } + + // Populate lambda and mu arrays on host then copy to device views + unsigned int const n_coefs = n_q_points_per_cell * _n_owned_cells; + std::vector lambda_host(n_coefs); + std::vector mu_host(n_coefs); + + for (auto const &cell : dealii::filter_iterators( + _matrix_free.get_dof_handler().active_cell_iterators(), + dealii::IteratorFilters::LocallyOwnedCell(), + dealii::IteratorFilters::ActiveFEIndexEqualTo(0))) + { + // Cast to triangulation cell to obtain material id + typename dealii::Triangulation::active_cell_iterator cell_tria(cell); + double const lambda_val = _material_properties.get_cell_value( + cell_tria, StateProperty::lame_first_parameter); + double const mu_val = _material_properties.get_cell_value( + cell_tria, StateProperty::lame_second_parameter); + + for (unsigned int i = 0; i < n_q_points_per_cell; ++i) + { + unsigned int const pos = _cell_it_to_mf_pos[cell][i]; + lambda_host[pos] = lambda_val; + mu_host[pos] = mu_val; + } + } + + using kokkos_default = dealii::MemorySpace::Default::kokkos_space; + _lambda = Kokkos::View( + Kokkos::view_alloc("mech_lambda", Kokkos::WithoutInitializing), n_coefs); + _mu = Kokkos::View( + Kokkos::view_alloc("mech_mu", Kokkos::WithoutInitializing), n_coefs); + + auto lambda_host_view = + Kokkos::create_mirror_view(Kokkos::WithoutInitializing, _lambda); + auto mu_host_view = + Kokkos::create_mirror_view(Kokkos::WithoutInitializing, _mu); + for (unsigned int i = 0; i < n_coefs; ++i) + { + lambda_host_view(i) = lambda_host[i]; + mu_host_view(i) = mu_host[i]; + } + Kokkos::deep_copy(_lambda, lambda_host_view); + Kokkos::deep_copy(_mu, mu_host_view); +} + +template +void MechanicalOperatorDevice:: + vmult(dealii::LA::distributed::Vector + &dst, + dealii::LA::distributed::Vector< + double, dealii::MemorySpace::Default> const &src) const +{ + dst = 0.; + vmult_add(dst, src); +} + +template +void MechanicalOperatorDevice:: + vmult_add( + dealii::LA::distributed::Vector + &dst, + dealii::LA::distributed::Vector< + double, dealii::MemorySpace::Default> const &src) const +{ + LocalMechanicalOperatorDevice local_operator(_lambda, _mu); + _matrix_free.cell_loop(local_operator, src, dst); + _matrix_free.copy_constrained_values(src, dst); +} + +template +void MechanicalOperatorDevice:: + initialize_dof_vector( + dealii::LA::distributed::Vector + &vector) const +{ + _matrix_free.initialize_dof_vector(vector); +} + +} // namespace adamantine + +#endif diff --git a/source/MechanicalPhysics.cc b/source/MechanicalPhysics.cc index c0c2e295..16f24fb1 100644 --- a/source/MechanicalPhysics.cc +++ b/source/MechanicalPhysics.cc @@ -12,6 +12,7 @@ #include #include #include +#include #include #include @@ -68,7 +69,13 @@ MechanicalPhysics:: } // Create the mechanical operator - _mechanical_operator = + if constexpr (std::is_same_v) + _mechanical_operator = std::make_unique< + MechanicalOperatorDevice>( + communicator, _material_properties); + else + Kokkos::abort("Not implemented"); + _mechanical_operator_host = std::make_unique>( communicator, _material_properties, reference_temperatures); @@ -124,8 +131,9 @@ void MechanicalPhysicsreinit(_dof_handler, _affine_constraints, _q_collection, - body_forces); + _mechanical_operator->reinit(_dof_handler, _affine_constraints); + _mechanical_operator_host->reinit(_dof_handler, _affine_constraints, + _q_collection, body_forces); } template :: update_rhs(std::vector>> const &body_forces) { - _mechanical_operator->assemble_rhs(body_forces); + _mechanical_operator_host->assemble_rhs(body_forces); } template const &has_melted, std::vector>> const &body_forces) { - _mechanical_operator->update_temperature(thermal_dof_handler, temperature, - has_melted); - _mechanical_operator->assemble_rhs(body_forces); + _mechanical_operator_host->update_temperature(thermal_dof_handler, + temperature, has_melted); + _mechanical_operator_host->assemble_rhs(body_forces); } template const &has_melted, std::vector>> const &body_forces) { - _mechanical_operator->update_temperature(thermal_dof_handler, temperature, - has_melted); + _mechanical_operator_host->update_temperature(thermal_dof_handler, + temperature, has_melted); + // Update the active fe indices, the plastic variables, and the displacement. unsigned int const n_quad_pts = _q_collection.max_n_quadrature_points(); unsigned int cell_id = 0; @@ -466,13 +475,17 @@ MechanicalPhysicsrhs().get_mpi_communicator()); - TrilinosVectorType rhs_device( - locally_owned_dofs, _mechanical_operator->rhs().get_mpi_communicator()); + dealii::LinearAlgebra::distributed::Vector + displacement(locally_owned_dofs, + _mechanical_operator_host->rhs().get_mpi_communicator()); + dealii::LinearAlgebra::distributed::Vector + rhs_device(locally_owned_dofs, + _mechanical_operator_host->rhs().get_mpi_communicator()); dealii::LinearAlgebra::ReadWriteVector rw_vector(locally_owned_dofs); - rw_vector.import_elements(_mechanical_operator->rhs(), + rw_vector.import_elements(_mechanical_operator_host->rhs(), dealii::VectorOperation::insert); rhs_device.import_elements(rw_vector, dealii::VectorOperation::insert); @@ -480,16 +493,19 @@ MechanicalPhysicsrhs().l2_norm(); + double const tol = 1e-12 * _mechanical_operator_host->rhs().l2_norm(); dealii::SolverControl solver_control(max_iter, tol); - dealii::SolverCG cg(solver_control); - cg.solve(_mechanical_operator->system_matrix(), displacement, rhs_device, - _mechanical_operator->preconditioner()); + dealii::SolverCG> + cg(solver_control); + cg.solve(*_mechanical_operator, displacement, rhs_device, + dealii::PreconditionIdentity{}); rw_vector.import_elements(displacement, dealii::VectorOperation::insert); dealii::LA::distributed::Vector - displacement_host(locally_owned_dofs, locally_relevant_dofs, - _mechanical_operator->rhs().get_mpi_communicator()); + displacement_host( + locally_owned_dofs, locally_relevant_dofs, + _mechanical_operator_host->rhs().get_mpi_communicator()); displacement_host.import_elements(rw_vector, dealii::VectorOperation::insert); _affine_constraints.distribute(displacement_host); @@ -500,7 +516,7 @@ MechanicalPhysics incremental_displacement( locally_owned_dofs, locally_relevant_dofs, - _mechanical_operator->rhs().get_mpi_communicator()); + _mechanical_operator_host->rhs().get_mpi_communicator()); incremental_displacement = displacement_host; if (_old_displacement.size() > 0) { diff --git a/source/MechanicalPhysics.hh b/source/MechanicalPhysics.hh index 62a36543..1110f949 100644 --- a/source/MechanicalPhysics.hh +++ b/source/MechanicalPhysics.hh @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -34,7 +35,7 @@ public: /** * Setup the DoFHandler, the AffineConstraints, and the - * MechanicalOperator. + * MechanicalOperatorDevice. */ void setup_dofs(std::vector>> const &body_forces = @@ -140,11 +141,14 @@ private: */ dealii::hp::QCollection _q_collection; /** - * Pointer to the MechanicalOperator + * Pointer to the MechanicalOperatorDevice */ + std::unique_ptr< + MechanicalOperatorDevice> + _mechanical_operator; std::unique_ptr> - _mechanical_operator; + _mechanical_operator_host; /** * Whether to include a gravitional body force in the calculation. */