"""Compiled C backend for cell models generated by gotranx
``beat`` advances the cell model through a function ``fun(states, t, parameters, dt)`` that acts on
a ``(num_states, num_points)`` array. This module compiles the C code that gotranx generates for
one cell (``gotranx ode2c``, or :class:`gotranx.codegen.CCodeGenerator`) together with a small
driver that loops over the points, and wraps the library in a :class:`CScheme` with that
signature, so that it can be passed as ``fun`` to :class:`beat.odesolver.DolfinODESolver` and
:class:`beat.odesolver.DolfinMultiODESolver`.
Rank 0 compiles the library into a cache directory, keyed by a hash of the source, the compiler
and the flags, and the other ranks load it from there. The cache directory must therefore be on a
file system that all ranks can see.
"""
from __future__ import annotations
import ctypes
import hashlib
import logging
import os
import platform
import subprocess
import tempfile
from collections.abc import Callable, Sequence
from pathlib import Path
from typing import Any, TypeVar
from mpi4py import MPI
import numpy as np
import numpy.typing as npt
logger = logging.getLogger(__name__)
T = TypeVar("T")
DRIVER_PATH = Path(__file__).with_name("c_backend_driver.c")
LIBRARY_NAME = "libbeat_scheme.so"
DEFAULT_CFLAGS: tuple[str, ...] = ("-O3", "-march=native")
OPENMP_FLAGS: tuple[str, ...] = ("-fopenmp", "-DBEAT_OPENMP")
[docs]
def default_cache_dir() -> Path:
"""``$BEAT_C_CACHE`` if it is set, else ``~/.cache/beat/c_backend``"""
if env := os.environ.get("BEAT_C_CACHE"):
return Path(env)
return Path.home() / ".cache" / "beat" / "c_backend"
[docs]
def default_compiler() -> str:
"""``$CC`` if it is set, else ``cc``"""
return os.environ.get("CC", "cc")
[docs]
def full_source(c_code: str, scheme: str) -> str:
"""The generated C code followed by the driver that loops over the points"""
return f"{c_code}\n\n#define BEAT_SCHEME {scheme}\n{DRIVER_PATH.read_text()}"
def _run_compiler(cmd: list[str]) -> subprocess.CompletedProcess:
return subprocess.run(cmd, capture_output=True, text=True, check=False)
[docs]
def compiler_version(cc: str) -> str:
"""The first line of ``cc --version``"""
try:
result = _run_compiler([cc, "--version"])
except FileNotFoundError as e:
raise RuntimeError(
f"The C compiler {cc!r} was not found. Install one, or set CC or pass cc=...",
) from e
return result.stdout.splitlines()[0] if result.stdout else ""
[docs]
def target_signature(cc: str, flags: Sequence[str]) -> str:
"""The compiler's predefined macros for ``flags`` if they target the native CPU, else ""
With ``-march=native`` (or ``-mcpu=native``) the library depends on the features of the CPU
it is compiled on, which ``platform.machine()`` does not tell apart (e.g. with and without
AVX-512), so the macros that the compiler defines for these flags go into the cache key.
"""
if not any("native" in flag for flag in flags):
return ""
result = _run_compiler([cc, *flags, "-dM", "-E", "-x", "c", "/dev/null"])
return result.stdout
[docs]
def cache_key(source: str, cc: str, cc_version: str, flags: Sequence[str]) -> str:
"""Hash of everything that changes the compiled library"""
h = hashlib.sha256()
for part in (source, cc, cc_version, *flags, platform.machine()):
h.update(part.encode())
h.update(b"\0")
return h.hexdigest()[:32]
def _compile(source: str, cc: str, flags: Sequence[str], directory: Path) -> Path:
library = directory / LIBRARY_NAME
if library.is_file():
return library
directory.mkdir(parents=True, exist_ok=True)
# Unique names and an atomic rename, since another job may compile the same key at once
fd, c_tmp = tempfile.mkstemp(dir=directory, suffix=".c")
with os.fdopen(fd, "w") as f:
f.write(source)
fd, lib_tmp = tempfile.mkstemp(dir=directory, suffix=".so")
os.close(fd)
# An unknown scheme name must fail here, not as an undefined symbol when the library is loaded
cmd = [cc, *flags, "-Werror=implicit-function-declaration", "-shared", "-fPIC"]
cmd += ["-o", lib_tmp, c_tmp, "-lm"]
logger.debug("Compiling the C scheme: %s", " ".join(cmd))
try:
result = _run_compiler(cmd)
except FileNotFoundError as e:
Path(c_tmp).unlink(missing_ok=True)
Path(lib_tmp).unlink(missing_ok=True)
raise RuntimeError(f"The C compiler {cc!r} was not found") from e
if result.returncode != 0:
Path(c_tmp).unlink(missing_ok=True)
Path(lib_tmp).unlink(missing_ok=True)
hint = ""
if "-fopenmp" in flags:
hint = "\nOpenMP was requested: check that the compiler supports -fopenmp."
raise RuntimeError(
f"Compiling the C scheme failed:\n{' '.join(cmd)}\n{result.stderr}{hint}",
)
if library.is_file():
# Another job finished the same key first; keep its files
Path(c_tmp).unlink(missing_ok=True)
Path(lib_tmp).unlink(missing_ok=True)
return library
os.replace(c_tmp, directory / "scheme.c")
os.replace(lib_tmp, library)
return library
def _check_visible(library: Path, comm: MPI.Intracomm) -> None:
"""Raise on all ranks if any rank cannot see the library that rank 0 built"""
if not comm.allreduce(library.is_file(), op=MPI.LAND):
raise RuntimeError(
f"Not all ranks can see {library}. The cache directory must be on a file system "
"that all ranks share; set BEAT_C_CACHE or pass cache_dir=...",
)
def _run_on_rank0(comm: MPI.Intracomm, func: Callable[[], T]) -> T:
"""Run ``func`` on rank 0 and broadcast its result; raise its error on every rank"""
result: T | None = None
error: str | None = None
if comm.rank == 0:
try:
result = func()
except Exception as e: # noqa: BLE001 - re-raised on every rank below
error = f"{type(e).__name__}: {e}"
result, error = comm.bcast((result, error), root=0)
if error is not None:
raise RuntimeError(error)
return result # type: ignore[return-value]
[docs]
def build_library(
source: str,
*,
cc: str,
cflags: tuple[str, ...],
cache_dir: Path,
comm: MPI.Intracomm,
) -> Path:
"""Compile ``source`` on rank 0 (unless it is cached) and return the library on all ranks"""
def build() -> Path:
key = cache_key(
source,
cc,
compiler_version(cc),
(*cflags, target_signature(cc, cflags)),
)
return _compile(source, cc, cflags, cache_dir / key)
library = _run_on_rank0(comm, build)
_check_visible(library, comm)
return library
def _check_array(name: str, array: Any) -> npt.NDArray[np.float64]:
if not isinstance(array, np.ndarray):
raise TypeError(f"{name} must be a numpy array, got {type(array).__name__}")
if array.dtype != np.float64:
raise TypeError(f"{name} must have dtype float64, got {array.dtype}")
if array.ndim not in (1, 2):
raise ValueError(f"{name} must be 1D or 2D, got shape {array.shape}")
if not array.flags.c_contiguous:
raise ValueError(f"{name} must be C-contiguous (use np.ascontiguousarray)")
return array
[docs]
class CScheme:
"""A compiled cell-model scheme, callable as ``fun(states, t, parameters, dt)``
Parameters
----------
library_path : Path
The shared library built by :func:`compile_scheme`
scheme : str
The name of the per-cell C function
num_states : int | None, optional
The number of states of the model; if given, ``states`` is checked against it
num_parameters : int | None, optional
The number of parameters of the model; if given, ``parameters`` is checked against it
openmp : bool, optional
Whether the library was compiled with OpenMP
num_threads : int | None, optional
The number of OpenMP threads, by default the OpenMP default
"""
def __init__(
self,
library_path: Path,
scheme: str,
*,
num_states: int | None = None,
num_parameters: int | None = None,
openmp: bool = False,
num_threads: int | None = None,
) -> None:
self.library_path = Path(library_path)
self.scheme = scheme
self.num_states = num_states
self.num_parameters = num_parameters
self.openmp = openmp
self.num_threads = num_threads
self._lib = ctypes.CDLL(str(self.library_path))
vec = self._lib.beat_scheme_vec
vec.restype = None
vec.argtypes = [
ctypes.c_int,
ctypes.c_long,
ctypes.c_void_p,
ctypes.c_double,
ctypes.c_double,
ctypes.c_void_p,
ctypes.c_int,
ctypes.c_long,
ctypes.c_void_p,
ctypes.c_int,
]
self._vec = vec
def __repr__(self) -> str:
return f"CScheme(scheme={self.scheme!r}, library_path={str(self.library_path)!r})"
def _index(self, function: str, name: str) -> int:
lookup = getattr(self._lib, function)
lookup.restype = ctypes.c_int
lookup.argtypes = [ctypes.c_char_p]
index = lookup(name.encode())
if index < 0:
raise KeyError(name)
return index
[docs]
def state_index(self, name: str) -> int:
"""The index of the state ``name``"""
return self._index("state_index", name)
[docs]
def parameter_index(self, name: str) -> int:
"""The index of the parameter ``name``"""
return self._index("parameter_index", name)
def _init_values(self, function: str, count: int | None, what: str) -> npt.NDArray[np.float64]:
if count is None:
raise ValueError(f"{what} is unknown; pass {what}=... to compile_scheme")
values = np.zeros(count)
init = getattr(self._lib, function)
init.restype = None
init.argtypes = [ctypes.c_void_p]
init(values.ctypes.data)
return values
[docs]
def init_state_values(self) -> npt.NDArray[np.float64]:
"""The default initial states of the model"""
return self._init_values("init_state_values", self.num_states, "num_states")
[docs]
def init_parameter_values(self) -> npt.NDArray[np.float64]:
"""The default parameters of the model"""
return self._init_values("init_parameter_values", self.num_parameters, "num_parameters")
def __call__(
self,
states: npt.NDArray[np.float64],
t: float,
parameters: npt.NDArray[np.float64],
dt: float,
out: npt.NDArray[np.float64] | None = None,
) -> npt.NDArray[np.float64]:
states = _check_array("states", states)
parameters = _check_array("parameters", parameters)
n_states = states.shape[0]
n_points = 1 if states.ndim == 1 else states.shape[1]
n_params = parameters.shape[0]
if self.num_states is not None and n_states != self.num_states:
raise ValueError(f"states has {n_states} states, but the model has {self.num_states}")
if self.num_parameters is not None and n_params != self.num_parameters:
raise ValueError(
f"parameters has {n_params} parameters, but the model has {self.num_parameters}",
)
if parameters.ndim == 1:
params_stride = 0
elif parameters.shape[1] == n_points:
params_stride = n_points
else:
raise ValueError(
f"parameters has {parameters.shape[1]} points, but states has {n_points}",
)
if out is None:
out = np.empty_like(states)
else:
out = _check_array("out", out)
if not out.flags.writeable:
raise ValueError("out must be writeable")
if out.shape != states.shape:
raise ValueError(f"out has shape {out.shape}, but states has {states.shape}")
self._vec(
n_states,
n_points,
states.ctypes.data,
float(t),
float(dt),
parameters.ctypes.data,
n_params,
params_stride,
out.ctypes.data,
self.num_threads or 0,
)
return out
[docs]
def max_threads(self) -> int:
"""The number of threads that one call uses (1 without OpenMP)"""
function = self._lib.beat_max_threads
function.restype = ctypes.c_int
function.argtypes = [ctypes.c_int]
return function(self.num_threads or 0)
[docs]
def compile_scheme(
c_code: str,
scheme: str = "generalized_rush_larsen",
*,
num_states: int | None = None,
num_parameters: int | None = None,
cache_dir: str | Path | None = None,
cc: str | None = None,
cflags: Sequence[str] = DEFAULT_CFLAGS,
openmp: bool = False,
num_threads: int | None = None,
comm: MPI.Intracomm = MPI.COMM_WORLD,
) -> CScheme:
"""Compile gotranx-generated C code into a ``fun(states, t, parameters, dt)``
Parameters
----------
c_code : str
C code with the per-cell function ``scheme``, as generated by gotranx
scheme : str, optional
The name of the per-cell function, by default "generalized_rush_larsen"
num_states : int | None, optional
The number of states, used to check the arrays that the scheme is called with.
The generated C code has no state or parameter counts, so without ``num_states`` and
``num_parameters`` the arrays are not checked against the model, and arrays that are too
small make the C code read out of bounds. :func:`from_ode` fills them in.
num_parameters : int | None, optional
The number of parameters, used to check the arrays that the scheme is called with
(see ``num_states``)
cache_dir : str | Path | None, optional
Where to put the library, by default :func:`default_cache_dir`. Must be visible to all ranks
cc : str | None, optional
The C compiler, by default :func:`default_compiler`
cflags : Sequence[str], optional
The compiler flags, by default ``("-O3", "-march=native")``
openmp : bool, optional
Compile with -fopenmp and run the point loop in parallel, by default False
num_threads : int | None, optional
The number of OpenMP threads, by default OMP_NUM_THREADS or the OpenMP default
comm : MPI.Intracomm, optional
The communicator; rank 0 compiles, by default MPI.COMM_WORLD
Returns
-------
CScheme
The compiled scheme
"""
if isinstance(cflags, str):
raise TypeError(f"cflags must be a sequence of strings, not a str: {cflags!r}")
if num_threads is not None and num_threads < 1:
raise ValueError(f"num_threads must be at least 1, got {num_threads}")
if num_threads is not None and not openmp:
raise ValueError("num_threads is only used with OpenMP; pass openmp=True")
if (num_states is None or num_parameters is None) and comm.rank == 0:
logger.warning(
"compile_scheme was called without num_states and num_parameters, so the arrays are "
"not checked against the model and too-small arrays make the C code read out of bounds",
)
cc = cc or default_compiler()
cache = Path(cache_dir) if cache_dir is not None else default_cache_dir()
flags = tuple(cflags) + (OPENMP_FLAGS if openmp else ())
library = build_library(
full_source(c_code, scheme),
cc=cc,
cflags=flags,
cache_dir=cache.resolve(),
comm=comm,
)
return CScheme(
library,
scheme,
num_states=num_states,
num_parameters=num_parameters,
openmp=openmp,
num_threads=num_threads,
)
[docs]
def generate_c_code(
ode: Any,
scheme: str = "generalized_rush_larsen",
codegen_kwargs: dict[str, Any] | None = None,
) -> str:
"""The C code for one cell, with index and initial value helpers and ``scheme``
Parameters
----------
ode : gotranx.ode.ODE
The cell model
scheme : str, optional
The name of a gotranx scheme, by default "generalized_rush_larsen"
codegen_kwargs : dict[str, Any] | None, optional
Extra keyword arguments for :meth:`gotranx.codegen.CCodeGenerator.scheme`
"""
from gotranx.codegen import CCodeGenerator
from gotranx.codegen.c import Format
from gotranx.schemes import get_scheme
codegen = CCodeGenerator(ode, format=Format.none)
return "\n".join(
[
codegen.imports(),
codegen.parameter_index(),
codegen.state_index(),
codegen.initial_parameter_values(),
codegen.initial_state_values(),
codegen.scheme(get_scheme(scheme), **(codegen_kwargs or {})),
],
)
[docs]
def from_ode(
ode: Any,
scheme: str = "generalized_rush_larsen",
*,
codegen_kwargs: dict[str, Any] | None = None,
comm: MPI.Intracomm = MPI.COMM_WORLD,
**compile_kwargs: Any,
) -> CScheme:
"""Generate C for a gotranx ODE and compile it with :func:`compile_scheme`
The code is generated on rank 0 only. ``compile_kwargs`` are passed on to
:func:`compile_scheme` (``cache_dir``, ``cc``, ``cflags``, ...).
Parameters
----------
ode : gotranx.ode.ODE
The cell model
scheme : str, optional
The name of a gotranx scheme, by default "generalized_rush_larsen"
codegen_kwargs : dict[str, Any] | None, optional
Extra keyword arguments for :meth:`gotranx.codegen.CCodeGenerator.scheme`
comm : MPI.Intracomm, optional
The communicator, by default MPI.COMM_WORLD
"""
c_code = _run_on_rank0(comm, lambda: generate_c_code(ode, scheme, codegen_kwargs))
return compile_scheme(
c_code,
scheme,
num_states=ode.num_states,
num_parameters=ode.num_parameters,
comm=comm,
**compile_kwargs,
)
[docs]
def openmp_available(cc: str | None = None) -> bool:
"""Whether the C compiler can build an OpenMP shared library"""
cc = cc or default_compiler()
with tempfile.TemporaryDirectory() as tmp:
source = Path(tmp) / "omp.c"
source.write_text("#include <omp.h>\nint f(void) { return omp_get_max_threads(); }\n")
try:
result = _run_compiler(
[cc, "-fopenmp", "-shared", "-fPIC", "-o", str(Path(tmp) / "omp.so"), str(source)],
)
except OSError:
return False
return result.returncode == 0