17789800893

Committed 17 Sep 2025 07:14AM UTC coverage: 82.981% (+0.001%) from 82.98%

Build # 17789800893

Build Type

Pull #783

github

Committed by

web-flow

Commit Message

Last revert

Pull Request Pull Request #783: Diagonal assembly of matrices

Coverage Stats

58 of 68 new or added lines in 4 files covered. (85.29%)

3725 of 4489 relevant lines covered (82.98%)

0.83 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

79.46

/ffcx/codegeneration/jit.py

# Copyright (C) 2004-2019 Garth N. Wells
#
# This file is part of FFCx.(https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later
"""Just-in-time compilation."""

from __future__ import annotations

import importlib
import io
import logging
import os
import re
import sys
import sysconfig
import tempfile
import time
from contextlib import redirect_stdout
from pathlib import Path

import cffi
import numpy as np
import numpy.typing as npt
import ufl

import ffcx
import ffcx.naming
from ffcx.codegeneration.C.file_template import libraries as _libraries

logger = logging.getLogger("ffcx")
root_logger = logging.getLogger()

# Get declarations directly from ufcx.h
file_dir = os.path.dirname(os.path.abspath(__file__))
with open(file_dir + "/ufcx.h") as f:
    ufcx_h = "".join(f.readlines())

# Emulate C preprocessor on __STDC_NO_COMPLEX__
if sys.platform.startswith("win32"):
    # Remove macro statements and content
    ufcx_h = re.sub(
        r"\#ifndef __STDC_NO_COMPLEX__.*?\#endif // __STDC_NO_COMPLEX__",
        "",
        ufcx_h,
        flags=re.DOTALL,
    )
else:
    # Remove only macros keeping content
    ufcx_h = ufcx_h.replace("#ifndef __STDC_NO_COMPLEX__", "")
    ufcx_h = ufcx_h.replace("#endif // __STDC_NO_COMPLEX__", "")

header = ufcx_h.split("<HEADER_DECL>")[1].split("</HEADER_DECL>")[0].strip(" /\n")
header = header.replace("{", "{{").replace("}", "}}")
UFC_HEADER_DECL = header + "\n"

UFC_FORM_DECL = "\n".join(re.findall("typedef struct ufcx_form.*?ufcx_form;", ufcx_h, re.DOTALL))

UFC_INTEGRAL_DECL = "\n".join(
    re.findall(r"typedef void ?\(ufcx_tabulate_tensor_float32\).*?\);", ufcx_h, re.DOTALL)
)
UFC_INTEGRAL_DECL += "\n".join(
    re.findall(r"typedef void ?\(ufcx_tabulate_tensor_float64\).*?\);", ufcx_h, re.DOTALL)
)
UFC_INTEGRAL_DECL += "\n".join(
    re.findall(r"typedef void ?\(ufcx_tabulate_tensor_complex64\).*?\);", ufcx_h, re.DOTALL)
)
UFC_INTEGRAL_DECL += "\n".join(
    re.findall(r"typedef void ?\(ufcx_tabulate_tensor_complex128\).*?\);", ufcx_h, re.DOTALL)
)

UFC_INTEGRAL_DECL += "\n".join(
    re.findall("typedef struct ufcx_integral.*?ufcx_integral;", ufcx_h, re.DOTALL)
)

UFC_EXPRESSION_DECL = "\n".join(
    re.findall("typedef struct ufcx_expression.*?ufcx_expression;", ufcx_h, re.DOTALL)
)


def _compute_option_signature(options):
    """Return options signature (some options should not affect signature)."""
    return str(sorted(options.items()))


