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 1fd8d778..3daef970 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/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}" 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 new file mode 100644 index 00000000..5fe30499 --- /dev/null +++ b/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"]