Skip to content

generated_kernel_abi.md

Latest commit

 

History

History
458 lines (386 loc) · 22.4 KB

generated_kernel_abi.md

File metadata and controls

458 lines (386 loc) · 22.4 KB
Raw

Tensorium Generated Kernel ABI (v3)

This document freezes the ABI contract used by generated functions:

  • tensorium_init
  • tensorium_rhs
  • tensorium_entry
  • tensorium_init_point
  • tensorium_init_grid_scf
  • tensorium_init_grid_affine
  • tensorium_rhs_grid_scf
  • tensorium_rhs_grid_affine
  • tensorium_rhs_grid_parallel
  • tensorium_residual_grid_scf
  • tensorium_residual_grid_affine
  • tensorium_residual_grid_parallel
  • tensorium_spectral_residual_<target>
  • tensorium_spectral_residual_jvp_<target>
  • tensorium_spectral_residual_grid_<target>

Source of truth for ABI constants: include/tensorium_mlir/Target/MLIRGen/GeneratedKernelABI.h.

Source of truth for the host-callable kernel descriptor: include/tensorium_mlir/Target/MLIRGen/MLIRGenHostABI.h.

Source of truth for the runtime spectral residual callback ABI: include/tensorium_mlir/Runtime/SpectralResidualKernel.h. This remains the public umbrella include; the runtime implementation is split under SpectralResidualTypes.h, SpectralResidualAssembly.h, SpectralResidualJVP.h, and SpectralEllipticSolver.h.

Architectural context: docs/language_mlir_abi_architecture.md describes how this ABI fits between Tensorium MLIR/LLVM kernels, generated host glue, standalone execution, and optional AMReX integration.

Versioning

  • ABI version attribute: tensorium.abi.version
  • Current value: 3
  • Memory layout attribute: tensorium.abi.memory_layout = "soa_component_major"
  • Memref ABI attribute: tensorium.abi.memref_abi = "strided_memref_rank1_f64"

Every generated function above carries these attrs, plus a stable tensorium.abi.kind.

Metadata attrs per function

Generated functions expose argument-order metadata:

  • tensorium.abi.param_names: runtime scalar parameter order (f64 values)
  • tensorium.abi.coord_names: coordinate buffer order (x,y,z or r,theta,phi)
  • tensorium.abi.field_names: field buffer order (matching function args)
  • tensorium.abi.output_names: written output field names
  • tensorium.abi.read_arg_indices: absolute argument indices read by the kernel
  • tensorium.abi.write_arg_indices: absolute argument indices (in function signature) written by the kernel
  • tensorium.abi.stencil_radius: required interior ghost/radius width for RHS and residual grid kernels, derived from lowered stencil reads
  • tensorium.abi.residual_kernel: internal marker on source residual modules that asks grid lowering to expose residual aliases

C/C++ low-level memref contract

After LLVM lowering, each memref<?xf64> or memref<Nxf64> argument is lowered to the 5-value descriptor:

  1. double *allocated
  2. double *aligned
  3. int64_t offset
  4. int64_t size
  5. int64_t stride

The helper host-side shape is provided as:

tensorium_mlir::abi::StridedMemRef1DF64

in include/tensorium_mlir/Target/MLIRGen/GeneratedKernelABI.h.

Generated C Host Header

The driver can emit a C header that mirrors the lowered ABI and adds thin buffer wrappers:

Tensorium_cc --emit-host-header tensorium_generated_host.h <file.tn>

The header contains:

  • raw extern void tensorium_* prototypes with expanded memref descriptors,
  • tensorium_memref1d_f64 for callers that need descriptor-level access,
  • convenience wrappers such as tensorium_call_init_grid_affine(...) and tensorium_call_rhs_grid_affine(...).

Convenience wrappers accept plain double * buffers and compute descriptor sizes from ABI metadata and field tensor ranks. For example, a rank-2 spatial field uses 9 * n_points in 3D, and a rank-3 spatial field uses 27 * n_points.

The internal host module descriptor also exposes simulation metadata (dimension, coordinate system, resolution, spatial scheme/order), field descriptors (name, variance, rank, component count), and per-kernel read/write/stencil metadata. This descriptor is the intended source for future C++ runtime and AMReX wrappers.

Runtime buffer contract

