tenferro_xla/
error.rs

1use std::error::Error as StdError;
2use std::path::PathBuf;
3
4use tenferro_tensor::{DType, ErrorKind, ValidationKind};
5
6/// Erased typed source used only at the XLA plugin boundary.
7pub type BoxError = Box<dyn StdError + Send + Sync + 'static>;
8
9/// Error type for StableHLO lowering and runtime PJRT plugin loading.
10///
11/// # Examples
12///
13/// ```
14/// use tenferro_runtime::{DType, GraphCompiler, TracedTensor};
15/// use tenferro_xla::{lower_compiled_to_stablehlo, Error};
16///
17/// let x = TracedTensor::input_symbolic_shape(DType::I64, 1).unwrap();
18/// let mut compiler = GraphCompiler::new();
19/// let y = x.neg().unwrap();
20/// let program = compiler
21///     .compile_with_input_specs(&y, &[(&x, DType::I64, &[2])])
22///     .unwrap();
23/// let err = lower_compiled_to_stablehlo(&program).unwrap_err();
24/// assert!(matches!(err, Error::UnsupportedDType { .. }));
25/// ```
26#[derive(Debug, thiserror::Error)]
27pub enum Error {
28    #[error("XLA lowering does not support dtype {dtype:?} in {context}")]
29    UnsupportedDType { dtype: DType, context: &'static str },
30    #[error("XLA lowering does not support ExecOp::{op}: {reason}")]
31    UnsupportedOp {
32        op: &'static str,
33        reason: &'static str,
34    },
35    #[error(
36        "XLA lowering supports only exact static shapes; ExecOp::{op} output {output_index} axis {axis} is {kind}"
37    )]
38    NonStaticShape {
39        op: &'static str,
40        output_index: usize,
41        axis: usize,
42        kind: &'static str,
43    },
44    #[error("invalid XLA program: {message}")]
45    InvalidProgram { message: String },
46    #[error("XLA tensor input/output error: {0}")]
47    Tensor(#[from] tenferro_tensor::Error),
48    #[error("XLA extension standard-op lowering failed: {source}")]
49    ExtensionLowering {
50        #[source]
51        source: tenferro_ops::ext_op::ExtensionLoweringError,
52    },
53    #[error("PJRT support requires enabling the tenferro-xla `pjrt` feature")]
54    PjrtFeatureDisabled,
55    #[error("PJRT execution requires an executor created from a loaded plugin")]
56    PjrtPluginNotLoaded,
57    #[error("PJRT call {call} failed: {message}")]
58    PjrtCall { call: &'static str, message: String },
59    #[error("environment variable {var} is not set; set it to a PJRT plugin .so path")]
60    MissingEnv { var: &'static str },
61    #[error("failed to load PJRT plugin from {path}: {source}")]
62    PluginLoad {
63        path: PathBuf,
64        #[source]
65        source: BoxError,
66    },
67}
68
69impl Error {
70    /// Return the stable coarse classification for this XLA failure.
71    ///
72    /// # Examples
73    ///
74    /// ```
75    /// use tenferro_tensor::ErrorKind;
76    /// use tenferro_xla::Error;
77    ///
78    /// assert_eq!(
79    ///     Error::UnsupportedOp { op: "Maximum", reason: "not lowered" }.kind(),
80    ///     ErrorKind::Unsupported,
81    /// );
82    /// ```
83    #[must_use]
84    pub fn kind(&self) -> ErrorKind {
85        match self {
86            Self::UnsupportedDType { .. } | Self::UnsupportedOp { .. } => ErrorKind::Unsupported,
87            Self::NonStaticShape { .. } => ErrorKind::Validation(ValidationKind::ShapeMismatch),
88            Self::InvalidProgram { .. } => ErrorKind::Validation(ValidationKind::InvalidArgument),
89            Self::Tensor(source) => source.kind(),
90            Self::ExtensionLowering { source } => source.kind(),
91            Self::PjrtFeatureDisabled | Self::PjrtPluginNotLoaded | Self::MissingEnv { .. } => {
92                ErrorKind::RuntimeState
93            }
94            Self::PluginLoad { .. } => ErrorKind::Io,
95            Self::PjrtCall { .. } => ErrorKind::BackendFailure,
96        }
97    }
98}
99
100/// Result alias for `tenferro-xla`.
101///
102/// # Examples
103///
104/// ```
105/// use tenferro_xla::Result;
106///
107/// fn ok() -> Result<()> {
108///     Ok(())
109/// }
110///
111/// ok().unwrap();
112/// ```
113pub type Result<T> = std::result::Result<T, Error>;
114
115#[cfg(test)]
116mod tests {
117    use std::error::Error as StdError;
118    use std::path::PathBuf;
119
120    use tenferro_ops::ext_op::ExtensionLoweringError;
121    use tenferro_tensor::{Error as TensorError, ErrorKind, ValidationKind};
122
123    use super::Error;
124
125    #[test]
126    fn xla_error_kind_distinguishes_capability_state_io_and_backend_failures() {
127        let cases = [
128            (
129                Error::UnsupportedDType {
130                    dtype: tenferro_tensor::DType::I64,
131                    context: "constant",
132                },
133                ErrorKind::Unsupported,
134            ),
135            (
136                Error::UnsupportedOp {
137                    op: "Custom",
138                    reason: "not lowered",
139                },
140                ErrorKind::Unsupported,
141            ),
142            (
143                Error::NonStaticShape {
144                    op: "Reshape",
145                    output_index: 0,
146                    axis: 1,
147                    kind: "symbolic",
148                },
149                ErrorKind::Validation(ValidationKind::ShapeMismatch),
150            ),
151            (
152                Error::InvalidProgram {
153                    message: "missing output".into(),
154                },
155                ErrorKind::Validation(ValidationKind::InvalidArgument),
156            ),
157            (
158                Error::Tensor(TensorError::invalid_argument("xla", "input", "invalid")),
159                ErrorKind::Validation(ValidationKind::InvalidArgument),
160            ),
161            (
162                Error::ExtensionLowering {
163                    source: ExtensionLoweringError::new_with_kind(
164                        ErrorKind::Unsupported,
165                        "cannot lower",
166                    ),
167                },
168                ErrorKind::Unsupported,
169            ),
170            (Error::PjrtFeatureDisabled, ErrorKind::RuntimeState),
171            (Error::PjrtPluginNotLoaded, ErrorKind::RuntimeState),
172            (
173                Error::MissingEnv { var: "PJRT_PLUGIN" },
174                ErrorKind::RuntimeState,
175            ),
176            (
177                Error::PluginLoad {
178                    path: PathBuf::from("plugin.so"),
179                    source: Box::new(std::io::Error::other("not found")),
180                },
181                ErrorKind::Io,
182            ),
183            (
184                Error::PjrtCall {
185                    call: "pjrt_execute",
186                    message: "invalid status".into(),
187                },
188                ErrorKind::BackendFailure,
189            ),
190        ];
191
192        for (error, expected) in cases {
193            assert_eq!(error.kind(), expected, "classified {error:?}");
194        }
195    }
196
197    #[test]
198    fn xla_error_sources_remain_typed_at_boundary() {
199        let tensor = Error::Tensor(TensorError::backend_source(
200            "xla_input",
201            std::io::Error::other("device read failed"),
202        ));
203        assert!(StdError::source(&tensor).is_some());
204
205        let lowering = Error::ExtensionLowering {
206            source: ExtensionLoweringError::from_source_with_kind(
207                ErrorKind::BackendFailure,
208                std::io::Error::other("shape source"),
209            ),
210        };
211        assert_eq!(lowering.kind(), ErrorKind::BackendFailure);
212        let source = StdError::source(&lowering).expect("lowering source should be retained");
213        let typed_source = source
214            .source()
215            .expect("typed lowering source should remain in the chain");
216        assert!(typed_source.downcast_ref::<std::io::Error>().is_some());
217
218        let plugin = Error::PluginLoad {
219            path: PathBuf::from("plugin.so"),
220            source: Box::new(std::io::Error::other("dlopen failed")),
221        };
222        assert!(StdError::source(&plugin).is_some());
223    }
224
225    #[test]
226    fn xla_extension_lowering_preserves_non_validation_kinds() {
227        for expected in [
228            ErrorKind::NumericalFailure,
229            ErrorKind::BackendFailure,
230            ErrorKind::RuntimeState,
231        ] {
232            let error = Error::ExtensionLowering {
233                source: ExtensionLoweringError::new_with_kind(expected, "typed failure"),
234            };
235            assert_eq!(error.kind(), expected);
236        }
237    }
238}
tenferro_xla/error.rs

tenferro_xla/
error.rs
