Feature: extend the partition algorithm to handle weights #1889

dutkalex · 2025-10-09T20:13:50Z

Describe your changes here:
This PR makes it possible to define element weights for the partitioning algorithm. This is achieved via a callback pointer. In the null case, the old unweighted algorithm is used. This introduces one breaking change in the public API, in the t8_forest_set_partition function signature. All the client code in the repo has been updated, but this will require an update of external user code: t8_forest_set_partition(forest, set_from, set_for_coarsening) => t8_forest_set_partition(forest, set_from, set_for_coarsening, nullptr)

Design tradeoffs:

The new algorithm replaces the old one, even if no weight function is provided. Keeping the old algorithm around could possibly result in slightly better performance since it requires less work, but this would come at the cost of increased maintenance and the need for duplicated testing to cover both code paths. Given that the partition algorithm is not a performance bottleneck in practice, the single code path solution was preferred. The final design includes the unweighted case as a fast "early exit" case.
In the general case, a single pass algorithm is not possible because the total weight of the forest wan't be known a priori and must be computed (first traversal of the forest) to then be able to define the new partition boundaries (second traversal). This means that each element weight is computed twice with the current implementation. If the weight function is expensive to evaluate, this could cause performance problems. There are basically two ways to work around this: the results could either be cached during the first evaluation (this means that we need to allocate a temporary array), or the weights could alternatively be computed and stored in an array ahead-of-time by the user, before calling the partition function. In the second case the user then simply provides a weight function that looks up the precomputed weights. This second option was preferred because it gives the user greater flexibility, and because it offers optimal performance in the case where the weight function is cheap to evaluate which is probably the default case ("Don't pay for what you don't use" principle).

All these boxes must be checked by the AUTHOR before requesting review:

The PR is small enough to be reviewed easily. If not, consider splitting up the changes in multiple PRs.
The title starts with one of the following prefixes: Documentation:, Bugfix:, Feature:, Improvement: or Other:.
If the PR is related to an issue, make sure to link it.
The author made sure that, as a reviewer, he/she would check all boxes below.

All these boxes must be checked by the REVIEWERS before merging the pull request:

As a reviewer please read through all the code lines and make sure that the code is fully understood, bug free, well-documented and well-structured.

General

The reviewer executed the new code features at least once and checked the results manually.
The code follows the t8code coding guidelines.
New source/header files are properly added to the CMake files.
The code is well documented. In particular, all function declarations, structs/classes and their members have a proper doxygen documentation.
All new algorithms and data structures are sufficiently optimal in terms of memory and runtime (If this should be merged, but there is still potential for optimization, create a new issue).

Tests

The code is covered in an existing or new test case using Google Test.
The code coverage of the project (reported in the CI) should not decrease. If coverage is decreased, make sure that this is reasonable and acceptable.
Valgrind doesn't find any bugs in the new code. This script can be used to check for errors; see also this wiki article.

If the Pull request introduces code that is not covered by the github action (for example coupling with a new library):

Should this use case be added to the github action?
If not, does the specific use case compile and all tests pass (check manually).

Scripts and Wiki

If a new directory with source files is added, it must be covered by the script/find_all_source_files.scp to check the indentation of these files.
If this PR introduces a new feature, it must be covered in an example or tutorial and a Wiki article.

License

The author added a BSD statement to doc/ (or already has one).

holke · 2025-10-14T10:46:18Z

Thanks again for this addition!

Since 95% of our users do not use weighted partitioning but the classical one and the weighted does add overhead i would argue we should at least skip the weight computation loop in standard mode.
Or keep the standard version as is - but i also see your point of avoiding code duplication which is very important.

dutkalex · 2025-10-14T11:12:44Z

Since 95% of our users do not use weighted partitioning but the classical one and the weighted does add overhead i would argue we should at least skip the weight computation loop in standard mode.

That's an interesting idea. Maybe we can get the best of both worlds with this kind of "fast-forward" behavior in the simple case. I will look into that.

dutkalex · 2025-10-14T11:25:07Z

I wonder what should be the input of the weight function. Currently it takes a forest, a local tree id and an element id (within that tree). I'm not quite convinced this is really usable in practice however, especially if the forest has just been refined or coarsened prior to entering the partition algorithm. Do you have ideas on what would be the right function signature in the general case?
Since the weight function is always used on a committed forest, then all the non mutating public API functions can be used by this weight function to query for additional information if needed.

codecov · 2025-10-15T11:56:09Z

Codecov Report

❌ Patch coverage is 94.02985% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 76.65%. Comparing base (7f65c1a) to head (c6766ee).
⚠️ Report is 8 commits behind head on main.

Files with missing lines	Patch %	Lines
src/t8_forest/t8_forest_partition.cxx	95.16%	3 Missing ⚠️
src/t8_forest/t8_forest.cxx	75.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1889      +/-   ##
==========================================
+ Coverage   76.60%   76.65%   +0.05%     
==========================================
  Files         109      109              
  Lines       18721    18766      +45     