HostModuleABI now materializes the buffer-level contract that a runtime should consume directly:

  • each HostFieldABI records field name, variance (up/down), rank, and component count per grid point;
  • each HostKernelABI records raw scalar/memref arguments and a buffers table;
  • each HostBufferABI records logical buffer name, C-safe name, absolute function argument index, role (Coordinate, Field, Output), access (Read, Write, ReadWrite, None), variance/rank, and component count;
  • requiredBufferScalars(buffer, nPoints) returns the exact scalar allocation size required by SoA component-major layout;
  • validateHostModuleABI(abi) checks the descriptor before a runtime trusts it.
  • tensorium_mlir::runtime::HostFieldStorage builds a deduplicated storage plan from the ABI and a uniform grid shape. It allocates one contiguous scalar arena for all logical buffers, then exposes stable per-kernel binding plans and rank-1 memref descriptors into that arena.
  • tensorium_mlir::runtime::GeneratedHostStorage provides the same uniform-grid storage plan from generated C descriptor tables, so standalone runners can consume tensorium_host_kernels / tensorium_host_buffers without linking the compiler-side MLIR ABI builder.
  • Generated host headers also emit tensorium_host_kernel_adapters, a uniform invocation table for grid kernels. GeneratedHostStorage::invoke(...) uses those adapters to bind runtime-owned buffers by descriptor order instead of requiring callers to spell every lowered field argument manually.
  • constraints DSL blocks lower through the same scalarization path as RHS kernels but additionally expose tensorium_residual_grid_affine / tensorium_residual_grid_scf / tensorium_residual_grid_parallel host-callable symbols. These kernels compute residual buffers F(u); the solver runtime is responsible for choosing the update method that drives those residuals toward zero.
  • tensorium_rhs_grid_parallel and tensorium_residual_grid_parallel have the same low-level argument layout as the other field-grid kernels. Their MLIR body uses scf.parallel; LLVM emission lowers it through OpenMP runtime calls.
  • GeneratedHostStorage also exposes a descriptor-level Euler helper for standalone runtime experiments: eulerUpdatePairsFromDerivativePrefix() discovers writable derivative fields named dX and maps them to state field X, then applyEulerUpdate(...) performs X += dt * dX over the runtime arena. This is intentionally minimal and should be replaced by the target runtime's integrator once AMReX owns the storage.
  • Spectral initial-data experiments use tensorium_spectral_residual_point and tensorium_spectral_residual_kernel_fn as the pointwise residual callback ABI. The runtime constructs the derivative bundle on the selected spectral grid, applies an optional coordinate map, then calls the generated callback to compute one scalar F(u)=0 value per collocation point. This keeps NRPy/Kadath/TwoPunctures-style formulations above the generic spectral grid and solver machinery.
  • For constraints modules with spatial { scheme = spectral order = 0 }, the compiler emits tensorium_spectral_residual_<target> point kernels for scalar residuals. Every referenced scalar field receives its complete value, gradient, and Hessian bundle, so one residual may contain derivatives of several unknowns. Generated host headers expose these through tensorium_spectral_residual_kernels.
  • SpectralResidualProblem is the runtime-side assembly surface for these point kernels. It binds the grid, generated callback, scalar params, optional auxiliary fields, optional coordinate map, optional derivative map, and optional generated grid kernel. A derivative map transforms the complete first/second derivative bundle before either kernel path is evaluated; the TwoPunctures implementation uses it to expose Cartesian derivatives from the compact (A,B,phi) collocation domain. An independent SpectralUnknownMap may first transform the solver variable and its logical derivatives into the physical unknown consumed by the generated residual. makeLinearBoundaryFactorUnknownMap() implements the generic product u=scale*(q_axis-boundary)*v; the two-puncture runtime uses it for U=(A-1)v without embedding that representation in the DSL equation or nonlinear solver. assembleSpectralResidual(...) computes the global collocation vector F(u) plus L2/max norms. When SpectralResidualProblem::gridKernel is set from tensorium_spectral_residual_grid_kernels, assembly uses the generated MLIR/LLVM global kernel; otherwise it falls back to the pointwise callback loop. For ABI v3 kernels, evaluateSpectralJacobianVectorProduct(...) calls the generated forward-mode JVP kernel. Manually supplied kernels that leave that callback null retain the centered finite-difference path as a fallback.
  • solveSpectralNewton(...) is the first scalar spectral elliptic solve path. It uses the common JVP interface to assemble a dense Jacobian, solves the dense Newton system with pivoting, and performs a damped residual-decreasing line search. In SpectralLinearSolveKind::Auto, small problems use the dense Jacobian path and grids above denseJacobianMaxUnknowns use matrix-free GMRES over the same JVP interface. GMRES supports right preconditioning through Jacobi JVP diagonals, a dense collocation inverse-Laplacian oracle, and a modal laplacian + shift inverse approximation intended as the scalable runtime path.
  • SpectralResidualSystemProblem is the multi-residual assembly surface. It combines several scalar SpectralResidualProblem equations into one equation-major vector [F0(points), F1(points), ...]. Each equation has one primary unknown and may use the complete derivative bundles of other current unknowns through auxiliary mappings. The first differentiated field encountered in the residual expression is the primary unknown; formulations should therefore place their principal operator first. Generated host headers expose tensorium_spectral_residual_systems[] with the constraint-block name, field-major unknown names, per-equation residual names, unknown indices, point/grid kernel indices, parameter names, auxiliary field names, and auxiliary-to-unknown maps (-1 means static auxiliary). evaluateSpectralResidualSystemJacobianVectorProduct(...) returns equation-major Jv, resolving both primal and tangent auxiliary-unknown maps so coupled residuals receive the current direction. It uses generated JVP kernels only when every equation exposes one, otherwise it falls back to a consistent finite-difference perturbation of the whole field bundle. solveSpectralNewton(...) also accepts a system problem plus a mutable field-major unknown bundle; it solves the square multi-residual system with either dense JVP assembly or matrix-free GMRES over the same system JVP. SpectralPreconditionerKind::DiagonalJVP adds a Jacobi smoke preconditioner by estimating the residual/Jacobian diagonal from JVP columns. SpectralPreconditionerKind::DenseLaplacianShift applies a right preconditioner by solving dense collocation systems for laplacian + shift, giving the first runtime inverse-Laplacian path on the existing spectral derivative matrices. It is kept as a small-grid oracle. SpectralPreconditionerKind::ModalLaplacianShift applies the same laplacian + shift idea in spectral coefficient space using Chebyshev/Fourier modal transforms. For the current Chebyshev/Chebyshev/Fourier layout it solves dense Chebyshev 2D modal blocks per Fourier mode and per field; current nonlinear and system runtime solve tests use this path. SpectralPreconditionerKind::MappedFiniteDifferenceLaplacianShift instead builds at most seven mapped finite-difference coefficients per collocation row and approximately inverts them with symmetric relaxation. The physical TwoPunctures regression combines this O(N) preconditioner with bounded- memory restarted GMRES. SpectralPreconditionerKind::MappedFiniteDifferenceMultigrid extends the same mapped sparse operator with a recursive geometric hierarchy. Restriction and prolongation use tensor-product Chebyshev/Fourier interpolation; exact Galerkin coarsening is retained for small two-level problems, while deeper hierarchies use stable mapped rediscretization at the terminal level. Dense LU is limited to the configured small terminal grid and cached across Krylov applications. At each Newton state, the point JVP kernel also evaluates the local reaction derivative dF/du in O(N) work (or estimates it with the finite-difference fallback).
  • tests/fixtures/elliptic/spectral_hamiltonian_toy_nonlinear_3d.tn is a manufactured nonlinear spectral constraint. Its runtime test solves laplacian(psi) + mass * psi + alpha * psi^5 + source = 0 from a non-exact local guess with matrix-free GMRES preconditioned by laplacian + mass, validating nonlinear Newton/JVP behavior against an analytic solution.
  • tests/fixtures/elliptic/spectral_lichnerowicz_manufactured_3d.tn is the first manufactured conformal Hamiltonian fixture: laplacian(u) + 1/8 * A2 * (background + u)^(-7) + source = 0. Its runtime test keeps A2 and source as auxiliary fields, solves from a non-exact guess with modal GMRES preconditioning, and checks against the analytic manufactured solution.
  • tools/dev/test_generated_two_puncture_solve_ll.sh exercises the first mapped-domain vertical slice. It compiles a DSL Poisson residual to LLVM, transforms the runtime derivative bundle through the compact two-puncture geometry, and solves a manufactured U=(A-1)v correction with Newton.
  • tests/fixtures/elliptic/spectral_two_puncture_hamiltonian_3d.tn contains a physical, generated two-centre Bowen-York/Lichnerowicz residual. The DSL computes both puncture distances, arbitrary momentum and spin contributions, their full symmetric-tensor contraction, the singular conformal background, and the nonlinear Hamiltonian residual. Its runtime test checks the closed single-puncture momentum/spin contractions, the exact Brill-Lindquist limit, and a nonlinear equal-and-opposite-momentum solve on the compact map.
  • SpectralAxis::interpolate(...) and SpectralGrid3D::interpolate(...) evaluate collocation fields at arbitrary logical probes with Chebyshev barycentric/Fourier tensor-product interpolation. The two-puncture regression uses the compactified endpoint value v(A=1,B=0,phi=0) for the ADM energy m1+m2-4*b*v_infinity and compares fixed probes across three resolutions.
  • TwoPunctureHandoff.h analytically inverts the compact map at arbitrary Cartesian target points, applies the configured spectral unknown transform, and writes conformally flat BSSN fields into caller-owned SoA buffers. The physical regression exercises finite puncture limits and rechecks the Hamiltonian and momentum constraints after interpolation.
  • The compiler also emits tensorium_spectral_residual_grid_<target> MLIR/LLVM kernels. These consume the runtime-computed spectral derivative buffers, auxiliary field buffers, coordinate buffers, scalar params, and one residual output buffer, then call the point kernel inside an MLIR scf.for loop. This moves global F(u) evaluation into generated code while keeping spectral differentiation in the runtime grid layer for now. Generated host headers expose uniform descriptors through tensorium_spectral_residual_grid_kernels.

