diff --git a/examples/main.c b/examples/main.c
index a8abb938398dbcd7fa22578cbae29ed58e8c8ffe..1ab317d964e3dec09460c685e487599c951c406b 100644
--- a/examples/main.c
+++ b/examples/main.c
@@ -301,9 +301,9 @@ int main(int argc, char *argv[]) {
#endif
/* Prepare the domain decomposition scheme */
+ enum repartition_type reparttype = REPART_NONE;
#ifdef WITH_MPI
struct partition initial_partition;
- enum repartition_type reparttype;
partition_init(&initial_partition, &reparttype, params, nr_nodes);
/* Let's report what we did */
@@ -466,9 +466,10 @@ int main(int argc, char *argv[]) {
/* Initialize the engine with the space and policies. */
if (myrank == 0) clocks_gettime(&tic);
struct engine e;
- engine_init(&e, &s, params, nr_nodes, myrank, nr_threads, with_aff,
- engine_policies, talking, &us, &prog_const, &hydro_properties,
- &potential, &cooling_func, &sourceterms);
+ engine_init(&e, &s, params, nr_nodes, myrank, nr_threads, N_total[0],
+ N_total[1], with_aff, engine_policies, talking, &us,
+ &prog_const, &hydro_properties, &potential, &cooling_func,
+ &sourceterms);
if (myrank == 0) {
clocks_gettime(&toc);
message("engine_init took %.3f %s.", clocks_diff(&tic, &toc),
@@ -524,16 +525,11 @@ int main(int argc, char *argv[]) {
/* Main simulation loop */
for (int j = 0; !engine_is_done(&e) && e.step != nsteps; j++) {
-/* Repartition the space amongst the nodes? */
-#ifdef WITH_MPI
- if (j % 100 == 2) e.forcerepart = reparttype;
-#endif
-
/* Reset timers */
timers_reset(timers_mask_all);
/* Take a step. */
- engine_step(&e);
+ engine_step(&e, reparttype);
/* Dump the task data using the given frequency. */
if (dump_tasks && (dump_tasks == 1 || j % dump_tasks == 1)) {
diff --git a/src/engine.c b/src/engine.c
index 8de8bfe7116367502142df7d4bb0776c8bd2be4b..5bd03863b4f7955c11f012dd7949eaabf0ae5c46 100644
--- a/src/engine.c
+++ b/src/engine.c
@@ -2694,8 +2694,9 @@ void engine_init_particles(struct engine *e, int flag_entropy_ICs) {
* @brief Let the #engine loose to compute the forces.
*
* @param e The #engine.
- */
-void engine_step(struct engine *e) {
+ * @param reparttype type of repartitioning to use (MPI only)
+*/
+void engine_step(struct engine *e, enum repartition_type reparttype) {
double snapshot_drift_time = 0.;
@@ -2709,6 +2710,27 @@ void engine_step(struct engine *e) {
/* Recover the (integer) end of the next time-step */
engine_collect_timestep(e);
+ /* If the last step updated all particles then repartition the space around
+ * the nodes.
+ *
+ * XXX handle fixdt, that will repartition all the time.
+ *
+ * XXX Look at node balance, try to use that to decide if necessary.
+ */
+#ifdef WITH_MPI
+ if (e->nodeID == 0) {
+ if ((e->updates != 0 && e->updates == e->total_nr_parts) ||
+ (e->g_updates != 0 && e->g_updates == e->total_nr_gparts)) {
+ message("will repartition %ld %ld %ld %ld", e->updates,
+ e->total_nr_parts, e->g_updates, e->total_nr_gparts);
+ e->forcerepart = reparttype;
+ }
+ }
+
+ /* All nodes do this together. */
+ MPI_Bcast(&e->forcerepart, 1, MPI_INT, 0, MPI_COMM_WORLD);
+#endif
+
/* Check for output */
while (e->ti_end_min >= e->ti_nextSnapshot && e->ti_nextSnapshot > 0) {
@@ -3171,6 +3193,8 @@ void engine_unpin() {
* @param nr_nodes The number of MPI ranks.
* @param nodeID The MPI rank of this node.
* @param nr_threads The number of threads per MPI rank.
+ * @param Ngas total number of gas particles in the simulation.
+ * @param Ndm total number of gravity particles in the simulation.
* @param with_aff use processor affinity, if supported.
* @param policy The queuing policy to use.
* @param verbose Is this #engine talkative ?
@@ -3183,8 +3207,8 @@ void engine_unpin() {
*/
void engine_init(struct engine *e, struct space *s,
const struct swift_params *params, int nr_nodes, int nodeID,
- int nr_threads, int with_aff, int policy, int verbose,
- const struct UnitSystem *internal_units,
+ int nr_threads, int Ngas, int Ndm, int with_aff, int policy,
+ int verbose, const struct UnitSystem *internal_units,
const struct phys_const *physical_constants,
const struct hydro_props *hydro,
const struct external_potential *potential,
@@ -3201,6 +3225,8 @@ void engine_init(struct engine *e, struct space *s,
e->step = 0;
e->nr_nodes = nr_nodes;
e->nodeID = nodeID;
+ e->total_nr_parts = Ngas;
+ e->total_nr_gparts = Ndm;
e->proxy_ind = NULL;
e->nr_proxies = 0;
e->forcerebuild = 1;
diff --git a/src/engine.h b/src/engine.h
index ec9e16553c3172f22e9bfe60971163488947a47e..6f2345ee7e94318e87913d37d719df634b7de803 100644
--- a/src/engine.h
+++ b/src/engine.h
@@ -140,6 +140,9 @@ struct engine {
/* Number of particles updated */
size_t updates, g_updates;
+ /* Total numbers of particles in the system. */
+ size_t total_nr_parts, total_nr_gparts;
+
/* The internal system of units */
const struct UnitSystem *internalUnits;
@@ -223,8 +226,8 @@ void engine_drift(struct engine *e);
void engine_dump_snapshot(struct engine *e);
void engine_init(struct engine *e, struct space *s,
const struct swift_params *params, int nr_nodes, int nodeID,
- int nr_threads, int with_aff, int policy, int verbose,
- const struct UnitSystem *internal_units,
+ int nr_threads, int Ngas, int Ndm, int with_aff, int policy,
+ int verbose, const struct UnitSystem *internal_units,
const struct phys_const *physical_constants,
const struct hydro_props *hydro,
const struct external_potential *potential,
@@ -235,7 +238,7 @@ void engine_launch(struct engine *e, int nr_runners, unsigned int mask,
void engine_prepare(struct engine *e, int nodrift);
void engine_print(struct engine *e);
void engine_init_particles(struct engine *e, int flag_entropy_ICs);
-void engine_step(struct engine *e);
+void engine_step(struct engine *e, enum repartition_type reparttype);
void engine_maketasks(struct engine *e);
void engine_split(struct engine *e, struct partition *initial_partition);
void engine_exchange_strays(struct engine *e, size_t offset_parts,