def get_cached_module(module_name, object_names, cache_dir, timeout):
    """Look for an existing C file and wait for compilation, or if it does not exist, create it."""
    cache_dir = Path(cache_dir)
    c_filename = cache_dir.joinpath(module_name).with_suffix(".c")
    ready_name = c_filename.with_suffix(".c.cached")

    # Ensure cache dir exists
    cache_dir.mkdir(exist_ok=True, parents=True)

    try:
        # Create C file with exclusive access
        with open(c_filename, "x"):
            pass
        return None, None
    except FileExistsError:
        logger.info("Cached C file already exists: " + str(c_filename))
        finder = importlib.machinery.FileFinder(
            str(cache_dir),
            (importlib.machinery.ExtensionFileLoader, importlib.machinery.EXTENSION_SUFFIXES),
        )
        finder.invalidate_caches()

        # Now, wait for ready
        for i in range(timeout):
            if os.path.exists(ready_name):
                spec = finder.find_spec(module_name)
                if spec is None:
                    raise ModuleNotFoundError("Unable to find JIT module.")
                compiled_module = importlib.util.module_from_spec(spec)
                spec.loader.exec_module(compiled_module)

                compiled_objects = [getattr(compiled_module.lib, name) for name in object_names]
                return compiled_objects, compiled_module

            logger.info(f"Waiting for {ready_name} to appear.")
            time.sleep(1)
        raise TimeoutError(
            "JIT compilation timed out, probably due to a failed previous compile. "
            f"Try cleaning cache (e.g. remove {c_filename}) or increase timeout option."
        )


def _compilation_signature(cffi_extra_compile_args, cffi_debug):
    """Compute the compilation-inputs part of the signature.

    Used to avoid cache conflicts across Python versions, architectures, installs.

    - SOABI includes platform, Python version, debug flags
    - CFLAGS includes prefixes, arch targets
    """
    if sys.platform.startswith("win32"):
        # NOTE: SOABI not defined on win32, EXT_SUFFIX contains e.g. '.cp312-win_amd64.pyd'
        return (
            str(cffi_extra_compile_args)
            + str(cffi_debug)
            + str(sysconfig.get_config_var("EXT_SUFFIX"))
        )
    else:
        return (
            str(cffi_extra_compile_args)
            + str(cffi_debug)
            + str(sysconfig.get_config_var("CFLAGS"))
            + str(sysconfig.get_config_var("SOABI"))
        )


def compile_forms(
    forms: list[ufl.Form],
    options: dict = {},
    cache_dir: Path | None = None,
    timeout: int = 10,
    cffi_extra_compile_args: list[str] = [],
    cffi_verbose: bool = False,
    cffi_debug: bool = False,
    cffi_libraries: list[str] = [],
    visualise: bool = False,
):
    """Compile a list of UFL forms into UFC Python objects.

    Args:
        forms: List of ufl.form to compile.
        options: Options
        cache_dir: Cache directory
        timeout: Timeout
        cffi_extra_compile_args: Extra compilation args for CFFI
        cffi_verbose: Use verbose compile
        cffi_debug: Use compiler debug mode
        cffi_libraries: libraries to use with compiler
        visualise: Toggle visualisation
    """
    p = ffcx.options.get_options(options)

    # If requested, replace bi-linear forms by their diagonal part
    if p["part"] == "diagonal":
        for i, form in enumerate(forms):
            arguments = form.arguments()
            numbers = tuple(sorted(set(a.number() for a in arguments)))
            arity = len(numbers)
            if arity == 2:
                blocked_form = ufl.extract_blocks(form, replace_argument=False)
                if isinstance(blocked_form, ufl.form.Form):
                    # If there are no sub-elements, continue
                    continue
                diagonal_form = ufl.ZeroBaseForm(())
                for j in range(len(blocked_form)):
                    if blocked_form[j][j] is not None:
                        diagonal_form += blocked_form[j][j]
                if diagonal_form == 0:
                    raise RuntimeError("Diagonal form seems to be zero.")
                forms[i] = diagonal_form  # type: ignore

    # Get a signature for these forms
    module_name = "libffcx_forms_" + ffcx.naming.compute_signature(
        forms,
        _compute_option_signature(p) + _compilation_signature(cffi_extra_compile_args, cffi_debug),
    )

    form_names = [ffcx.naming.form_name(form, i, module_name) for i, form in enumerate(forms)]

    if cache_dir is not None:
        cache_dir = Path(cache_dir)
        obj, mod = get_cached_module(module_name, form_names, cache_dir, timeout)
        if obj is not None:
            return obj, mod, (None, None)
    else:
        cache_dir = Path(tempfile.mkdtemp())

    try:
        decl = (
            UFC_HEADER_DECL.format(np.dtype(p["scalar_type"]).name)  # type: ignore
            + UFC_INTEGRAL_DECL
            + UFC_FORM_DECL
        )

        form_template = "extern ufcx_form {name};\n"
        for name in form_names:
            decl += form_template.format(name=name)

        impl = _compile_objects(
            decl,
            forms,
            form_names,
            module_name,
            p,
            cache_dir,
            cffi_extra_compile_args,
            cffi_verbose,
            cffi_debug,
            cffi_libraries,
            visualise=visualise,
        )
    except Exception as e:
        try:
            # remove c file so that it will not timeout next time
            c_filename = cache_dir.joinpath(module_name + ".c")
            os.replace(c_filename, c_filename.with_suffix(".c.failed"))
        except Exception:
            pass
        raise e

    obj, module = _load_objects(cache_dir, module_name, form_names)
    return obj, module, (decl, impl)