Generated host headers also expose the same runtime contract in C-compatible tables:

  • TENSORIUM_HOST_KERNEL_COUNT / tensorium_host_kernels;
  • TENSORIUM_HOST_BUFFER_COUNT / tensorium_host_buffers;
  • TENSORIUM_HOST_KERNEL_ADAPTER_COUNT / tensorium_host_kernel_adapters;
  • tensorium_host_buffer_desc::component_count, role, access, and arg_index are enough for a lightweight runtime to deduplicate buffers and bind generated wrappers without reconstructing tensor sizes by hand.

Runtime code should use this contract instead of reconstructing argument order from names. For AMReX this means:

  • allocate one MultiFab component group per logical field or map component ranges according to componentCount;
  • allocate at least stencilRadius ghost cells for RHS kernels;
  • bind read/write buffers according to HostBufferABI::access;
  • reject descriptors with non-empty validateHostModuleABI diagnostics.

Generated grid loops should not perform per-point heap allocation. Scratch buffers used by fallback init-grid lowering are hoisted outside the generated loop and reused for every point. RHS old-state snapshots, when required by read/write overlap, are full-grid temporaries allocated once before the loop and released after it.

The development probe can print this contract for a fixture:

build/tools/runtime/Tensorium_abi_probe \
  tests/fixtures/gr/schwarzschild_bssn_constraints_analytic_3d.tn

