tenferro_xla/
executor.rs

1use std::fmt;
2
3use tenferro_runtime::program::SemanticProgram;
4use tenferro_runtime::CompiledGraph;
5use tenferro_tensor::Tensor;
6
7use crate::Error;
8use crate::{lower_to_stablehlo, Result, StableHloModule};
9
10/// Options for the experimental XLA executor.
11///
12/// # Examples
13///
14/// ```
15/// use tenferro_xla::XlaExecutorOptions;
16///
17/// let options = XlaExecutorOptions::default();
18/// assert_eq!(options, XlaExecutorOptions::default());
19/// ```
20#[derive(Clone, Copy, Default, PartialEq, Eq)]
21#[non_exhaustive]
22pub struct XlaExecutorOptions {
23    _private: (),
24}
25
26impl fmt::Debug for XlaExecutorOptions {
27    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
28        f.debug_struct("XlaExecutorOptions").finish()
29    }
30}
31
32/// Experimental peer executor for XLA/PJRT.
33///
34/// # Examples
35///
36/// ```
37/// use tenferro_runtime::{GraphCompiler, TracedTensor};
38/// use tenferro_xla::XlaExecutor;
39///
40/// let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
41/// let mut compiler = GraphCompiler::new();
42/// let y = x.neg().unwrap();
43/// let program = compiler.compile(&y).unwrap();
44/// let module = XlaExecutor::default()
45///     .lower_compiled_to_stablehlo(&program)
46///     .unwrap();
47/// assert!(module.as_str().contains("stablehlo.negate"));
48/// ```
49pub struct XlaExecutor {
50    options: XlaExecutorOptions,
51    #[cfg(feature = "pjrt")]
52    plugin: Option<crate::pjrt::PjrtPlugin>,
53}
54
55impl fmt::Debug for XlaExecutor {
56    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
57        f.debug_struct("XlaExecutor")
58            .field("options", &self.options)
59            .field("has_loaded_pjrt_plugin", &self.has_loaded_pjrt_plugin())
60            .finish()
61    }
62}
63
64impl XlaExecutor {
65    /// Create an executor with explicit options.
66    ///
67    /// # Examples
68    ///
69    /// ```
70    /// use tenferro_xla::{XlaExecutor, XlaExecutorOptions};
71    ///
72    /// let executor = XlaExecutor::new(XlaExecutorOptions::default());
73    /// assert_eq!(executor.options(), XlaExecutorOptions::default());
74    /// ```
75    pub fn new(options: XlaExecutorOptions) -> Self {
76        Self {
77            options,
78            #[cfg(feature = "pjrt")]
79            plugin: None,
80        }
81    }
82
83    /// Create an executor by loading PJRT configuration from environment variables.
84    ///
85    /// Without the `pjrt` feature this returns [`Error::PjrtFeatureDisabled`].
86    ///
87    /// # Examples
88    ///
89    /// ```
90    /// use tenferro_xla::{Error, XlaExecutor};
91    ///
92    /// if let Err(err) = XlaExecutor::from_env() {
93    ///     assert!(matches!(err, Error::PjrtFeatureDisabled | Error::MissingEnv { .. } | Error::PluginLoad { .. }));
94    /// }
95    /// ```
96    ///
97    /// # Errors
98    ///
99    /// Without `pjrt`, returns `Error::PjrtFeatureDisabled`. With `pjrt`,
100    /// returns `Error::MissingEnv` for an unset plugin path or
101    /// `Error::PluginLoad` with the typed dynamic-library source.
102    #[cfg(not(feature = "pjrt"))]
103    pub fn from_env() -> Result<Self> {
104        Err(Error::PjrtFeatureDisabled)
105    }
106
107    /// Create an executor by loading PJRT configuration from environment variables.
108    ///
109    /// # Examples
110    ///
111    /// ```
112    /// use tenferro_xla::XlaExecutor;
113    ///
114    /// let _ = XlaExecutor::from_env();
115    /// ```
116    ///
117    /// # Errors
118    ///
119    /// Returns `Error::MissingEnv` when the configured plugin-path variable is
120    /// unset, or `Error::PluginLoad` with the typed dynamic-library source
121    /// when the path cannot be loaded.
122    #[cfg(feature = "pjrt")]
123    pub fn from_env() -> Result<Self> {
124        Self::from_env_var(crate::TENFERRO_PJRT_PLUGIN_ENV)
125    }
126
127    /// Create an executor by loading a PJRT plugin path from a specific
128    /// environment variable.
129    ///
130    /// # Examples
131    ///
132    /// ```
133    /// use tenferro_xla::XlaExecutor;
134    ///
135    /// let _ = XlaExecutor::from_env_var("__TENFERRO_XLA_DOCS_UNSET");
136    /// ```
137    ///
138    /// # Errors
139    ///
140    /// Returns `Error::MissingEnv` when `var` is unset, or `Error::PluginLoad`
141    /// with the typed dynamic-library source when its value cannot be loaded.
142    #[cfg(feature = "pjrt")]
143    pub fn from_env_var(var: &'static str) -> Result<Self> {
144        let plugin = crate::pjrt::PjrtPlugin::load_from_env(var)?;
145        Ok(Self {
146            options: XlaExecutorOptions::default(),
147            plugin: Some(plugin),
148        })
149    }
150
151    /// Return the executor options.
152    ///
153    /// # Examples
154    ///
155    /// ```
156    /// use tenferro_xla::XlaExecutor;
157    ///
158    /// assert_eq!(XlaExecutor::default().options(), Default::default());
159    /// ```
160    pub fn options(&self) -> XlaExecutorOptions {
161        self.options
162    }
163
164    /// Return whether this executor owns a loaded PJRT plugin.
165    ///
166    /// # Examples
167    ///
168    /// ```
169    /// use tenferro_xla::XlaExecutor;
170    ///
171    /// assert!(!XlaExecutor::default().has_loaded_pjrt_plugin());
172    /// ```
173    pub fn has_loaded_pjrt_plugin(&self) -> bool {
174        #[cfg(feature = "pjrt")]
175        {
176            self.plugin.is_some()
177        }
178        #[cfg(not(feature = "pjrt"))]
179        {
180            false
181        }
182    }
183
184    /// Lower a graph program to StableHLO without executing it.
185    ///
186    /// # Examples
187    ///
188    /// ```
189    /// use tenferro_runtime::{GraphCompiler, TracedTensor};
190    /// use tenferro_xla::XlaExecutor;
191    ///
192    /// let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
193    /// let mut compiler = GraphCompiler::new();
194    /// let y = x.neg().unwrap();
195    /// let program = compiler.compile(&y).unwrap();
196    /// let module = XlaExecutor::default().lower_to_stablehlo(program.program()).unwrap();
197    /// assert!(module.as_str().contains("stablehlo.negate"));
198    /// ```
199    ///
200    /// # Errors
201    ///
202    /// Returns `Error::UnsupportedDType`, `Error::UnsupportedOp`, or
203    /// `Error::NonStaticShape` for unsupported graph content, and
204    /// `Error::InvalidProgram` for inconsistent graph metadata.
205    pub fn lower_to_stablehlo(&self, program: &SemanticProgram) -> Result<StableHloModule> {
206        lower_to_stablehlo(program)
207    }
208
209    /// Lower a compiled graph to StableHLO without executing it.
210    ///
211    /// This is the preferred public lowering boundary for graph users. It
212    /// consumes [`CompiledGraph`] directly and does not route through native
213    /// [`Runtime`](tenferro_runtime::Runtime) execution staging.
214    ///
215    /// # Examples
216    ///
217    /// ```
218    /// use tenferro_runtime::{GraphCompiler, TracedTensor};
219    /// use tenferro_xla::XlaExecutor;
220    ///
221    /// let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
222    /// let mut compiler = GraphCompiler::new();
223    /// let y = x.neg().unwrap();
224    /// let program = compiler.compile(&y).unwrap();
225    /// let module = XlaExecutor::default()
226    ///     .lower_compiled_to_stablehlo(&program)
227    ///     .unwrap();
228    /// assert!(module.as_str().contains("stablehlo.negate"));
229    /// ```
230    ///
231    /// # Errors
232    ///
233    /// Returns `Error::UnsupportedDType`, `Error::UnsupportedOp`, or
234    /// `Error::NonStaticShape` for unsupported graph content, and
235    /// `Error::InvalidProgram` for inconsistent graph metadata.
236    pub fn lower_compiled_to_stablehlo(&self, program: &CompiledGraph) -> Result<StableHloModule> {
237        crate::lower_compiled_to_stablehlo(program)
238    }
239
240    /// Execute a graph program through a loaded PJRT plugin and return all outputs.
241    ///
242    /// Inputs must match the ordered semantic-program input metadata exactly. This
243    /// experimental execution path supports the same exact-static-shape,
244    /// `F32`/`F64` subset as StableHLO lowering.
245    ///
246    /// # Examples
247    ///
248    /// ```
249    /// use tenferro_runtime::{GraphCompiler, TracedTensor};
250    /// use tenferro_tensor::Tensor;
251    /// use tenferro_xla::{Error, XlaExecutor};
252    ///
253    /// let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
254    /// let mut compiler = GraphCompiler::new();
255    /// let y = x.neg().unwrap();
256    /// let program = compiler.compile(&y).unwrap();
257    /// let input = Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
258    /// let err = XlaExecutor::default()
259    ///     .run_many_with_inputs(program.program(), &[&input])
260    ///     .unwrap_err();
261    /// assert!(matches!(err, Error::PjrtFeatureDisabled | Error::PjrtPluginNotLoaded));
262    /// ```
263    ///
264    /// # Errors
265    ///
266    /// Returns `Error::PjrtFeatureDisabled` or `Error::PjrtPluginNotLoaded`
267    /// when no PJRT executor is available, `Error::InvalidProgram` for input
268    /// count/dtype/shape mismatches, and `Error::PjrtCall` for vendor status
269    /// failures.
270    pub fn run_many_with_inputs(
271        &self,
272        program: &SemanticProgram,
273        inputs: &[&Tensor],
274    ) -> Result<Vec<Tensor>> {
275        #[cfg(feature = "pjrt")]
276        {
277            let Some(plugin) = self.plugin.as_ref() else {
278                return Err(Error::PjrtPluginNotLoaded);
279            };
280            crate::pjrt::run_many_with_inputs(plugin, program, inputs)
281        }
282        #[cfg(not(feature = "pjrt"))]
283        {
284            let _ = (program, inputs);
285            Err(Error::PjrtFeatureDisabled)
286        }
287    }
288
289    /// Execute a compiled graph through a loaded PJRT plugin and return all outputs.
290    ///
291    /// This wrapper accepts the same [`CompiledGraph`] that native runtime
292    /// execution consumes, while preserving the current XLA/PJRT exact-static
293    /// `F32`/`F64` subset. It delegates to the semantic-program PJRT executor
294    /// and does not route through native [`Runtime`](tenferro_runtime::Runtime)
295    /// execution.
296    ///
297    /// # Examples
298    ///
299    /// ```
300    /// use tenferro_runtime::{GraphCompiler, TracedTensor};
301    /// use tenferro_tensor::Tensor;
302    /// use tenferro_xla::{Error, XlaExecutor};
303    ///
304    /// let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
305    /// let mut compiler = GraphCompiler::new();
306    /// let y = x.neg().unwrap();
307    /// let program = compiler.compile(&y).unwrap();
308    /// let input = Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
309    /// let err = XlaExecutor::default()
310    ///     .run_compiled_many_with_inputs(&program, &[&input])
311    ///     .unwrap_err();
312    /// assert!(matches!(err, Error::PjrtFeatureDisabled | Error::PjrtPluginNotLoaded));
313    /// ```
314    ///
315    /// # Errors
316    ///
317    /// Returns `Error::PjrtFeatureDisabled` or `Error::PjrtPluginNotLoaded`
318    /// when no PJRT executor is available, `Error::InvalidProgram` for input
319    /// count/dtype/shape mismatches, and `Error::PjrtCall` for vendor status
320    /// failures.
321    pub fn run_compiled_many_with_inputs(
322        &self,
323        program: &CompiledGraph,
324        inputs: &[&Tensor],
325    ) -> Result<Vec<Tensor>> {
326        self.run_many_with_inputs(program.program(), inputs)
327    }
328
329    /// Execute a single-output graph program through a loaded PJRT plugin.
330    ///
331    /// # Examples
332    ///
333    /// ```
334    /// use tenferro_runtime::{GraphCompiler, TracedTensor};
335    /// use tenferro_tensor::Tensor;
336    /// use tenferro_xla::{Error, XlaExecutor};
337    ///
338    /// let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
339    /// let mut compiler = GraphCompiler::new();
340    /// let y = x.neg().unwrap();
341    /// let program = compiler.compile(&y).unwrap();
342    /// let input = Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
343    /// let err = XlaExecutor::default().run_with_inputs(program.program(), &[&input]).unwrap_err();
344    /// assert!(matches!(err, Error::PjrtFeatureDisabled | Error::PjrtPluginNotLoaded));
345    /// ```
346    ///
347    /// # Errors
348    ///
349    /// Propagates the `run_many_with_inputs` errors and returns
350    /// `Error::InvalidProgram` if the program does not have exactly one
351    /// output.
352    pub fn run_with_inputs(&self, program: &SemanticProgram, inputs: &[&Tensor]) -> Result<Tensor> {
353        single_output_tensor(self.run_many_with_inputs(program, inputs)?)
354    }
355
356    /// Execute a single-output compiled graph through a loaded PJRT plugin.
357    ///
358    /// # Examples
359    ///
360    /// ```
361    /// use tenferro_runtime::{GraphCompiler, TracedTensor};
362    /// use tenferro_tensor::Tensor;
363    /// use tenferro_xla::{Error, XlaExecutor};
364    ///
365    /// let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
366    /// let mut compiler = GraphCompiler::new();
367    /// let y = x.neg().unwrap();
368    /// let program = compiler.compile(&y).unwrap();
369    /// let input = Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
370    /// let err = XlaExecutor::default()
371    ///     .run_compiled_with_inputs(&program, &[&input])
372    ///     .unwrap_err();
373    /// assert!(matches!(err, Error::PjrtFeatureDisabled | Error::PjrtPluginNotLoaded));
374    /// ```
375    ///
376    /// # Errors
377    ///
378    /// Propagates the `run_compiled_many_with_inputs` errors and returns
379    /// `Error::InvalidProgram` if the program does not have exactly one output.
380    pub fn run_compiled_with_inputs(
381        &self,
382        program: &CompiledGraph,
383        inputs: &[&Tensor],
384    ) -> Result<Tensor> {
385        single_output_tensor(self.run_compiled_many_with_inputs(program, inputs)?)
386    }
387}
388
389impl Default for XlaExecutor {
390    fn default() -> Self {
391        Self::new(XlaExecutorOptions::default())
392    }
393}
394
395fn single_output_tensor(mut outputs: Vec<Tensor>) -> Result<Tensor> {
396    if outputs.len() != 1 {
397        return Err(crate::Error::InvalidProgram {
398            message: format!(
399                "PJRT single-output execution expected 1 output, got {}",
400                outputs.len()
401            ),
402        });
403    }
404    Ok(outputs.remove(0))
405}
406
407#[cfg(test)]
408mod tests {
409    use super::{single_output_tensor, XlaExecutor, XlaExecutorOptions};
410    use crate::Error;
411    use tenferro_runtime::{GraphCompiler, TracedTensor};
412    use tenferro_tensor::Tensor;
413
414    #[test]
415    fn executor_options_and_debug_are_directly_covered() {
416        let options = XlaExecutorOptions::default();
417        let executor = XlaExecutor::new(options);
418
419        assert_eq!(executor.options(), options);
420        assert!(!executor.has_loaded_pjrt_plugin());
421        assert_eq!(format!("{options:?}"), "XlaExecutorOptions");
422        assert!(format!("{executor:?}").contains("has_loaded_pjrt_plugin"));
423    }
424
425    #[test]
426    fn default_executor_reports_missing_pjrt_before_dispatch() {
427        let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
428        let mut compiler = GraphCompiler::new();
429        let program = compiler.compile(&x.neg().unwrap()).unwrap();
430        let input = Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
431
432        let err = XlaExecutor::default()
433            .run_many_with_inputs(program.program(), &[&input])
434            .unwrap_err();
435        assert!(matches!(
436            err,
437            Error::PjrtFeatureDisabled | Error::PjrtPluginNotLoaded
438        ));
439        let err = XlaExecutor::default()
440            .run_with_inputs(program.program(), &[&input])
441            .unwrap_err();
442        assert!(matches!(
443            err,
444            Error::PjrtFeatureDisabled | Error::PjrtPluginNotLoaded
445        ));
446    }
447
448    #[test]
449    fn single_output_tensor_rejects_zero_or_multiple_outputs() {
450        let tensor = Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap();
451        assert_eq!(
452            single_output_tensor(vec![tensor.clone()])
453                .unwrap()
454                .as_slice::<f64>()
455                .unwrap(),
456            &[1.0]
457        );
458
459        let err = single_output_tensor(Vec::new()).unwrap_err();
460        assert!(
461            err.to_string().contains("got 0"),
462            "expected zero-output error, got {err:?}"
463        );
464
465        let err = single_output_tensor(vec![tensor.clone(), tensor]).unwrap_err();
466        assert!(
467            err.to_string().contains("got 2"),
468            "expected multi-output error, got {err:?}"
469        );
470    }
471}
tenferro_xla/executor.rs

tenferro_xla/
executor.rs
