JIT: two constants of the same name compiled to one constants[] slot

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)
 

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
 

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

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}"
+    )