diff --git a/CMakeLists.txt b/CMakeLists.txt index 37c83594..0d677b42 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -78,6 +78,7 @@ option(POLYSOLVE_WITH_UNICODE "Use Unicode in logging messages" option(POLYSOLVE_WITH_CUDA "Enable cuda support" OFF) # Polysolve options for enabling/disabling optional libraries +option(POLYSOLVE_WITH_MPI "Enable MPI support" OFF) option(POLYSOLVE_WITH_ACCELERATE "Enable Apple Accelerate" ${POLYSOLVE_ON_APPLE_SILICON}) option(POLYSOLVE_WITH_CHOLMOD "Enable Cholmod library" ON) option(POLYSOLVE_WITH_UMFPACK "Enable UmfPack library" ON) @@ -169,6 +170,19 @@ if (POLYSOLVE_WITH_CUDA) set(CMAKE_CUDA_STANDARD_REQUIRED ON) endif() +################################################################################ +# MPI +################################################################################ + +# Check that MPI exists +if(POLYSOLVE_WITH_MPI) + find_package(MPI QUIET) + if (NOT MPI_CXX_FOUND) + message(WARNING "POLYSOLVE_WITH_MPI was requested but MPI was not found, proceeding without MPI dependent solvers.") + set(POLYSOLVE_WITH_MPI OFF) + endif() +endif() + ################################################################################ # PolySolve Library ################################################################################ @@ -247,6 +261,21 @@ if(POLYSOLVE_WITH_CUDA) target_compile_definitions(polysolve_linear PUBLIC POLYSOLVE_WITH_CUDA) endif() +if(POLYSOLVE_WITH_CUDA AND NOT POLYSOLVE_LARGE_INDEX) + target_compile_definitions(polysolve PUBLIC POLYSOLVE_WITH_CUDSS) + target_compile_definitions(polysolve_linear PUBLIC POLYSOLVE_WITH_CUDSS) +endif() + +if(POLYSOLVE_WITH_CUDA AND NOT POLYSOLVE_LARGE_INDEX) + target_compile_definitions(polysolve PUBLIC POLYSOLVE_WITH_GPU_HYBRID) + target_compile_definitions(polysolve_linear PUBLIC POLYSOLVE_WITH_GPU_HYBRID) +endif() + +if(POLYSOLVE_WITH_MPI AND NOT POLYSOLVE_LARGE_INDEX) + target_compile_definitions(polysolve PUBLIC POLYSOLVE_WITH_CPU_HYBRID) + target_compile_definitions(polysolve_linear PUBLIC POLYSOLVE_WITH_CPU_HYBRID) +endif() + # Graph partition library Kaminpar does not support windows. if(POLYSOLVE_WITH_CUDA AND NOT WIN32) target_compile_definitions(polysolve PUBLIC POLYSOLVE_WITH_MAS) @@ -257,6 +286,12 @@ endif() # Dependencies ################################################################################ +# cuDSS +if (POLYSOLVE_WITH_CUDA) + include(cudss) + target_link_libraries(polysolve_linear PRIVATE cudss) +endif() + # CCCL if(POLYSOLVE_WITH_CUDA) include(cccl) @@ -275,6 +310,16 @@ endif() # Linear # ------ +if(POLYSOLVE_WITH_MPI) + message(STATUS "MPI enabled for PolySolve.") + + # Link to the linear sub-library + target_link_libraries(polysolve_linear PUBLIC MPI::MPI_CXX) + + # Define POLYSOLVE_WITH_MPI for your C++ #ifdef checks + target_compile_definitions(polysolve_linear PUBLIC POLYSOLVE_WITH_MPI) +endif() + # Accelerate solver (Include before Eigen) if(POLYSOLVE_WITH_ACCELERATE) set(BLA_VENDOR Apple) @@ -483,4 +528,4 @@ if(POLYSOLVE_WITH_TESTS) include("${catch2_SOURCE_DIR}/contrib/Catch.cmake") add_subdirectory(tests) -endif() +endif() \ No newline at end of file diff --git a/cmake/recipes/cudss.cmake b/cmake/recipes/cudss.cmake new file mode 100644 index 00000000..8c9fba48 --- /dev/null +++ b/cmake/recipes/cudss.cmake @@ -0,0 +1,27 @@ +# cuDSS solver + +if(TARGET cudss) + return() +endif() + +message(STATUS "Third-party: creating target 'cudss'") + +set(CUDSS_URL + "https://developer.download.nvidia.com/compute/cudss/redist/libcudss/linux-x86_64/libcudss-linux-x86_64-0.7.1.4_cuda13-archive.tar.xz" + CACHE STRING "cuDSS download URL") +set(CUDSS_URL_SHA256 + "84b34ebe7fad40ec10f2aab2957a63b6070bd8ce16e3ada3e6bcac7317256347" + CACHE STRING "cuDSS download URL SHA256 checksum") + +include(CPM) +CPMAddPackage( + NAME cudss + URL ${CUDSS_URL} + URL_HASH SHA256=${CUDSS_URL_SHA256} + DOWNLOAD_ONLY ON +) + +find_package(cudss CONFIG REQUIRED + PATHS "${cudss_SOURCE_DIR}/lib/cmake/cudss" + NO_DEFAULT_PATH +) diff --git a/cmake/recipes/hypre.cmake b/cmake/recipes/hypre.cmake index 88ca1098..6873bac6 100644 --- a/cmake/recipes/hypre.cmake +++ b/cmake/recipes/hypre.cmake @@ -6,7 +6,6 @@ endif() message(STATUS "Third-party: creating target 'HYPRE::HYPRE'") -set(HYPRE_ENABLE_MPI OFF CACHE INTERNAL "" FORCE) set(HYPRE_ENABLE_PRINT_ERRORS ON CACHE INTERNAL "" FORCE) set(HYPRE_ENABLE_BIGINT OFF CACHE INTERNAL "" FORCE) set(HYPRE_ENABLE_MIXEDINT OFF CACHE BOOL "" FORCE) @@ -22,6 +21,12 @@ else() set(HYPRE_ENABLE_CUDA OFF CACHE INTERNAL "" FORCE) endif() +if (POLYSOLVE_WITH_MPI) + set(HYPRE_ENABLE_MPI ON CACHE INTERNAL "" FORCE) +else() + set(HYPRE_ENABLE_MPI OFF CACHE INTERNAL "" FORCE) +endif() + # HYPRE unconditionally defines an "uninstall" target, which conflicts with other buggy libraries # as modern cmake requires unique target name. This is a hacky workaround until upstream is fixed. macro(add_custom_target _target_name) diff --git a/cmake/recipes/polyfem-data.cmake b/cmake/recipes/polyfem-data.cmake index c7a3081e..4121a46d 100644 --- a/cmake/recipes/polyfem-data.cmake +++ b/cmake/recipes/polyfem-data.cmake @@ -13,7 +13,7 @@ ExternalProject_Add( SOURCE_DIR ${POLYFEM_DATA_ROOT} GIT_REPOSITORY https://github.com/polyfem/polyfem-data - GIT_TAG 9c1bdd5bd02215e80bc1668547e5dbeb5484a527 + GIT_TAG 8a2eff19a33ccfd395c9294342e2c5c2babf2626 CONFIGURE_COMMAND "" BUILD_COMMAND "" diff --git a/linear-solver-spec.json b/linear-solver-spec.json index 88162960..fbf698c4 100644 --- a/linear-solver-spec.json +++ b/linear-solver-spec.json @@ -16,7 +16,9 @@ "Pardiso", "Hypre", "AMGCL", - "MAS" + "MAS", + "CPUHybrid", + "GPUHybrid" ], "doc": "Settings for the linear solver." }, @@ -49,7 +51,9 @@ "Eigen::BiCGSTAB", "Eigen::GMRES", "Eigen::MINRES", - "MAS" + "MAS", + "CPUHybrid", + "GPUHybrid" ] }, { @@ -506,5 +510,217 @@ "default": false, "type": "bool", "doc": "Use preconditioned residual norm for termination check." + }, + { + "pointer": "/CPUHybrid", + "default": null, + "type": "object", + "optional": [ + "block_dim", + "max_iter", + "relative_tolerance", + "absolute_tolerance", + "theta", + "decompose_subdomains", + "min_subdomain_size", + "max_subdomain_size", + "expand_subdomains", + "gmm_jump_threshold", + "gmm_tol", + "max_gmm_iterations", + "conditioning_threshold", + "additive_mode" + ], + "doc": "Settings for the CPUHybrid solver." + }, + { + "pointer": "/CPUHybrid/block_dim", + "default": 1, + "type": "int", + "doc": "Block size of the system matrix. " + }, + { + "pointer": "/CPUHybrid/max_iter", + "default": 10000, + "type": "int", + "doc": "Maximum number of PCG iterations. " + }, + { + "pointer": "/CPUHybrid/relative_tolerance", + "default": 1e-10, + "type": "float", + "doc": "Relative tolerance for solver convergence." + }, + { + "pointer": "/CPUHybrid/absolute_tolerance", + "default": 0, + "type": "float", + "doc": "Absolute tolerance for solver convergence." + }, + { + "pointer": "/CPUHybrid/theta", + "default": 0.5, + "type": "float", + "doc": "Strong threshold for AMG." + }, + { + "pointer": "/CPUHybrid/decompose_subdomains", + "default": true, + "type": "bool", + "doc": "Whether to decompose the problematic subspace into independent clusters." + }, + { + "pointer": "/CPUHybrid/min_subdomain_size", + "default": 3, + "type": "int", + "doc": "Minimum number of degrees of freedom or elements in a subdomain." + }, + { + "pointer": "/CPUHybrid/max_subdomain_size", + "default": 1000000000, + "type": "int", + "doc": "Maximum number of degrees of freedom or elements in a subdomain." + }, + { + "pointer": "/CPUHybrid/expand_subdomains", + "default": true, + "type": "bool", + "doc": "Enables one layer BFS expansion of subdomains to improve convergence." + }, + { + "pointer": "/CPUHybrid/gmm_jump_threshold", + "default": 10.0, + "type": "float", + "doc": "Threshold ratio for detecting row norm jumps in GMM." + }, + { + "pointer": "/CPUHybrid/gmm_tol", + "default": 1e-3, + "type": "float", + "doc": "Tolerance for the Gaussian Mixture Model. " + }, + { + "pointer": "/CPUHybrid/max_gmm_iterations", + "default": 20, + "type": "int", + "doc": "Maximum number of iterations allowed for the GMM routine." + }, + { + "pointer": "/CPUHybrid/conditioning_threshold", + "default": 100, + "type": "float", + "doc": "Minimal worst-case condition number to include in subspace correction. " + }, + { + "pointer": "/CPUHybrid/additive_mode", + "default": false, + "type": "bool", + "doc": "Apply subspace correction additively. " + }, + { + "pointer": "/GPUHybrid", + "default": null, + "type": "object", + "optional": [ + "block_dim", + "max_iter", + "relative_tolerance", + "absolute_tolerance", + "theta", + "decompose_subdomains", + "min_subdomain_size", + "max_subdomain_size", + "expand_subdomains", + "gmm_jump_threshold", + "gmm_tol", + "max_gmm_iterations", + "conditioning_threshold", + "additive_mode" + ], + "doc": "Settings for the GPUHybrid solver." + }, + { + "pointer": "/GPUHybrid/block_dim", + "default": 1, + "type": "int", + "doc": "Block size of the system matrix. " + }, + { + "pointer": "/GPUHybrid/max_iter", + "default": 10000, + "type": "int", + "doc": "Maximum number of PCG iterations. " + }, + { + "pointer": "/GPUHybrid/relative_tolerance", + "default": 1e-10, + "type": "float", + "doc": "Relative tolerance for solver convergence." + }, + { + "pointer": "/GPUHybrid/absolute_tolerance", + "default": 0, + "type": "float", + "doc": "Absolute tolerance for solver convergence." + }, + { + "pointer": "/GPUHybrid/theta", + "default": 0.5, + "type": "float", + "doc": "Strong threshold for AMG." + }, + { + "pointer": "/GPUHybrid/decompose_subdomains", + "default": true, + "type": "bool", + "doc": "Whether to decompose the problematic subspace into independent clusters." + }, + { + "pointer": "/GPUHybrid/min_subdomain_size", + "default": 3, + "type": "int", + "doc": "Minimum number of degrees of freedom or elements in a subdomain." + }, + { + "pointer": "/GPUHybrid/max_subdomain_size", + "default": 1000000000, + "type": "int", + "doc": "Maximum number of degrees of freedom or elements in a subdomain." + }, + { + "pointer": "/GPUHybrid/expand_subdomains", + "default": true, + "type": "bool", + "doc": "Enables one layer BFS expansion of subdomains to improve convergence." + }, + { + "pointer": "/GPUHybrid/gmm_jump_threshold", + "default": 10.0, + "type": "float", + "doc": "Threshold ratio for detecting row norm jumps in GMM." + }, + { + "pointer": "/GPUHybrid/gmm_tol", + "default": 1e-3, + "type": "float", + "doc": "Tolerance for the Gaussian Mixture Model. " + }, + { + "pointer": "/GPUHybrid/max_gmm_iterations", + "default": 20, + "type": "int", + "doc": "Maximum number of iterations allowed for the GMM routine." + }, + { + "pointer": "/GPUHybrid/conditioning_threshold", + "default": 100, + "type": "float", + "doc": "Minimal worst-case condition number to include in subspace correction. " + }, + { + "pointer": "/GPUHybrid/additive_mode", + "default": false, + "type": "bool", + "doc": "Apply subspace correction additively. " } ] diff --git a/src/polysolve/linear/CMakeLists.txt b/src/polysolve/linear/CMakeLists.txt index 7a7d3eeb..0430402b 100644 --- a/src/polysolve/linear/CMakeLists.txt +++ b/src/polysolve/linear/CMakeLists.txt @@ -15,8 +15,26 @@ set(SOURCES Pardiso.hpp SaddlePointSolver.cpp SaddlePointSolver.hpp + hybrid_utils/DisjointSet.cpp + hybrid_utils/DisjointSet.hpp ) +if(POLYSOLVE_WITH_CUDA) + list(APPEND SOURCES + GPUHybridSolver.cu + GPUHybridSolver.hpp + cuDSS.hpp + cuDSS.cpp + ) +endif() + +if(POLYSOLVE_WITH_MPI) + list(APPEND SOURCES + CPUHybridSolver.cpp + CPUHybridSolver.hpp + ) +endif() + if(POLYSOLVE_WITH_CUDA AND NOT WIN32) list(APPEND SOURCES MASSolver.cu diff --git a/src/polysolve/linear/CPUHybridSolver.cpp b/src/polysolve/linear/CPUHybridSolver.cpp new file mode 100644 index 00000000..fccb026b --- /dev/null +++ b/src/polysolve/linear/CPUHybridSolver.cpp @@ -0,0 +1,1455 @@ + +//////////////////////////////////////////////////////////////////////////////// +#include "CPUHybridSolver.hpp" +//////////////////////////////////////////////////////////////////////////////// + +#include "hybrid_utils/DisjointSet.hpp" + +#include + +#include +#include + +#if POLYSOLVE_WITH_ACCELERATE +#include +#endif + +#if POLYSOLVE_WITH_MKL +#include +#endif + +#include +#include +#include + +#include "_hypre_IJ_mv.h" +#include "_hypre_parcsr_mv.h" + +namespace polysolve::linear +{ + + //////////////////////////////////////////////////////////////////////////////// + namespace + { + using clock = std::chrono::steady_clock; + + double elapsed_seconds(const std::chrono::time_point &begin) + { + return std::chrono::duration(clock::now() - begin).count(); + } + } // namespace + + CPUHybridSolver::CPUHybridSolver() + { + // check if MPI is initialized + int done_already; + MPI_Initialized(&done_already); + + if (!done_already) + { + MPI_Init(nullptr, nullptr); + } + + if (!HYPRE_Initialized()) + { + HYPRE_Initialize(); + } + + // get MPI rank information + MPI_Comm_rank(MPI_COMM_WORLD, &myid); + MPI_Comm_size(MPI_COMM_WORLD, &num_procs); + + if (is_running_worker_loop) + { + return; + } + + Eigen::setNbThreads(1); + HYPRE_SetMemoryLocation(HYPRE_MEMORY_HOST); + HYPRE_SetExecutionPolicy(HYPRE_EXEC_HOST); + + if (myid != 0) + { + spdlog::set_level(spdlog::level::off); + } + else + { + spdlog::flush_on(spdlog::level::info); + } + + if (myid != 0) + { + is_running_worker_loop = true; + + run_worker_loop(); + + int finalized; + MPI_Finalized(&finalized); + if (!finalized) + { + MPI_Finalize(); + } + if (!HYPRE_Finalized()) + { + HYPRE_Finalize(); + } + std::exit(0); + } + + solver_id = next_id++; + + SolverCmd cmd = CMD_CREATE; + MPI_Bcast(&cmd, 1, MPI_INT, 0, MPI_COMM_WORLD); + MPI_Bcast(&solver_id, 1, MPI_INT, 0, MPI_COMM_WORLD); + } + + // Set solver parameters + void CPUHybridSolver::set_parameters(const json ¶ms) + { + if (myid == 0) + { + SolverCmd cmd = CMD_SET_PARAMETERS; + MPI_Bcast(&cmd, 1, MPI_INT, 0, MPI_COMM_WORLD); + MPI_Bcast(&solver_id, 1, MPI_INT, 0, MPI_COMM_WORLD); + } + + json shared_params; + std::string json_str; + int str_size = 0; + + if (myid == 0) + { + shared_params = params; // Use the incoming params on root + json_str = shared_params.dump(); + str_size = static_cast(json_str.size()); + } + + MPI_Bcast(&str_size, 1, MPI_INT, 0, MPI_COMM_WORLD); + + if (myid != 0) + { + json_str.resize(str_size); + } + + MPI_Bcast(json_str.data(), str_size, MPI_CHAR, 0, MPI_COMM_WORLD); + shared_params = json::parse(json_str); + + if (shared_params.contains("CPUHybrid")) + { + if (shared_params["CPUHybrid"].contains("max_iter")) + { + max_iter_ = shared_params["CPUHybrid"]["max_iter"]; + } + if (shared_params["CPUHybrid"].contains("relative_tolerance")) + { + rel_conv_tol_ = shared_params["CPUHybrid"]["relative_tolerance"]; + } + if (shared_params["CPUHybrid"].contains("absolute_tolerance")) + { + abs_conv_tol_ = shared_params["CPUHybrid"]["absolute_tolerance"]; + } + if (shared_params["CPUHybrid"].contains("theta")) + { + theta = shared_params["CPUHybrid"]["theta"]; + } + if (shared_params["CPUHybrid"].contains("block_dim")) + { + dimension_ = shared_params["CPUHybrid"]["block_dim"]; + } + if (shared_params["CPUHybrid"].contains("decompose_subdomains")) + { + decompose_subdomains = shared_params["CPUHybrid"]["decompose_subdomains"]; + } + if (shared_params["CPUHybrid"].contains("min_subdomain_size")) + { + min_subdomain_size = shared_params["CPUHybrid"]["min_subdomain_size"]; + } + if (shared_params["CPUHybrid"].contains("max_subdomain_size")) + { + max_subdomain_size = shared_params["CPUHybrid"]["max_subdomain_size"]; + } + if (shared_params["CPUHybrid"].contains("expand_subdomains")) + { + expand_subdomains = shared_params["CPUHybrid"]["expand_subdomains"]; + } + if (shared_params["CPUHybrid"].contains("gmm_jump_threshold")) + { + gmm_jump_threshold = shared_params["CPUHybrid"]["gmm_jump_threshold"]; + } + if (shared_params["CPUHybrid"].contains("gmm_tol")) + { + gmm_tol = shared_params["CPUHybrid"]["gmm_tol"]; + } + if (shared_params["CPUHybrid"].contains("max_gmm_iterations")) + { + max_gmm_iterations = shared_params["CPUHybrid"]["max_gmm_iterations"]; + } + if (shared_params["CPUHybrid"].contains("conditioning_threshold")) + { + conditioning_threshold = shared_params["CPUHybrid"]["conditioning_threshold"]; + } + if (shared_params["CPUHybrid"].contains("additive_mode")) + { + additive_mode = shared_params["CPUHybrid"]["additive_mode"]; + } + } + } + + void CPUHybridSolver::check_settings() const + { + if (myid != 0) + { + return; + } + } + + void CPUHybridSolver::get_info(json ¶ms) const + { + params["num_iterations"] = num_iterations; + params["final_res_norm"] = final_res_norm; + } + + //////////////////////////////////////////////////////////////////////////////// + + void CPUHybridSolver::factorize(const StiffnessMatrix &Ain) + { + if (myid == 0) + { + SolverCmd cmd = CMD_FACTORIZE; + MPI_Bcast(&cmd, 1, MPI_INT, 0, MPI_COMM_WORLD); + MPI_Bcast(&solver_id, 1, MPI_INT, 0, MPI_COMM_WORLD); + } + + check_settings(); + + int rows, cols, nnz; + if (myid == 0) + { + SPDLOG_TRACE("[{}] [start_solve] [0.000000] [num_procs={}] [problem_size={}]", name(), num_procs, Ain.rows()); + rows = Ain.rows(); + cols = Ain.cols(); + nnz = Ain.nonZeros(); + } + + MPI_Bcast(&rows, 1, MPI_INT, 0, MPI_COMM_WORLD); + MPI_Bcast(&cols, 1, MPI_INT, 0, MPI_COMM_WORLD); + MPI_Bcast(&nnz, 1, MPI_INT, 0, MPI_COMM_WORLD); + + partition_ranks(rows); + + MPI_Win A_win; + double *values; + int *inner_indices; + int *outer_pointers; + + { + auto phase_begin = clock::now(); + + uint64_t val_bytes = nnz * sizeof(double); + uint64_t inner_bytes = nnz * sizeof(int); + uint64_t outer_bytes = (cols + 1) * sizeof(int); + uint64_t total_bytes = myid == 0 ? val_bytes + inner_bytes + outer_bytes : 0; + + void *A_ptr; + MPI_Win_allocate_shared(total_bytes, 1, MPI_INFO_NULL, MPI_COMM_WORLD, &A_ptr, &A_win); + + if (myid != 0) + { + int disp_unit; + MPI_Aint sz; + MPI_Win_shared_query(A_win, 0, &sz, &disp_unit, &A_ptr); + } + + values = (double *)A_ptr; + inner_indices = (int *)((char *)A_ptr + val_bytes); + outer_pointers = (int *)((char *)A_ptr + inner_bytes + val_bytes); + + MPI_Win_fence(0, A_win); + + if (myid == 0) + { + std::memcpy(values, Ain.valuePtr(), val_bytes); + std::memcpy(inner_indices, Ain.innerIndexPtr(), inner_bytes); + std::memcpy(outer_pointers, Ain.outerIndexPtr(), outer_bytes); + } + + MPI_Win_fence(0, A_win); + + SPDLOG_TRACE("[{}] [create_shared_matrix_window] [{:.6f}]", name(), elapsed_seconds(phase_begin)); + } + + SharedSparseMatrix shared_A(rows, cols, nnz, outer_pointers, inner_indices, values); + + auto phase_begin = clock::now(); + + bad_indices_sets.clear(); + bad_indices_arrays.clear(); + select_bad_dofs(shared_A); + + if (myid == 0) + { + if (decompose_subdomains) + { + filter_subdomains(shared_A); + } + + if (expand_subdomains) + { + expand_subdomains_to_strongly_connected(shared_A); + } + + if (decompose_subdomains) + { + decompose_subdomains_to_disjoint_subsets(shared_A); + } + else + { + bad_indices_sets.emplace_back(all_bad_dofs.begin(), all_bad_dofs.end()); + } + + load_balance_subdomains(); + } + + share_bad_subdomains(); + factorize_submatrix(shared_A); + + SPDLOG_TRACE("[{}] [setup_problematic_dof_precond] [{:.6f}]", name(), elapsed_seconds(phase_begin)); + + if (has_matrix_) + { + HYPRE_IJMatrixDestroy(A); + has_matrix_ = false; + A = nullptr; + } + + { + auto phase_begin = clock::now(); + copy_matrix_to_hypre(shared_A); + has_matrix_ = true; + SPDLOG_TRACE("[{}] [copy_matrix_to_hypre] [{:.6f}]", name(), elapsed_seconds(phase_begin)); + } + + MPI_Win_free(&A_win); + } + + //////////////////////////////////////////////////////////////////////////////// + + namespace + { + + void eigen_to_hypre_par_vec(HYPRE_ParVector &par_x, HYPRE_IJVector &ij_x, const Eigen::VectorXd &x, int start_i, int end_i) + { + HYPRE_IJVectorSetValues(ij_x, end_i - start_i + 1, nullptr, x.data() + start_i); + HYPRE_IJVectorAssemble(ij_x); + HYPRE_IJVectorGetObject(ij_x, (void **)&par_x); + } + + void hypre_vec_to_eigen(const HYPRE_IJVector &ij_x, Eigen::Ref x, int start_i, int end_i, int num_procs) + { + x.setZero(); + HYPRE_IJVectorGetValues(ij_x, end_i - start_i + 1, nullptr, x.data() + start_i); + + std::vector recv_counts(num_procs); + std::vector displs(num_procs); + + int local_size = end_i - start_i + 1; + MPI_Allgather(&local_size, 1, MPI_INT, recv_counts.data(), 1, MPI_INT, MPI_COMM_WORLD); + + displs[0] = 0; + for (int i = 1; i < num_procs; ++i) + { + displs[i] = displs[i - 1] + recv_counts[i - 1]; + } + + MPI_Allgatherv(x.data() + start_i, local_size, MPI_DOUBLE, + x.data(), recv_counts.data(), displs.data(), + MPI_DOUBLE, MPI_COMM_WORLD); + } + + void HypreBoomerAMG_SetDefaultOptions(HYPRE_Solver &amg_precond) + { + // AMG coarsening options: + int coarsen_type = 10; // 10 = HMIS, 8 = PMIS, 6 = Falgout, 0 = CLJP + int agg_levels = 1; // number of aggressive coarsening levels + double theta = 0.25; // strength threshold: 0.25, 0.5, 0.8 + + // AMG interpolation options: + int interp_type = 6; // 6 = extended+i, 0 = classical + int Pmax = 4; // max number of elements per row in P + + // AMG relaxation options: + int relax_type = 8; // 8 = l1-GS, 6 = symm. GS, 3 = GS, 18 = l1-Jacobi + int relax_sweeps = 1; // relaxation sweeps on each level + + // Additional options: + int print_level = 0; // print AMG iterations? 1 = no, 2 = yes + int max_levels = 25; // max number of levels in AMG hierarchy + + int min_coarse_size = 5; + + HYPRE_BoomerAMGSetCoarsenType(amg_precond, coarsen_type); + HYPRE_BoomerAMGSetAggNumLevels(amg_precond, agg_levels); + HYPRE_BoomerAMGSetRelaxType(amg_precond, relax_type); + + // relax_type = 88; + HYPRE_BoomerAMGSetMinCoarseSize(amg_precond, min_coarse_size); + // HYPRE_BoomerAMGSetCycleRelaxType(amg_precond, relax_type, 1); + // HYPRE_BoomerAMGSetCycleRelaxType(amg_precond, relax_type, 2); + HYPRE_BoomerAMGSetCycleRelaxType(amg_precond, relax_type, 3); + // HYPRE_BoomerAMGSetDebugFlag(amg_precond, 1); + // HYPRE_BoomerAMGSetNodal(amg_precond, 0); + // HYPRE_BoomerAMGSetNodalDiag(amg_precond, 0); + HYPRE_BoomerAMGSetNumSweeps(amg_precond, relax_sweeps); + HYPRE_BoomerAMGSetStrongThreshold(amg_precond, theta); + HYPRE_BoomerAMGSetInterpType(amg_precond, interp_type); + HYPRE_BoomerAMGSetPMaxElmts(amg_precond, Pmax); + // print_level = 3; + HYPRE_BoomerAMGSetPrintLevel(amg_precond, print_level); + HYPRE_BoomerAMGSetMaxLevels(amg_precond, max_levels); + + // Use as a preconditioner (one V-cycle, zero tolerance) + HYPRE_BoomerAMGSetMaxIter(amg_precond, 1); + HYPRE_BoomerAMGSetTol(amg_precond, 0.0); + } + + void HypreBoomerAMG_SetElasticityOptions(HYPRE_Solver &amg_precond, int dim, double theta) + { + // Make sure the systems AMG options are set + HYPRE_BoomerAMGSetNumFunctions(amg_precond, dim); + + // More robust options with respect to convergence + HYPRE_BoomerAMGSetAggNumLevels(amg_precond, 0); + HYPRE_BoomerAMGSetStrongThreshold(amg_precond, theta); + } + + } // anonymous namespace + + //////////////////////////////////////////////////////////////////////////////// + + void CPUHybridSolver::solve(const Eigen::Ref rhs, Eigen::Ref result) + { + if (myid == 0) + { + SolverCmd cmd = CMD_SOLVE; + MPI_Bcast(&cmd, 1, MPI_INT, 0, MPI_COMM_WORLD); + MPI_Bcast(&solver_id, 1, MPI_INT, 0, MPI_COMM_WORLD); + } + + int problem_size = rhs.size(); + MPI_Bcast(&problem_size, 1, MPI_INT, 0, MPI_COMM_WORLD); + + MPI_Win vec_win; + void *vec_ptr; + create_shared_vec(vec_win, vec_ptr, 3 * problem_size); + SharedVector shared_vec((double *)vec_ptr, 3 * problem_size); + + shared_rhs.resize(my_size()); + shared_result.resize(my_size()); + z1.resize(my_size()); + z2.resize(my_size()); + z3.resize(my_size()); + r.resize(my_size()); + p.resize(my_size()); + buffer.resize(my_size()); + + std::vector displs(num_procs); + std::vector recv_counts(num_procs); + + for (int i = 0; i < num_procs; ++i) + { + int local_size = ends[i] - starts[i] + 1; + recv_counts[i] = local_size; + displs[i] = starts[i]; + } + + int local_size = my_size(); + MPI_Scatterv( + rhs.data(), + recv_counts.data(), + displs.data(), + MPI_DOUBLE, + shared_rhs.data(), + local_size, + MPI_DOUBLE, + 0, + MPI_COMM_WORLD); + + MPI_Scatterv( + result.data(), + recv_counts.data(), + displs.data(), + MPI_DOUBLE, + shared_result.data(), + local_size, + MPI_DOUBLE, + 0, + MPI_COMM_WORLD); + + HYPRE_ParVector par_b; + HYPRE_ParVector par_x; + init_hypre_vectors(); + + /* AMG preconditioner */ + HYPRE_Solver precond; + + /* Now set up the AMG preconditioner and specify any parameters */ + { + auto phase_begin = clock::now(); + HYPRE_BoomerAMGCreate(&precond); + + HypreBoomerAMG_SetDefaultOptions(precond); + if (dimension_ > 1) + { + HypreBoomerAMG_SetElasticityOptions( + precond, + dimension_, + theta); + } + + MPI_Barrier(MPI_COMM_WORLD); + + HYPRE_IJVectorSetData(ij_b, shared_rhs.data()); + HYPRE_IJVectorAssemble(ij_b); + HYPRE_IJVectorGetObject(ij_b, (void **)&par_b); + + HYPRE_IJVectorSetData(ij_x, shared_result.data()); + HYPRE_IJVectorAssemble(ij_x); + HYPRE_IJVectorGetObject(ij_x, (void **)&par_x); + + HYPRE_BoomerAMGSetup(precond, parcsr_A, par_b, par_x); + SPDLOG_TRACE("[{}] [amg_setup] [{:.6f}]", name(), elapsed_seconds(phase_begin)); + } + + /* Now setup and solve! */ + { + auto phase_begin = clock::now(); + + pcg_solve(shared_rhs, shared_result, par_b, par_x, precond, shared_vec, vec_win); + + MPI_Gatherv( + shared_result.data(), + local_size, + MPI_DOUBLE, + result.data(), + recv_counts.data(), + displs.data(), + MPI_DOUBLE, + 0, + MPI_COMM_WORLD); + + Eigen::VectorXd A_times_result; + matmul(shared_result, buffer); + buffer = shared_rhs - buffer; + final_res_norm = sqrt(dot(buffer, buffer)); + SPDLOG_TRACE("[{}] [pcg_solve] [{:.6f}] [pcg_iters={}] [residual={}]", name(), elapsed_seconds(phase_begin), num_iterations, final_res_norm); + } + + /* Destroy preconditioner */ + { + HYPRE_BoomerAMGDestroy(precond); + HYPRE_IJVectorDestroy(ij_x); + HYPRE_IJVectorDestroy(ij_b); + MPI_Win_free(&vec_win); + } + } + + void CPUHybridSolver::pcg_solve(Eigen::VectorXd &rhs, Eigen::VectorXd &result, HYPRE_ParVector &par_b, HYPRE_ParVector &par_x, HYPRE_Solver &precond, SharedVector &vec, MPI_Win &vec_win) + { + MPI_Barrier(MPI_COMM_WORLD); + + double pre_loop_time; + double bi_prod, rel_eps, abs_eps, gamma, old_gamma; + + r.setZero(); + { + auto phase_begin = clock::now(); + + bi_prod = dot(rhs, rhs); + if (bi_prod > 0.0) + { + rel_eps = rel_conv_tol_ * rel_conv_tol_; + abs_eps = abs_conv_tol_ * abs_conv_tol_; + } + else + { + result.setZero(); + MPI_Barrier(MPI_COMM_WORLD); + return; + } + + matmul(result, buffer); + r += rhs - buffer; + + p.setZero(); + z1.setZero(); + + { + custom_mixed_precond_iter(precond, r, z1, vec, vec_win); + } + + p = z1; + + gamma = dot(r, z1); + old_gamma = gamma; + SPDLOG_TRACE("[{}] [pre_loop] [{:.6f}] [rhs_norm={}]", name(), elapsed_seconds(phase_begin), sqrt(bi_prod)); + } + + for (int k = 0; k < max_iter_; ++k) + { + auto phase_begin = clock::now(); + num_iterations = k + 1; + + matmul(p, buffer); + double sdotp = dot(p, buffer); + + if (sdotp == 0.0) + { + SPDLOG_TRACE("[{}] [err_zero_sdotp] [0.000000]", name()); + break; + } + + double alpha = gamma / sdotp; + + if (alpha <= 0.0) + { + SPDLOG_TRACE("[{}] [err_negative_alpha] [0.000000]", name()); + break; + } + else if (alpha < __DBL_MIN__) + { + SPDLOG_TRACE("[{}] [err_subnormal_alpha] [0.000000]", name()); + break; + } + + result += alpha * p; + r -= alpha * buffer; + + double i_prod = dot(r, r); + if (rel_eps > 0 && (i_prod / bi_prod) < rel_eps) + { + SPDLOG_TRACE("[{}] [converged_rel] [0.000000]", name()); + break; + } + + if (abs_eps > 0 && i_prod < abs_eps) + { + SPDLOG_TRACE("[{}] [converged_abs] [0.000000]", name()); + break; + } + + z1.setZero(); + + custom_mixed_precond_iter(precond, r, z1, vec, vec_win); + + gamma = dot(r, z1); + double beta = gamma / old_gamma; + old_gamma = gamma; + + p = z1 + beta * p; + SPDLOG_TRACE("[{}] [pcg_iter] [{:.6f}] [iter={}] [residual={}]", name(), elapsed_seconds(phase_begin), k, sqrt(i_prod)); + } + MPI_Barrier(MPI_COMM_WORLD); + } + + void CPUHybridSolver::custom_mixed_precond_iter(const HYPRE_Solver &precond, Eigen::VectorXd &r, Eigen::VectorXd &z, SharedVector &vec, MPI_Win &vec_win) + { + z1.setZero(); + z2.setZero(); + z3.setZero(); + + if (bad_indices_arrays.size() == 0 || bad_indices_arrays[0].size() == 0) + { + amg_precond_iter(precond, r, z1); + z = z1; + return; + } + + if (additive_mode) + { + amg_precond_iter(precond, r, z1); + dss_precond_iter(z3, r, z2, vec, vec_win); + z = z1 + z2; + } + else + { + amg_precond_iter(precond, r, z1); + dss_precond_iter(z1, r, z2, vec, vec_win); + matmul(z2, buffer); + Eigen::VectorXd curr_r = r - buffer; + amg_precond_iter(precond, curr_r, z3); + z = z2 + z3; + } + } + + void CPUHybridSolver::amg_precond_iter(const HYPRE_Solver &precond, Eigen::VectorXd &eigen_b, Eigen::VectorXd &eigen_x) + { + auto phase_begin = clock::now(); + HYPRE_ParVector par_x; + HYPRE_ParVector par_b; + + MPI_Barrier(MPI_COMM_WORLD); + + HYPRE_IJVectorSetData(ij_x, eigen_x.data()); + HYPRE_IJVectorSetData(ij_b, eigen_b.data()); + HYPRE_IJVectorAssemble(ij_x); + HYPRE_IJVectorGetObject(ij_x, (void **)&par_x); + HYPRE_IJVectorAssemble(ij_b); + HYPRE_IJVectorGetObject(ij_b, (void **)&par_b); + + HYPRE_BoomerAMGSolve(precond, parcsr_A, par_b, par_x); + MPI_Barrier(MPI_COMM_WORLD); + SPDLOG_TRACE("[{}] [amg_v_cycle] [{:.6f}]", name(), elapsed_seconds(phase_begin)); + } + + void CPUHybridSolver::dss_precond_iter(Eigen::VectorXd &z, Eigen::VectorXd &r, Eigen::VectorXd &next_z, SharedVector &vec, MPI_Win &vec_win) + { + auto phase_begin = clock::now(); + const int problem_size = vec.size() / 3; + + MPI_Win_fence(0, vec_win); + for (int i = 0; i < my_size(); ++i) + { + vec(starts[myid] + i) = z(i); + vec(problem_size + starts[myid] + i) = r(i); + vec(2 * problem_size + starts[myid] + i) = 0.0; + } + MPI_Win_fence(0, vec_win); + + int index_counter = 0; + for (int index : bad_subdomain_assignments[myid]) + { + auto &subdomain = bad_indices_arrays[index]; + Eigen::VectorXd sub_rhs(subdomain.size()); + Eigen::VectorXd sub_result(subdomain.size()); + + for (int i = 0; i < subdomain.size(); ++i) + { + sub_rhs(index_mappings[index_counter][subdomain[i]]) = vec(problem_size + subdomain[i]) - global_to_row[subdomain[i]].dot(vec.segment(0, problem_size)); + } + + { + sub_result = D_solvers[index_counter]->solve(sub_rhs); + } + + for (int i = 0; i < subdomain.size(); ++i) + { + vec(2 * problem_size + subdomain[i]) = sub_result(index_mappings[index_counter][subdomain[i]]); + } + ++index_counter; + } + MPI_Barrier(MPI_COMM_WORLD); + + MPI_Win_fence(0, vec_win); + for (int i = 0; i < my_size(); ++i) + { + next_z(i) = vec(starts[myid] + i) + vec(2 * problem_size + starts[myid] + i); + } + MPI_Win_fence(0, vec_win); + + SPDLOG_TRACE("[{}] [subdomain_solve] [{:.6f}]", name(), elapsed_seconds(phase_begin)); + } + + void CPUHybridSolver::select_bad_dofs(SharedSparseMatrix &sparse_A) + { + auto phase_begin = clock::now(); + + MPI_Win row_norm_win; + int local_alloc_size = (myid == 0) ? (sparse_A.rows() * sizeof(double)) : 0; + void *row_norm_ptr; + + MPI_Win_allocate_shared(local_alloc_size, 1, MPI_INFO_NULL, MPI_COMM_WORLD, &row_norm_ptr, &row_norm_win); + if (myid != 0) + { + int disp_unit; + MPI_Aint sz; + MPI_Win_shared_query(row_norm_win, 0, &sz, &disp_unit, &row_norm_ptr); + } + + SharedVector row_norms((double *)row_norm_ptr, sparse_A.rows()); + + MPI_Win_fence(0, row_norm_win); + for (int i = starts[myid]; i <= ends[myid]; ++i) + { + double row_sum = 0; + for (SharedSparseMatrix::InnerIterator it(sparse_A, i); it; ++it) + { + row_sum += std::abs(it.value()); + } + row_norms(i) = row_sum; + } + MPI_Win_fence(0, row_norm_win); + + double global_mean = row_norms.segment(starts[myid], my_size()).sum() / row_norms.size(); + MPI_Allreduce(MPI_IN_PLACE, &global_mean, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); + + double global_var = (row_norms.segment(starts[myid], my_size()).array() - global_mean).square().sum() / row_norms.size(); + MPI_Allreduce(MPI_IN_PLACE, &global_var, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); + + all_bad_dofs.clear(); + + double mean_0 = 0.0, mean_1 = 0.0; + double var_0 = 0.0, var_1 = 0.0; + int gmm_iter = 0; + + double max_dist = -1.0; + double max_jump = -1.0; + double min_cost = -1.0; + int split_idx = 0; + + MPI_Win gamma_win; + void *gamma_ptr; + MPI_Win_allocate_shared(2 * local_alloc_size, 1, MPI_INFO_NULL, MPI_COMM_WORLD, &gamma_ptr, &gamma_win); + if (myid != 0) + { + int disp_unit; + MPI_Aint sz; + MPI_Win_shared_query(gamma_win, 0, &sz, &disp_unit, &gamma_ptr); + } + + SharedVector gamma((double *)gamma_ptr, 2 * sparse_A.rows()); + + mean_0 = row_norms.minCoeff(); + var_0 = global_var; + mean_1 = row_norms.maxCoeff(); + var_1 = global_var; + double w0 = 0.5; + double w1 = 0.5; + double var_reg = 1e-6; + + for (gmm_iter = 0; gmm_iter < max_gmm_iterations; ++gmm_iter) + { + double log_w0 = std::log(w0); + double log_w1 = std::log(w1); + double log_norm_const_0 = -0.5 * std::log(2.0 * M_PI * var_0); + double log_norm_const_1 = -0.5 * std::log(2.0 * M_PI * var_1); + + auto x = row_norms.segment(starts[myid], my_size()).array(); + + Eigen::ArrayXd log_g0 = log_w0 + log_norm_const_0 - 0.5 * (x - mean_0).square() / var_0; + Eigen::ArrayXd log_g1 = log_w1 + log_norm_const_1 - 0.5 * (x - mean_1).square() / var_1; + Eigen::ArrayXd max_log_g = log_g0.cwiseMax(log_g1); + Eigen::ArrayXd log_total = max_log_g + ((log_g0 - max_log_g).exp() + (log_g1 - max_log_g).exp()).log(); + + MPI_Win_fence(0, gamma_win); + gamma.segment(starts[myid], my_size()).array() = (log_g0 - log_total).exp(); + gamma.segment(starts[myid] + row_norms.size(), my_size()).array() = (log_g1 - log_total).exp(); + MPI_Win_fence(0, gamma_win); + + w0 = 1.0 / row_norms.size() * gamma.segment(starts[myid], my_size()).sum(); + w1 = 1.0 / row_norms.size() * gamma.segment(starts[myid] + row_norms.size(), my_size()).sum(); + + MPI_Allreduce(MPI_IN_PLACE, &w0, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); + MPI_Allreduce(MPI_IN_PLACE, &w1, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); + + double old_mean_0 = mean_0; + double old_mean_1 = mean_1; + double old_var_0 = var_0; + double old_var_1 = var_1; + + mean_0 = (row_norms.segment(starts[myid], my_size()).array() * gamma.segment(starts[myid], my_size()).array()).sum() / (w0 * row_norms.size()); + mean_1 = (row_norms.segment(starts[myid], my_size()).array() * gamma.segment(starts[myid] + row_norms.size(), my_size()).array()).sum() / (w1 * row_norms.size()); + + MPI_Allreduce(MPI_IN_PLACE, &mean_0, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); + MPI_Allreduce(MPI_IN_PLACE, &mean_1, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); + + var_0 = (gamma.segment(starts[myid], my_size()).array() * (row_norms.segment(starts[myid], my_size()).array() - mean_0).square()).sum() / (w0 * row_norms.size()); + var_1 = (gamma.segment(starts[myid] + row_norms.size(), my_size()).array() * (row_norms.segment(starts[myid], my_size()).array() - mean_1).square()).sum() / (w1 * row_norms.size()); + + MPI_Allreduce(MPI_IN_PLACE, &var_0, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); + MPI_Allreduce(MPI_IN_PLACE, &var_1, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); + + var_0 += var_reg; + var_1 += var_reg; + + if (std::abs(mean_0 - old_mean_0) / std::abs(old_mean_0) < gmm_tol && std::abs(mean_1 - old_mean_1) / std::abs(old_mean_1) < gmm_tol && std::abs(var_0 - old_var_0) / std::abs(old_var_0) < gmm_tol && std::abs(var_1 - old_var_1) / std::abs(old_var_1) < gmm_tol) + { + break; + } + } + + if (myid == 0) + { + if (std::abs(mean_1) / std::abs(mean_0) > gmm_jump_threshold) + { + for (int i = 0; i < row_norms.size(); ++i) + { + if (gamma(i) < gamma(i + row_norms.size())) + { + all_bad_dofs.insert(i); + } + } + } + } + + MPI_Win_free(&gamma_win); + MPI_Win_free(&row_norm_win); + + SPDLOG_TRACE("[{}] [bad_dof_selection] [{}] [strategy=GMM] [global_mean={}] [global_var={}] [mean_0={}] [mean_1={}] [var_0={}] [var_1={}] [gmm_iters={}] [num_bad_dofs={}]", + name(), elapsed_seconds(phase_begin), global_mean, global_var, mean_0, mean_1, var_0, var_1, gmm_iter, all_bad_dofs.size()); + } + + void CPUHybridSolver::factorize_submatrix(SharedSparseMatrix &sparse_A) + { + auto phase_begin = clock::now(); + D_solvers.clear(); + + for (int i : bad_subdomain_assignments[myid]) + { + if (bad_indices_sets[i].size() > 1000) + { +#if POLYSOLVE_WITH_MKL + D_solvers.push_back(std::make_unique>>>()); +#elif POLYSOLVE_WITH_ACCELERATE + // Uses Apple's Accelerate framework wrapper in Eigen + D_solvers.push_back(std::make_unique>>>()); +#else + // Fallback if neither high-performance solver is compiled + D_solvers.push_back(std::make_unique>>>()); +#endif + } + else + { + // Stick to the lightweight solver for smaller subdomains + D_solvers.push_back(std::make_unique>>>()); + } + } + + build_index_mappings(); + + int i_counter = 0; + for (int i : bad_subdomain_assignments[myid]) + { + Eigen::SparseMatrix D; + assemble_D(i_counter, i, D, sparse_A); + D_solvers[i_counter]->compute(D); + ++i_counter; + } + + MPI_Barrier(MPI_COMM_WORLD); + SPDLOG_TRACE("[{}] [factorize_submatrix] [{}]", name(), elapsed_seconds(phase_begin)); + } + + void CPUHybridSolver::matmul(Eigen::VectorXd &x, Eigen::VectorXd &result) + { + auto phase_begin = clock::now(); + result.resize(x.size()); + result.setZero(); + HYPRE_ParVector par_x; + HYPRE_ParVector par_result; + HYPRE_IJVectorSetData(ij_x, x.data()); + HYPRE_IJVectorAssemble(ij_x); + HYPRE_IJVectorGetObject(ij_x, (void **)&par_x); + HYPRE_IJVectorSetData(ij_b, result.data()); + HYPRE_IJVectorAssemble(ij_b); + HYPRE_IJVectorGetObject(ij_b, (void **)&par_result); + HYPRE_ParCSRMatrixMatvec(1.0, parcsr_A, par_x, 0.0, par_result); + SPDLOG_TRACE("[{}] [matmul] [{:.6f}]", name(), elapsed_seconds(phase_begin)); + } + + double CPUHybridSolver::dot(Eigen::VectorXd &a, Eigen::VectorXd &b) + { + HYPRE_ParVector par_a; + HYPRE_ParVector par_b; + HYPRE_IJVectorSetData(ij_x, a.data()); + HYPRE_IJVectorAssemble(ij_x); + HYPRE_IJVectorGetObject(ij_x, (void **)&par_a); + HYPRE_IJVectorSetData(ij_b, b.data()); + HYPRE_IJVectorAssemble(ij_b); + HYPRE_IJVectorGetObject(ij_b, (void **)&par_b); + double result; + HYPRE_ParVectorInnerProd(par_a, par_b, &result); + return result; + } + + void CPUHybridSolver::partition_ranks(const int rows) + { + starts.clear(); + ends.clear(); + int local_size = rows / num_procs; + for (int i = 0; i < num_procs; ++i) + { + starts.push_back(i == 0 ? 0 : local_size * i + i); + ends.push_back(i == (num_procs - 1) ? rows - 1 : starts.back() + local_size); + } + } + + void CPUHybridSolver::copy_matrix_to_hypre(SharedSparseMatrix &sparse_A) + { + HYPRE_IJMatrixCreate(MPI_COMM_WORLD, starts[myid], ends[myid], starts[myid], ends[myid], &A); + HYPRE_IJMatrixSetObjectType(A, HYPRE_PARCSR); + HYPRE_IJMatrixInitialize(A); + + for (HYPRE_Int k = starts[myid]; k <= ends[myid]; ++k) + { + HYPRE_Int row[1]; + row[0] = k; + int counter = 0; + std::vector cols; + std::vector vals; + for (SharedSparseMatrix::InnerIterator it(sparse_A, k); it; ++it) + { + ++counter; + row[0] = it.col(); + cols.push_back((HYPRE_Int)it.row()); + vals.push_back(it.value()); + } + HYPRE_Int n_cols[1] = {counter}; + if (counter > 0) + { + HYPRE_IJMatrixSetValues(A, 1, n_cols, row, cols.data(), vals.data()); + } + } + HYPRE_IJMatrixAssemble(A); + HYPRE_IJMatrixGetObject(A, (void **)&parcsr_A); + } + + void CPUHybridSolver::copy_matrix_to_hypre(Eigen::SparseMatrix &sparse_A) + { + HYPRE_IJMatrixCreate(MPI_COMM_WORLD, starts[myid], ends[myid], starts[myid], ends[myid], &A); + HYPRE_IJMatrixSetObjectType(A, HYPRE_PARCSR); + HYPRE_IJMatrixInitialize(A); + + for (HYPRE_Int k = starts[myid]; k <= ends[myid]; ++k) + { + HYPRE_Int row[1]; + row[0] = k; + int counter = 0; + std::vector cols; + std::vector vals; + for (Eigen::SparseMatrix::InnerIterator it(sparse_A, k); it; ++it) + { + ++counter; + row[0] = it.col(); + cols.push_back((HYPRE_Int)it.row()); + vals.push_back(it.value()); + } + HYPRE_Int n_cols[1] = {counter}; + HYPRE_IJMatrixSetValues(A, 1, n_cols, row, cols.data(), vals.data()); + } + HYPRE_IJMatrixAssemble(A); + HYPRE_IJMatrixGetObject(A, (void **)&parcsr_A); + } + + void CPUHybridSolver::init_hypre_vectors() + { + HYPRE_IJVectorCreate(MPI_COMM_WORLD, starts[myid], ends[myid], &ij_x); + HYPRE_IJVectorSetObjectType(ij_x, HYPRE_PARCSR); + HYPRE_IJVectorInitializeShell(ij_x); + + hypre_ParVector *par_x = (hypre_ParVector *)hypre_IJVectorObject((hypre_IJVector *)ij_x); + if (par_x && hypre_ParVectorLocalVector(par_x)) + { + hypre_Vector *vec_x = hypre_ParVectorLocalVector(par_x); + + hypre_VectorNumVectors(vec_x) = 1; + hypre_VectorMultiVecStorageMethod(vec_x) = 0; + + hypre_VectorVectorStride(vec_x) = hypre_VectorSize(vec_x); + hypre_VectorIndexStride(vec_x) = 1; + } + + HYPRE_IJVectorCreate(MPI_COMM_WORLD, starts[myid], ends[myid], &ij_b); + HYPRE_IJVectorSetObjectType(ij_b, HYPRE_PARCSR); + HYPRE_IJVectorInitializeShell(ij_b); + + hypre_ParVector *par_b = (hypre_ParVector *)hypre_IJVectorObject((hypre_IJVector *)ij_b); + if (par_b && hypre_ParVectorLocalVector(par_b)) + { + hypre_Vector *vec_b = hypre_ParVectorLocalVector(par_b); + + hypre_VectorNumVectors(vec_b) = 1; + hypre_VectorMultiVecStorageMethod(vec_b) = 0; + + hypre_VectorVectorStride(vec_b) = hypre_VectorSize(vec_b); + hypre_VectorIndexStride(vec_b) = 1; + } + } + + void CPUHybridSolver::create_shared_vec(MPI_Win &win, void *&base_ptr, int size) + { + int local_alloc_size = myid == 0 ? size * sizeof(double) : 0; + MPI_Win_allocate_shared(local_alloc_size, 1, MPI_INFO_NULL, MPI_COMM_WORLD, &base_ptr, &win); + if (myid != 0) + { + int disp_unit; + MPI_Aint sz; + MPI_Win_shared_query(win, 0, &sz, &disp_unit, &base_ptr); + } + } + + void CPUHybridSolver::assemble_D(int bad_i, int i, Eigen::SparseMatrix &D, SharedSparseMatrix &sparse_A) + { + D.resize(bad_indices_sets[i].size(), bad_indices_sets[i].size()); + std::vector> triplets; + for (int k : bad_indices_sets[i]) + { + global_to_row[k] = sparse_A.col(k); + for (SharedSparseMatrix::InnerIterator it(sparse_A, k); it; ++it) + { + auto ind_it = index_mappings[bad_i].find(it.row()); + if (ind_it != index_mappings[bad_i].end()) + { + triplets.push_back(Eigen::Triplet(index_mappings[bad_i][it.row()], index_mappings[bad_i][it.col()], it.value())); + } + } + } + + D.setFromTriplets(triplets.begin(), triplets.end()); + } + + void CPUHybridSolver::build_index_mappings() + { + index_mappings.clear(); + index_mappings.resize(bad_subdomain_assignments[myid].size()); + + int i_counter = 0; + for (int i : bad_subdomain_assignments[myid]) + { + int j_counter = 0; + for (auto j : bad_indices_sets[i]) + { + index_mappings[i_counter][j] = j_counter; + ++j_counter; + } + ++i_counter; + } + } + + void CPUHybridSolver::filter_subdomains(SharedSparseMatrix &sparse_A) + { + auto phase_begin = clock::now(); + + int num_too_small = 0; + int num_too_large = 0; + int num_not_poorly_conditioned = 0; + int original_num_bad_dofs = all_bad_dofs.size(); + + int counter = 0; + std::vector global_to_local(sparse_A.rows(), -1); + for (auto index : all_bad_dofs) + { + global_to_local[index] = counter; + ++counter; + } + + hybrid::DisjointSet decomposed_subdomains(all_bad_dofs.size()); + + for (int k : all_bad_dofs) + { + for (SharedSparseMatrix::InnerIterator it(sparse_A, k); it; ++it) + { + if (global_to_local[it.row()] != -1) + { + decomposed_subdomains.union_set(global_to_local[it.row()], global_to_local[it.col()]); + } + } + } + + std::unordered_map> chosen_sets; + for (auto index : all_bad_dofs) + { + chosen_sets[decomposed_subdomains.find_set(global_to_local[index])].push_back(index); + } + + all_bad_dofs.clear(); + + for (auto &kv : chosen_sets) + { + if (kv.second.size() < min_subdomain_size) + { + ++num_too_small; + continue; + } + if (kv.second.size() > max_subdomain_size) + { + ++num_too_large; + continue; + } + + double lambda_min = std::numeric_limits::max(); + double lambda_max = 0.0; + for (int k : kv.second) + { + double diag_value = 0.0; + double abs_off_diag_sum = 0.0; + for (SharedSparseMatrix::InnerIterator it(sparse_A, k); it; ++it) + { + if (global_to_local[it.row()] != -1) + { + if (it.row() == it.col()) + { + diag_value = it.value(); + } + else + { + abs_off_diag_sum += abs(it.value()); + } + } + } + lambda_min = std::min(lambda_min, diag_value - abs_off_diag_sum); + lambda_max = std::max(lambda_max, diag_value + abs_off_diag_sum); + } + + if (lambda_min * lambda_max < 0.0 || lambda_max / lambda_min > conditioning_threshold) + { + all_bad_dofs.insert(kv.second.begin(), kv.second.end()); + continue; + } + ++num_not_poorly_conditioned; + } + + SPDLOG_TRACE("[{}] [subdomain_filtering] [{}] [total_dofs_before={}] [total_dofs_after={}] [num_too_small={}] [num_too_large={}] [num_not_poorly_conditioned={}]", + name(), elapsed_seconds(phase_begin), original_num_bad_dofs, all_bad_dofs.size(), num_too_small, num_too_large, num_not_poorly_conditioned); + } + + void CPUHybridSolver::expand_subdomains_to_strongly_connected(SharedSparseMatrix &sparse_A) + { + auto phase_begin = clock::now(); + int num_bad_dofs_before = all_bad_dofs.size(); + + std::set new_bad_dofs; + ; + + for (int k : all_bad_dofs) + { + for (SharedSparseMatrix::InnerIterator it(sparse_A, k); it; ++it) + { + new_bad_dofs.insert(it.row()); + } + } + all_bad_dofs = std::move(new_bad_dofs); + + SPDLOG_TRACE("[{}] [subdomain_expansion] [{}] [num_dofs_before={}] [num_dofs_after={}]", + name(), elapsed_seconds(phase_begin), num_bad_dofs_before, all_bad_dofs.size()); + } + + void CPUHybridSolver::decompose_subdomains_to_disjoint_subsets(SharedSparseMatrix &sparse_A) + { + auto phase_begin = clock::now(); + int counter = 0; + std::vector global_to_local(sparse_A.rows(), -1); + for (auto index : all_bad_dofs) + { + global_to_local[index] = counter; + ++counter; + } + + hybrid::DisjointSet decomposed_subdomains(all_bad_dofs.size()); + + for (int k : all_bad_dofs) + { + for (SharedSparseMatrix::InnerIterator it(sparse_A, k); it; ++it) + { + if (global_to_local[it.row()] != -1) + { + decomposed_subdomains.union_set(global_to_local[it.row()], global_to_local[it.col()]); + } + } + } + + std::unordered_map> chosen_sets; + for (auto index : all_bad_dofs) + { + chosen_sets[decomposed_subdomains.find_set(global_to_local[index])].push_back(index); + } + + bad_indices_sets.clear(); + + for (auto &kv : chosen_sets) + { + if (kv.second.size() > max_subdomain_size) + { + continue; + } + bad_indices_sets.emplace_back(kv.second.begin(), kv.second.end()); + } + SPDLOG_TRACE("[{}] [subdomain_decomposition] [{}] [num_subdomains={}] ", + name(), elapsed_seconds(phase_begin), bad_indices_sets.size()); + } + + void CPUHybridSolver::share_bad_subdomains() + { + auto phase_begin = clock::now(); + int num_subdomains; + if (myid == 0) + { + num_subdomains = bad_indices_sets.size(); + MPI_Bcast(&num_subdomains, 1, MPI_INT, 0, MPI_COMM_WORLD); + if (bad_indices_sets.size() > 0) + { + for (int i = 0; i < num_subdomains; ++i) + { + int num_indices = bad_indices_sets[i].size(); + MPI_Bcast(&num_indices, 1, MPI_INT, 0, MPI_COMM_WORLD); + std::vector subdomain_vec; + for (auto index : bad_indices_sets[i]) + { + subdomain_vec.push_back(index); + } + MPI_Bcast(subdomain_vec.data(), num_indices, MPI_INT, 0, MPI_COMM_WORLD); + MPI_Barrier(MPI_COMM_WORLD); + } + } + for (int i = 0; i < num_procs; ++i) + { + int local_size = bad_subdomain_assignments[i].size(); + MPI_Bcast(&local_size, 1, MPI_INT, 0, MPI_COMM_WORLD); + MPI_Bcast(bad_subdomain_assignments[i].data(), local_size, MPI_INT, 0, MPI_COMM_WORLD); + } + } + else + { + MPI_Bcast(&num_subdomains, 1, MPI_INT, 0, MPI_COMM_WORLD); + bad_indices_sets.clear(); + bad_indices_sets.resize(num_subdomains); + for (int i = 0; i < num_subdomains; ++i) + { + int num_indices; + MPI_Bcast(&num_indices, 1, MPI_INT, 0, MPI_COMM_WORLD); + std::vector subdomain_vec; + subdomain_vec.resize(num_indices); + MPI_Bcast(subdomain_vec.data(), num_indices, MPI_INT, 0, MPI_COMM_WORLD); + for (auto index : subdomain_vec) + { + bad_indices_sets[i].insert(index); + } + MPI_Barrier(MPI_COMM_WORLD); + } + bad_subdomain_assignments.resize(num_procs); + for (int i = 0; i < num_procs; ++i) + { + int local_size; + MPI_Bcast(&local_size, 1, MPI_INT, 0, MPI_COMM_WORLD); + bad_subdomain_assignments[i].resize(local_size); + MPI_Bcast(bad_subdomain_assignments[i].data(), local_size, MPI_INT, 0, MPI_COMM_WORLD); + } + } + + bad_indices_arrays.clear(); + bad_indices_arrays.resize(bad_indices_sets.size()); + for (int i = 0; i < bad_indices_sets.size(); ++i) + { + bad_indices_arrays[i].reserve(bad_indices_sets[i].size()); + for (auto index : bad_indices_sets[i]) + { + bad_indices_arrays[i].push_back(index); + } + } + SPDLOG_TRACE("[{}] [share_bad_subdomains] [{}] ", + name(), elapsed_seconds(phase_begin)); + } + + void CPUHybridSolver::load_balance_subdomains() + { + auto phase_begin = clock::now(); + bad_subdomain_assignments.clear(); + bad_subdomain_assignments.resize(num_procs); + + std::vector> subdomain_sizes; + subdomain_sizes.reserve(bad_indices_sets.size()); + + for (auto &subdomain : bad_indices_sets) + { + subdomain_sizes.push_back(std::make_pair(subdomain_sizes.size(), subdomain.size())); + } + + std::sort(subdomain_sizes.begin(), subdomain_sizes.end(), [](const std::pair &l, const std::pair &r) { return l.second > r.second; }); + std::vector assigned_sizes(num_procs, 0); + + int total_bad_dofs = 0; + for (auto [i, size] : subdomain_sizes) + { + int min_size = assigned_sizes[0]; + total_bad_dofs += size; + int chosen_proc = 0; + for (int pi = 1; pi < num_procs; ++pi) + { + if (assigned_sizes[pi] < min_size) + { + min_size = assigned_sizes[pi]; + chosen_proc = pi; + } + } + bad_subdomain_assignments[chosen_proc].push_back(i); + assigned_sizes[chosen_proc] += size; + } + + const int max_size = subdomain_sizes.size() > 0 ? subdomain_sizes.front().second : 0; + const int min_size = subdomain_sizes.size() > 0 ? subdomain_sizes.back().second : 0; + + SPDLOG_TRACE("[{}] [subdomain_load_balance] [{}] [max_size={}] [min_size={}] [total_dofs={}]", + name(), elapsed_seconds(phase_begin), max_size, min_size, total_bad_dofs); + } + + //////////////////////////////////////////////////////////////////////////////// + + void CPUHybridSolver::run_worker_loop() + { + bool running = true; + while (running) + { + SolverCmd cmd; + MPI_Bcast(&cmd, 1, MPI_INT, 0, MPI_COMM_WORLD); + + if (cmd == CMD_EXIT) + { + running = false; + break; + } + + int id; + MPI_Bcast(&id, 1, MPI_INT, 0, MPI_COMM_WORLD); + + switch (cmd) + { + case CMD_CREATE: + { + worker_registry[id] = std::make_unique(); + worker_registry[id]->solver_id = id; + break; + } + case CMD_SET_PARAMETERS: + { + json dummy_params; + worker_registry[id]->set_parameters(dummy_params); + break; + } + case CMD_FACTORIZE: + { + StiffnessMatrix dummy_A; + worker_registry[id]->factorize(dummy_A); + break; + } + case CMD_SOLVE: + { + Eigen::VectorXd dummy_b, dummy_x; + worker_registry[id]->solve(dummy_b, dummy_x); + break; + } + case CMD_DESTROY: + { + worker_registry.erase(id); + break; + } + } + } + } + + CPUHybridSolver::~CPUHybridSolver() + { + if (myid == 0) + { + SolverCmd cmd = CMD_DESTROY; + MPI_Bcast(&cmd, 1, MPI_INT, 0, MPI_COMM_WORLD); + MPI_Bcast(&solver_id, 1, MPI_INT, 0, MPI_COMM_WORLD); + } + + if (has_matrix_) + { + HYPRE_IJMatrixDestroy(A); + has_matrix_ = false; + A = nullptr; + } + } + +} // namespace polysolve::linear \ No newline at end of file diff --git a/src/polysolve/linear/CPUHybridSolver.hpp b/src/polysolve/linear/CPUHybridSolver.hpp new file mode 100644 index 00000000..8f77468f --- /dev/null +++ b/src/polysolve/linear/CPUHybridSolver.hpp @@ -0,0 +1,195 @@ +#pragma once + +//////////////////////////////////////////////////////////////////////////////// +#include "Solver.hpp" + +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include + +#include + +namespace polysolve::linear +{ + class AbstractSolver + { + + public: + virtual void compute(const Eigen::SparseMatrix &A) = 0; + virtual Eigen::VectorXd solve(const Eigen::VectorXd &b) = 0; + virtual ~AbstractSolver() = default; + }; + + template + class EigenWrapper : public AbstractSolver + { + EigenSolverT solver; + + public: + void compute(const Eigen::SparseMatrix &A) override + { + solver.compute(A); + } + + Eigen::VectorXd solve(const Eigen::VectorXd &b) override + { + return solver.solve(b); + } + }; + + class CPUHybridSolver : public Solver + { + + public: + CPUHybridSolver(); + ~CPUHybridSolver(); + + typedef Eigen::Map SharedSparseMatrix; + typedef Eigen::Map SharedVector; + + enum SolverCmd + { + CMD_CREATE, + CMD_SET_PARAMETERS, + CMD_FACTORIZE, + CMD_SOLVE, + CMD_DESTROY, + CMD_EXIT + }; + + private: + POLYSOLVE_DELETE_MOVE_COPY(CPUHybridSolver) + + int solver_id; + static inline int next_id = 0; + static inline bool is_running_worker_loop = false; + static inline std::unordered_map> worker_registry; + + static void run_worker_loop(); + + public: + ////////////////////// + // Public interface // + ////////////////////// + + // Set solver parameters + virtual void set_parameters(const json ¶ms) override; + + // Retrieve solve information + virtual void get_info(json ¶ms) const override; + + void check_settings() const; + + // Factorize system matrix + virtual void factorize(const StiffnessMatrix &A) override; + + // Solve the linear system Ax = b + virtual void solve(const Ref b, Ref x) override; + + // Name of the solver type (for debugging purposes) + virtual std::string name() const override + { + return "CPUHybrid"; + } + + protected: + // AMG settings + double theta = 0.5; + + // Hybrid preconditioner settings + bool decompose_subdomains = true; + int min_subdomain_size = 3; + int max_subdomain_size = 1e9; + double gmm_jump_threshold = 10.0; + double gmm_tol = 1e-3; + int max_gmm_iterations = 20; + bool expand_subdomains = true; + bool additive_mode = false; + + // General solver settings + int dimension_ = 1; // 1 = scalar (Laplace), 2 or 3 = vector (Elasticity) + int max_iter_ = 10000; + double rel_conv_tol_ = 1e-10; + double abs_conv_tol_ = 0.0; + double conditioning_threshold = 100.0; + + // solve information + HYPRE_Int num_iterations; + HYPRE_Complex final_res_norm; + + private: + bool has_matrix_ = false; + + // MPI rank distribution + int myid = 0; + int num_procs = 1; + std::vector starts; + std::vector ends; + + // problem-specific data + std::unordered_map> global_to_row; + + // Hypre variables + HYPRE_IJMatrix A; + HYPRE_ParCSRMatrix parcsr_A; + HYPRE_IJVector ij_x; + HYPRE_IJVector ij_b; + + // hybrid preconditioner data + std::deque> D_solvers; + std::set all_bad_dofs; + std::vector> bad_indices_sets; + std::vector> bad_indices_arrays; + std::vector> bad_subdomain_assignments; + std::vector> index_mappings; + + Eigen::VectorXd shared_rhs, shared_result; + Eigen::VectorXd z1, z2, z3; + Eigen::VectorXd r, p, buffer; + + // factorization helpers + void partition_ranks(const int rows); + void copy_matrix_to_hypre(SharedSparseMatrix &sparse_A); + void copy_matrix_to_hypre(Eigen::SparseMatrix &sparse_A); + + // solve helpers + void init_hypre_vectors(); + + // hybrid preconditioner helpers + void assemble_D(int bad_i, int i, Eigen::SparseMatrix &D, SharedSparseMatrix &sparse_A); + void build_index_mappings(); + void decompose_subdomains_to_disjoint_subsets(SharedSparseMatrix &sparse_A); + void filter_subdomains(SharedSparseMatrix &sparse_A); + void expand_subdomains_to_strongly_connected(SharedSparseMatrix &sparse_A); + void share_bad_subdomains(); + void load_balance_subdomains(); + void select_bad_dofs(SharedSparseMatrix &sparse_A); + void factorize_submatrix(SharedSparseMatrix &sparse_A); + + // matrix multiplication + void matmul(Eigen::VectorXd &x, Eigen::VectorXd &result); + double dot(Eigen::VectorXd &x, Eigen::VectorXd &y); + + // preconditioning functions + void custom_mixed_precond_iter(const HYPRE_Solver &precond, Eigen::VectorXd &r, Eigen::VectorXd &z, SharedVector &vec, MPI_Win &vec_win); + void amg_precond_iter(const HYPRE_Solver &precond, Eigen::VectorXd &b, Eigen::VectorXd &x); + void dss_precond_iter(Eigen::VectorXd &z, Eigen::VectorXd &r, Eigen::VectorXd &next_z, SharedVector &vec, MPI_Win &vec_win); + + // MPI helpers + void create_shared_vec(MPI_Win &win, void *&base_ptr, int size); + int my_size() { return ends[myid] - starts[myid] + 1; }; + + // Krylov solve methods + void pcg_solve(Eigen::VectorXd &rhs, Eigen::VectorXd &result, HYPRE_ParVector &par_b, HYPRE_ParVector &par_x, HYPRE_Solver &precond, SharedVector &vec, MPI_Win &vec_win); + }; + +} // namespace polysolve::linear \ No newline at end of file diff --git a/src/polysolve/linear/GPUHybridSolver.cu b/src/polysolve/linear/GPUHybridSolver.cu new file mode 100644 index 00000000..e34ded99 --- /dev/null +++ b/src/polysolve/linear/GPUHybridSolver.cu @@ -0,0 +1,1230 @@ + + +#include "GPUHybridSolver.hpp" + +#include "hybrid_utils/DisjointSet.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include + +#include + +#ifdef HYPRE_ENABLE_MPI +#include +#endif + +#define CHECK_CUDA(call) \ + do \ + { \ + cudaError_t status = call; \ + if (status != cudaSuccess) \ + { \ + std::cerr << "CUDA Error at " << __FILE__ << ":" << __LINE__ \ + << " - " << cudaGetErrorName(status) \ + << " (" << cudaGetErrorString(status) << ")" << std::endl; \ + exit(EXIT_FAILURE); \ + } \ + } while (0) + +#define CHECK_CUDSS(call) \ + do \ + { \ + cudssStatus_t status = call; \ + if (status != CUDSS_STATUS_SUCCESS) \ + { \ + std::cerr << "cuDSS Error at " << __FILE__ << ":" << __LINE__ \ + << " code " << (int)status << std::endl; \ + exit(EXIT_FAILURE); \ + } \ + } while (0) + +namespace polysolve::linear +{ + + namespace + { + using clock = std::chrono::steady_clock; + + double elapsed_seconds(const std::chrono::time_point &begin) + { + return std::chrono::duration(clock::now() - begin).count(); + } + } // namespace + + GPUHybridSolver::GPUHybridSolver() + { +#ifdef HYPRE_ENABLE_MPI + int done_already; + + MPI_Initialized(&done_already); + if (!done_already) + { + MPI_Init(nullptr, nullptr); + } +#endif + if (!HYPRE_Initialized()) + { + HYPRE_Initialize(); + } + + HYPRE_SetMemoryLocation(HYPRE_MEMORY_DEVICE); + HYPRE_SetExecutionPolicy(HYPRE_EXEC_DEVICE); + HYPRE_SetSpGemmUseCusparse(false); + HYPRE_SetUseGpuRand(true); + + CHECK_CUDSS(cudssCreate(&cudss_handle)); + } + + void GPUHybridSolver::set_parameters(const json ¶ms) + { + if (params.contains("GPUHybrid")) + { + if (params["GPUHybrid"].contains("max_iter")) + { + max_iter_ = params["GPUHybrid"]["max_iter"]; + } + if (params["GPUHybrid"].contains("relative_tolerance")) + { + rel_conv_tol_ = params["GPUHybrid"]["relative_tolerance"]; + } + if (params["GPUHybrid"].contains("absolute_tolerance")) + { + abs_conv_tol_ = params["GPUHybrid"]["absolute_tolerance"]; + } + if (params["GPUHybrid"].contains("theta")) + { + theta = params["GPUHybrid"]["theta"]; + } + if (params["GPUHybrid"].contains("block_dim")) + { + dimension_ = params["GPUHybrid"]["block_dim"]; + } + if (params["GPUHybrid"].contains("decompose_subdomains")) + { + decompose_subdomains = params["GPUHybrid"]["decompose_subdomains"]; + } + if (params["GPUHybrid"].contains("min_subdomain_size")) + { + min_subdomain_size = params["GPUHybrid"]["min_subdomain_size"]; + } + if (params["GPUHybrid"].contains("max_subdomain_size")) + { + max_subdomain_size = params["GPUHybrid"]["max_subdomain_size"]; + } + if (params["GPUHybrid"].contains("gmm_jump_threshold")) + { + gmm_jump_threshold = params["GPUHybrid"]["gmm_jump_threshold"]; + } + if (params["GPUHybrid"].contains("expand_subdomains")) + { + expand_subdomains = params["GPUHybrid"]["expand_subdomains"]; + } + if (params["GPUHybrid"].contains("gmm_tol")) + { + gmm_tol = params["GPUHybrid"]["gmm_tol"]; + } + if (params["GPUHybrid"].contains("max_gmm_iterations")) + { + max_gmm_iterations = params["GPUHybrid"]["max_gmm_iterations"]; + } + if (params["GPUHybrid"].contains("conditioning_threshold")) + { + conditioning_threshold = params["GPUHybrid"]["conditioning_threshold"]; + } + if (params["GPUHybrid"].contains("additive_mode")) + { + additive_mode = params["GPUHybrid"]["additive_mode"]; + } + } + } + + void GPUHybridSolver::get_info(json ¶ms) const + { + params["num_iterations"] = num_iterations; + params["final_res_norm"] = final_res_norm; + } + + void GPUHybridSolver::check_settings() const + { + } + + void GPUHybridSolver::factorize(const StiffnessMatrix &Ain) + { + check_settings(); + SPDLOG_TRACE("[{}] [start_solve] [0.000000] [problem_size={}]", name(), Ain.rows()); + + { + auto phase_begin = clock::now(); + + d_outer_indices.resize(Ain.rows() + 1); + d_inner_indices.resize(Ain.nonZeros()); + d_values.resize(Ain.nonZeros()); + + thrust::copy(Ain.outerIndexPtr(), Ain.outerIndexPtr() + d_outer_indices.size(), d_outer_indices.begin()); + thrust::copy(Ain.innerIndexPtr(), Ain.innerIndexPtr() + d_inner_indices.size(), d_inner_indices.begin()); + thrust::copy(Ain.valuePtr(), Ain.valuePtr() + d_values.size(), d_values.begin()); + + SPDLOG_TRACE("[{}] [copy_matrix_to_gpu] [{:.6f}]", name(), elapsed_seconds(phase_begin)); + } + + { + auto phase_begin = clock::now(); + + bad_indices_arrays.clear(); + select_bad_dofs(); + + if (decompose_subdomains) + { + filter_subdomains(Ain); + } + + if (expand_subdomains) + { + expand_subdomains_to_strongly_connected(Ain); + } + + if (decompose_subdomains) + { + decompose_subdomains_to_disjoint_subsets(Ain); + } + else + { + bad_indices_arrays.emplace_back(h_all_bad_dofs.begin(), h_all_bad_dofs.end()); + } + + d_all_bad_dofs.clear(); + h_subdomain_sizes.clear(); + d_subdomain_sizes.clear(); + + for (int i = 0; i < bad_indices_arrays.size(); ++i) + { + d_all_bad_dofs.insert(d_all_bad_dofs.end(), bad_indices_arrays[i].begin(), bad_indices_arrays[i].end()); + h_subdomain_sizes.push_back(bad_indices_arrays[i].size()); + } + + d_subdomain_sizes.insert(d_subdomain_sizes.end(), h_subdomain_sizes.begin(), h_subdomain_sizes.end()); + + factorize_submatrix(); + + SPDLOG_TRACE("[{}] [setup_problematic_dof_precond] [{:.6f}]", name(), elapsed_seconds(phase_begin)); + } + + if (has_matrix_) + { + HYPRE_IJMatrixDestroy(A); + has_matrix_ = false; + A = nullptr; + } + + copy_matrix_to_hypre(); + has_matrix_ = true; + + d_outer_indices.clear(); + d_inner_indices.clear(); + d_values.clear(); + } + + namespace + { + void HypreBoomerAMG_SetDefaultOptions(HYPRE_Solver &amg_precond) + { + // AMG coarsening options: + int coarsen_type = 8; // 10 = HMIS, 8 = PMIS, 6 = Falgout, 0 = CLJP + int agg_levels = 1; // number of aggressive coarsening levels + double theta = 0.25; // strength threshold: 0.25, 0.5, 0.8 + + // AMG interpolation options: + int interp_type = 6; // 6 = extended+i, 0 = classical + int Pmax = 4; // max number of elements per row in P + + // AMG relaxation options: + int relax_type = 18; // 8 = l1-GS, 6 = symm. GS, 3 = GS, 18 = l1-Jacobi + int relax_sweeps = 1; // relaxation sweeps on each level + + // Additional options: + int print_level = 0; // print AMG iterations? 1 = no, 2 = yes + int max_levels = 25; // max number of levels in AMG hierarchy + + int min_coarse_size = 5; + + HYPRE_BoomerAMGSetCoarsenType(amg_precond, coarsen_type); + HYPRE_BoomerAMGSetAggNumLevels(amg_precond, agg_levels); + HYPRE_BoomerAMGSetRelaxType(amg_precond, relax_type); + + HYPRE_BoomerAMGSetRelaxOrder(amg_precond, false); + HYPRE_BoomerAMGSetRAP2(amg_precond, true); + HYPRE_BoomerAMGSetKeepTranspose(amg_precond, true); + + HYPRE_BoomerAMGSetMinCoarseSize(amg_precond, min_coarse_size); + HYPRE_BoomerAMGSetCycleRelaxType(amg_precond, relax_type, 3); + HYPRE_BoomerAMGSetNumSweeps(amg_precond, relax_sweeps); + HYPRE_BoomerAMGSetStrongThreshold(amg_precond, theta); + HYPRE_BoomerAMGSetInterpType(amg_precond, interp_type); + HYPRE_BoomerAMGSetPMaxElmts(amg_precond, Pmax); + // print_level = 3; + HYPRE_BoomerAMGSetPrintLevel(amg_precond, print_level); + HYPRE_BoomerAMGSetMaxLevels(amg_precond, max_levels); + + // Use as a preconditioner (one V-cycle, zero tolerance) + HYPRE_BoomerAMGSetMaxIter(amg_precond, 1); + HYPRE_BoomerAMGSetTol(amg_precond, 0.0); + } + + void HypreBoomerAMG_SetElasticityOptions(HYPRE_Solver &amg_precond, int dim, double theta) + { + // Make sure the systems AMG options are set + HYPRE_BoomerAMGSetNumFunctions(amg_precond, dim); + + // HYPRE_BoomerAMGSetDofFunc(amg_precond, (HYPRE_Int*) dof_to_function.data()); + + // More robust options with respect to convergence + HYPRE_BoomerAMGSetAggNumLevels(amg_precond, 0); + HYPRE_BoomerAMGSetStrongThreshold(amg_precond, theta); + } + } // namespace + + void GPUHybridSolver::solve(const Ref b, Ref x) + { + thrust::device_vector d_x(x.size()); + thrust::device_vector d_b(b.size()); + + thrust::copy(x.data(), x.data() + x.size(), d_x.begin()); + thrust::copy(b.data(), b.data() + b.size(), d_b.begin()); + + HYPRE_ParVector par_b; + HYPRE_ParVector par_x; + init_hypre_vectors(b.size()); + + set_hypre_vec(ij_b, par_b, d_b); + set_hypre_vec(ij_x, par_x, d_x); + + HYPRE_Solver precond; + + { + auto phase_begin = clock::now(); + + HYPRE_BoomerAMGCreate(&precond); + HypreBoomerAMG_SetDefaultOptions(precond); + if (dimension_ > 1) + { + HypreBoomerAMG_SetElasticityOptions( + precond, + dimension_, + theta); + } + + HYPRE_BoomerAMGSetup(precond, parcsr_A, par_b, par_x); + CHECK_CUDA(cudaDeviceSynchronize()); + SPDLOG_TRACE("[{}] [amg_setup] [{:.6f}]", name(), elapsed_seconds(phase_begin)); + } + + { + auto phase_begin = clock::now(); + + pcg_solve(d_b, d_x, par_b, par_x, precond); + + thrust::device_vector buffer(d_x.size()); + matmul(d_x, buffer); + vector_add(-1.0, d_b, buffer); + final_res_norm = sqrt(dot(buffer, buffer)); + + thrust::copy(d_x.begin(), d_x.end(), x.data()); + + CHECK_CUDA(cudaDeviceSynchronize()); + SPDLOG_TRACE("[{}] [pcg_solve] [{:.6f}] [pcg_iters={}] [residual={}]", name(), elapsed_seconds(phase_begin), num_iterations, final_res_norm); + } + + { + HYPRE_BoomerAMGDestroy(precond); + HYPRE_IJVectorDestroy(ij_x); + HYPRE_IJVectorDestroy(ij_b); + } + } + + void GPUHybridSolver::copy_matrix_to_hypre() + { + auto phase_begin = clock::now(); + + const HYPRE_Int num_rows = d_outer_indices.size() - 1; + const HYPRE_Int nnz = d_values.size(); + +#ifdef HYPRE_ENABLE_MPI + HYPRE_IJMatrixCreate(MPI_COMM_WORLD, 0, num_rows - 1, 0, num_rows - 1, &A); +#else + HYPRE_IJMatrixCreate(0, 0, num_rows - 1, 0, num_rows - 1, &A); +#endif + HYPRE_IJMatrixSetObjectType(A, HYPRE_PARCSR); + HYPRE_IJMatrixInitialize(A); + + thrust::device_vector d_rows(num_rows); + thrust::sequence(d_rows.begin(), d_rows.end()); + + thrust::device_vector d_n_cols(num_rows); + const HYPRE_Int *raw_outer = thrust::raw_pointer_cast(d_outer_indices.data()); + HYPRE_Int *raw_n_cols = thrust::raw_pointer_cast(d_n_cols.data()); + + thrust::for_each(thrust::device, + thrust::make_counting_iterator(0), + thrust::make_counting_iterator(num_rows), + [=] __device__(int i) { + raw_n_cols[i] = raw_outer[i + 1] - raw_outer[i]; + }); + + HYPRE_Int *gpu_n_cols = thrust::raw_pointer_cast(d_n_cols.data()); + HYPRE_Int *gpu_rows = thrust::raw_pointer_cast(d_rows.data()); + + HYPRE_Int *gpu_cols = thrust::raw_pointer_cast(d_inner_indices.data()); + double *gpu_vals = thrust::raw_pointer_cast(d_values.data()); + + HYPRE_IJMatrixSetValues(A, num_rows, gpu_n_cols, gpu_rows, gpu_cols, gpu_vals); + + HYPRE_IJMatrixAssemble(A); + + void *temp_A = nullptr; + HYPRE_IJMatrixGetObject(A, &temp_A); + parcsr_A = static_cast(temp_A); + + SPDLOG_TRACE("[{}] [copy_matrix_to_hypre] [{:.6f}]", name(), elapsed_seconds(phase_begin)); + } + + void GPUHybridSolver::init_hypre_vectors(const int size) + { +#ifdef HYPRE_ENABLE_MPI + HYPRE_IJVectorCreate(MPI_COMM_WORLD, 0, size - 1, &ij_x); +#else + HYPRE_IJVectorCreate(0, 0, size - 1, &ij_x); +#endif + HYPRE_IJVectorSetObjectType(ij_x, HYPRE_PARCSR); + HYPRE_IJVectorInitializeShell(ij_x); +#ifdef HYPRE_ENABLE_MPI + HYPRE_IJVectorCreate(MPI_COMM_WORLD, 0, size - 1, &ij_b); +#else + HYPRE_IJVectorCreate(0, 0, size - 1, &ij_b); +#endif + HYPRE_IJVectorSetObjectType(ij_b, HYPRE_PARCSR); + HYPRE_IJVectorInitializeShell(ij_b); + } + + void GPUHybridSolver::matmul(const thrust::device_vector &x, thrust::device_vector &result) + { + auto phase_begin = clock::now(); + HYPRE_ParVector par_x; + HYPRE_ParVector par_result; + + set_hypre_vec(ij_x, par_x, x); + set_hypre_vec(ij_b, par_result, result); + + HYPRE_ParCSRMatrixMatvec(1.0, parcsr_A, par_x, 0.0, par_result); + CHECK_CUDA(cudaDeviceSynchronize()); + SPDLOG_TRACE("[{}] [matmul] [{:.6f}]", name(), elapsed_seconds(phase_begin)); + } + + double GPUHybridSolver::dot(const thrust::device_vector &a, const thrust::device_vector &b) + { + HYPRE_ParVector par_a; + HYPRE_ParVector par_b; + + set_hypre_vec(ij_x, par_a, a); + set_hypre_vec(ij_b, par_b, b); + + double result; + HYPRE_ParVectorInnerProd(par_a, par_b, &result); + return result; + } + + void GPUHybridSolver::vector_copy(const thrust::device_vector &x, thrust::device_vector &y) + { + HYPRE_ParVector par_x; + HYPRE_ParVector par_y; + + set_hypre_vec(ij_x, par_x, x); + set_hypre_vec(ij_b, par_y, y); + + HYPRE_ParVectorCopy(par_x, par_y); + } + + void GPUHybridSolver::vector_add(double alpha, const thrust::device_vector &x, thrust::device_vector &y) + { + HYPRE_ParVector par_x; + HYPRE_ParVector par_y; + + set_hypre_vec(ij_x, par_x, x); + set_hypre_vec(ij_b, par_y, y); + + hypre_ParVectorAxpy(alpha, par_x, par_y); + } + + void GPUHybridSolver::vector_scale(double alpha, thrust::device_vector &x) + { + HYPRE_ParVector par_x; + + set_hypre_vec(ij_x, par_x, x); + + HYPRE_ParVectorScale(alpha, par_x); + } + + void GPUHybridSolver::set_hypre_vec(HYPRE_IJVector &my_ij_x, HYPRE_ParVector &par_x, const thrust::device_vector &x) + { + double *raw_ptr = const_cast(thrust::raw_pointer_cast(x.data())); + + HYPRE_IJVectorSetData(my_ij_x, raw_ptr); + HYPRE_IJVectorAssemble(my_ij_x); + HYPRE_IJVectorGetObject(my_ij_x, (void **)&par_x); + } + + void GPUHybridSolver::custom_mixed_precond_iter(const HYPRE_Solver &precond, thrust::device_vector &r, thrust::device_vector &z, thrust::device_vector &buffer, thrust::device_vector &z2) + { + if (d_all_bad_dofs.size() == 0) + { + amg_precond_iter(precond, r, z); + } + else if (additive_mode) + { + thrust::fill(buffer.begin(), buffer.end(), 0.0); + thrust::fill(z2.begin(), z2.end(), 0.0); + amg_precond_iter(precond, r, z); + dss_precond_iter(buffer, r, z2); + vector_add(1.0, z2, z); + } + else + { + thrust::fill(buffer.begin(), buffer.end(), 0.0); + thrust::fill(z2.begin(), z2.end(), 0.0); + amg_precond_iter(precond, r, z); + dss_precond_iter(z, r, z2); + matmul(z2, z); + vector_copy(r, buffer); + vector_add(-1.0, z, buffer); + thrust::fill(z.begin(), z.end(), 0.0); + amg_precond_iter(precond, buffer, z); + vector_add(1.0, z2, z); + } + } + + void GPUHybridSolver::dss_precond_iter(thrust::device_vector &z, thrust::device_vector &r, thrust::device_vector &next_z) + { + auto phase_begin = clock::now(); + + matmul(z, next_z); + vector_scale(-1.0, next_z); + vector_add(1.0, r, next_z); + + if (sparse_batch_count > 0) + { + thrust::gather( + thrust::device, + d_sparse_dof_map.begin(), + d_sparse_dof_map.end(), + next_z.begin(), + d_sparse_b.begin()); + + CHECK_CUDSS(cudssExecute(cudss_handle, CUDSS_PHASE_SOLVE, cudss_config, cudss_solver_data, batch_A, batch_x, batch_b)); + } + + thrust::fill(next_z.begin(), next_z.end(), 0.0); + + if (sparse_batch_count > 0) + { + thrust::scatter( + thrust::device, + d_sparse_x.begin(), + d_sparse_x.begin() + d_sparse_dof_map.size(), + d_sparse_dof_map.begin(), + next_z.begin()); + } + + vector_add(1.0, z, next_z); + + CHECK_CUDA(cudaDeviceSynchronize()); + SPDLOG_TRACE("[{}] [subdomain_solve] [{:.6f}]", name(), elapsed_seconds(phase_begin)); + } + + void GPUHybridSolver::amg_precond_iter(const HYPRE_Solver &precond, thrust::device_vector &b, thrust::device_vector &x) + { + auto phase_begin = clock::now(); + HYPRE_ParVector par_x; + HYPRE_ParVector par_b; + + set_hypre_vec(ij_x, par_x, x); + set_hypre_vec(ij_b, par_b, b); + + HYPRE_BoomerAMGSolve(precond, parcsr_A, par_b, par_x); + CHECK_CUDA(cudaDeviceSynchronize()); + SPDLOG_TRACE("[{}] [amg_v_cycle] [{:.6f}]", name(), elapsed_seconds(phase_begin)); + } + + void GPUHybridSolver::decompose_subdomains_to_disjoint_subsets(const Eigen::SparseMatrix &sparse_A) + { + auto phase_begin = clock::now(); + + std::vector global_to_local(sparse_A.rows(), -1); + int counter = 0; + for (auto index : h_all_bad_dofs) + { + global_to_local[index] = counter; + ++counter; + } + + hybrid::DisjointSet decomposed_subdomains(h_all_bad_dofs.size()); + + for (int k : h_all_bad_dofs) + { + for (Eigen::SparseMatrix::InnerIterator it(sparse_A, k); it; ++it) + { + if (global_to_local[it.row()] != -1) + { + decomposed_subdomains.union_set(global_to_local[it.row()], global_to_local[it.col()]); + } + } + } + + std::unordered_map> chosen_sets; + for (auto index : h_all_bad_dofs) + { + chosen_sets[decomposed_subdomains.find_set(global_to_local[index])].push_back(index); + } + + bad_indices_arrays.clear(); + + for (auto &kv : chosen_sets) + { + if (kv.second.size() > max_subdomain_size) + { + continue; + } + bad_indices_arrays.emplace_back(kv.second.begin(), kv.second.end()); + } + + SPDLOG_TRACE("[{}] [subdomain_decomposition] [{}] [num_subdomains={}] ", + name(), elapsed_seconds(phase_begin), bad_indices_arrays.size()); + } + + void GPUHybridSolver::select_bad_dofs() + { + auto phase_begin = clock::now(); + + const int num_rows = d_outer_indices.size() - 1; + + const int *row_offsets = thrust::raw_pointer_cast(d_outer_indices.data()); + const double *values = thrust::raw_pointer_cast(d_values.data()); + + thrust::device_vector d_row_norms(num_rows); + double *row_norms = thrust::raw_pointer_cast(d_row_norms.data()); + + thrust::for_each(thrust::device, + thrust::make_counting_iterator(0), + thrust::make_counting_iterator(num_rows), + [=] __device__(int i) { + int start = row_offsets[i]; + int end = row_offsets[i + 1]; + double sum = 0.0; + for (int j = start; j < end; ++j) + { + sum += fabs(values[j]); + } + row_norms[i] = sum; + }); + + double sum_norms = thrust::reduce(d_row_norms.begin(), d_row_norms.end(), 0.0); + double global_mean = sum_norms / num_rows; + + double var_sum = thrust::transform_reduce(thrust::device, d_row_norms.begin(), d_row_norms.end(), [global_mean] __device__(double x) -> double { return (x - global_mean) * (x - global_mean); }, 0.0, thrust::plus()); + double global_var = var_sum / num_rows; + + auto minmax = thrust::minmax_element(d_row_norms.begin(), d_row_norms.end()); + double mean_0 = *minmax.first; + double mean_1 = *minmax.second; + double var_0 = global_var; + double var_1 = global_var; + double w0 = 0.5; + double w1 = 0.5; + + double var_reg = 1e-6; + + thrust::device_vector d_gamma0(num_rows); + thrust::device_vector d_gamma1(num_rows); + double *g0 = thrust::raw_pointer_cast(d_gamma0.data()); + double *g1 = thrust::raw_pointer_cast(d_gamma1.data()); + + int gmm_iter; + + for (gmm_iter = 0; gmm_iter < max_gmm_iterations; ++gmm_iter) + { + + double log_likelihood = thrust::transform_reduce(thrust::device, thrust::make_counting_iterator(0), thrust::make_counting_iterator(num_rows), [=] __device__(int i) -> double { + double x = row_norms[i]; + + double log_w0 = log(w0); + double log_w1 = log(w1); + + double log_N0 = -0.5 * log(2.0 * M_PI * var_0) - 0.5 * (x - mean_0) * (x - mean_0) / var_0; + double log_N1 = -0.5 * log(2.0 * M_PI * var_1) - 0.5 * (x - mean_1) * (x - mean_1) / var_1; + + double log_g0 = log_w0 + log_N0; + double log_g1 = log_w1 + log_N1; + + double max_log_g = max(log_g0, log_g1); + double log_total = max_log_g + log(exp(log_g0 - max_log_g) + exp(log_g1 - max_log_g)); + + g0[i] = exp(log_g0 - log_total); + g1[i] = exp(log_g1 - log_total); + + return log_total; }, 0.0, thrust::plus()); + + double sum_g0 = thrust::reduce(d_gamma0.begin(), d_gamma0.end(), 0.0); + double sum_g1 = thrust::reduce(d_gamma1.begin(), d_gamma1.end(), 0.0); + + w0 = sum_g0 / num_rows; + w1 = sum_g1 / num_rows; + + double old_mean_0 = mean_0, old_mean_1 = mean_1; + double old_var_0 = var_0, old_var_1 = var_1; + + mean_0 = thrust::inner_product(d_gamma0.begin(), d_gamma0.end(), d_row_norms.begin(), 0.0) / sum_g0; + mean_1 = thrust::inner_product(d_gamma1.begin(), d_gamma1.end(), d_row_norms.begin(), 0.0) / sum_g1; + + var_0 = thrust::transform_reduce(thrust::device, thrust::make_counting_iterator(0), thrust::make_counting_iterator(num_rows), [=] __device__(int i) -> double { return g0[i] * (row_norms[i] - mean_0) * (row_norms[i] - mean_0); }, 0.0, thrust::plus()) / sum_g0 + var_reg; + + var_1 = thrust::transform_reduce(thrust::device, thrust::make_counting_iterator(0), thrust::make_counting_iterator(num_rows), [=] __device__(int i) -> double { return g1[i] * (row_norms[i] - mean_1) * (row_norms[i] - mean_1); }, 0.0, thrust::plus()) / sum_g1 + var_reg; + + // Check Convergence + if (abs(mean_0 - old_mean_0) / abs(old_mean_0) < gmm_tol && abs(mean_1 - old_mean_1) / abs(old_mean_1) < gmm_tol && abs(var_0 - old_var_0) / abs(old_var_0) < gmm_tol && abs(var_1 - old_var_1) / abs(old_var_1) < gmm_tol) + { + break; + } + } + + int num_bad_dofs = 0; + if (abs(mean_1) / abs(mean_0) > gmm_jump_threshold) + { + d_all_bad_dofs.resize(num_rows); + auto end_it = thrust::copy_if(thrust::device, + thrust::make_counting_iterator(0), + thrust::make_counting_iterator(num_rows), + d_all_bad_dofs.begin(), + [=] __device__(int i) { return g0[i] < g1[i]; }); + + num_bad_dofs = thrust::distance(d_all_bad_dofs.begin(), end_it); + } + + d_all_bad_dofs.resize(num_bad_dofs); + + std::vector h_bad_dofs(num_bad_dofs); + thrust::copy(d_all_bad_dofs.begin(), d_all_bad_dofs.end(), h_bad_dofs.begin()); + + h_all_bad_dofs.clear(); + h_all_bad_dofs.insert(h_bad_dofs.begin(), h_bad_dofs.end()); + + SPDLOG_TRACE("[{}] [bad_dof_selection] [{:.6f}] [global_mean={}] [global_var={}] [mean_0={}] [mean_1={}] [var_0={}] [var_1={}] [gmm_iters={}] [num_bad_dofs={}]", + name(), elapsed_seconds(phase_begin), global_mean, global_var, mean_0, mean_1, var_0, var_1, gmm_iter, num_bad_dofs); + } + + void GPUHybridSolver::filter_subdomains(const Eigen::SparseMatrix &sparse_A) + { + auto phase_begin = clock::now(); + + int num_too_small = 0; + int num_too_large = 0; + int num_not_poorly_conditioned = 0; + int original_num_bad_dofs = h_all_bad_dofs.size(); + + int counter = 0; + std::vector global_to_local(sparse_A.rows(), -1); + for (auto index : h_all_bad_dofs) + { + global_to_local[index] = counter; + ++counter; + } + + hybrid::DisjointSet decomposed_subdomains(h_all_bad_dofs.size()); + + for (int k : h_all_bad_dofs) + { + for (Eigen::SparseMatrix::InnerIterator it(sparse_A, k); it; ++it) + { + if (global_to_local[it.row()] != -1) + { + decomposed_subdomains.union_set(global_to_local[it.row()], global_to_local[it.col()]); + } + } + } + + std::unordered_map> chosen_sets; + for (auto index : h_all_bad_dofs) + { + chosen_sets[decomposed_subdomains.find_set(global_to_local[index])].push_back(index); + } + + h_all_bad_dofs.clear(); + + for (auto &kv : chosen_sets) + { + if (kv.second.size() < min_subdomain_size) + { + ++num_too_small; + continue; + } + if (kv.second.size() > max_subdomain_size) + { + ++num_too_large; + continue; + } + + double lambda_min = std::numeric_limits::max(); + double lambda_max = 0.0; + for (int k : kv.second) + { + double diag_value = 0.0; + double abs_off_diag_sum = 0.0; + for (Eigen::SparseMatrix::InnerIterator it(sparse_A, k); it; ++it) + { + if (global_to_local[it.row()] != -1) + { + if (it.row() == it.col()) + { + diag_value = it.value(); + } + else + { + abs_off_diag_sum += abs(it.value()); + } + } + } + lambda_min = std::min(lambda_min, diag_value - abs_off_diag_sum); + lambda_max = std::max(lambda_max, diag_value + abs_off_diag_sum); + } + + if (lambda_min * lambda_max < 0.0 || lambda_max / lambda_min > conditioning_threshold) + { + h_all_bad_dofs.insert(kv.second.begin(), kv.second.end()); + continue; + } + ++num_not_poorly_conditioned; + } + + SPDLOG_TRACE("[{}] [subdomain_filtering] [{}] [total_dofs_before={}] [total_dofs_after={}] [num_too_small={}] [num_too_large={}] [num_not_poorly_conditioned={}]", + name(), elapsed_seconds(phase_begin), original_num_bad_dofs, h_all_bad_dofs.size(), num_too_small, num_too_large, num_not_poorly_conditioned); + } + + void GPUHybridSolver::expand_subdomains_to_strongly_connected(const Eigen::SparseMatrix &sparse_A) + { + auto phase_begin = clock::now(); + int num_bad_dofs_before = h_all_bad_dofs.size(); + + std::set new_bad_dofs; + ; + + for (int k : h_all_bad_dofs) + { + for (Eigen::SparseMatrix::InnerIterator it(sparse_A, k); it; ++it) + { + new_bad_dofs.insert(it.row()); + } + } + + h_all_bad_dofs = std::move(new_bad_dofs); + + SPDLOG_TRACE("[{}] [subdomain_expansion] [{}] [num_dofs_before={}] [num_dofs_after={}]", + name(), elapsed_seconds(phase_begin), num_bad_dofs_before, h_all_bad_dofs.size()); + } + + void GPUHybridSolver::factorize_submatrix() + { + auto phase_begin = clock::now(); + + if (d_all_bad_dofs.size() == 0) + { + return; + } + + free_device_memory(); + + CHECK_CUDSS(cudssConfigCreate(&cudss_config)); + CHECK_CUDSS(cudssDataCreate(cudss_handle, &cudss_solver_data)); + + int total_sparse_dofs = 0; + std::vector h_sparse_row_starts; + std::vector h_sparse_nnz_starts; + std::vector h_sparse_dof_starts; + h_sparse_row_starts.push_back(0); + h_sparse_dof_starts.push_back(0); + h_sparse_nrows.clear(); + h_sparse_ncols.clear(); + h_sparse_vec_ncols.clear(); + h_sparse_ld.clear(); + for (int size : h_subdomain_sizes) + { + h_sparse_nrows.push_back(size); + h_sparse_ncols.push_back(size); + h_sparse_vec_ncols.push_back(1); + h_sparse_ld.push_back(size); + + h_sparse_row_starts.push_back(h_sparse_row_starts.back() + size + 1); + h_sparse_dof_starts.push_back(h_sparse_dof_starts.back() + size); + + total_sparse_dofs += size; + } + + sparse_batch_count = h_sparse_nrows.size(); + + d_sparse_dof_map.clear(); + d_sparse_dof_map.reserve(total_sparse_dofs); + + if (total_sparse_dofs > 0) + { + std::vector h_sparse_dof_map; + + for (int i = 0; i < bad_indices_arrays.size(); ++i) + { + h_sparse_dof_map.insert(h_sparse_dof_map.end(), bad_indices_arrays[i].begin(), bad_indices_arrays[i].end()); + } + + d_sparse_dof_map.insert(d_sparse_dof_map.begin(), h_sparse_dof_map.begin(), h_sparse_dof_map.end()); + + h_sparse_nnz.clear(); + h_sparse_nnz.resize(sparse_batch_count); + + thrust::device_vector d_sparse_nnz(sparse_batch_count, 0); + thrust::device_vector d_sparse_batch_offsets(h_sparse_dof_starts.begin(), h_sparse_dof_starts.end()); + thrust::device_vector d_sparse_batch_sizes(h_sparse_nrows.begin(), h_sparse_nrows.end()); + + int total_sparse_dofs = d_sparse_dof_map.size(); + + int *raw_bad_dofs = thrust::raw_pointer_cast(d_sparse_dof_map.data()); + int *raw_batch_offsets = thrust::raw_pointer_cast(d_sparse_batch_offsets.data()); + int *raw_batch_sizes = thrust::raw_pointer_cast(d_sparse_batch_sizes.data()); + int *raw_outer = thrust::raw_pointer_cast(d_outer_indices.data()); + int *raw_inner = thrust::raw_pointer_cast(d_inner_indices.data()); + + thrust::device_vector d_row_nnz(total_sparse_dofs, 0); + int *raw_row_nnz = thrust::raw_pointer_cast(d_row_nnz.data()); + + thrust::fill(d_sparse_nnz.begin(), d_sparse_nnz.end(), 0); + int *raw_sparse_nnz = thrust::raw_pointer_cast(d_sparse_nnz.data()); + + int local_sparse_batch_count = sparse_batch_count; + + thrust::for_each(thrust::device, + thrust::make_counting_iterator(0), + thrust::make_counting_iterator(total_sparse_dofs), + [=] __device__(int i) { + int *batch_ptr = thrust::upper_bound(thrust::seq, raw_batch_offsets, raw_batch_offsets + local_sparse_batch_count, i); + int batch_idx = (batch_ptr - raw_batch_offsets) - 1; + + int offset = raw_batch_offsets[batch_idx]; + int size = raw_batch_sizes[batch_idx]; + + int global_row = raw_bad_dofs[i]; + int row_start = raw_outer[global_row]; + int row_end = raw_outer[global_row + 1]; + + int *sub_begin = raw_bad_dofs + offset; + int *sub_end = sub_begin + size; + + int row_nnz = 0; + for (int j = row_start; j < row_end; ++j) + { + int global_col = raw_inner[j]; + if (thrust::binary_search(thrust::seq, sub_begin, sub_end, global_col)) + { + row_nnz++; + } + } + + raw_row_nnz[i] = row_nnz; + atomicAdd(&raw_sparse_nnz[batch_idx], row_nnz); + }); + + thrust::copy(d_sparse_nnz.begin(), d_sparse_nnz.end(), h_sparse_nnz.begin()); + h_sparse_nnz_starts.resize(h_sparse_nnz.size()); + h_sparse_nnz_starts[0] = 0; + std::partial_sum(h_sparse_nnz.begin(), h_sparse_nnz.end() - 1, h_sparse_nnz_starts.begin() + 1); + int total_sparse_nnz = std::accumulate(h_sparse_nnz.begin(), h_sparse_nnz.end(), 0); + d_sparse_inner_indices.clear(); + d_sparse_outer_indices.clear(); + d_sparse_values.clear(); + d_sparse_x.clear(); + d_sparse_b.clear(); + + d_sparse_outer_indices.resize(total_sparse_dofs + sparse_batch_count); + d_sparse_inner_indices.resize(total_sparse_nnz); + + d_sparse_x.resize(total_sparse_dofs); + d_sparse_b.resize(total_sparse_dofs); + d_sparse_values.resize(total_sparse_nnz); + + thrust::device_vector d_row_nnz_starts(total_sparse_dofs, 0); + thrust::exclusive_scan(thrust::device, d_row_nnz.begin(), d_row_nnz.end(), d_row_nnz_starts.begin()); + + int *raw_row_nnz_starts = thrust::raw_pointer_cast(d_row_nnz_starts.data()); + int *raw_sparse_outer = thrust::raw_pointer_cast(d_sparse_outer_indices.data()); + int *raw_sparse_inner = thrust::raw_pointer_cast(d_sparse_inner_indices.data()); + + double *raw_global_values = thrust::raw_pointer_cast(d_values.data()); + double *raw_sparse_values = thrust::raw_pointer_cast(d_sparse_values.data()); + + thrust::for_each(thrust::device, + thrust::make_counting_iterator(0), + thrust::make_counting_iterator(total_sparse_dofs), + [=] __device__(int i) { + int *batch_ptr = thrust::upper_bound(thrust::seq, raw_batch_offsets, raw_batch_offsets + local_sparse_batch_count, i); + int batch_idx = (batch_ptr - raw_batch_offsets) - 1; + + int offset = raw_batch_offsets[batch_idx]; + int size = raw_batch_sizes[batch_idx]; + + int r = i - offset; + int outer_start = offset + batch_idx; + + if (r == 0) + { + raw_sparse_outer[outer_start] = 0; + } + + int batch_nnz_start = raw_row_nnz_starts[offset]; + + raw_sparse_outer[outer_start + r + 1] = (raw_row_nnz_starts[i] + raw_row_nnz[i]) - batch_nnz_start; + + int current_nnz = raw_row_nnz_starts[i]; + + int global_row = raw_bad_dofs[i]; + int row_start = raw_outer[global_row]; + int row_end = raw_outer[global_row + 1]; + + int *sub_begin = raw_bad_dofs + offset; + int *sub_end = sub_begin + size; + + for (int j = row_start; j < row_end; ++j) + { + int global_col = raw_inner[j]; + + int *ptr = thrust::lower_bound(thrust::seq, sub_begin, sub_end, global_col); + + if (ptr != sub_end && *ptr == global_col) + { + int local_col = ptr - sub_begin; + + raw_sparse_inner[current_nnz] = local_col; + raw_sparse_values[current_nnz] = raw_global_values[j]; + + current_nnz++; + } + } + }); + + std::vector h_sparse_outer_void; + std::vector h_sparse_inner_void; + std::vector h_sparse_values_void; + std::vector h_sparse_x_void; + std::vector h_sparse_b_void; + + for (int i = 0; i < sparse_batch_count; ++i) + { + h_sparse_outer_void.push_back(static_cast(thrust::raw_pointer_cast(d_sparse_outer_indices.data()) + h_sparse_row_starts[i])); + h_sparse_inner_void.push_back(static_cast(thrust::raw_pointer_cast(d_sparse_inner_indices.data()) + h_sparse_nnz_starts[i])); + h_sparse_x_void.push_back(static_cast(thrust::raw_pointer_cast(d_sparse_x.data()) + h_sparse_dof_starts[i])); + h_sparse_b_void.push_back(static_cast(thrust::raw_pointer_cast(d_sparse_b.data()) + h_sparse_dof_starts[i])); + h_sparse_values_void.push_back(static_cast(thrust::raw_pointer_cast(d_sparse_values.data()) + h_sparse_nnz_starts[i])); + } + + d_sparse_outer_void = h_sparse_outer_void; + d_sparse_inner_void = h_sparse_inner_void; + d_sparse_values_void = h_sparse_values_void; + d_sparse_x_void = h_sparse_x_void; + d_sparse_b_void = h_sparse_b_void; + + auto precision = CUDA_R_64F; + + CHECK_CUDSS(cudssMatrixCreateBatchDn( + &batch_x, sparse_batch_count, h_sparse_nrows.data(), h_sparse_vec_ncols.data(), h_sparse_ld.data(), + thrust::raw_pointer_cast(d_sparse_x_void.data()), CUDA_R_32I, precision, CUDSS_LAYOUT_COL_MAJOR)); + + CHECK_CUDSS(cudssMatrixCreateBatchDn( + &batch_b, sparse_batch_count, h_sparse_nrows.data(), h_sparse_vec_ncols.data(), h_sparse_ld.data(), + thrust::raw_pointer_cast(d_sparse_b_void.data()), CUDA_R_32I, precision, CUDSS_LAYOUT_COL_MAJOR)); + + CHECK_CUDSS(cudssMatrixCreateBatchCsr( + &batch_A, sparse_batch_count, h_sparse_nrows.data(), h_sparse_ncols.data(), h_sparse_nnz.data(), + thrust::raw_pointer_cast(d_sparse_outer_void.data()), nullptr, thrust::raw_pointer_cast(d_sparse_inner_void.data()), thrust::raw_pointer_cast(d_sparse_values_void.data()), + CUDA_R_32I, precision, CUDSS_MTYPE_SYMMETRIC, + CUDSS_MVIEW_FULL, CUDSS_BASE_ZERO)); + + CHECK_CUDSS(cudssExecute(cudss_handle, CUDSS_PHASE_ANALYSIS, cudss_config, cudss_solver_data, batch_A, nullptr, nullptr)); + CHECK_CUDSS(cudssExecute(cudss_handle, CUDSS_PHASE_FACTORIZATION, cudss_config, cudss_solver_data, batch_A, nullptr, nullptr)); + } + + CHECK_CUDA(cudaDeviceSynchronize()); + SPDLOG_TRACE("[{}] [factorize_submatrix] [{}] [n_sparse={}] [n_sparse_dofs={}]", + name(), elapsed_seconds(phase_begin), sparse_batch_count, total_sparse_dofs); + } + + void GPUHybridSolver::pcg_solve(thrust::device_vector &rhs, thrust::device_vector &result, HYPRE_ParVector &par_b, HYPRE_ParVector &par_x, HYPRE_Solver &precond) + { + double bi_prod, abs_eps, rel_eps, gamma, old_gamma; + + thrust::device_vector r(rhs.size()); + thrust::device_vector p(rhs.size()); + thrust::device_vector z(rhs.size()); + thrust::device_vector z2(rhs.size()); + thrust::device_vector buffer(rhs.size()); + + { + auto phase_begin = clock::now(); + + bi_prod = dot(rhs, rhs); + + if (bi_prod > 0.0) + { + rel_eps = rel_conv_tol_ * rel_conv_tol_; + abs_eps = abs_conv_tol_ * abs_conv_tol_; + } + else + { + thrust::fill(result.begin(), result.end(), 0.0); + num_iterations = 0; + final_res_norm = 0; + return; + } + + matmul(result, buffer); + + vector_copy(rhs, r); + vector_add(-1.0, buffer, r); + + custom_mixed_precond_iter(precond, r, z, buffer, z2); + + vector_copy(z, p); + + gamma = dot(r, z); + old_gamma = gamma; + SPDLOG_TRACE("[{}] [pre_loop] [{:.6f}] [rhs_norm={}]", name(), elapsed_seconds(phase_begin), sqrt(bi_prod)); + } + + for (int k = 0; k < max_iter_; ++k) + { + auto phase_begin = clock::now(); + num_iterations = k + 1; + + matmul(p, buffer); + double sdotp = dot(p, buffer); + + if (sdotp == 0.0) + { + SPDLOG_TRACE("[{}] [err_zero_sdotp] [0.000000]", name()); + break; + } + + double alpha = gamma / sdotp; + + if (alpha <= 0.0) + { + SPDLOG_TRACE("[{}] [err_negative_alpha] [0.000000]", name()); + break; + } + else if (alpha < __DBL_MIN__) + { + SPDLOG_TRACE("[{}] [err_subnormal_alpha] [0.000000]", name()); + break; + } + + vector_add(alpha, p, result); + vector_add(-1.0 * alpha, buffer, r); + double i_prod = dot(r, r); + + if (rel_eps > 0 && (i_prod / bi_prod) < rel_eps) + { + SPDLOG_TRACE("[{}] [converged_rel] [0.000000]", name()); + break; + } + + if (abs_eps > 0 && i_prod < abs_eps) + { + SPDLOG_TRACE("[{}] [converged_abs] [0.000000]", name()); + break; + } + + thrust::fill(z.begin(), z.end(), 0.0); + custom_mixed_precond_iter(precond, r, z, buffer, z2); + + gamma = dot(r, z); + + double beta = gamma / old_gamma; + old_gamma = gamma; + + vector_scale(beta, p); + vector_add(1.0, z, p); + + CHECK_CUDA(cudaDeviceSynchronize()); + SPDLOG_TRACE("[{}] [pcg_iter] [{:.6f}] [iter={}] [residual={}]", name(), elapsed_seconds(phase_begin), k, sqrt(i_prod)); + } + } + + GPUHybridSolver::~GPUHybridSolver() + { + if (has_matrix_) + { + HYPRE_IJMatrixDestroy(A); + has_matrix_ = false; + A = nullptr; + } + + free_device_memory(); + + if (cudss_handle) + { + cudssDestroy(cudss_handle); + cudss_handle = nullptr; + } + } + + void GPUHybridSolver::free_device_memory() + { + // Destroy cuDSS Opaque Structures + if (batch_A) + { + CHECK_CUDSS(cudssMatrixDestroy(batch_A)); + batch_A = nullptr; + } + if (batch_x) + { + CHECK_CUDSS(cudssMatrixDestroy(batch_x)); + batch_x = nullptr; + } + if (batch_b) + { + CHECK_CUDSS(cudssMatrixDestroy(batch_b)); + batch_b = nullptr; + } + if (cudss_solver_data) + { + CHECK_CUDSS(cudssDataDestroy(cudss_handle, cudss_solver_data)); + cudss_solver_data = nullptr; + } + if (cudss_config) + { + CHECK_CUDSS(cudssConfigDestroy(cudss_config)); + cudss_config = nullptr; + } + } +} // namespace polysolve::linear \ No newline at end of file diff --git a/src/polysolve/linear/GPUHybridSolver.hpp b/src/polysolve/linear/GPUHybridSolver.hpp new file mode 100644 index 00000000..b71d656a --- /dev/null +++ b/src/polysolve/linear/GPUHybridSolver.hpp @@ -0,0 +1,162 @@ +#pragma once + +//////////////////////////////////////////////////////////////////////////////// +#include "Solver.hpp" + +#include + +#include +#include + +#include +#include +#include + +#include +#include + +#include + +extern "C" +{ + HYPRE_Int hypre_ParVectorAxpy(HYPRE_Complex alpha, HYPRE_ParVector x, HYPRE_ParVector y); +} + +namespace polysolve::linear +{ + + class GPUHybridSolver : public Solver + { + + public: + GPUHybridSolver(); + ~GPUHybridSolver(); + + private: + POLYSOLVE_DELETE_MOVE_COPY(GPUHybridSolver) + + public: + ////////////////////// + // Public interface // + ////////////////////// + + // Set solver parameters + virtual void set_parameters(const json ¶ms) override; + + // Retrieve solve information + virtual void get_info(json ¶ms) const override; + + void check_settings() const; + + // Factorize system matrix + virtual void factorize(const StiffnessMatrix &A) override; + + // Solve the linear system Ax = b + virtual void solve(const Ref b, Ref x) override; + + // Name of the solver type (for debugging purposes) + virtual std::string name() const override + { + return "GPUHybrid"; + } + + protected: + // AMG settings + double theta = 0.5; + + // Hybrid preconditioner settings + bool decompose_subdomains = true; + int min_subdomain_size = 1; + int max_subdomain_size = 1e9; + double gmm_jump_threshold = 10.0; + double gmm_tol = 1e-3; + int max_gmm_iterations = 20; + bool expand_subdomains = true; + bool additive_mode = false; + + // General solver settings + int dimension_ = 1; // 1 = scalar (Laplace), 2 or 3 = vector (Elasticity) + int max_iter_ = 10000; + double rel_conv_tol_ = 1e-10; + double abs_conv_tol_ = 0.0; + double conditioning_threshold = 100.0; + + // solve information + HYPRE_Int num_iterations; + HYPRE_Complex final_res_norm; + + private: + bool has_matrix_ = false; + + // Hypre variables + HYPRE_IJMatrix A; + HYPRE_ParCSRMatrix parcsr_A; + HYPRE_IJVector ij_x; + HYPRE_IJVector ij_b; + + // hybrid preconditioner data + std::set h_all_bad_dofs; + std::vector h_subdomain_sizes; + thrust::device_vector d_subdomain_sizes; + thrust::device_vector d_all_bad_dofs; + std::vector> bad_indices_arrays; + + thrust::device_vector d_inner_indices; + thrust::device_vector d_outer_indices; + thrust::device_vector d_values; + + cudssHandle_t cudss_handle = nullptr; + cudssConfig_t cudss_config = nullptr; + cudssData_t cudss_solver_data = nullptr; + + cudssMatrix_t batch_A = nullptr; + cudssMatrix_t batch_x = nullptr; + cudssMatrix_t batch_b = nullptr; + + int sparse_batch_count = 0; + + thrust::device_vector d_sparse_dof_map; + + // sparse solve data + std::vector h_sparse_nrows, h_sparse_ncols, h_sparse_nnz, h_sparse_vec_ncols, h_sparse_ld; + + thrust::device_vector d_sparse_inner_indices, d_sparse_outer_indices; + thrust::device_vector d_sparse_values; + thrust::device_vector d_sparse_x, d_sparse_b; + thrust::device_vector d_sparse_inner_void, d_sparse_outer_void, d_sparse_values_void; + thrust::device_vector d_sparse_x_void, d_sparse_b_void; + + public: + void free_device_memory(); + + // factorization helpers + void copy_matrix_to_hypre(); + + // solve helpers + void init_hypre_vectors(const int size); + + // hybrid preconditioner helpers + void decompose_subdomains_to_disjoint_subsets(const Eigen::SparseMatrix &sparse_A); + void filter_subdomains(const Eigen::SparseMatrix &sparse_A); + void expand_subdomains_to_strongly_connected(const Eigen::SparseMatrix &sparse_A); + void select_bad_dofs(); + void factorize_submatrix(); + + // linear algebra + void set_hypre_vec(HYPRE_IJVector &ij_x, HYPRE_ParVector &par_x, const thrust::device_vector &x); + void matmul(const thrust::device_vector &x, thrust::device_vector &result); + double dot(const thrust::device_vector &x, const thrust::device_vector &y); + void vector_copy(const thrust::device_vector &x, thrust::device_vector &y); + void vector_add(double alpha, const thrust::device_vector &x, thrust::device_vector &y); + void vector_scale(double alpha, thrust::device_vector &x); + + // preconditioning functions + void custom_mixed_precond_iter(const HYPRE_Solver &precond, thrust::device_vector &r, thrust::device_vector &z, thrust::device_vector &buffer, thrust::device_vector &z2); + void amg_precond_iter(const HYPRE_Solver &precond, thrust::device_vector &b, thrust::device_vector &x); + void dss_precond_iter(thrust::device_vector &z, thrust::device_vector &r, thrust::device_vector &next_z); + + // Krylov solve methods + void pcg_solve(thrust::device_vector &rhs, thrust::device_vector &result, HYPRE_ParVector &par_b, HYPRE_ParVector &par_x, HYPRE_Solver &precond); + }; + +} // namespace polysolve::linear \ No newline at end of file diff --git a/src/polysolve/linear/HypreSolver.cpp b/src/polysolve/linear/HypreSolver.cpp index 6001ded5..974125d1 100644 --- a/src/polysolve/linear/HypreSolver.cpp +++ b/src/polysolve/linear/HypreSolver.cpp @@ -21,21 +21,15 @@ namespace polysolve::linear MPI_Initialized(&done_already); if (!done_already) { - /* Initialize MPI */ - int argc = 1; - char name[] = ""; - char *argv[] = {name}; - char **argvv = &argv[0]; - int myid, num_procs; - MPI_Init(&argc, &argvv); - MPI_Comm_rank(MPI_COMM_WORLD, &myid); - MPI_Comm_size(MPI_COMM_WORLD, &num_procs); + MPI_Init(nullptr, nullptr); } #endif if (!HYPRE_Initialized()) { HYPRE_Initialize(); } + HYPRE_SetMemoryLocation(HYPRE_MEMORY_HOST); + HYPRE_SetExecutionPolicy(HYPRE_EXEC_HOST); } // Set solver parameters @@ -135,7 +129,11 @@ namespace polysolve::linear void eigen_to_hypre_par_vec(HYPRE_ParVector &par_x, HYPRE_IJVector &ij_x, const Eigen::VectorXd &x) { +#ifdef HYPRE_ENABLE_MPI + HYPRE_IJVectorCreate(MPI_COMM_WORLD, 0, x.size() - 1, &ij_x); +#else HYPRE_IJVectorCreate(0, 0, x.size() - 1, &ij_x); +#endif HYPRE_IJVectorSetObjectType(ij_x, HYPRE_PARCSR); HYPRE_IJVectorInitialize(ij_x); @@ -155,7 +153,7 @@ namespace polysolve::linear // AMG coarsening options: int coarsen_type = 10; // 10 = HMIS, 8 = PMIS, 6 = Falgout, 0 = CLJP int agg_levels = 1; // number of aggressive coarsening levels - double theta = 0.5; // strength threshold: 0.25, 0.5, 0.8 + double theta = 0.5; // strength threshold: 0.25, 0.5, 0.8 // AMG interpolation options: int interp_type = 6; // 6 = extended+i, 0 = classical @@ -208,7 +206,7 @@ namespace polysolve::linear // refinement (this is generally applicable for any system) int interp_refine = 1; - if (nodal_coarsening) + if (nodal_coarsening) { HYPRE_BoomerAMGSetNodal(amg_precond, nodal); HYPRE_BoomerAMGSetNodalDiag(amg_precond, nodal_diag); @@ -243,16 +241,16 @@ namespace polysolve::linear for (int i = 0; i < positions.rows(); ++i) { - rbm_xy(0 + i*dim) = positions(i, 1); - rbm_xy(1 + i*dim) = -1 * positions(i, 0); + rbm_xy(0 + i * dim) = positions(i, 1); + rbm_xy(1 + i * dim) = -1 * positions(i, 0); if (dim == 3) { - rbm_zx(1 + i*dim) = positions(i, 2); - rbm_zx(2 + i*dim) = -1 * positions(i, 1); + rbm_zx(1 + i * dim) = positions(i, 2); + rbm_zx(2 + i * dim) = -1 * positions(i, 1); - rbm_yz(2 + i*dim) = positions(i, 0); - rbm_yz(0 + i*dim) = -1 * positions(i, 2); + rbm_yz(2 + i * dim) = positions(i, 0); + rbm_yz(0 + i * dim) = -1 * positions(i, 2); } } diff --git a/src/polysolve/linear/HypreSolver.hpp b/src/polysolve/linear/HypreSolver.hpp index 5196f164..e8b0c763 100644 --- a/src/polysolve/linear/HypreSolver.hpp +++ b/src/polysolve/linear/HypreSolver.hpp @@ -53,7 +53,7 @@ namespace polysolve::linear // Name of the solver type (for debugging purposes) virtual std::string name() const override { return "Hypre"; } - virtual void set_tolerance(const double tol) override {conv_tol_ = tol;} + virtual void set_tolerance(const double tol) override { conv_tol_ = tol; } protected: int dimension_ = 1; // 1 = scalar (Laplace), 2 or 3 = vector (Elasticity) diff --git a/src/polysolve/linear/Solver.cpp b/src/polysolve/linear/Solver.cpp index 4334ac6f..b327f057 100644 --- a/src/polysolve/linear/Solver.cpp +++ b/src/polysolve/linear/Solver.cpp @@ -62,6 +62,15 @@ namespace polysolve::linear #ifdef POLYSOLVE_WITH_MAS #include "MASSolver.hpp" #endif +#ifdef POLYSOLVE_WITH_CPU_HYBRID +#include "CPUHybridSolver.hpp" +#endif +#ifdef POLYSOLVE_WITH_GPU_HYBRID +#include "GPUHybridSolver.hpp" +#endif +#ifdef POLYSOLVE_WITH_CUDSS +#include "cuDSS.hpp" +#endif #include @@ -403,6 +412,24 @@ namespace polysolve::linear { return std::make_unique(); #endif +#ifdef POLYSOLVE_WITH_GPU_HYBRID + } + else if (solver == "GPUHybrid") + { + return std::make_unique(); +#endif +#ifdef POLYSOLVE_WITH_CPU_HYBRID + } + else if (solver == "CPUHybrid") + { + return std::make_unique(); +#endif +#ifdef POLYSOLVE_WITH_CUDSS + } + else if (solver == "cuDSS") + { + return std::make_unique(); +#endif #ifdef POLYSOLVE_WITH_HYPRE } else if (solver == "Hypre") @@ -538,6 +565,15 @@ namespace polysolve::linear #ifdef POLYSOLVE_WITH_MAS "MAS", #endif +#ifdef POLYSOLVE_WITH_GPU_HYBRID + "GPUHybrid", +#endif +#ifdef POLYSOLVE_WITH_CPU_HYBRID + "CPUHybrid", +#endif +#ifdef POLYSOLVE_WITH_CUDSS + "cuDSS", +#endif #ifdef POLYSOLVE_WITH_HYPRE "Hypre", #endif diff --git a/src/polysolve/linear/Solver.hpp b/src/polysolve/linear/Solver.hpp index 3df533fc..c6837056 100644 --- a/src/polysolve/linear/Solver.hpp +++ b/src/polysolve/linear/Solver.hpp @@ -113,7 +113,7 @@ namespace polysolve::linear /// If the problem is nullspace for multigrid solvers virtual void set_is_nullspace(const VectorXd &x) {} - /// Set solver tolerance + /// Set solver tolerance virtual void set_tolerance(const double tol) {} /// diff --git a/src/polysolve/linear/cuDSS.cpp b/src/polysolve/linear/cuDSS.cpp new file mode 100644 index 00000000..348ff771 --- /dev/null +++ b/src/polysolve/linear/cuDSS.cpp @@ -0,0 +1,210 @@ +#include "cuDSS.hpp" + +#include +#include +#include + +#include + +#ifndef CHECK_CUDA +#define CHECK_CUDA(func) \ + { \ + cudaError_t status = (func); \ + if (status != cudaSuccess) \ + { \ + std::cerr << "CUDA Error: " << cudaGetErrorString(status) \ + << " at " << __FILE__ << ":" << __LINE__ << std::endl; \ + throw std::runtime_error("CUDA Error"); \ + } \ + } +#endif + +#ifndef CHECK_CUDSS +#define CHECK_CUDSS(func) \ + { \ + cudssStatus_t status = (func); \ + if (status != CUDSS_STATUS_SUCCESS) \ + { \ + std::cerr << "cuDSS Error Code: " << status \ + << " at " << __FILE__ << ":" << __LINE__ << std::endl; \ + throw std::runtime_error("cuDSS Error"); \ + } \ + } +#endif + +namespace polysolve::linear +{ + + namespace + { + using clock = std::chrono::steady_clock; + + double elapsed_seconds(const std::chrono::time_point &begin) + { + return std::chrono::duration(clock::now() - begin).count(); + } + } // namespace + + cuDSSSolver::cuDSSSolver() + { + CHECK_CUDSS(cudssCreate(&cudss_handle)); + } + + cuDSSSolver::~cuDSSSolver() + { + free_device_memory(); + if (cudss_handle) + { + CHECK_CUDSS(cudssDestroy(cudss_handle)); + cudss_handle = nullptr; + } + } + + void cuDSSSolver::free_device_memory() + { + if (MatrixA) + { + CHECK_CUDSS(cudssMatrixDestroy(MatrixA)); + MatrixA = nullptr; + } + if (MatrixX) + { + CHECK_CUDSS(cudssMatrixDestroy(MatrixX)); + MatrixX = nullptr; + } + if (MatrixB) + { + CHECK_CUDSS(cudssMatrixDestroy(MatrixB)); + MatrixB = nullptr; + } + if (solverData) + { + CHECK_CUDSS(cudssDataDestroy(cudss_handle, solverData)); + solverData = nullptr; + } + if (config) + { + CHECK_CUDSS(cudssConfigDestroy(config)); + config = nullptr; + } + + if (d_csrRowOffsets) + { + CHECK_CUDA(cudaFree(d_csrRowOffsets)); + d_csrRowOffsets = nullptr; + } + if (d_csrColIndices) + { + CHECK_CUDA(cudaFree(d_csrColIndices)); + d_csrColIndices = nullptr; + } + if (d_csrValues) + { + CHECK_CUDA(cudaFree(d_csrValues)); + d_csrValues = nullptr; + } + if (d_x) + { + CHECK_CUDA(cudaFree(d_x)); + d_x = nullptr; + } + if (d_b) + { + CHECK_CUDA(cudaFree(d_b)); + d_b = nullptr; + } + } + + void cuDSSSolver::analyze_pattern(const StiffnessMatrix &A, const int precond_num) + { + free_device_memory(); + + { + auto phase_begin = clock::now(); + CHECK_CUDSS(cudssConfigCreate(&config)); + CHECK_CUDSS(cudssDataCreate(cudss_handle, &solverData)); + SPDLOG_TRACE("[cuDSS] [create_solver_objects] [{:.6f}]", elapsed_seconds(phase_begin)); + } + + { + auto phase_begin = clock::now(); + m_nrows = A.rows(); + m_ncols = A.cols(); + m_nnz = A.nonZeros(); + + CHECK_CUDA(cudaMalloc(&d_csrRowOffsets, (m_nrows + 1) * sizeof(int))); + CHECK_CUDA(cudaMalloc(&d_csrColIndices, m_nnz * sizeof(int))); + CHECK_CUDA(cudaMalloc(&d_csrValues, m_nnz * sizeof(double))); + + CHECK_CUDA(cudaMemcpy(d_csrRowOffsets, A.outerIndexPtr(), (m_nrows + 1) * sizeof(int), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(d_csrColIndices, A.innerIndexPtr(), m_nnz * sizeof(int), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(d_csrValues, A.valuePtr(), m_nnz * sizeof(double), cudaMemcpyHostToDevice)); + + CHECK_CUDSS(cudssMatrixCreateCsr( + &MatrixA, m_nrows, m_ncols, m_nnz, + d_csrRowOffsets, nullptr, d_csrColIndices, d_csrValues, + CUDA_R_32I, CUDA_R_64F, CUDSS_MTYPE_SYMMETRIC, + CUDSS_MVIEW_FULL, CUDSS_BASE_ZERO)); + + CHECK_CUDA(cudaDeviceSynchronize()); + SPDLOG_TRACE("[cuDSS] [copy_sparse_matrix] [{:.6f}]", elapsed_seconds(phase_begin)); + } + + { + auto phase_begin = clock::now(); + + CHECK_CUDSS(cudssExecute(cudss_handle, CUDSS_PHASE_REORDERING, config, solverData, + MatrixA, nullptr, nullptr)); + CHECK_CUDSS(cudssExecute(cudss_handle, CUDSS_PHASE_SYMBOLIC_FACTORIZATION, config, solverData, + MatrixA, nullptr, nullptr)); + CHECK_CUDA(cudaDeviceSynchronize()); + + SPDLOG_TRACE("[cuDSS] [pattern_analysis] [{:.6f}]", elapsed_seconds(phase_begin)); + } + } + + void cuDSSSolver::factorize(const StiffnessMatrix &A) + { + auto phase_begin = clock::now(); + CHECK_CUDSS(cudssExecute(cudss_handle, CUDSS_PHASE_FACTORIZATION, config, solverData, + MatrixA, nullptr, nullptr)); + CHECK_CUDA(cudaDeviceSynchronize()); + SPDLOG_TRACE("[cuDSS] [numerical_factorization] [{:.6f}]", elapsed_seconds(phase_begin)); + } + + void cuDSSSolver::solve(const Ref b, Ref x) + { + { + auto phase_begin = clock::now(); + if (d_x == nullptr || d_b == nullptr) + { + CHECK_CUDA(cudaMalloc(&d_x, m_nrows * sizeof(double))); + CHECK_CUDA(cudaMalloc(&d_b, m_nrows * sizeof(double))); + + CHECK_CUDSS(cudssMatrixCreateDn( + &MatrixX, m_nrows, 1, m_nrows, + d_x, CUDA_R_64F, CUDSS_LAYOUT_COL_MAJOR)); + + CHECK_CUDSS(cudssMatrixCreateDn( + &MatrixB, m_nrows, 1, m_nrows, + d_b, CUDA_R_64F, CUDSS_LAYOUT_COL_MAJOR)); + } + + CHECK_CUDA(cudaMemcpy(d_b, b.data(), m_nrows * sizeof(double), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(d_x, x.data(), m_nrows * sizeof(double), cudaMemcpyHostToDevice)); + + CHECK_CUDA(cudaDeviceSynchronize()); + SPDLOG_TRACE("[cuDSS] [copy_vectors] [{:.6f}]", elapsed_seconds(phase_begin)); + } + + { + auto phase_begin = clock::now(); + CHECK_CUDSS(cudssExecute(cudss_handle, CUDSS_PHASE_SOLVE, config, solverData, + MatrixA, MatrixX, MatrixB)); + CHECK_CUDA(cudaDeviceSynchronize()); + SPDLOG_TRACE("[cuDSS] [solve] [{:.6f}]", elapsed_seconds(phase_begin)); + } + + CHECK_CUDA(cudaMemcpy(x.data(), d_x, m_nrows * sizeof(double), cudaMemcpyDeviceToHost)); + } +} // namespace polysolve::linear \ No newline at end of file diff --git a/src/polysolve/linear/cuDSS.hpp b/src/polysolve/linear/cuDSS.hpp new file mode 100644 index 00000000..3e01f36c --- /dev/null +++ b/src/polysolve/linear/cuDSS.hpp @@ -0,0 +1,48 @@ +#pragma once + +#include "Solver.hpp" +#include +#include +#include +#include +#include + +namespace polysolve::linear +{ + class cuDSSSolver : public Solver + { + public: + cuDSSSolver(); + ~cuDSSSolver(); + + void analyze_pattern(const StiffnessMatrix &A, const int precond_num) override; + void factorize(const StiffnessMatrix &A) override; + void solve(const Ref b, Ref x) override; + + std::string name() const override { return "cuDSS"; } + + private: + void free_device_memory(); + + double pattern_analysis_time, factorization_time, solve_time; + double reordering_time, symbolic_time; + + cudssHandle_t cudss_handle = nullptr; + cudssConfig_t config = nullptr; + cudssData_t solverData = nullptr; + + cudssMatrix_t MatrixA = nullptr; + cudssMatrix_t MatrixX = nullptr; + cudssMatrix_t MatrixB = nullptr; + + int m_nrows = 0; + int m_ncols = 0; + int m_nnz = 0; + + int *d_csrRowOffsets = nullptr; + int *d_csrColIndices = nullptr; + double *d_csrValues = nullptr; + double *d_x = nullptr; + double *d_b = nullptr; + }; +} // namespace polysolve::linear \ No newline at end of file diff --git a/src/polysolve/linear/hybrid_utils/DisjointSet.cpp b/src/polysolve/linear/hybrid_utils/DisjointSet.cpp new file mode 100644 index 00000000..9b67c080 --- /dev/null +++ b/src/polysolve/linear/hybrid_utils/DisjointSet.cpp @@ -0,0 +1,50 @@ +#include "DisjointSet.hpp" + +namespace polysolve::linear::hybrid +{ + DisjointSet::DisjointSet(int n) + { + rank.assign(n, 0); + parent.reserve(n); + for (int i = 0; i < n; i++) + { + parent.push_back(i); + } + } + + int DisjointSet::find_set(int v) + { + if (parent[v] != v) + { + parent[v] = find_set(parent[v]); // Path compression + } + return parent[v]; + } + + void DisjointSet::union_set(int x, int y) + { + // Find the absolute root representatives first + int root_x = find_set(x); + int root_y = find_set(y); + + // If they already belong to the same set, do nothing + if (root_x == root_y) + { + return; + } + + // Union by rank optimization + if (rank[root_x] > rank[root_y]) + { + parent[root_y] = root_x; + } + else + { + parent[root_x] = root_y; + if (rank[root_x] == rank[root_y]) + { + rank[root_y]++; + } + } + } +} // namespace polysolve::linear::hybrid \ No newline at end of file diff --git a/src/polysolve/linear/hybrid_utils/DisjointSet.hpp b/src/polysolve/linear/hybrid_utils/DisjointSet.hpp new file mode 100644 index 00000000..d04471c0 --- /dev/null +++ b/src/polysolve/linear/hybrid_utils/DisjointSet.hpp @@ -0,0 +1,23 @@ +#pragma once + +#include + +namespace polysolve::linear::hybrid +{ + class DisjointSet + { + public: + // Initializes a disjoint set of size n + explicit DisjointSet(int n); + + // Finds the representative root of the set containing v (with path compression) + int find_set(int v); + + // Unites the sets containing x and y (by rank) + void union_set(int x, int y); + + private: + std::vector parent; + std::vector rank; + }; +} // namespace polysolve::linear::hybrid \ No newline at end of file diff --git a/tests/test_linear_solver.cpp b/tests/test_linear_solver.cpp index 70deee92..318b6e89 100644 --- a/tests/test_linear_solver.cpp +++ b/tests/test_linear_solver.cpp @@ -134,6 +134,12 @@ TEST_CASE("all", "[solver]") params[s]["use_preconditioned_residual_norm"] = false; solver->set_parameters(params); } + else if (s == "GPUHybrid" || s == "CPUHybrid") + { + params[s]["relative_tolerance"] = 0.0; + params[s]["absolute_tolerance"] = 1e-8; + solver->set_parameters(params); + } Eigen::VectorXd b(A.rows()); b.setRandom(); Eigen::VectorXd x(b.size()); @@ -238,6 +244,51 @@ TEST_CASE("mas_block_dim", "[.][solver]") } } +TEST_CASE("hybrid_convergence", "[.][solver]") +{ + const std::string path = POLYFEM_DATA_DIR; + Eigen::SparseMatrix A; + const bool ok = loadMarket(A, path + "/A_contact.mtx"); + REQUIRE(ok); + json solver_info; + + std::vector solvers; +#ifdef POLYSOLVE_WITH_CPU_HYBRID + solvers.push_back("CPUHybrid"); +#endif +#ifdef POLYSOLVE_WITH_GPU_HYBRID + solvers.push_back("GPUHybrid"); +#endif + + for (const auto &s : solvers) + { + auto solver = Solver::create(s, ""); + json params; + params[s]["block_size"] = 3; + params[s]["relative_tolerance"] = 0; + params[s]["absolute_tolerance"] = 1e-12; + + solver->set_parameters(params); + Eigen::VectorXd b(A.rows()); + b.setRandom(); + Eigen::VectorXd x(b.size()); + x.setZero(); + + solver->analyze_pattern(A, A.rows()); + solver->factorize(A); + solver->solve(b, x); + + REQUIRE(solver->name() == s); + + solver->get_info(solver_info); + + const double err = (A * x - b).norm(); + INFO("solver: " + s); + REQUIRE(err < 1e-4); + REQUIRE(solver_info["num_iterations"] < 500); + } +} + TEST_CASE("pre_factor", "[solver]") { const std::string path = POLYFEM_DATA_DIR;