Signature contracts

tensorium_init_point

MLIR-level:

  • (params..., coords..., alpha: memref<1xf64>, gamma: memref<9xf64>, gammaU: memref<9xf64>) -> ()

LLVM-level:

  • scalar f64 params/coords first,
  • then 3 memref descriptors (alpha, gamma, gammaU), each expanded to 5 C arguments.

tensorium_init_grid_{scf,affine}

MLIR-level:

  • (params..., coord_buffers..., alpha, gamma, gammaU) -> ()
  • coord/output buffers are memref<?xf64>.

LLVM-level:

  • scalar params first,
  • then one descriptor per coord/output buffer.

tensorium_rhs_grid_{scf,affine}

MLIR-level:

  • (nx:index, ny:index, nz:index, dx:f64, dy:f64, dz:f64, params..., fields...) -> ()
  • each field is memref<?xf64>.

LLVM-level:

  • prefix: i64,i64,i64,double,double,double,
  • then scalar params (double),
  • then one 5-argument memref descriptor per field buffer.

tensorium_residual_grid_{scf,affine}

Same low-level signature shape as tensorium_rhs_grid_{scf,affine}:

  • prefix: i64,i64,i64,double,double,double,
  • scalar params,
  • one 5-argument memref descriptor per participating field buffer.

