-
- wraps `partition_hint` struct for python - `partition_hint`struct in `load_balance.hpp` is adjusted in a way that in case of zero unsigned int `cpu_group_size`/`gpu_group_size` it is set to the default value (via setter/getter) - adds documentation for `partition_hint` - adds testing for `domain_decomposition` including `partition_hint` (in unit and unit_distributed) - adds `partition_hint` in example `ring.py` - corrects `config.cpp` (and doc) to test for `ARB_GPU_ENABLED` instead of `ARB_WITH_GPU` Fixes #776 Addresses #799, #769
953d5007
domain_decomposition.cpp 5.87 KiB
#include <limits>
#include <string>
#include <sstream>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <arbor/context.hpp>
#include <arbor/domain_decomposition.hpp>
#include <arbor/load_balance.hpp>
#include "context.hpp"
#include "recipe.hpp"
#include "strprintf.hpp"
namespace pyarb {
std::string gd_string(const arb::group_description& g) {
return util::pprintf(
"<arbor.group_description: num_cells {}, gids [{}], {}, {}>",
g.gids.size(), util::csv(g.gids, 5), g.kind, g.backend);
}
std::string dd_string(const arb::domain_decomposition& d) {
return util::pprintf(
"<arbor.domain_decomposition: domain_id {}, num_domains {}, num_local_cells {}, num_global_cells {}, groups {}>",
d.domain_id, d.num_domains, d.num_local_cells, d.num_global_cells, d.groups.size());
}
std::string ph_string(const arb::partition_hint& h) {
return util::pprintf(
"<arbor.partition_hint: cpu_group_size {}, gpu_group_size {}, prefer_gpu {}>",
h.cpu_group_size, h.gpu_group_size, (h.prefer_gpu == 1) ? "True" : "False");
}
void register_domain_decomposition(pybind11::module& m) {
using namespace pybind11::literals;
// Group description
pybind11::class_<arb::group_description> group_description(m, "group_description",
"The indexes of a set of cells of the same kind that are grouped together in a cell group.");
group_description
.def(pybind11::init<arb::cell_kind, std::vector<arb::cell_gid_type>, arb::backend_kind>(),
"Construct a group description with cell kind, list of gids, and backend kind.",
"kind"_a, "gids"_a, "backend"_a)
.def_readonly("kind", &arb::group_description::kind,
"The type of cell in the cell group.")
.def_readonly("gids", &arb::group_description::gids,
"The list of gids of the cells in the group.")
.def_readonly("backend", &arb::group_description::backend,
"The hardware backend on which the cell group will run.")
.def("__str__", &gd_string)
.def("__repr__", &gd_string);
// Partition hint
pybind11::class_<arb::partition_hint> partition_hint(m, "partition_hint",
"Provide a hint on how the cell groups should be partitioned.");
partition_hint
.def(pybind11::init<std::size_t, std::size_t, bool>(),
"cpu_group_size"_a = 1, "gpu_group_size"_a = std::numeric_limits<std::size_t>::max(), "prefer_gpu"_a = true,
"Construct a partition hint with arguments:\n"
" cpu_group_size: The size of cell group assigned to CPU, each cell in its own group by default.\n"
" Must be positive, else set to default value.\n"
" gpu_group_size: The size of cell group assigned to GPU, all cells in one group by default.\n"
" Must be positive, else set to default value.\n"
" prefer_gpu: Whether GPU is preferred, True by default.")
.def_readwrite("cpu_group_size", &arb::partition_hint::cpu_group_size,
"The size of cell group assigned to CPU.")
.def_readwrite("gpu_group_size", &arb::partition_hint::gpu_group_size,
"The size of cell group assigned to GPU.")
.def_readwrite("prefer_gpu", &arb::partition_hint::prefer_gpu,
"Whether GPU usage is preferred.")
.def_property_readonly_static("max_size", [](pybind11::object) { return arb::partition_hint::max_size; },
"Get the maximum size of cell groups.")
.def("__str__", &ph_string)
.def("__repr__", &ph_string);
// Domain decomposition
pybind11::class_<arb::domain_decomposition> domain_decomposition(m, "domain_decomposition",
"The domain decomposition is responsible for describing the distribution of cells across cell groups and domains.");
domain_decomposition
.def(pybind11::init<>())
.def("gid_domain",
[](const arb::domain_decomposition& d, arb::cell_gid_type gid) {
return d.gid_domain(gid);
},
"Query the domain id that a cell assigned to (using global identifier gid).",
"gid"_a)
.def_readonly("num_domains", &arb::domain_decomposition::num_domains,
"Number of domains that the model is distributed over.")
.def_readonly("domain_id", &arb::domain_decomposition::domain_id,
"The index of the local domain.\n"
"Always 0 for non-distributed models, and corresponds to the MPI rank for distributed runs.")
.def_readonly("num_local_cells", &arb::domain_decomposition::num_local_cells,
"Total number of cells in the local domain.")
.def_readonly("num_global_cells", &arb::domain_decomposition::num_global_cells,
"Total number of cells in the global model (sum of num_local_cells over all domains).")
.def_readonly("groups", &arb::domain_decomposition::groups,
"Descriptions of the cell groups on the local domain.")
.def("__str__", &dd_string)
.def("__repr__", &dd_string);
// Partition load balancer
// The Python recipe has to be shimmed for passing to the function that takes a C++ recipe.
m.def("partition_load_balance",
[](std::shared_ptr<py_recipe>& recipe, const context_shim& ctx, arb::partition_hint_map hint_map) {
return arb::partition_load_balance(py_recipe_shim(recipe), ctx.context, std::move(hint_map));
},
"Construct a domain_decomposition that distributes the cells in the model described by recipe\n"
"over the distributed and local hardware resources described by context.\n"
"Optionally, provide a dictionary of partition hints for certain cell kinds, by default empty.",
"recipe"_a, "context"_a, "hints"_a=arb::partition_hint_map{});
}
} // namespace pyarb