JIT: compile kernels with -fno-math-errno by default (#834)

docs/developer/subsystems/jit-cache.md

@@ -125,6 +125,7 @@ performs the cold compile but only rank 0 publishes the result.
 | `UW_JIT_CACHE`       | Set to `0`/`false`/`no` to disable disk cache           |
 | `UW_JIT_CACHE_DIR`   | Override the cache directory location                   |
 | `XDG_CACHE_HOME`     | Used when `UW_JIT_CACHE_DIR` is unset                   |
+| `UW3_JIT_CFLAGS`     | Replaces the kernels' default compile flags, `-O3 -g0 -fno-math-errno`; `-std=c99` is always kept. Use a lower level for a huge expression whose `-O3` compile is slow or runs out of memory (`-O1 -g0 -fno-math-errno`), or drop a flag the compiler rejects (nvc rejects `-g0` and `-fno-math-errno`). Without `-fno-math-errno`, gcc and clang on Linux cannot merge repeated `sqrt`, `pow` and `exp` calls (#834). The flags are part of the cache key. |
 
 ## Code references
 

src/underworld3/utilities/_jitextension.py

@@ -1570,20 +1570,25 @@ def _patch_coords(expr):
 
     MODNAME = "fn_ptr_ext_" + str(name)
 
-    # JIT compile flags for the generated kernels. Default keeps -O3 (kernel
-    # runtime speed) but adds -g0 to drop the debug info that the base Python
-    # CFLAGS injects via sysconfig -- pure overhead, and a memory hog on huge
-    # expressions, for these generated kernels. For very large expressions
-    # whose gcc -O3 compile is slow or OOM-killed, set UW3_JIT_CFLAGS to a
-    # lower optimisation level, e.g. UW3_JIT_CFLAGS="-O1 -g0". -std=c99 is
-    # always prepended (the generated code relies on it).
-    _default_jit_cflags = ["-O3", "-g0"]
+    # JIT compile flags for the generated kernels. -std=c99 is always first (the
+    # generated code relies on it); UW3_JIT_CFLAGS replaces the rest.
+    #   -O3   kernel run-time speed.
+    #   -g0   drops the debug info the base Python CFLAGS inject via sysconfig:
+    #         overhead, and a memory hog on huge expressions.
+    #   -fno-math-errno  the kernels never read errno. gcc and clang on Linux
+    #         default to -fmath-errno, which makes each sqrt, pow and exp call a
+    #         side effect that cannot be merged with a repeat of the same call; the
+    #         flag cut a viscoplastic Jacobian's callbacks to a third on gcc 14, with
+    #         bit-identical results for the C SymPy prints (#834). clang targeting
+    #         macOS already defaults to it.
+    # For a very large expression whose -O3 compile is slow or OOM-killed, use a
+    # lower level, e.g. UW3_JIT_CFLAGS="-O1 -g0 -fno-math-errno". The override
+    # replaces all three, so a compiler that rejects one (nvc rejects -g0 and
+    # -fno-math-errno) can still be used.
+    default_flags = ["-O3", "-g0", "-fno-math-errno"]
     _jit_cflags_env = os.environ.get("UW3_JIT_CFLAGS")
-    extra_compile_args = (
-        ["-std=c99", *_jit_cflags_env.split()]
-        if _jit_cflags_env is not None
-        else ["-std=c99", *_default_jit_cflags]
-    )
+    chosen_flags = _jit_cflags_env.split() if _jit_cflags_env is not None else default_flags
+    extra_compile_args = ["-std=c99", *chosen_flags]
     if verbose:
         print(
             f"JIT compile flags: {extra_compile_args}"

tests/conftest.py

@@ -191,6 +191,24 @@ def pytest_make_collect_report(collector):
     report.result = []
 
 
+@pytest.fixture
+def jit_bundles(monkeypatch):
+    """The source bundles the JIT generates while a test runs, one dict per bundle
+    (``setup.py``, ``cy_ext.h``, ``cy_ext.pyx``), captured from ``generate_c_source``."""
+    import underworld3.utilities._jitextension as jx
+
+    bundles = []
+    generate = jx.generate_c_source
+
+    def capture(*args, **kwargs):
+        result = generate(*args, **kwargs)
+        bundles.append(dict(result[1]))
+        return result
+
+    monkeypatch.setattr(jx, "generate_c_source", capture)
+    return bundles
+
+
 @pytest.fixture(scope="module", autouse=True)
 def isolate_module_state():
     """Reset the global model BEFORE a module's own fixtures are built.

tests/test_0021_jit_finds_petsc_headers.py

@@ -33,20 +33,10 @@ def test_petscs_include_dirs_hold_the_headers_the_callbacks_include():
     assert any(os.path.exists(os.path.join(d, "petscconf.h")) for d in dirs), dirs
 
 
-def test_the_generated_setup_py_passes_petscs_include_dirs(monkeypatch):
+def test_the_generated_setup_py_passes_petscs_include_dirs(jit_bundles):
     """Fails wherever the include directories are left out, including CI, whose
     conda-forge PETSc would compile without them."""
-    import underworld3.utilities._jitextension as jx
-
-    bundles = []
-    generate = jx.generate_c_source
-
-    def capture(*args, **kwargs):
-        result = generate(*args, **kwargs)
-        bundles.append(dict(result[1]))
-        return result
-
-    monkeypatch.setattr(jx, "generate_c_source", capture)
+    bundles = jit_bundles
     mesh = uw.meshing.UnstructuredSimplexBox(
         minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=0.5)
     u = uw.discretisation.MeshVariable("U0021s", mesh, 1, degree=1)

tests/test_0025_jit_compile_flags.py

@@ -0,0 +1,40 @@
+"""The flags the JIT compiles its kernels with (#834).
+
+gcc, and clang on Linux, default to ``-fmath-errno``: ``sqrt``, ``pow`` and ``exp`` may
+set ``errno``, so each call is a side effect that cannot be merged with a repeat of the
+same call. The kernels never read ``errno``; with ``-fno-math-errno`` the viscoplastic
+box's Jacobian callbacks cost a third as much on gcc 14, and every assembled value is
+bit-identical. ``UW3_JIT_CFLAGS`` replaces the default flags, this one included, so a
+compiler that rejects it can still be used.
+"""
+import ast
+import re
+
+import pytest
+
+import underworld3 as uw
+
+pytestmark = [pytest.mark.level_1, pytest.mark.tier_c]
+
+
+def _compile_args(bundles):
+    """The ``extra_compile_args`` of the bundle a small integral generates."""
+    mesh = uw.meshing.UnstructuredSimplexBox(
+        minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=0.5)
+    # PETSc integrates only on a mesh that carries a field
+    uw.discretisation.MeshVariable("U0025", mesh, 1, degree=1)
+    uw.maths.Integral(mesh, 1.0 + mesh.X[0]).evaluate()
+    assert bundles, "no JIT bundle was generated"
+    match = re.search(r"extra_compile_args=(\[.*?\])", bundles[-1]["setup.py"])
+    assert match, bundles[-1]["setup.py"]
+    return ast.literal_eval(match.group(1))
+
+
+def test_the_default_flags(jit_bundles, monkeypatch):
+    monkeypatch.delenv("UW3_JIT_CFLAGS", raising=False)
+    assert _compile_args(jit_bundles) == ["-std=c99", "-O3", "-g0", "-fno-math-errno"]
+
+
+def test_an_override_replaces_every_flag_but_the_language_standard(jit_bundles, monkeypatch):
+    monkeypatch.setenv("UW3_JIT_CFLAGS", "-O1")
+    assert _compile_args(jit_bundles) == ["-std=c99", "-O1"]