// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced // at the Lawrence Livermore National Laboratory. All Rights reserved. See files // LICENSE and NOTICE for details. LLNL-CODE-806117. // // This file is part of the MFEM library. For more information and source code // availability visit https://mfem.org. // // MFEM is free software; you can redistribute it and/or modify it under the // terms of the BSD-3 license. We welcome feedback and contributions, see file // CONTRIBUTING.md for details. // // -------------------------------------------- // Elasticity LOR Block Preconditioning Miniapp // -------------------------------------------- // // Description: // The purpose of this miniapp is to demonstrate how to effectively // precondition vector valued PDEs, such as elasticity, on GPUs. // // Using 3D elasticity as an example, the linear operator can be broken into // vector components and has the form: // // A = [A_00 A_01 A_02] // [A_10 A_11 A_12] // [A_20 A_21 A_22] // // Traditional AMG requires having the entire matrix in memory which can be // prohibitive on GPUs. An effective preconditioning strategy for materials // which are not nearly incompressible is to use // P^{-1} = diag(AMG(A_00), AMG(A_11), AMG(A_22)) where AMG(A) is the AMG // approximation inv(A) [1]. This requires storing 3 blocks instead of 9, // but this is still prohibitive for high order discretizations on GPUs. This // is alleviated here by performing AMG on the low-order refined // operators instead. // // This miniapp solves the same beam problem described in Example 2. Run // times of the new solver (partial assembly with block diagonal LOR-AMG) // can be compared with the LEGACY approach of assembling the full matrix // and preconditioning with AMG. // // For the partial assembly approach, the operator actions and component // matrix assembly are supported on GPUs. // // The LEGACY approach should be performed with "-vdim" ordering while PARTIAL // requires "-nodes". // // There is also an option "-ss" or "--sub-solve" for the partial assembly // version which replaces P^{-1} with diag(inv(A_00), inv(A_11), inv(A_22)) // where the action of inv(A_ii) is performed with a CG inner solve that // is preconditioned with AMG(A_ii). This seems to give order independent // conditioning of the outer CG solve, but is much slower than performing // a single AMG iteration per block. // // This miniapp supports beam-tri.mesh, beam-quad.mesh, and beam-hex.mesh. // beam-tet.mesh can be run if MFEM is build with // ElementRestriction::MaxNbNbr set to 32 instead of 16. // // This miniapp shows how to test if the derived component integrators are // correct using BlockFESpaceOperator. If "-ca" (for componentwise action) // is used, a block operator where each block is a component of the // elasticity operator is used for A rather than the vector version. // This yields the same answer, but is less efficient. "-ca" can be called // with "-pa" for a version where each component is partially assembled, or // without where each component is called with full assembly, although the // latter may only work for order 1 on GPUs. // // Sample runs: // lor_elast -m ../../data/beam-tri.mesh // lor_elast -m ../../data/beam-quad.mesh // lor_elast -m ../../data/beam-hex.mesh // mpirun -np 4 lor_elast -m ../../data/beam-hex.mesh -l 3 -vdim // mpirun -np 4 lor_elast -m ../../data/beam-hex.mesh -l 3 -vdim -elast // // Device sample runs: // lor_elast --device cuda -m ../../data/beam-hex.mesh -l 3 -o 2 -pa // lor_elast --device cuda -m ../../data/beam-hex.mesh -l 3 -o 2 -pa -ss // lor_elast --device cuda -m ../../data/beam-hex.mesh -l 3 -o 2 -pa -ca // lor_elast --device cuda -m ../../data/beam-hex.mesh -l 3 -ca // // References: // [1] Mihajlović, M.D. and Mijalković, S., "A component decomposition // preconditioning for 3D stress analysis problems", Numerical Linear // Algebra with Applications, 2002. // #include "mfem.hpp" #include #include #include "block_fespace_operator.hpp" using namespace std; using namespace mfem; int main(int argc, char *argv[]) { // 1. Initialize MPI and HYPRE. Mpi::Init(argc, argv); Hypre::Init(); // 2. Parse command-line options. const char *mesh_file = "../../data/beam-tri.mesh"; int order = 1; bool pa = false; bool visualization = false; bool amg_elast = 0; bool reorder_space = true; const char *device_config = "cpu"; int ref_levels = 0; bool sub_solve = false; bool componentwise_action = false; OptionsParser args(argc, argv); args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use."); args.AddOption(&order, "-o", "--order", "Finite element order (polynomial degree)."); args.AddOption(&amg_elast, "-elast", "--amg-for-elasticity", "-sys", "--amg-for-systems", "Use the special AMG elasticity solver (GM/LN approaches), " "or standard AMG for systems (unknown approach)."); args.AddOption(&sub_solve, "-ss", "--sub-solve", "-no-ss", "--no-sub-solve", "Blocks are solved with a few CG iterations instead of a single AMG application."); args.AddOption(&pa, "-pa", "--partial-assembly", "-no-pa", "--no-partial-assembly", "Enable Partial Assembly."); args.AddOption(&reorder_space, "-nodes", "--by-nodes", "-vdim", "--by-vdim", "Use byNODES ordering of vector space instead of byVDIM"); args.AddOption(&device_config, "-d", "--device", "Device configuration string, see Device::Configure()."); args.AddOption(&ref_levels, "-l","--reflevels", "How many mesh refinements to perform."); args.AddOption(&componentwise_action, "-ca", "--component-action", "-no-ca", "--no-component-action", "Uses partial assembly with a block operator of components instead of the monolithic vector integrator."); args.AddOption(&visualization, "-vis", "--visualization", "-no-vis", "--no-visualization", "Enable or disable ParaView and GLVis output."); args.ParseCheck(); // 3. Enable hardware devices such as GPUs, and programming models such as // CUDA, OCCA, RAJA and OpenMP based on command line options. Device device(device_config); if (Mpi::Root()) { device.Print(); } // 4. Read the (serial) mesh from the given mesh file on all processors. We // can handle triangular, quadrilateral, tetrahedral, hexahedral, surface // and volume meshes with the same code. Mesh mesh(mesh_file, 1, 1); int dim = mesh.Dimension(); if (mesh.attributes.Max() < 2 || mesh.bdr_attributes.Max() < 2) { if (Mpi::Root()) { cerr << "\nInput mesh should have at least two materials and " << "two boundary attributes! (See schematic in ex2.cpp)\n" << endl; } return 3; } // 5. Refine the serial mesh on all processors to increase the resolution. for (int l = 0; l < ref_levels; l++) { mesh.UniformRefinement(); } // 6. Define a parallel mesh by a partitioning of the serial mesh. ParMesh pmesh(MPI_COMM_WORLD, mesh); // 7. Define a parallel finite element spaces on the parallel mesh. Here we // use vector finite elements, i.e. dim copies of a scalar finite element // space. If using partial assembly, also assemble the low order refined // (LOR) fespace. H1_FECollection fec(order, dim); const Ordering::Type fes_ordering = reorder_space ? Ordering::byNODES : Ordering::byVDIM; ParFiniteElementSpace fespace(&pmesh, &fec, dim, fes_ordering); ParFiniteElementSpace scalar_fespace(&pmesh, &fec, 1, fes_ordering); unique_ptr lor_disc; unique_ptr scalar_lor_fespace; if (pa || componentwise_action) { lor_disc.reset(new ParLORDiscretization(fespace)); ParFiniteElementSpace &lor_space = lor_disc->GetParFESpace(); const FiniteElementCollection &lor_fec = *lor_space.FEColl(); ParMesh &lor_mesh = *lor_space.GetParMesh(); scalar_lor_fespace.reset( new ParFiniteElementSpace(&lor_mesh, &lor_fec, 1, fes_ordering)); } HYPRE_BigInt size = fespace.GlobalTrueVSize(); if (Mpi::Root()) { cout << "Number of finite element unknowns: " << size << endl << "Assembling: " << flush; } // 8. Determine the list of true (i.e. parallel conforming) essential // boundary dofs. In this example, the boundary conditions are defined by // marking only boundary attribute 1 from the mesh as essential and // converting it to a list of true dofs. Array ess_tdof_list, ess_bdr(pmesh.bdr_attributes.Max()); ess_bdr = 0; ess_bdr[0] = 1; fespace.GetEssentialTrueDofs(ess_bdr, ess_tdof_list); // 9. Set up the parallel linear form b(.) which corresponds to the // right-hand side of the FEM linear system. In this case, b_i equals the // boundary integral of f*phi_i where f represents a "pull down" force on // the Neumann part of the boundary and phi_i are the basis functions in // the finite element fespace. The force is defined by the object f, which // is a vector of Coefficient objects. The fact that f is non-zero on // boundary attribute 2 is indicated by the use of piece-wise constants // coefficient for its last component. VectorArrayCoefficient f(dim); for (int i = 0; i < dim-1; i++) { f.Set(i, new ConstantCoefficient(0.0)); } { Vector pull_force(pmesh.bdr_attributes.Max()); pull_force = 0.0; pull_force(1) = -1.0e-2; f.Set(dim-1, new PWConstCoefficient(pull_force)); } ParLinearForm b(&fespace); b.AddBoundaryIntegrator(new VectorBoundaryLFIntegrator(f)); if (Mpi::Root()) { cout << "r.h.s. ... " << flush; } b.Assemble(); // 10. Define the solution vector x as a parallel finite element grid // function corresponding to fespace. Initialize x with initial guess of // zero, which satisfies the boundary conditions. ParGridFunction x(&fespace); x = 0.0; // 11. Set up the parallel bilinear form a(.,.) on the finite element space // corresponding to the linear elasticity integrator with piece-wise // constants coefficient lambda and mu. Vector lambda(pmesh.attributes.Max()); lambda = 1.0; lambda(0) = lambda(1)*50; PWConstCoefficient lambda_func(lambda); Vector mu(pmesh.attributes.Max()); mu = 1.0; mu(0) = mu(1)*50; PWConstCoefficient mu_func(mu); ElasticityIntegrator integrator(lambda_func, mu_func); ParBilinearForm a(&fespace); if (pa || componentwise_action) { a.SetAssemblyLevel( AssemblyLevel::PARTIAL); } a.AddDomainIntegrator(&integrator); a.UseExternalIntegrators(); // 12. Assemble the parallel bilinear form and the corresponding linear // system, applying any necessary transformations such as: parallel // assembly, eliminating boundary conditions, applying conforming // constraints for non-conforming AMR, static condensation, etc. if (Mpi::Root()) { cout << "matrix ... " << flush; } StopWatch total_timer; StopWatch assembly_timer; assembly_timer.Start(); total_timer.Start(); a.Assemble(); OperatorPtr A; Vector B, X; Operator *a_lhs = componentwise_action ? nullptr : &a; if (!componentwise_action) { a.FormLinearSystem(ess_tdof_list, x, b, A, X, B); } if (Mpi::Root()) { cout << "done." << endl; cout << "Size of linear system: " << fespace.GlobalTrueVSize() << endl; } Array block_offsets(dim + 1); block_offsets[0] = 0; BlockDiagonalPreconditioner *blockDiag; unique_ptr prec = nullptr; // 13. For partial assembly, assemble forms on LOR space. Construct the block // diagonal preconditioner by fully assembling the component bilinear // on the LOR space. If additionally "-ss" is enabled, create the // block CG solvers and the high order, partially assembled components. vector> bilinear_forms; vector> lor_block; // amg_blocks stores preconditioners of lor_block. vector> amg_blocks; // cg_blocks only gets used if -ss is enabled. vector> cg_blocks; // diag_ho only used if -hoa enabled. The high order partial assembled operators // with the essential dofs eliminated and constrained to one. vector> ho_bilinear_form_blocks; vector> diag_ho; // If -ca is used, component bilinear forms are stored in pa_components, and // pointers to fespaces. vector> pa_components; vector fespaces; // get block essential boundary info. // need to allocate here since constrained operator will not own essential dofs. Array ess_tdof_list_block_ho, ess_bdr_block_ho(pmesh.bdr_attributes.Max()); ess_bdr_block_ho = 0; ess_bdr_block_ho[0] = 1; ElasticityIntegrator lor_integrator(lambda_func, mu_func); if (pa || componentwise_action) { // 13(a) Create the diagonal LOR matrices and corresponding AMG preconditioners. lor_integrator.AssemblePA(lor_disc->GetParFESpace()); for (int j = 0; j < dim; j++) { ElasticityComponentIntegrator *block = new ElasticityComponentIntegrator( lor_integrator, j, j); // create the LOR matrix and corresponding AMG preconditioners. bilinear_forms.emplace_back(new ParBilinearForm(scalar_lor_fespace.get())); bilinear_forms[j]->SetAssemblyLevel(AssemblyLevel::FULL); bilinear_forms[j]->EnableSparseMatrixSorting(Device::IsEnabled()); bilinear_forms[j]->AddDomainIntegrator(block); bilinear_forms[j]->Assemble(); // get block essential boundary info Array ess_tdof_list_block, ess_bdr_block(pmesh.bdr_attributes.Max()); ess_bdr_block = 0; ess_bdr_block[0] = 1; scalar_lor_fespace->GetEssentialTrueDofs(ess_bdr_block, ess_tdof_list_block); lor_block.emplace_back(bilinear_forms[j]->ParallelAssemble()); lor_block[j]->EliminateBC(ess_tdof_list_block, Operator::DiagonalPolicy::DIAG_ONE); amg_blocks.emplace_back(new HypreBoomerAMG); amg_blocks[j]->SetStrengthThresh(0.25); amg_blocks[j]->SetRelaxType(16); // Chebyshev amg_blocks[j]->SetOperator(*lor_block[j]); block_offsets[j+1] = amg_blocks[j]->Height(); // 13(b) If needed, create the block components for operator action. if (componentwise_action) { for (int i = 0; i < dim; i++) { ElasticityComponentIntegrator *action_block = new ElasticityComponentIntegrator( integrator, i, j); if (i == j) { fespaces.emplace_back(&scalar_fespace); } pa_components.emplace_back(new ParBilinearForm(&scalar_fespace)); pa_components[i + dim*j]->SetAssemblyLevel(pa ? AssemblyLevel::PARTIAL : AssemblyLevel::FULL); pa_components[i + dim*j]->EnableSparseMatrixSorting(Device::IsEnabled()); pa_components[i + dim*j]->AddDomainIntegrator(action_block); pa_components[i + dim*j]->Assemble(); } } } block_offsets.PartialSum(); // 13(c) If needed, create CG solvers for diagonal sub-systems. if (sub_solve) { // create diagonal high order partial assembly operators for (int i = 0; i < dim; i++) { ElasticityComponentIntegrator *block = new ElasticityComponentIntegrator( integrator, i, i); scalar_fespace.GetEssentialTrueDofs(ess_bdr_block_ho, ess_tdof_list_block_ho); ho_bilinear_form_blocks.emplace_back(new ParBilinearForm(&scalar_fespace)); ho_bilinear_form_blocks[i]->SetAssemblyLevel(AssemblyLevel::PARTIAL); ho_bilinear_form_blocks[i]->AddDomainIntegrator(block); ho_bilinear_form_blocks[i]->Assemble(); const auto *prolong = scalar_fespace.GetProlongationMatrix(); auto *rap = new RAPOperator(*prolong, *ho_bilinear_form_blocks[i], *prolong); diag_ho.emplace_back(new ConstrainedOperator(rap, ess_tdof_list_block_ho, true, Operator::DiagonalPolicy::DIAG_ONE)); } // create CG solvers for (int i = 0; i < dim; i++) { cg_blocks.emplace_back(new CGSolver(MPI_COMM_WORLD)); cg_blocks[i]->iterative_mode = false; cg_blocks[i]->SetOperator(*diag_ho[i]); cg_blocks[i]->SetPreconditioner(*amg_blocks[i]); cg_blocks[i]->SetMaxIter(30); cg_blocks[i]->SetRelTol(1e-8); } } blockDiag = new BlockDiagonalPreconditioner(block_offsets); for (int i = 0; i < dim; i++) { if (sub_solve) { blockDiag->SetDiagonalBlock(i, cg_blocks[i].get()); } else { blockDiag->SetDiagonalBlock(i, amg_blocks[i].get()); } } prec.reset(blockDiag); } else { // 13(d) If not using PA, configure preconditioner on global matrix. auto *amg = new HypreBoomerAMG(*A.As()); if (amg_elast && !a.StaticCondensationIsEnabled()) { amg->SetElasticityOptions(&fespace); } else { amg->SetSystemsOptions(dim, reorder_space); } prec.reset(amg); } // 13(e) For componentwise action, create block operator and form linear system. unique_ptr A_components = nullptr; unique_ptr pa_blocks; if (componentwise_action) { pa_blocks.reset(new BlockFESpaceOperator(fespaces)); for (int j = 0; j < dim; j++) { for (int i = 0; i < dim; i++) { pa_blocks->SetBlock(i,j,pa_components[i + dim*j].get()); } } Operator *A_temp; pa_blocks->FormLinearSystem(ess_tdof_list, x, b, A_temp, X, B); A_components.reset(A_temp); a_lhs = pa_blocks.get(); } assembly_timer.Stop(); // 14. Create the global CG solver, solve, and recover solution. CGSolver solver(MPI_COMM_WORLD); solver.SetRelTol(1e-8); solver.SetMaxIter(2500); solver.SetPrintLevel(1); if (prec) { solver.SetPreconditioner(*prec); } solver.SetOperator(A_components ? *A_components : *A); StopWatch linear_solve_timer; linear_solve_timer.Start(); solver.Mult(B, X); linear_solve_timer.Stop(); a_lhs->RecoverFEMSolution(X, b, x); total_timer.Stop(); // Print run times if (Mpi::Root()) { cout << "Elapsed Times\n"; cout << "Assembly (s) = " << assembly_timer.RealTime() << endl; cout << "Linear Solve (s) = " << linear_solve_timer.RealTime() << endl; cout << "Total Solve (s) " << total_timer.RealTime() << endl; } // 15. For non-NURBS meshes, make the mesh curved based on the finite element // space. This means that we define the mesh elements through a fespace // based transformation of the reference element. This allows us to save // the displaced mesh as a curved mesh when using high-order finite // element displacement field. We assume that the initial mesh (read from // the file) is not higher order curved mesh compared to the chosen FE // space. pmesh.SetNodalFESpace(&fespace); // 16. If visualization is enabled, Save in parallel the displaced mesh and // the inverted solution (which gives the backward displacements to the // original grid). This output can be viewed later using GLVis: "glvis // -np -m mesh -g sol". // // Also, save the displacement, with displaced mesh, to VTK. if (visualization) { GridFunction *nodes = pmesh.GetNodes(); *nodes += x; x *= -1; ostringstream mesh_name, sol_name; mesh_name << "mesh." << setfill('0') << setw(6) << Mpi::WorldRank(); sol_name << "sol." << setfill('0') << setw(6) << Mpi::WorldRank(); ofstream mesh_ofs(mesh_name.str().c_str()); mesh_ofs.precision(8); pmesh.Print(mesh_ofs); ofstream sol_ofs(sol_name.str().c_str()); sol_ofs.precision(8); x.Save(sol_ofs); ParaViewDataCollection pd("LOR_Elasticity", &pmesh); pd.SetPrefixPath("ParaView"); pd.RegisterField("displacement", &x); pd.SetLevelsOfDetail(order); pd.SetDataFormat(VTKFormat::BINARY); pd.SetHighOrderOutput(true); pd.SetCycle(0); pd.SetTime(0.0); pd.Save(); } return 0; }