Code generation¶

def render_c_module(
  fun: Function,
  *,
  header_name: str | None = None,
  source_name: str | None = None,
  lang: str = "c",
  casadi: bool = False,
  lanes: LaneCount = "auto",
  dialect: CDialect = "gnu",
  vector_libm: VectorLibm = "none",
  reciprocal: bool = False,
  cpu: CpuLevel = "generic",
) -> CModule:
  """Render ``fun`` into its header / ``.c`` pair from a single lowering. The kernel is always C.
  ``lang`` picks the header a caller includes (``f.h`` or ``f.hpp``) and ``casadi`` adds the
  CasADi 3.8 compatible symbols to both."""
  _check_lang(lang)
  ctx, body = _render_observed(
    fun,
    casadi=casadi,
    recipe=BuildRecipe(cpu=cpu, lanes=lanes, dialect=dialect, vector_libm=vector_libm, reciprocal=reciprocal),
    source_name=source_name,
  )
  symbol = c_ident(fun.name)
  return CModule(
    fun=fun,
    header_name=f"{symbol}.{'hpp' if lang == 'cpp' else 'h'}" if header_name is None else header_name,
    source_name=f"{symbol}.c" if source_name is None else source_name,
    body=body,
    program=ctx.prog,
    workspace_size=entry_workspace(fun, ctx.workspace_size, casadi=casadi),
    backends=ctx.backends,
    lang=lang,
    casadi=casadi,
    recipe=ctx.recipe,
  )

def render_c_source(
  fun: Function,
  *,
  casadi: bool = False,
  lanes: LaneCount = "auto",
  dialect: CDialect = "gnu",
  vector_libm: VectorLibm = "none",
  reciprocal: bool = False,
  cpu: CpuLevel = "generic",
) -> str:
  """Render a standalone pointer-ABI C implementation of ``fun`` and its callees. ``casadi`` adds
  the CasADi 3.8 compatible symbols.

  A ``LoweringError`` for a function that cannot be lowered propagates. There is no fallback.
  """
  recipe = BuildRecipe(cpu=cpu, lanes=lanes, dialect=dialect, vector_libm=vector_libm, reciprocal=reciprocal)
  return _render_observed(fun, casadi=casadi, recipe=recipe)[1]

def render_c_api_header(
  fun: Function,
  *,
  lang: str = "c",
  casadi: bool = False,
  lanes: LaneCount = "auto",
  dialect: CDialect = "gnu",
  vector_libm: VectorLibm = "none",
  reciprocal: bool = False,
  cpu: CpuLevel = "generic",
) -> str:
  """Render the public header for ``fun``: the ABI declarations, ``SZ_*`` constants, the typed
  buffers of the chosen ``lang``, the
  sparse-output tables, and with ``casadi`` the CasADi query prototypes."""
  _check_lang(lang)
  if casadi:
    check_casadi_layout(fun)
  ctx = _lower(fun, recipe=BuildRecipe(cpu=cpu, lanes=lanes, dialect=dialect, vector_libm=vector_libm, reciprocal=reciprocal))
  return ctx.recipe.comment(f"{c_ident(fun.name)}.c") + _render_header(
    ctx.fun, ctx.backends, entry_workspace(fun, ctx.workspace_size, casadi=casadi), lang=lang, casadi=casadi
  )

def write_module(
  fun: Function | Solver,
  out_dir: Path,
  *,
  lang: str = "c",
  casadi: bool = False,
  lanes: LaneCount = "auto",
  dialect: CDialect = "gnu",
  vector_libm: VectorLibm = "none",
  reciprocal: bool = False,
  cpu: CpuLevel = "generic",
) -> CModule:
  """Write ``fun``'s header / ``.c`` into ``out_dir`` and return the module. A ``Solver`` renders its ``function``."""
  if isinstance(fun, Solver):
    fun = fun.function
  module = render_c_module(fun, lang=lang, casadi=casadi, lanes=lanes, dialect=dialect, vector_libm=vector_libm, reciprocal=reciprocal, cpu=cpu)
  out_dir.mkdir(parents=True, exist_ok=True)
  (out_dir / module.header_name).write_text(module.header)
  (out_dir / module.source_name).write_text(module.source)
  return module

