onnx_runtime_ep_cpu/
lib.rs

1//! # `onnx-runtime-ep-cpu`
2//!
3//! The CPU execution provider for the ORT 2.0 runtime (see `docs/ORT2.md` §4.4
4//! and §54 Phase 1). It implements [`onnx_runtime_ep_api::ExecutionProvider`]
5//! and hosts pure-Rust reference kernels for the Phase-1 op set (`MatMul`,
6//! `Add`, `Relu`, `Reshape`, `Transpose`, `Gather`, `LayerNormalization`).
7//!
8//! ## Backends: correctness baseline + SIMD fast path
9//!
10//! The GEMM hot spot is served through [`backend::CpuBackend`] (`docs/ORT2.md`
11//! §25.2). The **default** backend is a pure-Rust blocked, register-tiled,
12//! rayon-parallelized f32 GEMM — the portable, offline correctness baseline that
13//! compiles anywhere with no C++/FFI. On supported x86 hosts, the built-in
14//! `SimdX86` implementation provides the default fast path. Every backend lives behind the
15//! [`onnx_runtime_ep_api::Kernel`] trait, so neither the EP contract nor the
16//! session observes which one ran. See [`kernels::matmul`] for the hot spot.
17//!
18//! ## `unsafe`
19//!
20//! The default (Generic) path is `unsafe`-minimal: the only `unsafe` is the raw
21//! device-buffer access the ep-api contract forces (aligned host
22//! `alloc`/`dealloc`, `memcpy`, and strided element reads/writes), each isolated
23//! and `SAFETY`-documented. The blocked rayon GEMM itself contains no `unsafe`;
24//! all kernel arithmetic is safe Rust operating on
25//! dense `Vec<f32>` buffers produced by the two audited accessors in [`kernels`].
26
27// Kernel entry points mirror ONNX operator schemas, whose independent tensors and
28// dimensions often exceed Clippy's generic argument-count threshold.
29#![allow(clippy::too_many_arguments)]
30
31pub mod backend;
32pub mod dtype;
33pub mod kernels;
34pub mod optimizer;
35pub mod provider;
36pub mod strided;
37mod trace;
38pub mod weight_offload;
39
40pub use backend::CpuBackend;
41pub use kernels::qmoe::WeightOffloadHostCache;
42pub use optimizer::{ProjectionFusion, cpu_optimization_passes};
43pub use provider::CpuExecutionProvider;
44pub use weight_offload::placement::{
45    ArbitrationAction, GpuLayersOverrideReport, HostFallbackReason, IqFormat, KvAdmissionDecision,
46    KvAdmissionLimitingFactor, LayerPlacement, LayerWeightRegions, Placement, PlacementError,
47    PlacementPlan, QuantTileFormat, RegionPlacement, SnappedTileSize, TileSizeError,
48    VramArbitrationConfig, VramArbitrationError, VramArbitrationOutcome, VramArbitrationState,
49    VramDemand, VramSubBudgets, arbitrate_vram, decide_kv_admission, plan_placement,
50    snap_transfer_tile_bytes,
51};
52pub use weight_offload::weight_handle::{
53    ExecutionProviderCapabilities, LazyDeviceWeightBinder, LazyWeight, LazyWeightBoundary,
54    NXRT_WEIGHT_PAGING_CAPABILITY, NegotiatedWeight, Phase3aHostOnlyBinder, ResidentWeight,
55    ResidentWeightMaterializer, WeightHandle, WeightHandleError,
56};
57pub use weight_offload::{
58    LinuxProcessMemoryStats, WEIGHT_OFFLOAD_ENV, WEIGHT_OFFLOAD_HOST_BYTES_ENV,
59    WeightOffloadLayerStats, WeightOffloadStats, set_weight_offload_host_budget,
60    weight_offload_stats,
61};
62
63pub use kernels::slice::{SliceAxisPlan, slice_axes_steps, slice_plan};
64
65pub use kernels::matmul_nbits::with_decode_pool_scope;
onnx_runtime_ep_cpu/lib.rs

onnx_runtime_ep_cpu/
lib.rs