def compile_expressions(
    expressions: list[tuple[ufl.Expr, npt.NDArray[np.floating]]],  # type: ignore
    options: dict = {},
    cache_dir: Path | None = None,
    timeout: int = 10,
    cffi_extra_compile_args: list[str] = [],
    cffi_verbose: bool = False,
    cffi_debug: bool = False,
    cffi_libraries: list[str] = [],
    visualise: bool = False,
):
    """Compile a list of UFL expressions into UFC Python objects.

    Args:
        expressions: List of (UFL expression, evaluation points).
        options: Options
        cache_dir: Cache directory
        timeout: Timeout
        cffi_extra_compile_args: Extra compilation args for CFFI
        cffi_verbose: Use verbose compile
        cffi_debug: Use compiler debug mode
        cffi_libraries: libraries to use with compiler
        visualise: Toggle visualisation
    """
    p = ffcx.options.get_options(options)

    module_name = "libffcx_expressions_" + ffcx.naming.compute_signature(
        expressions,
        _compute_option_signature(p) + _compilation_signature(cffi_extra_compile_args, cffi_debug),
    )
    expr_names = [
        ffcx.naming.expression_name(expression, module_name) for expression in expressions
    ]

    if cache_dir is not None:
        cache_dir = Path(cache_dir)
        obj, mod = get_cached_module(module_name, expr_names, cache_dir, timeout)
        if obj is not None:
            return obj, mod, (None, None)
    else:
        cache_dir = Path(tempfile.mkdtemp())

    try:
        decl = (
            UFC_HEADER_DECL.format(np.dtype(p["scalar_type"]).name)  # type: ignore
            + UFC_INTEGRAL_DECL
            + UFC_FORM_DECL
            + UFC_EXPRESSION_DECL
        )

        expression_template = "extern ufcx_expression {name};\n"
        for name in expr_names:
            decl += expression_template.format(name=name)

        impl = _compile_objects(
            decl,
            expressions,
            expr_names,
            module_name,
            p,
            cache_dir,
            cffi_extra_compile_args,
            cffi_verbose,
            cffi_debug,
            cffi_libraries,
            visualise=visualise,
        )
    except Exception as e:
        try:
            # remove c file so that it will not timeout next time
            c_filename = cache_dir.joinpath(module_name + ".c")
            os.replace(c_filename, c_filename.with_suffix(".c.failed"))
        except Exception:
            pass
        raise e

    obj, module = _load_objects(cache_dir, module_name, expr_names)
    return obj, module, (decl, impl)