==========================================
+ Hits        14341    14385      +44     
- Misses       4380     4381       +1

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

dutkalex · 2025-10-29T18:04:34Z

I have refactored to use t8_shmem_array_allgatherv in the weighted case as you suggested @holke, but it seems like I'm still missing something to make this work...
https://github.com/dutkalex/t8code/blob/weighted-partition/src/t8_forest/t8_forest_partition.cxx#L519

src/t8_forest/t8_forest_partition.cxx

Co-authored-by: spenke91 <[email protected]>

spenke91 · 2025-11-19T09:25:57Z

Thanks a lot for your contribution @dutkalex ! I am out of office this week, but looking forward to reviewing it next week 👍

spenke91

Thank you so much @dutkalex for your contribuition! I like the new feature a lot and as discussed, it seems to work flawlessly. As you can see, I only have minor remarks, which are mostly on how to pass the weight function to the forest.

Don't hesitate to contact me in Mattermost if you want to discuss something! :-)

spenke91 · 2025-11-25T10:41:08Z

src/t8_forest/t8_forest_partition.h

 */
 void
-t8_forest_partition (t8_forest_t forest);
+t8_forest_partition (t8_forest_t forest, t8_weight_fcn_t *weight_callback);


Does it make sense to add the weight_callback as an extra argument here? To me, it seems like an unnecessary duplication, given that the forest already has the new member weight_function.
My suggestion would be to omit the extra argument here and then use forest->weight_function within t8_forest_partition rather than the function argument? Otherwise my impression is that we are bound to sooner or later run into hard-to-debug problems of the argument and the member being out of sync.

What do you think?

spenke91 · 2025-11-25T10:42:12Z

src/t8_forest/t8_forest_partition.h

 */
 void
-t8_forest_partition (t8_forest_t forest);
+t8_forest_partition (t8_forest_t forest, t8_weight_fcn_t *weight_callback);


Does it make sense to add the weight_callback as an extra argument here? To me, it seems like an unnecessary duplication, given that the forest already has the new member weight_function.
My suggestion would be to omit the extra argument here and then use forest->weight_function within t8_forest_partition rather than the function argument? Otherwise my impression is that we are bound to sooner or later run into hard-to-debug problems of the argument and the member being out of sync.

What do you think?

Suggested change

t8_forest_partition (t8_forest_t forest, t8_weight_fcn_t *weight_callback);

t8_forest_partition (t8_forest_t forest);

I'm actually on the fence regarding this, and it probably shows in the code...

You definitely have a point here, and this problem should be addressed.

I was actually tempted at some point to simply remove the additional parameter and rely on forest->weight_function as you suggest, but from a design perspective this feels just wrong. Conceptually, the weight function is not a property of the forest data structure and should not be stored in the forest. The weight function is a user customization point of the partitioning algorithm. It therefore makes much more sense to have the callback be passed explicitly to the partition algorithm (like you would do when calling std::transform or std::reduce for exemple), and get rid of forest->weight_function instead. However, I could not find a way to make this work because the current API for forest mutation is fundamentally designed around this idea of attaching these customization points to the forest.

This is obviously a problem that goes beyond the scope of this feature and that will have to be addressed at some point. In the meantime, we can either:

Go with the simple and consistent approach, but this means doubling down on this tech debt.

Implement a proper separation of concerns in new code to avoid making things worse, at the cost of some inconsistencies across the codebase, and eventually some breaking changes in the APIs.

For reference, here is a possible solution to make these APIs extensible and future-proof:

// public_api.h struct t8_partition_options { t8_weight_fcn_t* weight_fcn; int set_for_coarsening; }; extern const t8_partition_options default_values; t8forest_t t8_partition(t8forest_t forest_from, t8_partition_options options); // client_code.c t8_partition_options partition_options = default_values; partition_options.weight_fcn = /*...*/ ; t8forest_t new_forest = t8_partition( old_forest, partition_options );

Thank you for the detailed explanation. I discussed with @holke about this today. We do see your point about the design flaws of storing information like the weight function as member data of the forest; however, we would still suggest to stick to it in this PR and go for a function t8_forest_set_partition_weights for the following reasons:

it would be consistent with how this kind of data is handled so far in t8code (whether we like it or not);

it would render / keep this PR rather non-intrusive, in particular by not changing the function interfaces;

