onnx-runtime-ep-api 0.1.0-dev.5

Execution Provider API for the ORT 2.0 runtime: ExecutionProvider and Kernel traits, OpRegistry, and the ORT graph ABI bridge
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

//! # `onnx-runtime-ep-api`
//!
//! The Execution Provider (EP) interface for the ORT 2.0 runtime
//! (see `docs/ORT2.md` §4). Every backend — CPU, CUDA, MLX, or a legacy ORT
//! plugin loaded via `dlopen` — implements the same [`ExecutionProvider`]
//! trait; only the loading mechanism differs.
//!
//! This is a **Phase 1 skeleton**: trait and type *signatures* are real and
//! reference the [`onnx_runtime_ir`] contract, but method bodies live in the
//! concrete EP crates (e.g. `onnx-runtime-ep-cpu`). Zero-copy tensor views and
//! the ORT ABI bridge require `unsafe` FFI and are stubbed here.
//!
//! ## Modules
//! * [`provider`] — [`ExecutionProvider`], [`EpConfig`], device buffers/fences.
//! * [`kernel`] — [`Kernel`] trait, [`KernelMatch`], [`Cost`].
//! * [`registry`] — [`OpRegistry`], [`OpKey`], [`KernelFactory`], [`EpRegistry`].
//! * [`epcontext`] — [`EpContext`] runtime form + `source`-keyed [`EpContextRegistry`] (§55).
//! * [`tensor`] — [`TensorView`] / [`TensorMut`] zero-copy device views.
//! * [`weight`] — capability-negotiated lazy [`WeightHandle`] delivery.
//! * [`abi`] — ORT graph ABI bridge for legacy plugin EPs (Phase 2).

pub mod abi;
pub mod epcontext;
pub mod kernel;
pub mod provider;
pub mod registry;
pub mod tensor;
pub mod weight;

pub use epcontext::{EpContext, EpContextRegistry, build_ep_context_registry};
pub use error::{EpError, Result};
pub use kernel::{
    ARG_BYTES, ARG_DEVICE, ARG_FLOPS, ARG_KERNEL_VARIANT, ARG_KERNEL_VARIANT_REASON,
    CAT_KERNEL_WORKER, CaptureSupport, Cost, Kernel, KernelInput, KernelMatch,
    KernelVariantSelection, ViewOutput, kernel_variant_tracing_enabled, kernel_worker_span,
    record_kernel_metrics, record_kernel_variant_selection, record_kernel_variant_stage_selection,
};
pub use onnx_runtime_optimizer::OptimizationPass as OptimizerPass;
pub use provider::{
    CaptureRegionShapeStatus, DeviceBuffer, EpConfig, EpId, ExecutionProvider, Fence,
    OrtPluginExport, StructuralCaptureDecline,
};
pub use registry::{EpRegistry, KernelFactory, OpKey, OpRegistry};
pub use tensor::{
    DevicePtr, DevicePtrMut, ExternalMmapRegion, TensorBacking, TensorMut, TensorView,
};
pub use weight::{
    ExecutionProviderCapabilities, LazyDeviceWeightBinder, LazyWeight, LazyWeightBoundary,
    NXRT_WEIGHT_PAGING_CAPABILITY, NegotiatedWeight, Phase3aHostOnlyBinder, ResidentWeight,
    ResidentWeightMaterializer, WeightHandle, WeightHandleError,
};

// Re-export the device vocabulary from the IR so EP authors have one import.
pub use onnx_runtime_ir::{DeviceId, DeviceType};

mod error {
    use std::path::PathBuf;

    /// A convenience `Result` alias for EP operations.
    pub type Result<T> = std::result::Result<T, EpError>;

    /// Errors produced by execution providers and kernels (subset of the
    /// runtime top-level `Error`, `docs/ORT2.md` §22).
    #[derive(Debug, thiserror::Error)]
    pub enum EpError {
        #[error("no EP handler for {domain}::{op_type} at opset {opset} on any available device")]
        NoEpForOp {
            domain: String,
            op_type: String,
            opset: u64,
        },

        #[error("kernel execution failed: {0}")]
        KernelFailed(String),

        #[error("EP panicked during execution")]
        EpPanicked,

        #[error("EP plugin load failed: {path}: {reason}")]
        EpLoadFailed { path: PathBuf, reason: String },

        #[error("device OOM: requested {requested} bytes, available {available}")]
        OutOfMemory { requested: usize, available: usize },

        #[error("allocation alignment mismatch")]
        AlignmentError,

        #[error("invalid tensor view: {reason}")]
        InvalidTensorView { reason: String },

        #[error("EP not initialized")]
        NotInitialized,

        #[error("no EP is registered for EPContext source key {source_key:?}")]
        NoEpForContext { source_key: Option<String> },

        #[error("EP {ep} does not support EPContext save/load")]
        UnsupportedContext { ep: String },

        #[error(
            "duplicate EPContext source key {source_key:?}: already registered to {existing:?}, \
             cannot re-register to {new:?}"
        )]
        DuplicateContextSource {
            source_key: String,
            existing: crate::provider::EpId,
            new: crate::provider::EpId,
        },

        #[error(transparent)]
        Ir(#[from] onnx_runtime_ir::IrError),

        #[error(transparent)]
        WeightHandle(#[from] crate::weight::WeightHandleError),
    }
}