The semantic difference is the ABI kind:

  • tensorium.abi.kind = "residual_grid_scf" or "residual_grid_affine".

Residual grid kernels write the declared residual targets from a constraints block. Host wrappers and descriptor tables expose those outputs in the same tensorium.abi.output_names and access metadata used by RHS kernels.

tensorium_{rhs,residual}_grid_parallel

Same low-level signature shape as tensorium_rhs_grid_affine:

  • prefix: i64,i64,i64,double,double,double,
  • scalar params,
  • one 5-argument memref descriptor per participating field buffer.

The semantic difference is the ABI kind:

  • tensorium.abi.kind = "rhs_grid_parallel" or "residual_grid_parallel".

The Tensorium MLIR body is a three-dimensional scf.parallel over the interior stencil domain. The LLVM lowering pipeline converts this path to OpenMP runtime calls, so executables that link the generated object need an OpenMP runtime.

tensorium_spectral_residual_<target>

MLIR-level:

  • (field0_value, field0_d1, ..., field0_d33, field1_value, field1_d1, ..., fieldN_d33, x1, x2, x3, params...) -> f64

Host callback-level:

  • generated headers define a tensorium_spectral_residual_kernel_desc entry;
  • the callback receives tensorium_spectral_residual_point, scalar params, and optional user data;
  • the first tensorium.abi.field_names entry is the primary unknown;
  • subsequent scalar fields remain available by value through point.aux_values[] for manual callback compatibility and expose their full bundles through point.aux_derivatives[] in the same field-name order;
  • point.physical[] is populated by the runtime coordinate map before the generated residual is called.

Residual-system descriptors can bind an auxiliary field to another current unknown. Assembly then propagates that unknown's unknown-map, projector, coordinate-derivative map, and complete derivative bundle into both point and grid kernels.

tensorium_spectral_residual_jvp_<target>

MLIR-level:

  • (field0_bundle, ..., fieldN_bundle, direction_field0_bundle, ..., direction_fieldN_bundle, x1, x2, x3, params...) -> f64, where every bundle is (value, d1, d2, d3, d11, d12, d13, d22, d23, d33).

Host callback-level:

  • the ABI v3 residual descriptor pairs evaluate_jvp and jvp_symbol_name with the primal callback;
  • the callback receives a primal tensorium_spectral_residual_point and a direction point with the same derivative/auxiliary layout;
  • parameters and coordinates have zero tangent, while auxiliary tangents are populated from coupled unknown directions when the system descriptor maps them to another unknown.

The compiler emits the primal and tangent MLIR together from the DSL residual. The forward-mode rules currently cover scalar +, -, *, /, sqrt, sin, first/second derivatives of any referenced scalar field, contractions that form a Laplacian, contracted scalar gradients, and laplacian(). Constants, parameters, and coordinates have zero tangent. Opaque external scalar calls without a registered derivative remain unsupported in this generated path. The runtime automatically falls back to centered finite differences when evaluate_jvp is null.

tensorium_spectral_residual_grid_<target>

MLIR-level:

  • (n_points:index, params..., field0_bundle..., fieldN_bundle..., x1, x2, x3, residual_out) -> ()
  • every bundle component, coordinate, and output buffer is memref<?xf64>.

Generated host wrapper:

  • tensorium_call_spectral_residual_grid_<target>(n_points, params..., field0_value, ..., fieldN_d33, x1, x2, x3, residual_out)
  • tensorium_spectral_residual_grid_kernels[] exposes the same generated kernel through a uniform callback consumed by SpectralResidualProblem; its auxiliary SoA bundles use tensorium_spectral_residual_derivative_fields.

This kernel does not compute spectral derivatives itself. The runtime supplies the derivative buffers from its selected spectral basis and coordinate mapping, then the generated MLIR loop evaluates the residual at every collocation point.

Memory layout contract (SoA, component-major)

Tensor field components are flattened as:

  • flat = component * nPoints + pointLinearIndex

where:

  • component is row-major over tensor indices,
  • pointLinearIndex is row-major over (x,y,z) grid index.

Examples:

  • covariant/convariant 2-tensor uses 9 components.
  • rank-3 tensor uses 27 components.

RHS read/write semantics

tensorium_rhs_grid_* snapshots all field buffers before stencil reads, then writes dt_assign targets to original field buffers. This guarantees read consistency within one RHS sweep and avoids write-after-read hazards.