def workspace_size(
  fun: Function,
  *,
  casadi: bool = False,
  lanes: LaneCount = "auto",
  dialect: CDialect = "gnu",
  vector_libm: VectorLibm = "none",
  reciprocal: bool = False,
  cpu: CpuLevel = "generic",
) -> int:
  """Doubles of scratch ``fun`` needs in ``w[]``, the value its header's ``SZ_W`` quotes.
  ``CModule.workspace_size`` is the same number without a second lowering, so prefer it when the
  module is already in hand."""
  return entry_workspace(
    fun,
    _lower(fun, recipe=BuildRecipe(cpu=cpu, lanes=lanes, dialect=dialect, vector_libm=vector_libm, reciprocal=reciprocal)).workspace_size,
    casadi=casadi,
  )

def c_api_signature(symbol: str = "f") -> str:
  """The universal ABI entry signature, spelled for a given symbol name."""
  return C_API_SIGNATURE.replace(" f(", f" {symbol}(")

def render_cpp_header(fun: Function, backends: tuple[str, ...], sz_w: int, *, casadi: bool, sparsities: tuple[SparsityPattern | None, ...]) -> str:
  """The ``.hpp`` for ``fun``. The kernel symbols are declared ``extern "C"`` inside the function's
  namespace, since a namespace and a function cannot share the global name. C linkage keeps the
  symbol unmangled, so ``f::f`` is the same entry a C caller reaches as ``f``."""
  symbol = c_ident(fun.name)
  inputs, outputs = buffer_idents(fun)
  params = [f"const {i}_t& {i}" for i in inputs] + [f"{o}_t& {o}" for o in outputs] + ["workspace_t& workspace"]
  arg_init = ", ".join(f"{i}.ptr()" for i in inputs) or "nullptr"
  res_init = ", ".join(f"{o}.ptr()" for o in outputs) or "nullptr"
  lines = [
    "#pragma once",
    "",
    "#include <array>",
    "#include <cassert>",
    "#include <cstddef>",
    *(["#include <cstdint>", "", *stats_c_defs()] if backends else []),
    "",
    *abi_status_defines(guarded=True),
    *(["", *casadi_defines()] if casadi else []),
    "",
    f"#define {symbol}_SZ_ARG {len(fun.inputs)}",
    f"#define {symbol}_SZ_RES {len(fun.outputs)}",
    f"#define {symbol}_SZ_IW 0",
    f"#define {symbol}_SZ_W {sz_w}",
    "",
    BUFFER_TEMPLATE,
    "",
    f"namespace {symbol} {{",
    f"// The pointer ABI for {fun.name}; the same C symbols the C header declares.",
    'extern "C" ' + c_api_signature(symbol) + ";",
    *(f'extern "C" int {s}_stats(scaly_solver_stats* out);' for s in solver_stats_symbols(fun)),
    *(f'extern "C" {decl}' for decl in (casadi_declarations(symbol) if casadi else [])),
    "",
    *(_buffer_alias(i, e.shape) for i, e in zip(inputs, fun.inputs, strict=True)),
    *(_buffer_alias(o, e.shape) for o, e in zip(outputs, fun.outputs, strict=True)),
    f"using workspace_t = Buffer<double, {symbol}_SZ_W>;",
    f"constexpr int sz_arg = {symbol}_SZ_ARG;",
    f"constexpr int sz_res = {symbol}_SZ_RES;",
    f"constexpr int sz_iw = {symbol}_SZ_IW;",
    f"constexpr int sz_w = {symbol}_SZ_W;",
    "",
    f"inline int call({', '.join(params)}) {{",
    f"  const double* arg[sz_arg > 0 ? sz_arg : 1] = {{{arg_init}}};",
    f"  double* res[sz_res > 0 ? sz_res : 1] = {{{res_init}}};",
    f"  return {symbol}(arg, res, nullptr, sz_w ? workspace.ptr() : nullptr, 0);",
    "}",
  ]
  for name, sp in zip(fun.output_names, sparsities, strict=True):
    if sp is not None:
      lines += ["", *_sparse_namespace(name, sp)]
  lines += [f"}}  // namespace {symbol}"]
  return "\n".join(lines) + "\n"

def check_casadi_layout(fun: Function) -> None:
  """CasADi buffers are column-major and scaly's are row-major. Scalars, vectors and compact sparse
  outputs agree. A dense matrix with both dimensions above one would need a transpose, so it is
  rejected here rather than hand the caller a transposed matrix."""
  for kind, names, exprs, sparsities in (
    ("input", fun.input_names, fun.inputs, (None,) * len(fun.inputs)),
    ("output", fun.output_names, fun.outputs, fun.output_sparsities),
  ):
    for name, expr, sp in zip(names, exprs, sparsities, strict=True):
      if sp is None and (len(expr.shape) > 2 or (len(expr.shape) == 2 and min(expr.shape) > 1)):
        raise ValueError(
          f"casadi=True: {kind} {name!r} of {fun.name!r} has dense matrix shape {expr.shape}; CasADi is column-major, so only scalars, vectors and compact sparse outputs are supported"
        )