we might change this forest-based design in the future, but we do not have immediate plans to do so (because we also don't have the resources right now). Once we change it, it would be a quite deep change that won't get significantly more extensive by the weighted load balancing.

Would it be okay with you if we go forward with this approach?
Speaking of "we": If you prefer that I can also take over the PR and change it accordingly...

Would it be okay with you if we go forward with this approach?

That makes perfect sense, and I understand the consistency and backwards compatibility argument. I just wanted to make sure that this specific design aspect was taken into account before moving forward with this. At the end of the day, you guys have the broader and more insightful vision into this matter, and I trust your jugement.

Speaking of "we": If you prefer that I can also take over the PR and change it accordingly...

I guess it depends on how fast you'd like this to get merged: I'm a little bit over-subscribed at the moment so it will probably have to wait a little bit on my side. If you want to get this merged now, feel free to take over from here.

Ok then I try to free some time this week to finalize it 👍

spenke91 · 2025-11-25T10:46:15Z

src/t8_forest/t8_forest_partition.cxx

+// Compute forest->element_offsets according to the weight function, if provided
 static void
-t8_forest_partition_compute_new_offset (t8_forest_t forest)
+t8_forest_partition_compute_new_offset (t8_forest_t forest, t8_weight_fcn_t *weight_fcn)


same as in my comment below for t8_forest_partition: How about omitting the new argument here and using forest->weight_function within this function instead?

src/t8_forest/t8_forest_partition.cxx

spenke91 · 2025-11-26T15:30:10Z

src/t8_forest/t8_forest_general.h

+t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening,
+                         t8_weight_fcn_t *weight_callback);


After discussing with my colleagues, I suggest not to introduce the remove the weight function as an argument of t8_forest_set_partition, but to instead introduce a separate function t8_forest_set_partition_weights that assigns weight function to a forest.
This way, we can avoid having to update all examples, tutorials. etc and do not need to deal with nullptr that much.

I 100% hear the backwards compatibility argument, but this would result in a user API that is quite counter-intuitive. Moreover, I would argue that this API will have to be reworked at some point anyway, because introducing a new function for each new customization point is not sustainable long-term

As touched upon above, you are right that this strategy does not scale well with an increasing number of options. For this specific feature, however, we still think it would be preferable to go along with it, because it is so simple.

spenke91 · 2025-11-26T15:40:49Z

test/t8_forest/t8_gtest_forest_commit.cxx


-  /* partition the forest */
+  /* partition the forest with a unit weight function */
+  t8_forest_t forest_partition;


I would suggest to revert this test case to its original state, since it is now really about testing the weighted load balancing. Do you happen to already have some other test case you used locally for validation? Otherwise, I can also add a suitable test case in a subsequent PR.

I agree that this is not the best approach to test this feature, and adding a dedicated test suite for the partitioning algorithm is the right way to do things. However, beyond comparing the result of the unweighted algorithm with the weighted one with a constant weight callback, I don't know how to test this in a robust way. For development purposes, I have used a lot of intrusive whitebox testing manually, but this would not be very helpful in a proper automated test suite...

Ok, then I suggest you leave the test case as it currently is in this PR and I will move the weighted partitioning to another test case in a later PR.

Sharing a (probably half-baked) thought I had this morning about how to test this feature in the general case in a non-intrusive and hopefully robust way:
Given a non-trivial weight function, we could manually compute a posteriori each partition's weight, find the global min and max, and assert that the difference between the two is of the order of magnitude of a single element's weight.

Interesting idea, thanks for sharing! I thought about maybe having some simple weight function that is for example 1 for the first half of the elements and 2 for the other, or that is maybe just equal to the element number. I'm optimistic to find something useful here 👍

src/t8_forest/t8_forest_partition.cxx

Co-authored-by: spenke91 <[email protected]>

spenke91

Thank you for the quick replies! I discussed with @holke today about this PR and here's what we think.

spenke91 · 2025-11-28T11:51:17Z

test/t8_forest/t8_gtest_forest_commit.cxx


-  /* partition the forest */
+  /* partition the forest with a unit weight function */
+  t8_forest_t forest_partition;


Ok, then I suggest you leave the test case as it currently is in this PR and I will move the weighted partitioning to another test case in a later PR.

spenke91 · 2025-11-28T12:01:42Z

src/t8_forest/t8_forest_partition.h

 */
 void
-t8_forest_partition (t8_forest_t forest);
+t8_forest_partition (t8_forest_t forest, t8_weight_fcn_t *weight_callback);


Thank you for the detailed explanation. I discussed with @holke about this today. We do see your point about the design flaws of storing information like the weight function as member data of the forest; however, we would still suggest to stick to it in this PR and go for a function t8_forest_set_partition_weights for the following reasons:

it would be consistent with how this kind of data is handled so far in t8code (whether we like it or not);

it would render / keep this PR rather non-intrusive, in particular by not changing the function interfaces;

we might change this forest-based design in the future, but we do not have immediate plans to do so (because we also don't have the resources right now). Once we change it, it would be a quite deep change that won't get significantly more extensive by the weighted load balancing.

Would it be okay with you if we go forward with this approach?
Speaking of "we": If you prefer that I can also take over the PR and change it accordingly...

spenke91 · 2025-11-28T12:04:08Z

src/t8_forest/t8_forest_general.h

+t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening,
+                         t8_weight_fcn_t *weight_callback);


