From dba87db313b64f0b08e1bed34457cf16563674fa Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 9 Sep 2026 21:26:06 -0700 Subject: [PATCH 1/3] Two constants of the same name were one symbol in the generated C A layered viscosity built from two ViscousFlowModels returned the UNIFORM-viscosity answer: L2 2.819e-1 against the exact layered Couette profile, and 1.502e-10 against plain linear shear. The composed flux was correct and so was the constants manifest ((0, \eta = 1.0), (1, \eta = 1000.0)); only the emitted C was wrong. _JITConstant is a plain sympy.Symbol subclass that never adopted the disambiguation UWexpression uses (docs/developer/design/ SYMBOL_DISAMBIGUATION_2025-12.md), and it has neither half of it. It constructed through Symbol.__new__, which is cached by NAME, so the second placeholder was literally the first object and assigning its _ccodestr overwrote the first one's: a, b = _JITConstant(0, name='same'), _JITConstant(1, name='same') a is b -> True a._ccodestr, b._ccodestr -> constants[1], constants[1] Every occurrence then rendered as one constants[] slot, so the blend collapsed to (phi_0 + phi_1) * constants[k]. Every ViscousFlowModel calls its viscosity \eta, so any model with two of them was affected, not just multi-material. Construct via Symbol.__xnew__ to bypass the cache and put the slot index in _hashable_content, so identity is the slot rather than the name. The same solve now gives 1.807e-7 against the layered exact solution. Also: the manifest sort tie-broke on str(expr), which for a UWexpression is its current VALUE, so two same-named constants could swap slots when a parameter changed. Tie-break on instance_number instead. Regression tests at both levels in test_0103_jit_rampable_constants.py. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01MSGAFeA7qYXgkuw9ud8F2G --- src/underworld3/constitutive_models.py | 9 ++ src/underworld3/utilities/_jitextension.py | 44 ++++++++- tests/test_0103_jit_rampable_constants.py | 103 +++++++++++++++++++++ 3 files changed, 154 insertions(+), 2 deletions(-) diff --git a/src/underworld3/constitutive_models.py b/src/underworld3/constitutive_models.py index 1c807d4a7..8c8e38cdc 100644 --- a/src/underworld3/constitutive_models.py +++ b/src/underworld3/constitutive_models.py @@ -4518,6 +4518,15 @@ def __init__( Set to True if IndexSwarmVariable does not maintain partition of unity. Default: False (assumes IndexSwarmVariable maintains partition of unity) """ + # NB (2026-09-09): this class used to return the UNIFORM-viscosity + # answer (layered Couette L2 2.8e-1; identical to plain linear shear to + # 1.5e-10) while its composed flux and the constants manifest both + # looked correct. The cause was in the JIT, not here: every constituent + # names its viscosity \eta, and the constants[] placeholder was named + # for the expression alone, so sympy treated two slots as one symbol + # and the blend collapsed to a single viscosity. Fixed in + # utilities/_jitextension.py::_extract_constants; regression in + # tests/test_0103_jit_rampable_constants.py. # Validate compatibility before initialization self._validate_model_compatibility(constitutive_models) diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index 5301b43fb..83adffc5c 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -334,16 +334,44 @@ class _JITConstant(sympy.Symbol): Used by the JIT compiler to route constant UWexpressions through PETSc's PetscDSSetConstants() mechanism instead of baking values as C literals. + + Identity is the ``constants[]`` INDEX, not the name. Two constants may + legitimately share a display name — every ``ViscousFlowModel`` calls its + viscosity :math:`\eta`, so a two-material model has two of them — and each + needs its own slot. This follows the same mechanism ``UWexpression`` uses + to keep same-named symbols on different meshes apart (see + ``docs/developer/design/SYMBOL_DISAMBIGUATION_2025-12.md``): construct via + ``Symbol.__xnew__`` to bypass SymPy's name-keyed instance cache, and add + the discriminator to ``_hashable_content``. + + Without both halves, ``sympy.Symbol.__new__`` returns the *cached instance + for that name*: the second placeholder was literally the first object, and + assigning its ``_ccodestr`` overwrote the first one's. Every occurrence + then rendered as one ``constants[]`` slot, so a layered viscosity + assembled as ``(phi_0 + phi_1) * constants[k]`` — a single uniform + viscosity — while the manifest and the symbolic expression both looked + correct. """ + __slots__ = ("_const_index", "_ccodestr") + def __new__(cls, index, name=None): if name is None: name = f"_jit_const_{index}" - obj = super().__new__(cls, name) + # __xnew__, not __new__: the latter is cached by (name, assumptions), + # which knows nothing about _const_index. + obj = sympy.Symbol.__xnew__(cls, name) obj._const_index = index obj._ccodestr = f"constants[{index}]" return obj + def _hashable_content(self): + """Two placeholders differ if their constants[] slot differs.""" + return sympy.Symbol._hashable_content(self) + (self._const_index,) + + def __getnewargs_ex__(self): + return ((self._const_index, self.name), {}) + def _ccode(self, printer): return self._ccodestr @@ -394,13 +422,25 @@ def _extract_constants(all_fns, mesh): # Sort by the user-given symbol name, not ``str(expr)`` — ``__str__`` on a # UWexpression returns the current *value*, which shuffles the index # assignment whenever a value changes. ``.name`` is stable. - sorted_constants = sorted(constant_exprs, key=lambda e: (e.name, _stable_sort_key(e))) + # + # Two constants can legitimately SHARE a name: every ViscousFlowModel calls + # its viscosity \eta, so a model with two of them has two \eta constants. + # ``instance_number`` (creation order, identical on every rank running the + # same script) breaks that tie without reintroducing the value into the key. + sorted_constants = sorted( + constant_exprs, + key=lambda e: (e.name, getattr(e, "instance_number", -1), _stable_sort_key(e)), + ) manifest = [] subs_map = {} for i, expr in enumerate(sorted_constants): # Use ``expr.name`` (stable) instead of ``str(expr)`` (= current value) # so the placeholder symbol's identity is independent of parameter value. + # + # Same-named constants stay apart through _JITConstant's own + # disambiguation (its _hashable_content carries the slot index), not + # through a unique name — see the class docstring. jit_const = _JITConstant(i, name=f"_jit_const_{expr.name}") manifest.append((i, expr)) subs_map[expr] = jit_const diff --git a/tests/test_0103_jit_rampable_constants.py b/tests/test_0103_jit_rampable_constants.py index 180df0683..e4abe436a 100644 --- a/tests/test_0103_jit_rampable_constants.py +++ b/tests/test_0103_jit_rampable_constants.py @@ -89,3 +89,106 @@ def test_ramped_solution_matches_rebuilt_solution(): assert np.isclose(ramped, rebuilt, rtol=1e-10), ( f"ramped {ramped} != rebuilt {rebuilt}") + + +def test_two_constants_with_the_same_name_do_not_collapse(): + """Two constants may legitimately share a symbol name — every + ViscousFlowModel calls its viscosity \\eta — and they must reach the + compiled kernel as two different ``constants[]`` slots. + + The JIT placeholder is a plain ``sympy.Symbol`` subclass, so a placeholder + named for the expression alone made the two the SAME symbol whatever their + index. A two-material Stokes solve then assembled + ``(phi_0 + phi_1) * constants[k]`` — one uniform viscosity — and returned + exactly the linear-shear answer while the manifest and the symbolic + expression both looked correct. + + NB the two constants have to be built the way a solver builds them (through + ``Parameters``, which tags each one). Two BARE expressions of the same name + are the same symbol to sympy by design — ``2*a + 3*b`` is ``5*\\eta`` — so + they never reach the JIT as two things in the first place. + """ + import sympy + from underworld3.utilities._jitextension import _extract_constants + + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.5, qdegree=2) + v = uw.discretisation.MeshVariable("vjc", mesh, mesh.dim, degree=2) + p = uw.discretisation.MeshVariable("pjc", mesh, 1, degree=1) + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + + lower = uw.constitutive_models.ViscousFlowModel(stokes.Unknowns, material_name="lo") + lower.Parameters.shear_viscosity_0 = 1.0 + upper = uw.constitutive_models.ViscousFlowModel(stokes.Unknowns, material_name="up") + upper.Parameters.shear_viscosity_0 = 1000.0 + a = lower.Parameters.shear_viscosity_0 + b = upper.Parameters.shear_viscosity_0 + assert a != b, "the Parameters route no longer distinguishes same-named constants" + + phi0, phi1, xs = sympy.symbols("phi0 phi1 xs") + manifest, subs_map = _extract_constants((a * phi0 * xs + b * phi1 * xs,), mesh) + + assert len(manifest) == 2, manifest + placeholders = {subs_map[a], subs_map[b]} + assert len(placeholders) == 2, "the two placeholders are the same symbol" + + lowered = sympy.expand((a * phi0 * xs + b * phi1 * xs).xreplace(subs_map)) + assert len(lowered.free_symbols & placeholders) == 2, lowered + + +def test_two_materials_solve_the_layered_problem_not_the_uniform_one(): + """The end-to-end form of the same defect. + + Two ViscousFlowModels composed into a layered viscosity must give the + layered answer. Before the fix this returned exactly linear shear (L2 + 2.819e-1 against the layered profile, 1.502e-10 against uniform shear) + because the two \\eta constants shared a constants[] slot. The check that + does not depend on the mesh or the level-set representation is that the + composed model and the equivalent single-model blend now agree. + """ + import sympy + + eta_top, h = 1.0e3, 0.5 + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.1, qdegree=2, regular=True) + swarm = uw.swarm.Swarm(mesh) + material = uw.swarm.IndexSwarmVariable("Mjc", swarm, indices=2, proxy_degree=1) + swarm.populate(fill_param=3) + X = np.asarray(swarm._particle_coordinates.data) + with uw.synchronised_array_update(): + material.data[:, 0] = (X[:, 1] > h).astype(int) + + def _solve(tag, build_model): + v = uw.discretisation.MeshVariable(f"vjc{tag}", mesh, mesh.dim, degree=2) + p = uw.discretisation.MeshVariable(f"pjc{tag}", mesh, 1, degree=1) + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + build_model(stokes) + stokes.add_dirichlet_bc((1.0, 0.0), "Top") + stokes.add_dirichlet_bc((0.0, 0.0), "Bottom") + stokes.add_dirichlet_bc((sympy.oo, 0.0), "Left") + stokes.add_dirichlet_bc((sympy.oo, 0.0), "Right") + stokes.tolerance = 1e-8 + stokes.solve() + return np.asarray(v.data[:, 0]), np.asarray(v.coords) + + def _composed(stokes): + lower = uw.constitutive_models.ViscousFlowModel( + stokes.Unknowns, material_name="lower") + lower.Parameters.shear_viscosity_0 = 1.0 + upper = uw.constitutive_models.ViscousFlowModel( + stokes.Unknowns, material_name="upper") + upper.Parameters.shear_viscosity_0 = eta_top + stokes.constitutive_model = uw.MultiMaterialConstitutiveModel( + stokes.Unknowns, material, [lower, upper]) + + def _blended(stokes): + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = ( + material.createMask([1.0, eta_top])) + + composed, coords = _solve("c", _composed) + blended, _ = _solve("b", _blended) + + # the two routes are the same problem, and must give the same answer + assert np.sqrt(np.mean((composed - blended) ** 2)) < 1e-9 + + # and it is NOT the uniform-viscosity answer (which is linear shear) + assert np.sqrt(np.mean((composed - coords[:, 1]) ** 2)) > 1e-2 From 174bb0875175dc79a88ea3285320a404afcd1adf Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 10 Sep 2026 11:35:12 -0700 Subject: [PATCH 2/3] Review (#717): normalise the sort key, and shorten the incident note MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The manifest sort key called instance_number directly. That is an int for every UWexpression built today, because uw_object.__init__ assigns it — so the TypeError Copilot describes is not reachable at present. But `_uw_id = None` is a state UWexpression deliberately supports (its _hashable_content branches on it), so a set mixing None with an int would sort-crash the moment that state became reachable. Normalise instead of relying on the current initialisation order. The NB block in MultiMaterialConstitutiveModel narrated the whole incident — debug numbers, root cause, file paths — in a hot runtime class. The style charter asks for intent and constraints, not history, and the history already lives in the commit message and the PR. Reduced to the constraint that matters to a reader of that class: constituents that share a parameter name depend on _JITConstant keeping its slots distinct, and here is the test. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01MSGAFeA7qYXgkuw9ud8F2G --- src/underworld3/constitutive_models.py | 13 ++++--------- src/underworld3/utilities/_jitextension.py | 14 ++++++++++---- 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/src/underworld3/constitutive_models.py b/src/underworld3/constitutive_models.py index 8c8e38cdc..02a1a565c 100644 --- a/src/underworld3/constitutive_models.py +++ b/src/underworld3/constitutive_models.py @@ -4518,15 +4518,10 @@ def __init__( Set to True if IndexSwarmVariable does not maintain partition of unity. Default: False (assumes IndexSwarmVariable maintains partition of unity) """ - # NB (2026-09-09): this class used to return the UNIFORM-viscosity - # answer (layered Couette L2 2.8e-1; identical to plain linear shear to - # 1.5e-10) while its composed flux and the constants manifest both - # looked correct. The cause was in the JIT, not here: every constituent - # names its viscosity \eta, and the constants[] placeholder was named - # for the expression alone, so sympy treated two slots as one symbol - # and the blend collapsed to a single viscosity. Fixed in - # utilities/_jitextension.py::_extract_constants; regression in - # tests/test_0103_jit_rampable_constants.py. + # Constituents that share a parameter name (every ViscousFlowModel + # calls its viscosity \eta) rely on _JITConstant keeping its + # constants[] slots distinct; see utilities/_jitextension.py and the + # regression in tests/test_0103_jit_rampable_constants.py. # Validate compatibility before initialization self._validate_model_compatibility(constitutive_models) diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index 83adffc5c..d53323551 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -427,10 +427,16 @@ def _extract_constants(all_fns, mesh): # its viscosity \eta, so a model with two of them has two \eta constants. # ``instance_number`` (creation order, identical on every rank running the # same script) breaks that tie without reintroducing the value into the key. - sorted_constants = sorted( - constant_exprs, - key=lambda e: (e.name, getattr(e, "instance_number", -1), _stable_sort_key(e)), - ) + def _order_key(e): + # instance_number is an int for every UWexpression built today + # (uw_object.__init__ assigns it), but ``_uw_id = None`` is a state the + # class deliberately supports — _hashable_content branches on it — so + # normalise rather than let a mixed set raise TypeError mid-sort. + instance = getattr(e, "instance_number", None) + return (e.name, -1 if instance is None else int(instance), + _stable_sort_key(e)) + + sorted_constants = sorted(constant_exprs, key=_order_key) manifest = [] subs_map = {} From 143061261848de2e17f0ce8aa7278674b26a7594 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 10 Sep 2026 14:13:34 -0700 Subject: [PATCH 3/3] Review: the fix needed BOTH halves, and the tests had to prove it The first fix put the constants[] slot index in the placeholder NAME. The style review pushed toward the "structural" version -- Symbol.__xnew__ plus the index in _hashable_content -- and I replaced one with the other. That was wrong: they fix two different properties. identity _hashable_content + __xnew__ without it, two same-named placeholders are ONE object and the blend silently collapses ordering the index in the NAME without it, Symbol.sort_key() ties, Add term order follows the hash seed, ranks emit different C, and getext aborts the run sort_key() is derived from the name, so _hashable_content does nothing for it. Identity without ordering is a parallel abort; ordering without identity is a silently wrong answer. Both are now in, and the docstring says so. Measured with only the identity half: three distinct compiled modules from one model across four hash seeds (a6cb2e83, a6cb2e83, df7e63dd, 2f8aadf7), and `mpirun -np 2` aborting on 4 of 6 launches. With both halves, one module across five seeds and 0 of 6 launches aborting. Three regression tests, each verified to FAIL on the specific bug it guards: - distinct sort_key, and a sum that canonicalises identically written either way. Fails with only the name fix reverted, while the identity test still passes -- so the pre-existing tests could not have caught this. - slot assignment unmoved by a value ramp. NB the first version of this test used 1000 -> 1e-6 and was VACUOUS, because "1.00000000000000e-6" sorts as a prefix of "1.00000000000000"; 1000 -> 0.5 does permute. - a hermetic hash-seed sweep (UW_JIT_CACHE_DIR, cold cache per seed) asserting one compiled module. This replaces an mpirun-based test written first and DELETED: it caught the bug 0 times in 6 with the defect present and confirmed, because whether two ranks disagree depends on their seeds. A coin-flip assertion is not a regression test. Also dropped the instance_number None guard added in the previous round: the charter forbids guarding states that cannot occur, and the comment justifying it was factually wrong (UWexpression.__init__ assigns _uw_id directly; it never calls uw_object.__init__). Made the _JITConstant docstring raw -- it contained \e, an invalid escape that is scheduled to become a SyntaxError. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01MSGAFeA7qYXgkuw9ud8F2G --- src/underworld3/utilities/_jitextension.py | 82 +++++++------- tests/test_0103_jit_rampable_constants.py | 73 ++++++++++++ .../test_0105_jit_source_seed_independence.py | 106 ++++++++++++++++++ 3 files changed, 220 insertions(+), 41 deletions(-) create mode 100644 tests/test_0105_jit_source_seed_independence.py diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index d53323551..2d1ebbe11 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -329,38 +329,44 @@ def prepare_for_cache_key(fn, constants_subs_map): # ============================================================================ class _JITConstant(sympy.Symbol): - """Symbol subclass that renders as constants[i] in generated C code. - - Used by the JIT compiler to route constant UWexpressions through - PETSc's PetscDSSetConstants() mechanism instead of baking values - as C literals. - - Identity is the ``constants[]`` INDEX, not the name. Two constants may - legitimately share a display name — every ``ViscousFlowModel`` calls its - viscosity :math:`\eta`, so a two-material model has two of them — and each - needs its own slot. This follows the same mechanism ``UWexpression`` uses - to keep same-named symbols on different meshes apart (see - ``docs/developer/design/SYMBOL_DISAMBIGUATION_2025-12.md``): construct via - ``Symbol.__xnew__`` to bypass SymPy's name-keyed instance cache, and add - the discriminator to ``_hashable_content``. - - Without both halves, ``sympy.Symbol.__new__`` returns the *cached instance - for that name*: the second placeholder was literally the first object, and - assigning its ``_ccodestr`` overwrote the first one's. Every occurrence - then rendered as one ``constants[]`` slot, so a layered viscosity - assembled as ``(phi_0 + phi_1) * constants[k]`` — a single uniform - viscosity — while the manifest and the symbolic expression both looked - correct. + r"""Symbol subclass that renders as ``constants[i]`` in generated C code. + + Used by the JIT compiler to route constant UWexpressions through PETSc's + ``PetscDSSetConstants()`` mechanism instead of baking values as C literals. + + Two constants may legitimately share a display name — every + ``ViscousFlowModel`` calls its viscosity :math:`\eta`, so a two-material + model has two of them — and each needs its own ``constants[]`` slot. Two + separate SymPy properties have to hold for that to work, and they are not + the same property: + + **Identity** — the slot index is in ``_hashable_content``, and the symbol + is built with ``Symbol.__xnew__`` to bypass SymPy's ``(cls, name)`` + instance cache. Without both, ``Symbol.__new__`` hands back the cached + instance for that name: the second placeholder IS the first object, and + setting its ``_ccodestr`` overwrites the first one's, so every occurrence + renders as one slot. + + **Ordering** — the slot index is also in the NAME. ``_hashable_content`` + does nothing for ``Symbol.sort_key()``, which is derived from the name, so + two same-named placeholders sort equal; term order inside an ``Add`` then + falls back to hash order, which is randomised per process. The generated C + then differs between MPI ranks and ``getext``'s cross-rank hash check + aborts the run — intermittently, since it depends on the hash seed. + + Identity without ordering is a parallel abort; ordering without identity is + a silently wrong answer. Keep both. ``tests/test_0103_jit_rampable_constants.py`` + pins each one separately. """ __slots__ = ("_const_index", "_ccodestr") def __new__(cls, index, name=None): - if name is None: - name = f"_jit_const_{index}" - # __xnew__, not __new__: the latter is cached by (name, assumptions), - # which knows nothing about _const_index. - obj = sympy.Symbol.__xnew__(cls, name) + # The index leads the name so that sort_key() orders placeholders by + # slot; see the class docstring on why the name alone is not enough + # and _hashable_content alone is not either. + suffix = "" if name is None else f"_{name}" + obj = sympy.Symbol.__xnew__(cls, f"_jit_const_{index}{suffix}") obj._const_index = index obj._ccodestr = f"constants[{index}]" return obj @@ -427,16 +433,13 @@ def _extract_constants(all_fns, mesh): # its viscosity \eta, so a model with two of them has two \eta constants. # ``instance_number`` (creation order, identical on every rank running the # same script) breaks that tie without reintroducing the value into the key. - def _order_key(e): - # instance_number is an int for every UWexpression built today - # (uw_object.__init__ assigns it), but ``_uw_id = None`` is a state the - # class deliberately supports — _hashable_content branches on it — so - # normalise rather than let a mixed set raise TypeError mid-sort. - instance = getattr(e, "instance_number", None) - return (e.name, -1 if instance is None else int(instance), - _stable_sort_key(e)) - - sorted_constants = sorted(constant_exprs, key=_order_key) + # Creation order breaks a name tie. It is identical on every rank of an + # SPMD run, and unlike the value it does not move when a parameter is + # ramped — a slot permutation between two solves of the same model would + # invalidate the JIT cache for no reason. + sorted_constants = sorted( + constant_exprs, key=lambda e: (e.name, e.instance_number, _stable_sort_key(e)) + ) manifest = [] subs_map = {} @@ -444,10 +447,7 @@ def _order_key(e): # Use ``expr.name`` (stable) instead of ``str(expr)`` (= current value) # so the placeholder symbol's identity is independent of parameter value. # - # Same-named constants stay apart through _JITConstant's own - # disambiguation (its _hashable_content carries the slot index), not - # through a unique name — see the class docstring. - jit_const = _JITConstant(i, name=f"_jit_const_{expr.name}") + jit_const = _JITConstant(i, name=expr.name) manifest.append((i, expr)) subs_map[expr] = jit_const diff --git a/tests/test_0103_jit_rampable_constants.py b/tests/test_0103_jit_rampable_constants.py index e4abe436a..5691da513 100644 --- a/tests/test_0103_jit_rampable_constants.py +++ b/tests/test_0103_jit_rampable_constants.py @@ -135,6 +135,79 @@ def test_two_constants_with_the_same_name_do_not_collapse(): assert len(lowered.free_symbols & placeholders) == 2, lowered +def test_same_named_placeholders_order_deterministically(): + """The ORDERING half of the same-name problem, pinned without MPI. + + ``_hashable_content`` gives two placeholders separate identities but does + nothing for ``Symbol.sort_key()``, which comes from the name. Two + placeholders that sort equal leave term order inside an ``Add`` to hash + order, which is randomised per process — so the generated C differs + between MPI ranks and the cross-rank hash check aborts the run, on roughly + half of launches. This test is the deterministic proxy: distinct sort keys, + and a sum that canonicalises the same however it is written. + """ + import sympy + from underworld3.utilities._jitextension import _JITConstant + + c0 = _JITConstant(0, name=r"\eta") + c1 = _JITConstant(1, name=r"\eta") + + assert c0.sort_key() != c1.sort_key(), ( + "same-named placeholders sort equal; Add term order will follow the " + "hash seed and the generated C will differ between ranks") + assert c0 != c1 and c0 is not c1 # the identity half + + p0, p1 = sympy.symbols("p0 p1") + written_one_way = sympy.printing.ccode(p0 / c0 + p1 / c1) + written_the_other = sympy.printing.ccode(p1 / c1 + p0 / c0) + assert written_one_way == written_the_other, ( + written_one_way, written_the_other) + assert "constants[0]" in written_one_way and "constants[1]" in written_one_way + + +def test_the_manifest_order_does_not_move_when_a_value_changes(): + """Slot assignment follows creation order, not value. + + ``_stable_sort_key`` falls through to ``str(expr)``, which for a + UWexpression is its CURRENT VALUE, so tie-breaking two same-named + constants on it permutes their ``constants[]`` slots the moment one is + ramped past the other lexically. Ramping 1000 -> 0.5 does exactly that. + Slots that move for a reason unrelated to the model invalidate the JIT + cache needlessly, and would swap the two values outright for anything that + keyed on the manifest rather than on the generated source. + + NB the values matter: 1000 -> 1e-6 does NOT expose it, because + ``"1.00000000000000e-6"`` sorts after ``"1.00000000000000"`` as a prefix. + """ + import sympy + from underworld3.utilities._jitextension import _extract_constants + + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.5, qdegree=2) + v = uw.discretisation.MeshVariable("vmo", mesh, mesh.dim, degree=2) + p = uw.discretisation.MeshVariable("pmo", mesh, 1, degree=1) + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + + lower = uw.constitutive_models.ViscousFlowModel(stokes.Unknowns, material_name="lo") + lower.Parameters.shear_viscosity_0 = 1.0 + upper = uw.constitutive_models.ViscousFlowModel(stokes.Unknowns, material_name="up") + upper.Parameters.shear_viscosity_0 = 1000.0 + a = lower.Parameters.shear_viscosity_0 + b = upper.Parameters.shear_viscosity_0 + assert a != b and a.name == b.name, "expected two distinct same-named constants" + + x = sympy.Symbol("xmo") + slot_of = lambda manifest: {id(e): i for i, e in manifest} + + before = slot_of(_extract_constants((a * x + b * x,), mesh)[0]) + assert len(before) == 2, before + + b.sym = 0.5 # in place: the same object, a new value + after = slot_of(_extract_constants((a * x + b * x,), mesh)[0]) + + assert before == after, ( + "ramping a value permuted the constants[] slots", before, after) + + def test_two_materials_solve_the_layered_problem_not_the_uniform_one(): """The end-to-end form of the same defect. diff --git a/tests/test_0105_jit_source_seed_independence.py b/tests/test_0105_jit_source_seed_independence.py new file mode 100644 index 000000000..0e5463162 --- /dev/null +++ b/tests/test_0105_jit_source_seed_independence.py @@ -0,0 +1,106 @@ +"""The generated C must not depend on the Python hash seed. + +Every rank of an MPI job generates the C source independently and `getext` +allgathers a hash of it, so anything that leaks Python's per-process hash +randomisation into the emitted code aborts the run: + + RuntimeError: JIT C-source hash differs across MPI ranks: {...} + +Two constants that share a display name are the way in. `Symbol.sort_key()` is +derived from the NAME, so two same-named `constants[]` placeholders sort equal +and term order inside an `Add` falls back to hash order. + +**This is deliberately not an MPI test.** Whether two ranks happen to disagree +depends on their seeds, so an mpirun-based check passes or fails at random — +measured on the bug: 4 of 6 launches diverged on one occasion and 0 of 6 on +another, with the defect present and confirmed both times. A seed sweep in +subprocesses is the same property, tested deterministically. With the bug: + + seed 0 -> a6cb2e83 seed 2 -> df7e63dd + seed 1 -> a6cb2e83 seed 3 -> 2f8aadf7 + +Three distinct modules from one model. Fixed, all seeds give one hash. + +The in-process halves — distinct `sort_key`, distinct identity, canonical +ordering — are pinned in `test_0103_jit_rampable_constants.py`, which is where +to look first when this fails. +""" + +import os +import pathlib +import subprocess +import sys +import textwrap + +import pytest + +pytestmark = [pytest.mark.level_2, pytest.mark.tier_b] + +# A two-material Stokes: both ViscousFlowModels name their viscosity \eta, so +# the residual carries two same-named constants. +_CHILD = textwrap.dedent( + """ + import pathlib, sys + import numpy as np, sympy, underworld3 as uw + + mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.25, qdegree=2, regular=True) + v = uw.discretisation.MeshVariable("vs", mesh, mesh.dim, degree=2) + p = uw.discretisation.MeshVariable("ps", mesh, 1, degree=1) + + swarm = uw.swarm.Swarm(mesh) + material = uw.swarm.IndexSwarmVariable("Ms", swarm, indices=2, proxy_degree=1) + swarm.populate(fill_param=2) + coords = np.asarray(swarm._particle_coordinates.data) + with uw.synchronised_array_update(): + material.data[:, 0] = (coords[:, 1] > 0.5).astype(int) + + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + lower = uw.constitutive_models.ViscousFlowModel(stokes.Unknowns, material_name="lo") + lower.Parameters.shear_viscosity_0 = 1.0 + upper = uw.constitutive_models.ViscousFlowModel(stokes.Unknowns, material_name="up") + upper.Parameters.shear_viscosity_0 = 1000.0 + stokes.constitutive_model = uw.MultiMaterialConstitutiveModel( + stokes.Unknowns, material, [lower, upper]) + stokes.add_dirichlet_bc((1.0, 0.0), "Top") + stokes.add_dirichlet_bc((0.0, 0.0), "Bottom") + stokes._build() + + # The compiled module is named for the hash of the canonical C source. + cache = pathlib.Path(sys.argv[1]) + names = sorted({q.name.split(".")[0] for q in cache.glob("*.so")}) + print("MODULES:" + ",".join(names)) + """ +) + + +def _module_hash(seed, cache_dir, tmp_path): + child = tmp_path / "child.py" + child.write_text(_CHILD) + env = dict(os.environ) + env["PYTHONHASHSEED"] = str(seed) + env["UW_JIT_CACHE_DIR"] = str(cache_dir) + result = subprocess.run( + [sys.executable, str(child), str(cache_dir)], + capture_output=True, text=True, env=env, timeout=900, + ) + line = [ln for ln in result.stdout.splitlines() if ln.startswith("MODULES:")] + assert line, ( + f"child failed under PYTHONHASHSEED={seed}\n" + f"stdout tail:\n{result.stdout[-2000:]}\nstderr tail:\n{result.stderr[-2000:]}" + ) + return line[0].removeprefix("MODULES:") + + +@pytest.mark.parametrize("seeds", [(0, 1, 2)]) +def test_the_generated_module_is_the_same_under_every_hash_seed(seeds, tmp_path): + hashes = {} + for seed in seeds: + cache = tmp_path / f"cache_{seed}" # a cold cache per seed + cache.mkdir() + hashes[seed] = _module_hash(seed, cache, tmp_path) + + distinct = set(hashes.values()) + assert len(distinct) == 1, ( + "the emitted C depends on the Python hash seed, so MPI ranks will " + f"disagree and getext will abort: {hashes}" + )