diff --git a/src/underworld3/constitutive_models.py b/src/underworld3/constitutive_models.py index 1c807d4a..02a1a565 100644 --- a/src/underworld3/constitutive_models.py +++ b/src/underworld3/constitutive_models.py @@ -4518,6 +4518,10 @@ def __init__( Set to True if IndexSwarmVariable does not maintain partition of unity. Default: False (assumes IndexSwarmVariable maintains partition of unity) """ + # 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 5301b43f..2d1ebbe1 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -329,21 +329,55 @@ 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. + 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}" - obj = super().__new__(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 + 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,14 +428,26 @@ 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. + # 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 = {} 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. - 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 180df068..5691da51 100644 --- a/tests/test_0103_jit_rampable_constants.py +++ b/tests/test_0103_jit_rampable_constants.py @@ -89,3 +89,179 @@ 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_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. + + 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 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 00000000..0e546316 --- /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}" + )