def _compile_objects(
    decl,
    ufl_objects,
    object_names,
    module_name,
    options,
    cache_dir,
    cffi_extra_compile_args,
    cffi_verbose,
    cffi_debug,
    cffi_libraries,
    visualise: bool = False,
):
    import ffcx.compiler

    libraries = _libraries + cffi_libraries if cffi_libraries is not None else _libraries

    # JIT uses module_name as prefix, which is needed to make names of all struct/function
    # unique across modules
    _, code_body = ffcx.compiler.compile_ufl_objects(
        ufl_objects, prefix=module_name, options=options, visualise=visualise
    )

    # Raise error immediately prior to compilation if no support for C99
    # _Complex. Doing this here allows FFCx to be used for complex codegen on
    # Windows.
    if sys.platform.startswith("win32"):
        if np.issubdtype(options["scalar_type"], np.complexfloating):
            raise NotImplementedError("win32 platform does not support C99 _Complex numbers")
        elif isinstance(options["scalar_type"], str) and "complex" in options["scalar_type"]:
            raise NotImplementedError("win32 platform does not support C99 _Complex numbers")

    # Compile in C17 mode
    if sys.platform.startswith("win32"):
        cffi_base_compile_args = ["-std:c17"]
    else:
        cffi_base_compile_args = ["-std=c17"]

    cffi_final_compile_args = cffi_base_compile_args + cffi_extra_compile_args

    ffibuilder = cffi.FFI()

    ffibuilder.set_source(
        module_name,
        code_body,
        include_dirs=[ffcx.codegeneration.get_include_path()],
        extra_compile_args=cffi_final_compile_args,
        libraries=libraries,
    )

    ffibuilder.cdef(decl)

    c_filename = cache_dir.joinpath(module_name + ".c")
    ready_name = c_filename.with_suffix(".c.cached")

    # Compile (ensuring that compile dir exists)
    cache_dir.mkdir(exist_ok=True, parents=True)

    logger.info(79 * "#")
    logger.info("Calling JIT C compiler")
    logger.info(79 * "#")

    t0 = time.time()
    f = io.StringIO()
    # Temporarily set root logger handlers to string buffer only
    # since CFFI logs into root logger
    old_handlers = root_logger.handlers.copy()
    root_logger.handlers = [logging.StreamHandler(f)]
    with redirect_stdout(f):
        ffibuilder.compile(tmpdir=cache_dir, verbose=True, debug=cffi_debug)
    s = f.getvalue()
    if cffi_verbose:
        print(s)

    logger.info(f"JIT C compiler finished in {time.time() - t0:.4f}")

    # Create a "status ready" file. If this fails, it is an error,
    # because it should not exist yet.
    # Copy the stdout verbose output of the build into the ready file
    fd = open(ready_name, "x")
    fd.write(s)
    fd.close()

    # Copy back the original handlers (in case someone is logging into
    # root logger and has custom handlers)
    root_logger.handlers = old_handlers

    return code_body


def _load_objects(cache_dir, module_name, object_names):
    # Create module finder that searches the compile path
    finder = importlib.machinery.FileFinder(
        str(cache_dir),
        (importlib.machinery.ExtensionFileLoader, importlib.machinery.EXTENSION_SUFFIXES),
    )

    # Find module. Clear search cache to be sure dynamically created
    # (new) modules are found
    finder.invalidate_caches()
    spec = finder.find_spec(module_name)
    if spec is None:
        raise ModuleNotFoundError("Unable to find JIT module.")

    # Load module
    compiled_module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(compiled_module)

    compiled_objects = []
    for name in object_names:
        obj = getattr(compiled_module.lib, name)
        compiled_objects.append(obj)

    return compiled_objects, compiled_module