def casadi_sparsity(shape: tuple[int, ...], sp: SparsityPattern | None) -> tuple[int, ...]:
  """CasADi's compressed encoding: ``{nrow, ncol, 1}`` for dense, ``{nrow, ncol, colind..., row...}``
  for sparse. ``check_casadi_layout`` has already ruled out dense matrices."""
  if sp is None:
    return (*_dense_dims(shape), 1)
  col_ptr, row_ind, _ = sp.to_csc()
  return (*sp.shape, *col_ptr, *row_ind)

def casadi_scratch(fun: Function) -> int:
  """Doubles added to ``SZ_W`` for the compressed-column gather: ``nnz`` per compact sparse output
  whose native order is not already compressed-column (the QP path's patterns are, so they add 0)."""
  return sum(sp.nnz for sp in fun.output_sparsities if _needs_gather(sp) and sp is not None)

def load_library(path: Path, *, isolated: bool) -> ctypes.CDLL:
  """Load a shared library. ``isolated`` keeps Linux solver dependencies out of the host process linker namespace."""
  global _SOLVER_NAMESPACE, _SOLVER_NAMESPACE_ANCHOR
  if not isolated or sys.platform != "linux":
    return ctypes.CDLL(str(path))
  libc = ctypes.CDLL(None)
  dlmopen = libc.dlmopen
  dlmopen.argtypes = [ctypes.c_long, ctypes.c_char_p, ctypes.c_int]
  dlmopen.restype = ctypes.c_void_p
  with _SOLVER_NAMESPACE_LOCK:
    handle = dlmopen(_LM_ID_NEWLM if _SOLVER_NAMESPACE is None else _SOLVER_NAMESPACE, os.fsencode(path), os.RTLD_NOW | os.RTLD_LOCAL)
    if not handle:
      raise OSError(f"dlmopen failed for {path}")
    # Python 3.14 reopens a path even when CDLL receives a handle, so bind it directly.
    lib = ctypes.CDLL(None)
    lib._handle = handle
    lib._name = str(path)
    if _SOLVER_NAMESPACE is None:
      dlinfo = libc.dlinfo
      dlinfo.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p]
      dlinfo.restype = ctypes.c_int
      namespace = ctypes.c_long()
      if dlinfo(handle, _RTLD_DI_LMID, ctypes.byref(namespace)) != 0:
        raise OSError(f"dlinfo failed for {path}")
      _SOLVER_NAMESPACE = namespace.value
      _SOLVER_NAMESPACE_ANCHOR = lib
  return lib

def invalidate_cache(fun: Function) -> None:
  """Drop both the in-memory artifact entry and the on-disk cache directory for ``fun``.

  Safe to call when nothing is cached yet. Codegen failures (``NotImplementedError``) are
  swallowed, since there cannot be a corresponding cache entry to remove.
  """
  compiler = find_c_compiler()
  if compiler is None:
    return
  try:
    module = _render_native(fun, compiler.cc)
  except NotImplementedError:
    return
  key = _compute_cache_key(module.body, fun_name=fun.name, compile_flags=(*compile_flags(module.recipe), *module.link_flags))
  with _artifact_lock:
    _artifact_cache.pop(key, None)
  cache_dir = cache_root() / key
  if cache_dir.exists():
    shutil.rmtree(cache_dir, ignore_errors=True)

def find_c_compiler() -> Compiler | None:
  """Find the C compiler the JIT uses: ``SCALY_CC``, else ``CC``, else ``cc`` on ``PATH``.

  Returns the compiler path and which of those three supplied it, or ``None`` if nothing is found.
  """
  for key in ("SCALY_CC", "CC"):
    override = os.environ.get(key)
    if override:
      found = shutil.which(override)
      return Compiler(found, key) if found is not None else None
  found = shutil.which("cc")
  return Compiler(found, "PATH") if found is not None else None

def cache_root() -> Path:
  """Return the JIT cache directory: ``SCALY_CACHE_DIR``, else ``$XDG_CACHE_HOME/scaly/jit``, else ``~/.cache/scaly/jit``."""
  override = env_path("SCALY_CACHE_DIR")
  if override is not None:
    return override
  xdg = env_path("XDG_CACHE_HOME")
  base = xdg if xdg is not None else Path.home() / ".cache"
  return base / "scaly" / "jit"