As touched upon above, you are right that this strategy does not scale well with an increasing number of options. For this specific feature, however, we still think it would be preferable to go along with it, because it is so simple.

spenke91

Thank you for your replies and once again for your contribution 👍

spenke91 · 2025-11-28T13:24:30Z

src/t8_forest/t8_forest_partition.h

 */
 void
-t8_forest_partition (t8_forest_t forest);
+t8_forest_partition (t8_forest_t forest, t8_weight_fcn_t *weight_callback);


Ok then I try to free some time this week to finalize it 👍

spenke91 · 2025-12-01T09:45:13Z

test/t8_forest/t8_gtest_forest_commit.cxx


-  /* partition the forest */
+  /* partition the forest with a unit weight function */
+  t8_forest_t forest_partition;


Interesting idea, thanks for sharing! I thought about maybe having some simple weight function that is for example 1 for the first half of the elements and 2 for the other, or that is maybe just equal to the element number. I'm optimistic to find something useful here 👍

weight extension with dummy weights

3c7cb6a

dutkalex marked this pull request as draft October 9, 2025 20:14

dutkalex added 5 commits October 9, 2025 22:44

cleanup and add comments

ebe78e5

fix scan bug

8c4170b

indent

27a54b3

expose the weight function as an optional parameter

3d5c2b6

update comment and indent

0d6ab35

dutkalex added 6 commits October 14, 2025 13:47

add nullptr fast-forwards

9075a91

extract first pass into helper function

b17975f

use forest_from instead of forest

c7f3817

indent

3bfc468

restore optimized path for uniform weights

56b51b1

remove obfuscated while loop

c3ef88d

dutkalex marked this pull request as ready for review October 15, 2025 11:23

dutkalex added 7 commits October 28, 2025 11:08

expose the weight function through t8_forest_set_partition

0f3f711

indent

7773c4c

Merge branch 'main' into weighted-partition

02cea3a

bugfix: replace forest with forest_from

032622c

include weighted path in tests

6eebd25

indent

4431478

refactor to use allgatherv in the weighted case

9557cbe

spenke91 reviewed Nov 13, 2025

View reviewed changes

src/t8_forest/t8_forest_partition.cxx Outdated Show resolved Hide resolved

dutkalex and others added 3 commits November 13, 2025 13:51

Update src/t8_forest/t8_forest_partition.cxx

0b94ce9

Co-authored-by: spenke91 <[email protected]>

Merge branch 'main' into weighted-partition

f81300f

bugfix

2fa7769

indent

349a881

dutkalex requested a review from spenke91 November 18, 2025 11:02

spenke91 self-assigned this Nov 25, 2025

spenke91 requested changes Nov 26, 2025

View reviewed changes

spenke91 assigned dutkalex and unassigned spenke91 Nov 26, 2025

dutkalex and others added 4 commits November 26, 2025 18:10

Add suggested comment

8b65012

Co-authored-by: spenke91 <[email protected]>

Add suggested comment

95287fd

Co-authored-by: spenke91 <[email protected]>

Add suggested comment

0cf3d08

Co-authored-by: spenke91 <[email protected]>

Merge branch 'main' into weighted-partition

c6766ee

spenke91 reviewed Nov 28, 2025

View reviewed changes

spenke91 reviewed Dec 1, 2025

View reviewed changes

	t8_forest_partition (t8_forest_t forest, t8_weight_fcn_t *weight_callback);
	t8_forest_partition (t8_forest_t forest);

		t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening,
		t8_weight_fcn_t *weight_callback);

Feature: extend the partition algorithm to handle weights #1889

Are you sure you want to change the base?

Feature: extend the partition algorithm to handle weights #1889

Conversation

dutkalex commented Oct 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

General

Tests

Scripts and Wiki

License

Uh oh!

holke commented Oct 14, 2025

Uh oh!

dutkalex commented Oct 14, 2025

Uh oh!

dutkalex commented Oct 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov bot commented Oct 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

dutkalex commented Oct 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

spenke91 commented Nov 19, 2025

Uh oh!

spenke91 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dutkalex Nov 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dutkalex Nov 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dutkalex Nov 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

spenke91 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

dutkalex commented Oct 9, 2025 •

edited

Loading

dutkalex commented Oct 14, 2025 •

edited

Loading

codecov bot commented Oct 15, 2025 •

edited

Loading

dutkalex commented Oct 29, 2025 •

edited

Loading

dutkalex Nov 26, 2025 •

edited

Loading

dutkalex Nov 26, 2025 •

edited

Loading

dutkalex Nov 28, 2025 •

edited

Loading