compiler: Write the complex arithmetic header atomically

devito/passes/iet/dtypes.py

@@ -1,3 +1,6 @@
+import os
+from hashlib import sha1
+
 import numpy as np
 
 from devito.arch.compiler import Compiler
@@ -28,10 +31,17 @@ def _complex_includes(iet: Callable, langbb: type[LangBB], compiler: Compiler,
 
     # Some languages such as c++11 need some extra arithmetic definitions
     if langbb.get('def-complex'):
+        code = str(langbb['def-complex'])
         dest = compiler.get_jit_dir()
-        hfile = dest.joinpath('complex_arith.h')
-        with open(str(hfile), 'w') as ff:
-            ff.write(str(langbb['def-complex']))
+        hfile = dest.joinpath(f'complex_arith_{sha1(code.encode()).hexdigest()}.h')
+        # The JIT dir is shared by all processes compiling as the same user, e.g.
+        # the workers of a parallel test run. The header is keyed on its content
+        # and renamed into place once fully written, so that a concurrent compiler
+        # never reads a truncated header or one meant for another language
+        if not hfile.exists():
+            tmpfile = hfile.with_name(f'{hfile.name}.{os.getpid()}.tmp')
+            tmpfile.write_text(code)
+            os.replace(tmpfile, hfile)
         lib += (str(hfile),)
 
     metadata['includes'] = lib

tests/test_dtypes.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 import numpy as np
 import pytest
 import sympy
@@ -147,6 +149,30 @@ def test_complex_headers(dtype: np.dtype[np.inexact], kwargs: dict[str, str]) ->
         assert header not in op._includes
 
 
+@pytest.mark.parametrize('kwargs', _configs, ids=kw_id)
+def test_complex_arith_header(kwargs: dict[str, str]) -> None:
+    """
+    Tests that the extra complex arithmetic definitions some languages need are
+    written in full to a single content-keyed header, shared by all Operators.
+    """
+    code = _get_language(**kwargs).get('def-complex')
+    if code is None:
+        pytest.skip(f"`{kwargs['language']}` needs no complex arithmetic header")
+
+    # Set up two operators
+    grid = Grid(shape=(3, 3))
+    x, y = grid.dimensions
+
+    u = Function(name='u', grid=grid, dtype=np.complex64)
+    ops = [Operator(Eq(u, x*y + i*1j), **kwargs) for i in (1, 2)]
+
+    # Both include the same header, and its content is intact
+    headers = {i for op in ops for i in op._includes if 'complex_arith' in i}
+    assert len(headers) == 1
+    hfile, = headers
+    assert Path(hfile).read_text() == str(code)
+
+
 @pytest.mark.parametrize('dtype', [np.complex64, np.complex128])
 @pytest.mark.parametrize('kwargs', _configs, ids=kw_id)
 def test_imag_unit(dtype: np.complexfloating, kwargs: dict[str, str]) -> None: