tenferro_xla/
lib.rs

1//! Experimental StableHLO lowering and runtime PJRT plugin loading for tenferro.
2//!
3//! This crate is an optional peer executor over `tenferro-runtime`
4//! [`SemanticProgram`](tenferro_runtime::program::SemanticProgram) values. It does not
5//! implement `TensorBackend` and it does not change native CPU, CUDA, or
6//! WebGPU execution.
7//!
8//! # Examples
9//!
10//! ```
11//! use tenferro_runtime::{GraphCompiler, TracedTensor};
12//! use tenferro_xla::lower_compiled_to_stablehlo;
13//!
14//! let x = TracedTensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
15//! let y = (&x + &x).unwrap();
16//! let mut compiler = GraphCompiler::new();
17//! let program = compiler.compile(&y).unwrap();
18//! let module = lower_compiled_to_stablehlo(&program).unwrap();
19//! assert!(module.as_str().contains("stablehlo.add"));
20//! ```
21
22mod error;
23mod executor;
24mod layout;
25mod lowering;
26mod stablehlo;
27
28#[cfg(feature = "pjrt")]
29mod pjrt;
30
31pub use error::{Error, Result};
32pub use executor::{XlaExecutor, XlaExecutorOptions};
33#[cfg(feature = "pjrt")]
34pub use pjrt::PjrtPlugin;
35pub use stablehlo::{StableHloModule, StableHloModuleFingerprint};
36
37/// Environment variable used for the default PJRT plugin path.
38///
39/// # Examples
40///
41/// ```
42/// use tenferro_xla::TENFERRO_PJRT_PLUGIN_ENV;
43///
44/// assert_eq!(TENFERRO_PJRT_PLUGIN_ENV, "TENFERRO_PJRT_PLUGIN");
45/// ```
46pub const TENFERRO_PJRT_PLUGIN_ENV: &str = "TENFERRO_PJRT_PLUGIN";
47
48/// Environment variable used for a GPU-specific PJRT plugin path.
49///
50/// # Examples
51///
52/// ```
53/// use tenferro_xla::TENFERRO_PJRT_GPU_PLUGIN_ENV;
54///
55/// assert_eq!(TENFERRO_PJRT_GPU_PLUGIN_ENV, "TENFERRO_PJRT_GPU_PLUGIN");
56/// ```
57pub const TENFERRO_PJRT_GPU_PLUGIN_ENV: &str = "TENFERRO_PJRT_GPU_PLUGIN";
58
59/// Lower a static-shaped graph program to StableHLO MLIR text.
60///
61/// # Examples
62///
63/// ```
64/// use tenferro_runtime::{GraphCompiler, TracedTensor};
65/// use tenferro_xla::lower_to_stablehlo;
66///
67/// let x = TracedTensor::from_vec_col_major(vec![1], vec![3.0_f64]).unwrap();
68/// let mut compiler = GraphCompiler::new();
69/// let y = x.neg().unwrap();
70/// let program = compiler.compile(&y).unwrap();
71/// let module = lower_to_stablehlo(program.program()).unwrap();
72/// assert!(module.as_str().contains("stablehlo.negate"));
73/// ```
74///
75/// # Errors
76///
77/// Returns `Error::UnsupportedDType`, `Error::UnsupportedOp`, or
78/// `Error::NonStaticShape` for unsupported graph content, and
79/// `Error::InvalidProgram` for inconsistent graph metadata.
80pub fn lower_to_stablehlo(
81    program: &tenferro_runtime::program::SemanticProgram,
82) -> Result<StableHloModule> {
83    lowering::lower_to_stablehlo(program)
84}
85
86/// Lower a compiled graph to StableHLO MLIR text.
87///
88/// This is the preferred public lowering boundary for graph users. It consumes
89/// the backend-neutral [`CompiledGraph`](tenferro_runtime::CompiledGraph) view
90/// directly and does not route through native [`Runtime`](tenferro_runtime::Runtime)
91/// execution staging.
92///
93/// # Examples
94///
95/// ```
96/// use tenferro_runtime::{GraphCompiler, TracedTensor};
97/// use tenferro_xla::lower_compiled_to_stablehlo;
98///
99/// let x = TracedTensor::from_vec_col_major(vec![1], vec![3.0_f64]).unwrap();
100/// let mut compiler = GraphCompiler::new();
101/// let y = x.neg().unwrap();
102/// let program = compiler.compile(&y).unwrap();
103/// let module = lower_compiled_to_stablehlo(&program).unwrap();
104/// assert!(module.as_str().contains("stablehlo.negate"));
105/// ```
106///
107/// # Errors
108///
109/// Returns `Error::UnsupportedDType`, `Error::UnsupportedOp`, or
110/// `Error::NonStaticShape` for unsupported graph content, and
111/// `Error::InvalidProgram` for inconsistent graph metadata.
112pub fn lower_compiled_to_stablehlo(
113    program: &tenferro_runtime::CompiledGraph,
114) -> Result<StableHloModule> {
115    lowering::lower_to_stablehlo(program.program())
116}
tenferro_xla/lib.rs

tenferro_xla/
lib.rs
