tenferro_xla/
executor.rs

1use std::fmt;
2
3use tenferro_runtime::GraphProgram;
4use tenferro_tensor::Tensor;
5
6use crate::Error;
7use crate::{lower_to_stablehlo, Result, StableHloModule};
8
9/// Options for the experimental XLA executor.
10///
11/// # Examples
12///
13/// ```
14/// use tenferro_xla::XlaExecutorOptions;
15///
16/// let options = XlaExecutorOptions::default();
17/// assert_eq!(options, XlaExecutorOptions::default());
18/// ```
19#[derive(Clone, Copy, Default, PartialEq, Eq)]
20#[non_exhaustive]
21pub struct XlaExecutorOptions {
22    _private: (),
23}
24
25impl fmt::Debug for XlaExecutorOptions {
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        f.debug_struct("XlaExecutorOptions").finish()
28    }
29}
30
31/// Experimental peer executor for XLA/PJRT.
32///
33/// # Examples
34///
35/// ```
36/// use tenferro_runtime::{GraphCompiler, TracedTensor};
37/// use tenferro_xla::XlaExecutor;
38///
39/// let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
40/// let mut compiler = GraphCompiler::new();
41/// let y = x.neg().unwrap();
42/// let program = compiler.compile(&y).unwrap();
43/// let module = XlaExecutor::default().lower_to_stablehlo(&program).unwrap();
44/// assert!(module.as_str().contains("stablehlo.negate"));
45/// ```
46pub struct XlaExecutor {
47    options: XlaExecutorOptions,
48    #[cfg(feature = "pjrt")]
49    plugin: Option<crate::pjrt::PjrtPlugin>,
50}
51
52impl fmt::Debug for XlaExecutor {
53    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
54        f.debug_struct("XlaExecutor")
55            .field("options", &self.options)
56            .field("has_loaded_pjrt_plugin", &self.has_loaded_pjrt_plugin())
57            .finish()
58    }
59}
60
61impl XlaExecutor {
62    /// Create an executor with explicit options.
63    ///
64    /// # Examples
65    ///
66    /// ```
67    /// use tenferro_xla::{XlaExecutor, XlaExecutorOptions};
68    ///
69    /// let executor = XlaExecutor::new(XlaExecutorOptions::default());
70    /// assert_eq!(executor.options(), XlaExecutorOptions::default());
71    /// ```
72    pub fn new(options: XlaExecutorOptions) -> Self {
73        Self {
74            options,
75            #[cfg(feature = "pjrt")]
76            plugin: None,
77        }
78    }
79
80    /// Create an executor by loading PJRT configuration from environment variables.
81    ///
82    /// Without the `pjrt` feature this returns [`Error::PjrtFeatureDisabled`].
83    ///
84    /// # Examples
85    ///
86    /// ```
87    /// use tenferro_xla::{Error, XlaExecutor};
88    ///
89    /// if let Err(err) = XlaExecutor::from_env() {
90    ///     assert!(matches!(err, Error::PjrtFeatureDisabled | Error::MissingEnv { .. } | Error::PluginLoad { .. }));
91    /// }
92    /// ```
93    #[cfg(not(feature = "pjrt"))]
94    pub fn from_env() -> Result<Self> {
95        Err(Error::PjrtFeatureDisabled)
96    }
97
98    /// Create an executor by loading PJRT configuration from environment variables.
99    ///
100    /// # Examples
101    ///
102    /// ```
103    /// use tenferro_xla::XlaExecutor;
104    ///
105    /// let _ = XlaExecutor::from_env();
106    /// ```
107    #[cfg(feature = "pjrt")]
108    pub fn from_env() -> Result<Self> {
109        Self::from_env_var(crate::TENFERRO_PJRT_PLUGIN_ENV)
110    }
111
112    /// Create an executor by loading a PJRT plugin path from a specific
113    /// environment variable.
114    ///
115    /// # Examples
116    ///
117    /// ```
118    /// use tenferro_xla::XlaExecutor;
119    ///
120    /// let _ = XlaExecutor::from_env_var("__TENFERRO_XLA_DOCS_UNSET");
121    /// ```
122    #[cfg(feature = "pjrt")]
123    pub fn from_env_var(var: &'static str) -> Result<Self> {
124        let plugin = crate::pjrt::PjrtPlugin::load_from_env(var)?;
125        Ok(Self {
126            options: XlaExecutorOptions::default(),
127            plugin: Some(plugin),
128        })
129    }
130
131    /// Return the executor options.
132    ///
133    /// # Examples
134    ///
135    /// ```
136    /// use tenferro_xla::XlaExecutor;
137    ///
138    /// assert_eq!(XlaExecutor::default().options(), Default::default());
139    /// ```
140    pub fn options(&self) -> XlaExecutorOptions {
141        self.options
142    }
143
144    /// Return whether this executor owns a loaded PJRT plugin.
145    ///
146    /// # Examples
147    ///
148    /// ```
149    /// use tenferro_xla::XlaExecutor;
150    ///
151    /// assert!(!XlaExecutor::default().has_loaded_pjrt_plugin());
152    /// ```
153    pub fn has_loaded_pjrt_plugin(&self) -> bool {
154        #[cfg(feature = "pjrt")]
155        {
156            self.plugin.is_some()
157        }
158        #[cfg(not(feature = "pjrt"))]
159        {
160            false
161        }
162    }
163
164    /// Lower a graph program to StableHLO without executing it.
165    ///
166    /// # Examples
167    ///
168    /// ```
169    /// use tenferro_runtime::{GraphCompiler, TracedTensor};
170    /// use tenferro_xla::XlaExecutor;
171    ///
172    /// let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
173    /// let mut compiler = GraphCompiler::new();
174    /// let y = x.neg().unwrap();
175    /// let program = compiler.compile(&y).unwrap();
176    /// let module = XlaExecutor::default().lower_to_stablehlo(&program).unwrap();
177    /// assert!(module.as_str().contains("stablehlo.negate"));
178    /// ```
179    pub fn lower_to_stablehlo(&self, program: &GraphProgram) -> Result<StableHloModule> {
180        lower_to_stablehlo(program)
181    }
182
183    /// Execute a graph program through a loaded PJRT plugin and return all outputs.
184    ///
185    /// Inputs must match [`GraphProgram::input_specs`] exactly. This
186    /// experimental execution path supports the same exact-static-shape,
187    /// `F32`/`F64` subset as StableHLO lowering.
188    ///
189    /// # Examples
190    ///
191    /// ```
192    /// use tenferro_runtime::{GraphCompiler, TracedTensor};
193    /// use tenferro_tensor::Tensor;
194    /// use tenferro_xla::{Error, XlaExecutor};
195    ///
196    /// let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
197    /// let mut compiler = GraphCompiler::new();
198    /// let y = x.neg().unwrap();
199    /// let program = compiler.compile(&y).unwrap();
200    /// let input = Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
201    /// let err = XlaExecutor::default()
202    ///     .run_many_with_inputs(&program, &[&input])
203    ///     .unwrap_err();
204    /// assert!(matches!(err, Error::PjrtFeatureDisabled | Error::PjrtPluginNotLoaded));
205    /// ```
206    pub fn run_many_with_inputs(
207        &self,
208        program: &GraphProgram,
209        inputs: &[&Tensor],
210    ) -> Result<Vec<Tensor>> {
211        #[cfg(feature = "pjrt")]
212        {
213            let Some(plugin) = self.plugin.as_ref() else {
214                return Err(Error::PjrtPluginNotLoaded);
215            };
216            crate::pjrt::run_many_with_inputs(plugin, program, inputs)
217        }
218        #[cfg(not(feature = "pjrt"))]
219        {
220            let _ = (program, inputs);
221            Err(Error::PjrtFeatureDisabled)
222        }
223    }
224
225    /// Execute a single-output graph program through a loaded PJRT plugin.
226    ///
227    /// # Examples
228    ///
229    /// ```
230    /// use tenferro_runtime::{GraphCompiler, TracedTensor};
231    /// use tenferro_tensor::Tensor;
232    /// use tenferro_xla::{Error, XlaExecutor};
233    ///
234    /// let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
235    /// let mut compiler = GraphCompiler::new();
236    /// let y = x.neg().unwrap();
237    /// let program = compiler.compile(&y).unwrap();
238    /// let input = Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
239    /// let err = XlaExecutor::default().run_with_inputs(&program, &[&input]).unwrap_err();
240    /// assert!(matches!(err, Error::PjrtFeatureDisabled | Error::PjrtPluginNotLoaded));
241    /// ```
242    pub fn run_with_inputs(&self, program: &GraphProgram, inputs: &[&Tensor]) -> Result<Tensor> {
243        single_output_tensor(self.run_many_with_inputs(program, inputs)?)
244    }
245}
246
247impl Default for XlaExecutor {
248    fn default() -> Self {
249        Self::new(XlaExecutorOptions::default())
250    }
251}
252
253fn single_output_tensor(mut outputs: Vec<Tensor>) -> Result<Tensor> {
254    if outputs.len() != 1 {
255        return Err(crate::Error::InvalidProgram {
256            message: format!(
257                "PJRT single-output execution expected 1 output, got {}",
258                outputs.len()
259            ),
260        });
261    }
262    Ok(outputs.remove(0))
263}
264
265#[cfg(test)]
266mod tests {
267    use super::{single_output_tensor, XlaExecutor, XlaExecutorOptions};
268    use crate::Error;
269    use tenferro_runtime::{GraphCompiler, TracedTensor};
270    use tenferro_tensor::Tensor;
271
272    #[test]
273    fn executor_options_and_debug_are_directly_covered() {
274        let options = XlaExecutorOptions::default();
275        let executor = XlaExecutor::new(options);
276
277        assert_eq!(executor.options(), options);
278        assert!(!executor.has_loaded_pjrt_plugin());
279        assert_eq!(format!("{options:?}"), "XlaExecutorOptions");
280        assert!(format!("{executor:?}").contains("has_loaded_pjrt_plugin"));
281    }
282
283    #[test]
284    fn default_executor_reports_missing_pjrt_before_dispatch() {
285        let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
286        let mut compiler = GraphCompiler::new();
287        let program = compiler.compile(&x.neg().unwrap()).unwrap();
288        let input = Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
289
290        let err = XlaExecutor::default()
291            .run_many_with_inputs(&program, &[&input])
292            .unwrap_err();
293        assert!(matches!(
294            err,
295            Error::PjrtFeatureDisabled | Error::PjrtPluginNotLoaded
296        ));
297        let err = XlaExecutor::default()
298            .run_with_inputs(&program, &[&input])
299            .unwrap_err();
300        assert!(matches!(
301            err,
302            Error::PjrtFeatureDisabled | Error::PjrtPluginNotLoaded
303        ));
304    }
305
306    #[test]
307    fn single_output_tensor_rejects_zero_or_multiple_outputs() {
308        let tensor = Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap();
309        assert_eq!(
310            single_output_tensor(vec![tensor.clone()])
311                .unwrap()
312                .as_slice::<f64>()
313                .unwrap(),
314            &[1.0]
315        );
316
317        let err = single_output_tensor(Vec::new()).unwrap_err();
318        assert!(
319            err.to_string().contains("got 0"),
320            "expected zero-output error, got {err:?}"
321        );
322
323        let err = single_output_tensor(vec![tensor.clone(), tensor]).unwrap_err();
324        assert!(
325            err.to_string().contains("got 2"),
326            "expected multi-output error, got {err:?}"
327        );
328    }
329}
tenferro_xla/executor.rs

tenferro_xla/
executor.rs
