firedrakeproject / firedrake

Firedrake is an automated system for the portable solution of partial differential equations using the finite element method (FEM)
https://firedrakeproject.org
Other
482 stars 157 forks source link

`utils._new_uid` should be collective over a communicator #3631

Closed wence- closed 1 week ago

wence- commented 2 weeks ago

As noted in #3628, the _new_uid facility can get out of sync across processes when multiple communicators are involved. Although most things in firedrake don't care about the uid, writing data does, so we should make sure that if the user doesn't provide names to functions, then _new_uid is safe even in a scenario with divergent processes as long as their program is correctly logically collective.

The way to handle this is to adapt _new_uid https://github.com/firedrakeproject/firedrake/blob/cb77d32ac8559d8d57fe0ef0efe251d9cbaf2309/firedrake/utils.py#L23

It should take a communicator (since it's only ever called in a context where a communicator is available), pull out the firedrake "internal" communicator (cc @JDBetteridge), and use that to key an id generation. This can be done by attaching an attribute to the communicator that is the uid counter. Since the calls to _new_uid are collective over a communicator, pulling the uid out of the communicator attribute and updating it will be safe in this scenario.