diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 36e5ea32b..14350e624 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -6,6 +6,7 @@ add_subdirectory(serac) add_subdirectory(drivers) +add_subdirectory(tools) if (ENABLE_DOCS) add_subdirectory(docs) diff --git a/src/tools/CMakeLists.txt b/src/tools/CMakeLists.txt new file mode 100644 index 000000000..d4d572ef6 --- /dev/null +++ b/src/tools/CMakeLists.txt @@ -0,0 +1,9 @@ +# Copyright (c) 2019-2024, Lawrence Livermore National Security, LLC and +# other Serac Project Developers. See the top-level LICENSE file for +# details. +# +# SPDX-License-Identifier: (BSD-3-Clause) + +blt_add_executable(NAME partitioner + SOURCES partitioner.cpp + DEPENDS_ON serac_mesh) diff --git a/src/tools/partitioner.cpp b/src/tools/partitioner.cpp new file mode 100644 index 000000000..28343f998 --- /dev/null +++ b/src/tools/partitioner.cpp @@ -0,0 +1,116 @@ +// Copyright (c) 2019-2024, Lawrence Livermore National Security, LLC and +// other Serac Project Developers. See the top-level LICENSE file for +// details. +// +// SPDX-License-Identifier: (BSD-3-Clause) + +#include +#include + +#include "mfem.hpp" + +#include "axom/CLI11.hpp" +#include "axom/core/utilities/Timer.hpp" + +#include "serac/mesh/mesh_utils.hpp" +#include "serac/serac_config.hpp" + +using timer = axom::utilities::Timer; + +// partition the range {0, 1, 2, ... , n - 1} into +// `num_blocks` roughly equal-sized contiguous chunks +std::vector partition_range(uint32_t n, uint32_t num_blocks) +{ + uint32_t quotient = n / num_blocks; + uint32_t remainder = n % num_blocks; + + std::vector blocks(num_blocks + 1); + if (blocks.size() > 0) { + blocks[0] = 0; + for (uint32_t i = 1; i < num_blocks + 1; i++) { + if (remainder > 0) { + blocks[i] = blocks[i - 1] + quotient + 1; + remainder--; + } else { + blocks[i] = blocks[i - 1] + quotient; + } + } + } + return blocks; +} + +//------------------------------------------------------------------------------ + +int main(int argc, char* argv[]) +{ + axom::CLI::App app{"a tool for partitioning large meshes into smaller parts suitable for MPI"}; + + std::string input_mesh; + app.add_option("-i, --input-mesh", input_mesh, "Input file to use")->required()->check(axom::CLI::ExistingFile); + + std::string output_prefix; + app.add_option("-o, --output-prefix", output_prefix, "Prefix for output meshes")->required(); + + uint32_t num_parts = 4; + app.add_option("-n, --num_parts", num_parts, "number of partitions to generate"); + + uint32_t num_threads = 1; + // this is disabled temporarily, as it seems mfem's implementation is not thread safe + // app.add_option("-j, --num_threads", num_threads, "number of partitions to generate"); + + CLI11_PARSE(app, argc, argv); + + timer stopwatch; + + std::cout << "reading in mesh file ... "; + stopwatch.start(); + std::ifstream infile(input_mesh); + mfem::Mesh mesh(infile); + stopwatch.stop(); + std::cout << "completed after " << stopwatch.elapsed() * 1000.0 << "ms" << std::endl; + + std::cout << "partitioning mesh into " << num_parts << " pieces ... "; + stopwatch.start(); + int* partitioning = mesh.GeneratePartitioning(int(num_parts)); + stopwatch.stop(); + std::cout << "completed after " << stopwatch.elapsed() * 1000.0 << "ms" << std::endl; + + mfem::MeshPartitioner partitioner(mesh, int(num_parts), partitioning); + + timer fileio; + fileio.start(); + std::vector chunks = partition_range(num_parts, num_threads); + + std::vector threads; + for (uint32_t k = 0; k < num_threads; k++) { + threads.push_back(std::thread( + [&](uint32_t tid) { + mfem::MeshPart mesh_part; + for (uint32_t i = chunks[tid]; i < chunks[tid + 1]; i++) { + std::string filename = mfem::MakeParFilename(output_prefix + ".mesh.", int(i)); + if (num_threads == 1) { + std::cout << "extracting part " << i << " and writing it to " << filename << " ... "; + } + stopwatch.start(); + partitioner.ExtractPart(int(i), mesh_part); + std::ofstream f(filename); + f.precision(16); + mesh_part.Print(f); + stopwatch.stop(); + if (num_threads == 1) { + std::cout << "completed after " << stopwatch.elapsed() * 1000.0 << "ms" << std::endl; + } + } + }, + k)); + } + + for (auto& thread : threads) { + thread.join(); + } + + fileio.stop(); + std::cout << "writing files to disk took " << fileio.elapsed() * 1000.0 << "ms total" << std::endl; + + delete partitioning; +}