1	# Copyright (C) 2004-2019 Garth N. Wells
2	#
3	# This file is part of FFCx.(https://www.fenicsproject.org)
4	#
5	# SPDX-License-Identifier: LGPL-3.0-or-later
6	"""Just-in-time compilation."""
7
8	from __future__ import annotations	1✔
9
10	import importlib	1✔
11	import io	1✔
12	import logging	1✔
13	import os	1✔
14	import re	1✔
15	import sys	1✔
16	import sysconfig	1✔
17	import tempfile	1✔
18	import time	1✔
19	from contextlib import redirect_stdout	1✔
20	from pathlib import Path	1✔
21
22	import cffi	1✔
23	import numpy as np	1✔
24	import numpy.typing as npt	1✔
25	import ufl	1✔
26
27	import ffcx	1✔
28	import ffcx.naming	1✔
29	from ffcx.codegeneration.C.file_template import libraries as _libraries	1✔
30
31	logger = logging.getLogger("ffcx")	1✔
32	root_logger = logging.getLogger()	1✔
33
34	# Get declarations directly from ufcx.h
35	file_dir = os.path.dirname(os.path.abspath(__file__))	1✔
36	with open(file_dir + "/ufcx.h") as f:	1✔
37	ufcx_h = "".join(f.readlines())	1✔
38
39	# Emulate C preprocessor on __STDC_NO_COMPLEX__
40	if sys.platform.startswith("win32"):	1✔
41	# Remove macro statements and content
42	ufcx_h = re.sub(	×
43	r"\#ifndef __STDC_NO_COMPLEX__.*?\#endif // __STDC_NO_COMPLEX__",
44	"",
45	ufcx_h,
46	flags=re.DOTALL,
47	)
48	else:
49	# Remove only macros keeping content
50	ufcx_h = ufcx_h.replace("#ifndef __STDC_NO_COMPLEX__", "")	1✔
51	ufcx_h = ufcx_h.replace("#endif // __STDC_NO_COMPLEX__", "")	1✔
52
53	header = ufcx_h.split("<HEADER_DECL>")[1].split("</HEADER_DECL>")[0].strip(" /\n")	1✔
54	header = header.replace("{", "{{").replace("}", "}}")	1✔
55	UFC_HEADER_DECL = header + "\n"	1✔
56
57	UFC_FORM_DECL = "\n".join(re.findall("typedef struct ufcx_form.*?ufcx_form;", ufcx_h, re.DOTALL))	1✔
58
59	UFC_INTEGRAL_DECL = "\n".join(	1✔
60	re.findall(r"typedef void ?\(ufcx_tabulate_tensor_float32\).*?\);", ufcx_h, re.DOTALL)
61	)
62	UFC_INTEGRAL_DECL += "\n".join(	1✔
63	re.findall(r"typedef void ?\(ufcx_tabulate_tensor_float64\).*?\);", ufcx_h, re.DOTALL)
64	)
65	UFC_INTEGRAL_DECL += "\n".join(	1✔
66	re.findall(r"typedef void ?\(ufcx_tabulate_tensor_complex64\).*?\);", ufcx_h, re.DOTALL)
67	)
68	UFC_INTEGRAL_DECL += "\n".join(	1✔
69	re.findall(r"typedef void ?\(ufcx_tabulate_tensor_complex128\).*?\);", ufcx_h, re.DOTALL)
70	)
71
72	UFC_INTEGRAL_DECL += "\n".join(	1✔
73	re.findall("typedef struct ufcx_integral.*?ufcx_integral;", ufcx_h, re.DOTALL)
74	)
75
76	UFC_EXPRESSION_DECL = "\n".join(	1✔
77	re.findall("typedef struct ufcx_expression.*?ufcx_expression;", ufcx_h, re.DOTALL)
78	)
79
80
81	def _compute_option_signature(options):	1✔
82	"""Return options signature (some options should not affect signature)."""
83	return str(sorted(options.items()))	1✔
84
85
86	def get_cached_module(module_name, object_names, cache_dir, timeout):	1✔
87	"""Look for an existing C file and wait for compilation, or if it does not exist, create it."""
88	cache_dir = Path(cache_dir)	1✔
89	c_filename = cache_dir.joinpath(module_name).with_suffix(".c")	1✔
90	ready_name = c_filename.with_suffix(".c.cached")	1✔
91
92	# Ensure cache dir exists
93	cache_dir.mkdir(exist_ok=True, parents=True)	1✔
94
95	try:	1✔
96	# Create C file with exclusive access
97	with open(c_filename, "x"):	1✔
98	pass	1✔
99	return None, None	1✔
100	except FileExistsError:	×
101	logger.info("Cached C file already exists: " + str(c_filename))	×
102	finder = importlib.machinery.FileFinder(	×
103	str(cache_dir),
104	(importlib.machinery.ExtensionFileLoader, importlib.machinery.EXTENSION_SUFFIXES),
105	)
106	finder.invalidate_caches()	×
107
108	# Now, wait for ready
109	for i in range(timeout):	×
110	if os.path.exists(ready_name):	×
111	spec = finder.find_spec(module_name)	×
112	if spec is None:	×
113	raise ModuleNotFoundError("Unable to find JIT module.")	×
114	compiled_module = importlib.util.module_from_spec(spec)	×
115	spec.loader.exec_module(compiled_module)	×
116
117	compiled_objects = [getattr(compiled_module.lib, name) for name in object_names]	×
118	return compiled_objects, compiled_module	×
119
120	logger.info(f"Waiting for {ready_name} to appear.")	×
121	time.sleep(1)	×
122	raise TimeoutError(	×
123	"JIT compilation timed out, probably due to a failed previous compile. "
124	f"Try cleaning cache (e.g. remove {c_filename}) or increase timeout option."
125	)
126
127
128	def _compilation_signature(cffi_extra_compile_args, cffi_debug):	1✔
129	"""Compute the compilation-inputs part of the signature.
130
131	Used to avoid cache conflicts across Python versions, architectures, installs.
132
133	- SOABI includes platform, Python version, debug flags
134	- CFLAGS includes prefixes, arch targets
135	"""
136	if sys.platform.startswith("win32"):	1✔
137	# NOTE: SOABI not defined on win32, EXT_SUFFIX contains e.g. '.cp312-win_amd64.pyd'
138	return (	×
139	str(cffi_extra_compile_args)
140	+ str(cffi_debug)
141	+ str(sysconfig.get_config_var("EXT_SUFFIX"))
142	)
143	else:
144	return (	1✔
145	str(cffi_extra_compile_args)
146	+ str(cffi_debug)
147	+ str(sysconfig.get_config_var("CFLAGS"))
148	+ str(sysconfig.get_config_var("SOABI"))
149	)
150
151
152	def compile_forms(	1✔
153	forms: list[ufl.Form],
154	options: dict = {},
155	cache_dir: Path \| None = None,
156	timeout: int = 10,
157	cffi_extra_compile_args: list[str] = [],
158	cffi_verbose: bool = False,
159	cffi_debug: bool = False,
160	cffi_libraries: list[str] = [],
161	visualise: bool = False,
162	):
163	"""Compile a list of UFL forms into UFC Python objects.
164
165	Args:
166	forms: List of ufl.form to compile.
167	options: Options
168	cache_dir: Cache directory
169	timeout: Timeout
170	cffi_extra_compile_args: Extra compilation args for CFFI
171	cffi_verbose: Use verbose compile
172	cffi_debug: Use compiler debug mode
173	cffi_libraries: libraries to use with compiler
174	visualise: Toggle visualisation
175	"""
176	p = ffcx.options.get_options(options)	1✔
177
178	# If requested, replace bi-linear forms by their diagonal part
179	if p["part"] == "diagonal":	1✔
180	for i, form in enumerate(forms):	1✔
181	arguments = form.arguments()	1✔
182	numbers = tuple(sorted(set(a.number() for a in arguments)))	1✔
183	arity = len(numbers)	1✔
184	if arity == 2:	1✔
185	blocked_form = ufl.extract_blocks(form, replace_argument=False)	1✔
186	if isinstance(blocked_form, ufl.form.Form):	1✔
187	# If there are no sub-elements, continue
188	continue	1✔
189	diagonal_form = ufl.ZeroBaseForm(())	1✔
190	for j in range(len(blocked_form)):	1✔
191	if blocked_form[j][j] is not None:	1✔
192	diagonal_form += blocked_form[j][j]	1✔
193	if diagonal_form == 0:	1✔
NEW 194	raise RuntimeError("Diagonal form seems to be zero.")	×
195	forms[i] = diagonal_form # type: ignore	1✔
196
197	# Get a signature for these forms
198	module_name = "libffcx_forms_" + ffcx.naming.compute_signature(	1✔
199	forms,
200	_compute_option_signature(p) + _compilation_signature(cffi_extra_compile_args, cffi_debug),
201	)
202
203	form_names = [ffcx.naming.form_name(form, i, module_name) for i, form in enumerate(forms)]	1✔
204
205	if cache_dir is not None:	1✔
206	cache_dir = Path(cache_dir)	1✔
207	obj, mod = get_cached_module(module_name, form_names, cache_dir, timeout)	1✔
208	if obj is not None:	1✔
209	return obj, mod, (None, None)	×
210	else:
211	cache_dir = Path(tempfile.mkdtemp())	1✔
212
213	try:	1✔
214	decl = (	1✔
215	UFC_HEADER_DECL.format(np.dtype(p["scalar_type"]).name) # type: ignore
216	+ UFC_INTEGRAL_DECL
217	+ UFC_FORM_DECL
218	)
219
220	form_template = "extern ufcx_form {name};\n"	1✔
221	for name in form_names:	1✔
222	decl += form_template.format(name=name)	1✔
223
224	impl = _compile_objects(	1✔
225	decl,
226	forms,
227	form_names,
228	module_name,
229	p,
230	cache_dir,
231	cffi_extra_compile_args,
232	cffi_verbose,
233	cffi_debug,
234	cffi_libraries,
235	visualise=visualise,
236	)
237	except Exception as e:	1✔
238	try:	1✔
239	# remove c file so that it will not timeout next time
240	c_filename = cache_dir.joinpath(module_name + ".c")	1✔
241	os.replace(c_filename, c_filename.with_suffix(".c.failed"))	1✔
242	except Exception:	1✔
243	pass	1✔
244	raise e	1✔
245
246	obj, module = _load_objects(cache_dir, module_name, form_names)	1✔
247	return obj, module, (decl, impl)	1✔
248
249
250	def compile_expressions(	1✔
251	expressions: list[tuple[ufl.Expr, npt.NDArray[np.floating]]], # type: ignore
252	options: dict = {},
253	cache_dir: Path \| None = None,
254	timeout: int = 10,
255	cffi_extra_compile_args: list[str] = [],
256	cffi_verbose: bool = False,
257	cffi_debug: bool = False,
258	cffi_libraries: list[str] = [],
259	visualise: bool = False,
260	):
261	"""Compile a list of UFL expressions into UFC Python objects.
262
263	Args:
264	expressions: List of (UFL expression, evaluation points).
265	options: Options
266	cache_dir: Cache directory
267	timeout: Timeout
268	cffi_extra_compile_args: Extra compilation args for CFFI
269	cffi_verbose: Use verbose compile
270	cffi_debug: Use compiler debug mode
271	cffi_libraries: libraries to use with compiler
272	visualise: Toggle visualisation
273	"""
274	p = ffcx.options.get_options(options)	1✔
275
276	module_name = "libffcx_expressions_" + ffcx.naming.compute_signature(	1✔
277	expressions,
278	_compute_option_signature(p) + _compilation_signature(cffi_extra_compile_args, cffi_debug),
279	)
280	expr_names = [	1✔
281	ffcx.naming.expression_name(expression, module_name) for expression in expressions
282	]
283
284	if cache_dir is not None:	1✔
285	cache_dir = Path(cache_dir)	×
286	obj, mod = get_cached_module(module_name, expr_names, cache_dir, timeout)	×
287	if obj is not None:	×
288	return obj, mod, (None, None)	×
289	else:
290	cache_dir = Path(tempfile.mkdtemp())	1✔
291
292	try:	1✔
293	decl = (	1✔
294	UFC_HEADER_DECL.format(np.dtype(p["scalar_type"]).name) # type: ignore
295	+ UFC_INTEGRAL_DECL
296	+ UFC_FORM_DECL
297	+ UFC_EXPRESSION_DECL
298	)
299
300	expression_template = "extern ufcx_expression {name};\n"	1✔
301	for name in expr_names:	1✔
302	decl += expression_template.format(name=name)	1✔
303
304	impl = _compile_objects(	1✔
305	decl,
306	expressions,
307	expr_names,
308	module_name,
309	p,
310	cache_dir,
311	cffi_extra_compile_args,
312	cffi_verbose,
313	cffi_debug,
314	cffi_libraries,
315	visualise=visualise,
316	)
317	except Exception as e:	×
318	try:	×
319	# remove c file so that it will not timeout next time
320	c_filename = cache_dir.joinpath(module_name + ".c")	×
321	os.replace(c_filename, c_filename.with_suffix(".c.failed"))	×
322	except Exception:	×
323	pass	×
324	raise e	×
325
326	obj, module = _load_objects(cache_dir, module_name, expr_names)	1✔
327	return obj, module, (decl, impl)	1✔
328
329
330	def _compile_objects(	1✔
331	decl,
332	ufl_objects,
333	object_names,
334	module_name,
335	options,
336	cache_dir,
337	cffi_extra_compile_args,
338	cffi_verbose,
339	cffi_debug,
340	cffi_libraries,
341	visualise: bool = False,
342	):
343	import ffcx.compiler	1✔
344
345	libraries = _libraries + cffi_libraries if cffi_libraries is not None else _libraries	1✔
346
347	# JIT uses module_name as prefix, which is needed to make names of all struct/function
348	# unique across modules
349	_, code_body = ffcx.compiler.compile_ufl_objects(	1✔
350	ufl_objects, prefix=module_name, options=options, visualise=visualise
351	)
352
353	# Raise error immediately prior to compilation if no support for C99
354	# _Complex. Doing this here allows FFCx to be used for complex codegen on
355	# Windows.
356	if sys.platform.startswith("win32"):	1✔
357	if np.issubdtype(options["scalar_type"], np.complexfloating):	×
358	raise NotImplementedError("win32 platform does not support C99 _Complex numbers")	×
359	elif isinstance(options["scalar_type"], str) and "complex" in options["scalar_type"]:	×
360	raise NotImplementedError("win32 platform does not support C99 _Complex numbers")	×
361
362	# Compile in C17 mode
363	if sys.platform.startswith("win32"):	1✔
364	cffi_base_compile_args = ["-std:c17"]	×
365	else:
366	cffi_base_compile_args = ["-std=c17"]	1✔
367
368	cffi_final_compile_args = cffi_base_compile_args + cffi_extra_compile_args	1✔
369
370	ffibuilder = cffi.FFI()	1✔
371
372	ffibuilder.set_source(	1✔
373	module_name,
374	code_body,
375	include_dirs=[ffcx.codegeneration.get_include_path()],
376	extra_compile_args=cffi_final_compile_args,
377	libraries=libraries,
378	)
379
380	ffibuilder.cdef(decl)	1✔
381
382	c_filename = cache_dir.joinpath(module_name + ".c")	1✔
383	ready_name = c_filename.with_suffix(".c.cached")	1✔
384
385	# Compile (ensuring that compile dir exists)
386	cache_dir.mkdir(exist_ok=True, parents=True)	1✔
387
388	logger.info(79 * "#")	1✔
389	logger.info("Calling JIT C compiler")	1✔
390	logger.info(79 * "#")	1✔
391
392	t0 = time.time()	1✔
393	f = io.StringIO()	1✔
394	# Temporarily set root logger handlers to string buffer only
395	# since CFFI logs into root logger
396	old_handlers = root_logger.handlers.copy()	1✔
397	root_logger.handlers = [logging.StreamHandler(f)]	1✔
398	with redirect_stdout(f):	1✔
399	ffibuilder.compile(tmpdir=cache_dir, verbose=True, debug=cffi_debug)	1✔
400	s = f.getvalue()	1✔
401	if cffi_verbose:	1✔
402	print(s)	×
403
404	logger.info(f"JIT C compiler finished in {time.time() - t0:.4f}")	1✔
405
406	# Create a "status ready" file. If this fails, it is an error,
407	# because it should not exist yet.
408	# Copy the stdout verbose output of the build into the ready file
409	fd = open(ready_name, "x")	1✔
410	fd.write(s)	1✔
411	fd.close()	1✔
412
413	# Copy back the original handlers (in case someone is logging into
414	# root logger and has custom handlers)
415	root_logger.handlers = old_handlers	1✔
416
417	return code_body	1✔
418
419
420	def _load_objects(cache_dir, module_name, object_names):	1✔
421	# Create module finder that searches the compile path
422	finder = importlib.machinery.FileFinder(	1✔
423	str(cache_dir),
424	(importlib.machinery.ExtensionFileLoader, importlib.machinery.EXTENSION_SUFFIXES),
425	)
426
427	# Find module. Clear search cache to be sure dynamically created
428	# (new) modules are found
429	finder.invalidate_caches()	1✔
430	spec = finder.find_spec(module_name)	1✔
431	if spec is None:	1✔
432	raise ModuleNotFoundError("Unable to find JIT module.")	×
433
434	# Load module
435	compiled_module = importlib.util.module_from_spec(spec)	1✔
436	spec.loader.exec_module(compiled_module)	1✔
437
438	compiled_objects = []	1✔
439	for name in object_names:	1✔
440	obj = getattr(compiled_module.lib, name)	1✔
441	compiled_objects.append(obj)	1✔
442
443	return compiled_objects, compiled_module	1✔

FEniCS / ffcx / 17789800893

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous