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 decode_affinity;
33pub mod decode_numa;
34pub mod decode_spmd;
35pub mod dtype;
36pub mod kernels;
37#[cfg(feature = "mlas")]
38pub mod nchwc_layout;
39pub mod optimizer;
40pub mod provider;
41pub mod strided;
42mod trace;
43pub mod weight_offload;
44
45pub use backend::CpuBackend;
46pub use kernels::qmoe::WeightOffloadHostCache;
47pub use optimizer::{
48    ConvBatchNormActivationFusion, MatMulNBitsBiasFusion, ProjectionFusion, cpu_optimization_passes,
49};
50pub use provider::CpuExecutionProvider;
51pub use weight_offload::placement::{
52    ArbitrationAction, GpuLayersOverrideReport, HostFallbackReason, IqFormat, KvAdmissionDecision,
53    KvAdmissionLimitingFactor, LayerPlacement, LayerWeightRegions, Placement, PlacementError,
54    PlacementPlan, QuantTileFormat, RegionPlacement, SnappedTileSize, TileSizeError,
55    VramArbitrationConfig, VramArbitrationError, VramArbitrationOutcome, VramArbitrationState,
56    VramDemand, VramSubBudgets, arbitrate_vram, decide_kv_admission, plan_placement,
57    snap_transfer_tile_bytes,
58};
59pub use weight_offload::weight_handle::{
60    ExecutionProviderCapabilities, LazyDeviceWeightBinder, LazyWeight, LazyWeightBoundary,
61    NXRT_WEIGHT_PAGING_CAPABILITY, NegotiatedWeight, Phase3aHostOnlyBinder, ResidentWeight,
62    ResidentWeightMaterializer, WeightHandle, WeightHandleError,
63};
64pub use weight_offload::{
65    LinuxProcessMemoryStats, WEIGHT_OFFLOAD_ENV, WEIGHT_OFFLOAD_HOST_BYTES_ENV,
66    WeightOffloadLayerStats, WeightOffloadStats, set_weight_offload_host_budget,
67    weight_offload_stats,
68};
69
70pub use kernels::selection::non_max_suppression;
71pub use kernels::slice::{SliceAxisPlan, slice_axes_steps, slice_plan};
72
73pub use kernels::matmul_nbits::set_decode_thread_budget;
74pub use kernels::matmul_nbits::bound_process_to_decode_budget;
75pub use kernels::matmul_nbits::with_decode_pool_scope;
onnx_runtime_ep_cpu/lib.rs

onnx_runtime_ep_cpu/
lib.rs
