brainevent.register_ffi_target

brainevent.register_ffi_target#

brainevent.register_ffi_target(target_name, module, func_name, *, platform='CUDA', replace=False, content_id=None)[source]#

Register a compiled function as a JAX FFI target.

After registration, the function can be invoked inside @jax.jit via jax.ffi.ffi_call(target_name, ...).

The whole check-and-register sequence is guarded by a module-level lock and is idempotent: re-registering the same target_name with an equivalent module (identical shared-library content, function name, and platform) is a no-op and does not disturb the live keep-alives.

A different module (e.g. an edited-and-rebuilt kernel) under an already-registered name is refused deterministically on this JAX version, regardless of ``replace`` (see Notes): the installed JAX cannot re-point a live FFI target to new code in a way that can be performed or verified. The error directs the caller to register the rebuild under a distinct target name. replace is retained for API stability and reserved for a future JAX that supports verifiable re-pointing.

Parameters:

  • target_name (str) – Globally unique FFI target identifier.

  • module (CompiledModule) – The loaded module containing the function.

  • func_name (str) – Function name within the module.

  • platform (str) – Target platform ("CUDA" or "cpu").

  • replace (bool) – Reserved. On the installed JAX a content change under an existing target name raises whether or not this is set, because a live re-point cannot be verified (see Notes). Kept for API stability / forward compatibility; force_rebuild=True passes it through.

  • content_id (str | None) – Deterministic content identity for the registration, overriding the default .so byte hash. The compilation pipeline passes its cache key (a digest of source, ABI specs, headers, and build options) so that recompiling unchanged source — whose .so bytes may still differ because compilers embed build paths and timestamps — is recognised as the same registration (an idempotent no-op) rather than refused, while any real source/spec change still raises.

Raises:

KernelRegistrationError – If target_name is already registered to different .so content. This is deterministic on all platforms (CPU and CUDA alike) — see Notes for why a live re-point is refused rather than attempted.

Notes

Return type:

None

Registration is process-global and intentionally has no unload path: every registered .so is pinned in _LIVE_MODULES (a list, so future re-pointing support can append keep-alives without dropping the old image) for the lifetime of the process, so no XLA FFI target ever dangles.

Why a live re-point is refused. Probed on the installed JAX (0.10.2), XLA binds the FFI handler pointer into each compiled executable at compile time and resolves it by name only once, so already-traced callables keep dispatching to the original handler regardless. Worse, the two platforms disagree on re-registration and neither offers a lookup to confirm success:

  • On the CPU/”Host” platform, XLA’s registry rejects a re-registration whose handler pointer differs (“Duplicate FFI handler registration … with different bundle addresses”) — a raise.

  • On CUDA, XLA accepts the duplicate registration but silently keeps the old handler — so blindly re-registering would report success while the stale kernel keeps executing (this is audit finding 4).

Because success cannot be positively verified on this JAX, this function refuses deterministically instead of guessing. The reliable way to run edited code in a live process is a distinct target name (a new name= / target_prefix=), which the error messages recommend; a version check can enable true replace once some JAX supports it.