From e91f1bebb3306b697a511c3f5a3beedd514ccea7 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 8 Oct 2026 13:39:42 +1100 Subject: [PATCH 1/2] JIT: always compile kernels with -fno-math-errno (#834) gcc's default -fmath-errno makes each sqrt, pow and exp call a side effect, so a repeated call with the same argument is computed again. The kernels never read errno. With the flag, a viscoplastic Jacobian's callbacks cost about a third as much on gcc 14 (8.2 to 2.8 microseconds per quadrature point), and the assembled residual and Jacobian, the solutions and the iteration counts are bit-identical (measured on Linux for #823). Apple clang never sets errno, so this is the macOS behaviour already. The flag is prepended with -std=c99, so a UW3_JIT_CFLAGS override, which replaces the optimisation flags, keeps it. The flags are in the generated setup.py, which is part of the cache key: each cached module recompiles once. test_0025 pins the compile arguments, with and without an override. Underworld development team with AI support from Claude Code --- src/underworld3/utilities/_jitextension.py | 18 +++++--- tests/test_0025_jit_compile_flags.py | 50 ++++++++++++++++++++++ 2 files changed, 61 insertions(+), 7 deletions(-) create mode 100644 tests/test_0025_jit_compile_flags.py diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index 1fd8d778..d78ad33f 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -1575,15 +1575,19 @@ def _patch_coords(expr): # 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). + # lower optimisation level, e.g. UW3_JIT_CFLAGS="-O1 -g0". Two flags are + # always prepended: -std=c99 (the generated code relies on it) and + # -fno-math-errno. The kernels never read errno, and gcc's default + # -fmath-errno makes each sqrt, pow and exp a side effect it cannot merge + # with a repeat of the same call: a third of the cost of a viscoplastic + # Jacobian's callbacks with it, bit-identical results (#834). Apple clang + # never sets errno, so this is already the macOS behaviour. _default_jit_cflags = ["-O3", "-g0"] _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] - ) + extra_compile_args = [ + "-std=c99", "-fno-math-errno", + *(_jit_cflags_env.split() if _jit_cflags_env is not None else _default_jit_cflags), + ] if verbose: print( f"JIT compile flags: {extra_compile_args}" diff --git a/tests/test_0025_jit_compile_flags.py b/tests/test_0025_jit_compile_flags.py new file mode 100644 index 00000000..08597ad6 --- /dev/null +++ b/tests/test_0025_jit_compile_flags.py @@ -0,0 +1,50 @@ +"""The flags the JIT compiles its kernels with (#834). + +gcc defaults to ``-fmath-errno``: ``sqrt``, ``pow`` and ``exp`` may set ``errno``, so it +treats each call as a side effect and cannot merge repeated calls. 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. The flag must survive a +``UW3_JIT_CFLAGS`` override, which replaces only the optimisation flags. +""" +import ast +import re + +import pytest + +import underworld3 as uw + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + + +def _compile_args(monkeypatch): + """The ``extra_compile_args`` of the next JIT bundle's generated setup.py.""" + 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) + 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(monkeypatch): + monkeypatch.delenv("UW3_JIT_CFLAGS", raising=False) + assert _compile_args(monkeypatch) == ["-std=c99", "-fno-math-errno", "-O3", "-g0"] + + +def test_an_override_replaces_only_the_optimisation_flags(monkeypatch): + monkeypatch.setenv("UW3_JIT_CFLAGS", "-O1 -g0") + assert _compile_args(monkeypatch) == ["-std=c99", "-fno-math-errno", "-O1", "-g0"] From 9298328d9fe11bcebea37bbb7b284c640496a8ee Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 8 Oct 2026 13:58:33 +1100 Subject: [PATCH 2/2] Review fixes for #834: the flag is a default UW3_JIT_CFLAGS can drop - nvc rejects -fno-math-errno (and -g0). Always prepending the flag left a JIT on such a compiler with no way to compile; it is now one of the default flags, which UW3_JIT_CFLAGS replaces. -std=c99 stays fixed. - The comment names clang on Linux as well as gcc and states the bit-identity for the C SymPy prints. - jit-cache.md documents UW3_JIT_CFLAGS. - The bundle capture shared with test_0021 is a conftest fixture, jit_bundles. - test_0025 starts at tier C, as TESTING-RELIABILITY-SYSTEM.md asks of a new test. Underworld development team with AI support from Claude Code --- docs/developer/subsystems/jit-cache.md | 1 + src/underworld3/utilities/_jitextension.py | 35 ++++++++++---------- tests/conftest.py | 18 ++++++++++ tests/test_0021_jit_finds_petsc_headers.py | 14 ++------ tests/test_0025_jit_compile_flags.py | 38 ++++++++-------------- 5 files changed, 53 insertions(+), 53 deletions(-) diff --git a/docs/developer/subsystems/jit-cache.md b/docs/developer/subsystems/jit-cache.md index 795b9c43..d2fca725 100644 --- a/docs/developer/subsystems/jit-cache.md +++ b/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 diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index d78ad33f..3daef970 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -1570,24 +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". Two flags are - # always prepended: -std=c99 (the generated code relies on it) and - # -fno-math-errno. The kernels never read errno, and gcc's default - # -fmath-errno makes each sqrt, pow and exp a side effect it cannot merge - # with a repeat of the same call: a third of the cost of a viscoplastic - # Jacobian's callbacks with it, bit-identical results (#834). Apple clang - # never sets errno, so this is already the macOS behaviour. - _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", "-fno-math-errno", - *(_jit_cflags_env.split() if _jit_cflags_env is not None else _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}" diff --git a/tests/conftest.py b/tests/conftest.py index d104c7f1..faed91b1 100644 --- a/tests/conftest.py +++ b/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. diff --git a/tests/test_0021_jit_finds_petsc_headers.py b/tests/test_0021_jit_finds_petsc_headers.py index 0641c672..f8b53209 100644 --- a/tests/test_0021_jit_finds_petsc_headers.py +++ b/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) diff --git a/tests/test_0025_jit_compile_flags.py b/tests/test_0025_jit_compile_flags.py index 08597ad6..5fe30499 100644 --- a/tests/test_0025_jit_compile_flags.py +++ b/tests/test_0025_jit_compile_flags.py @@ -1,10 +1,11 @@ """The flags the JIT compiles its kernels with (#834). -gcc defaults to ``-fmath-errno``: ``sqrt``, ``pow`` and ``exp`` may set ``errno``, so it -treats each call as a side effect and cannot merge repeated calls. 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. The flag must survive a -``UW3_JIT_CFLAGS`` override, which replaces only the optimisation flags. +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 @@ -13,22 +14,11 @@ import underworld3 as uw -pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] +pytestmark = [pytest.mark.level_1, pytest.mark.tier_c] -def _compile_args(monkeypatch): - """The ``extra_compile_args`` of the next JIT bundle's generated setup.py.""" - 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) +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 @@ -40,11 +30,11 @@ def capture(*args, **kwargs): return ast.literal_eval(match.group(1)) -def test_the_default_flags(monkeypatch): +def test_the_default_flags(jit_bundles, monkeypatch): monkeypatch.delenv("UW3_JIT_CFLAGS", raising=False) - assert _compile_args(monkeypatch) == ["-std=c99", "-fno-math-errno", "-O3", "-g0"] + assert _compile_args(jit_bundles) == ["-std=c99", "-O3", "-g0", "-fno-math-errno"] -def test_an_override_replaces_only_the_optimisation_flags(monkeypatch): - monkeypatch.setenv("UW3_JIT_CFLAGS", "-O1 -g0") - assert _compile_args(monkeypatch) == ["-std=c99", "-fno-math-errno", "-O1", "-g0"] +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"]