tenferro_ad/
extension.rs

1//! Eager AD support for out-of-tree extension primitives.
2
3use std::sync::Arc;
4
5use computegraph::GraphOperation;
6use tenferro_ops::std_tensor_op::StdTensorOp;
7use tenferro_runtime::{
8    Error, ErrorPhase, ExtensionModule, InputSignature, PrepareCapability, PrepareError,
9    PreparedOperationExecutorHandle, Result, Runtime, RuntimeConfigError,
10};
11use tenferro_tensor::{Tensor, TensorRead, TensorValue};
12
13use crate::eager::{
14    eager_capture_active, eager_grad_recording_enabled, execute_prepared_extension_on,
15    EagerRuntime, EagerSession, EagerTensor,
16};
17use tenferro_runtime::ExecutionEngine;
18
19pub use tenferro_runtime::extension::{
20    apply, ExtensionCacheKey, ExtensionCacheLimits, ExtensionCacheSelector, ExtensionCacheStore,
21    ExtensionExecutionContext, ExtensionFamilyId, ExtensionOp,
22};
23
24/// Closed backend kind selected by the eager runtime owner for an extension.
25///
26/// # Examples
27///
28/// ```rust
29/// use tenferro_ad::extension::{EagerExtensionBackendKind, EagerExtensionTarget};
30/// use tenferro_runtime::EngineId;
31///
32/// let target = EagerExtensionTarget {
33///     engine_id: EngineId::new("example.engine")?,
34///     backend_kind: EagerExtensionBackendKind::Cpu,
35/// };
36/// assert!(matches!(
37///     target.backend_kind,
38///     EagerExtensionBackendKind::Cpu
39/// ));
40/// assert_eq!(target.engine_id.as_str(), "example.engine");
41/// # Ok::<(), Box<dyn std::error::Error>>(())
42/// ```
43#[doc(hidden)]
44#[derive(Clone, Copy, Debug, Eq, PartialEq)]
45pub enum EagerExtensionBackendKind {
46    /// The eager runtime owns a CPU backend.
47    Cpu,
48    /// The eager runtime owns a CUDA backend.
49    #[cfg(feature = "cuda")]
50    Cuda,
51}
52
53/// Exact engine target selected by the eager runtime owner.
54///
55/// # Examples
56///
57/// ```rust
58/// use tenferro_ad::extension::{EagerExtensionBackendKind, EagerExtensionTarget};
59/// use tenferro_runtime::EngineId;
60///
61/// let target = EagerExtensionTarget {
62///     engine_id: EngineId::new("example.engine")?,
63///     backend_kind: EagerExtensionBackendKind::Cpu,
64/// };
65/// assert_eq!(target.backend_kind, EagerExtensionBackendKind::Cpu);
66/// # Ok::<(), Box<dyn std::error::Error>>(())
67/// ```
68#[doc(hidden)]
69#[derive(Clone, Debug, Eq, PartialEq)]
70pub struct EagerExtensionTarget {
71    /// Exact runtime engine selected for this eager context.
72    pub engine_id: tenferro_runtime::EngineId,
73    /// Closed backend kind selected for this eager context.
74    pub backend_kind: EagerExtensionBackendKind,
75}
76
77#[cfg(test)]
78mod tests;
79
80/// Validate an untracked eager input for in-place execution on `engine` and
81/// install the extension module the in-place kernel needs.
82///
83/// # Errors
84///
85/// Returns a runtime-state error for a tracked input or inside trace capture,
86/// and the runtime's typed error when the engine cannot accept the input or
87/// the module cannot be installed.
88#[doc(hidden)]
89pub fn prepare_eager_in_place_input(
90    engine: &mut dyn ExecutionEngine,
91    input: &EagerTensor,
92    family_id: &'static str,
93    module_factory: impl FnOnce(EagerExtensionTarget) -> Result<Arc<dyn ExtensionModule>>,
94) -> Result<()> {
95    if input.requires_grad || input.trace.is_some() || eager_capture_active() {
96        return Err(Error::runtime_state(
97            "eager in-place",
98            ErrorPhase::Execution,
99            "in-place execution requires an untracked value outside trace capture",
100        ));
101    }
102    let target = input.ctx.eager_extension_target(engine)?;
103    validate_eager_extension_input_signature(&input.ctx, &target, &[input.tensor_read()])?;
104    let module = module_factory(target.clone())?;
105    input
106        .ctx
107        .ensure_extension_module_for_engine(module, family_id, &target.engine_id)?;
108    Ok(())
109}
110
111/// Adopt a runtime value produced by an extension kernel as an untracked eager
112/// tensor of `ctx`.
113///
114/// # Examples
115///
116/// ```
117/// use tenferro_ad::{EagerRuntime, Tensor};
118/// use tenferro_tensor::TensorValue;
119///
120/// let ctx = EagerRuntime::new()?;
121/// let value = TensorValue::from_tensor(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?);
122/// let tensor = tenferro_ad::extension::adopt_untracked_eager_value(ctx, value)?;
123/// assert!(!tensor.tracks_grad());
124/// # Ok::<(), Box<dyn std::error::Error>>(())
125/// ```
126///
127/// # Errors
128///
129/// Returns [`Error::RuntimeState`] when the [`TensorValue`] cannot be
130/// transferred into its allocation group.
131#[must_use = "the adopted eager tensor carries the runtime value"]
132pub fn adopt_untracked_eager_value(
133    ctx: Arc<EagerRuntime>,
134    value: TensorValue,
135) -> Result<EagerTensor> {
136    EagerTensor::new_untracked_value_result(ctx, value)
137}
138
139/// Apply an extension op to eager tensors on `engine`, recording it for AD.
140///
141/// # Errors
142///
143/// Returns an input-count or context-mismatch error, a runtime error when no
144/// installed extension module serves the op on this engine, or the extension
145/// kernel's own typed error.
146pub fn apply_eager(
147    engine: &mut dyn ExecutionEngine,
148    op: Arc<dyn ExtensionOp>,
149    inputs: &[&EagerTensor],
150) -> Result<Vec<EagerTensor>> {
151    let ctx = validate_eager_extension_inputs(op.as_ref(), inputs)?;
152    ctx.with_eager_session(engine, |session| {
153        apply_eager_in_session(session, op, inputs)
154    })
155}
156
157/// Apply an extension op inside an eager session, preferring a prepared
158/// session executor and falling back to a compiled one-op program.
159fn apply_eager_in_session(
160    session: &mut EagerSession<'_>,
161    op: Arc<dyn ExtensionOp>,
162    inputs: &[&EagerTensor],
163) -> Result<Vec<EagerTensor>> {
164    let ctx = validate_eager_extension_inputs(op.as_ref(), inputs)?;
165    ensure_session_runtime(session, &ctx)?;
166    let std_op = StdTensorOp::Extension(Arc::clone(&op));
167    let input_reads: Vec<_> = inputs.iter().map(|tensor| tensor.tensor_read()).collect();
168    let outputs = match try_prepared_eager_extension(session, &std_op, &input_reads)? {
169        Some(outputs) => outputs,
170        None => ctx.exec_extension_outputs_read(session.engine(), &op, &input_reads)?,
171    };
172    finish_eager_extension_outputs(ctx, std_op, inputs, outputs, session)
173}
174
175fn ensure_session_runtime(session: &EagerSession<'_>, ctx: &Arc<EagerRuntime>) -> Result<()> {
176    if !Arc::ptr_eq(session.runtime(), ctx) {
177        return Err(Error::ContextMismatch {
178            lhs: session.runtime().id(),
179            rhs: ctx.id(),
180        });
181    }
182    Ok(())
183}
184
185/// Run a prepared extension executor on `engine`: through the session-capable
186/// entry when it has one, otherwise through the typed engine escape.
187fn execute_prepared_executor(
188    engine: &mut dyn ExecutionEngine,
189    executor: &dyn tenferro_runtime::PreparedOperationExecutor,
190    input_reads: &[TensorRead<'_>],
191) -> Result<Vec<Tensor>> {
192    if executor.supports_session() {
193        return execute_prepared_extension_on(engine, executor, input_reads);
194    }
195    let identity = engine.target().context_identity();
196    crate::eager::with_lent_extension_caches(engine, |engine, caches| {
197        let mut erased =
198            tenferro_runtime::ErasedExecutionContext::from_dyn(identity, engine.as_any_mut());
199        executor.execute(&mut erased, caches, input_reads)
200    })
201}
202
203fn try_prepared_eager_extension(
204    session: &mut EagerSession<'_>,
205    op: &StdTensorOp,
206    input_reads: &[TensorRead<'_>],
207) -> Result<Option<Vec<Tensor>>> {
208    let ctx = Arc::clone(session.runtime());
209    // An engine kind without an eager extension target falls back to the
210    // compiled one-op program.
211    let Ok(target) = ctx.eager_extension_target(session.engine()) else {
212        return Ok(None);
213    };
214    let Some(executor) = prepared_eager_extension_executor(&ctx, &target, op, input_reads)? else {
215        return Ok(None);
216    };
217    Ok(Some(execute_prepared_executor(
218        session.engine(),
219        executor.as_ref(),
220        input_reads,
221    )?))
222}
223
224fn prepared_eager_extension_executor(
225    ctx: &EagerRuntime,
226    target: &EagerExtensionTarget,
227    op: &StdTensorOp,
228    input_reads: &[TensorRead<'_>],
229) -> Result<Option<PreparedOperationExecutorHandle>> {
230    let StdTensorOp::Extension(ext) = op else {
231        return Ok(None);
232    };
233    let signature = InputSignature::from_reads(input_reads).map_err(|source| {
234        Error::runtime_state_source("extension::apply_eager", ErrorPhase::Execution, source)
235    })?;
236    let PrepareCapability::Prepared(plan) =
237        ctx.runtime()
238            .prepare_extension_immediate(&target.engine_id, ext.as_ref(), &signature)?
239    else {
240        return Ok(None);
241    };
242    Ok(plan.executor().cloned())
243}
244
245/// Install `module` and apply `op` on `engine`.
246///
247/// # Errors
248///
249/// As [`apply_eager`], plus the runtime's error when the module cannot be
250/// installed.
251#[doc(hidden)]
252pub fn apply_eager_with_extension_session(
253    engine: &mut dyn ExecutionEngine,
254    op: Arc<dyn ExtensionOp>,
255    inputs: &[&EagerTensor],
256    module: Arc<dyn ExtensionModule>,
257) -> Result<Vec<EagerTensor>> {
258    let ctx = validate_eager_extension_inputs(op.as_ref(), inputs)?;
259    ctx.install_extension_module(module)?;
260    apply_eager(engine, op, inputs)
261}
262
263/// Install the module `module_factory` builds for `engine`'s target and apply
264/// `op` on `engine`.
265///
266/// # Errors
267///
268/// As [`apply_eager`], plus input-ingress and module-installation errors.
269#[doc(hidden)]
270pub fn apply_eager_with_targeted_extension_session(
271    engine: &mut dyn ExecutionEngine,
272    op: Arc<dyn ExtensionOp>,
273    inputs: &[&EagerTensor],
274    module_factory: impl FnOnce(
275        EagerExtensionTarget,
276    ) -> tenferro_runtime::Result<Arc<dyn ExtensionModule>>,
277) -> Result<Vec<EagerTensor>> {
278    let ctx = validate_eager_extension_inputs(op.as_ref(), inputs)?;
279    let target = ctx.eager_extension_target(engine)?;
280    let input_reads: Vec<_> = inputs.iter().map(|tensor| tensor.tensor_read()).collect();
281    validate_eager_extension_input_signature(&ctx, &target, &input_reads)?;
282    let module = module_factory(target.clone())?;
283    ctx.ensure_extension_module_for_engine(module, op.family_id(), &target.engine_id)?;
284    apply_eager(engine, op, inputs)
285}
286
287/// [`apply_eager_with_targeted_extension_session`] inside an eager session.
288///
289/// # Errors
290///
291/// As [`apply_eager_with_targeted_extension_session`], plus a context mismatch
292/// when the inputs belong to another runtime.
293#[doc(hidden)]
294pub fn apply_eager_with_targeted_extension_in_session(
295    session: &mut EagerSession<'_>,
296    op: Arc<dyn ExtensionOp>,
297    inputs: &[&EagerTensor],
298    module_factory: impl FnOnce(
299        EagerExtensionTarget,
300    ) -> tenferro_runtime::Result<Arc<dyn ExtensionModule>>,
301) -> Result<Vec<EagerTensor>> {
302    let ctx = validate_eager_extension_inputs(op.as_ref(), inputs)?;
303    ensure_session_runtime(session, &ctx)?;
304    let target = ctx.eager_extension_target(session.engine())?;
305    let input_reads: Vec<_> = inputs.iter().map(|tensor| tensor.tensor_read()).collect();
306    validate_eager_extension_input_signature(&ctx, &target, &input_reads)?;
307    let module = module_factory(target.clone())?;
308    ctx.ensure_extension_module_for_engine(module, op.family_id(), &target.engine_id)?;
309    apply_eager_in_session(session, op, inputs)
310}
311
312pub(crate) fn validate_eager_extension_target(
313    runtime: &Runtime,
314    target: &EagerExtensionTarget,
315) -> Result<()> {
316    let snapshot = runtime.snapshot().map_err(|source| {
317        Error::runtime_state_source(
318            "extension::apply_eager_with_extension_session",
319            ErrorPhase::Execution,
320            source,
321        )
322    })?;
323    if snapshot.engine(&target.engine_id).is_none() {
324        return Err(Error::runtime_state_source(
325            "extension::apply_eager_with_extension_session",
326            ErrorPhase::Execution,
327            RuntimeConfigError::MissingEngine {
328                engine_id: target.engine_id.clone(),
329            },
330        ));
331    }
332    Ok(())
333}
334
335fn validate_eager_extension_input_signature(
336    ctx: &EagerRuntime,
337    target: &EagerExtensionTarget,
338    input_reads: &[TensorRead<'_>],
339) -> Result<()> {
340    let signature = InputSignature::from_reads(input_reads).map_err(|source| {
341        Error::runtime_state_source(
342            "extension::apply_eager_with_extension_session",
343            ErrorPhase::Execution,
344            source,
345        )
346    })?;
347    let snapshot = ctx.runtime().snapshot().map_err(|source| {
348        Error::runtime_state_source(
349            "extension::apply_eager_with_extension_session",
350            ErrorPhase::Execution,
351            source,
352        )
353    })?;
354    let engine = snapshot.engine(&target.engine_id).ok_or_else(|| {
355        Error::runtime_state_source(
356            "extension::apply_eager_with_extension_session",
357            ErrorPhase::Execution,
358            RuntimeConfigError::MissingEngine {
359                engine_id: target.engine_id.clone(),
360            },
361        )
362    })?;
363    for (input_index, entry) in signature.entries().iter().enumerate() {
364        if !engine.accepts_input_signature(entry) {
365            return Err(Error::runtime_state_source(
366                "extension::apply_eager_with_extension_session",
367                ErrorPhase::Execution,
368                PrepareError::NoInputIngress {
369                    input_index,
370                    placement: entry.placement().clone(),
371                },
372            ));
373        }
374    }
375    Ok(())
376}
377
378fn validate_eager_extension_inputs(
379    op: &dyn ExtensionOp,
380    inputs: &[&EagerTensor],
381) -> Result<Arc<EagerRuntime>> {
382    let Some(first) = inputs.first() else {
383        return Err(Error::invalid_argument(
384            "extension::apply_eager",
385            ErrorPhase::Execution,
386            "inputs",
387            "at least one input tensor is required",
388        ));
389    };
390    if inputs.len() != op.input_count() {
391        return Err(Error::invalid_argument(
392            "extension::apply_eager",
393            ErrorPhase::Execution,
394            "inputs",
395            format!(
396                "op family {:?} expects {} inputs, got {}",
397                op.family_id(),
398                op.input_count(),
399                inputs.len()
400            ),
401        ));
402    }
403
404    let ctx = Arc::clone(&first.ctx);
405    for tensor in inputs.iter().skip(1) {
406        if !first.same_context(tensor) {
407            return Err(Error::ContextMismatch {
408                lhs: first.ctx_id(),
409                rhs: tensor.ctx_id(),
410            });
411        }
412    }
413    Ok(ctx)
414}
415
416fn finish_eager_extension_outputs(
417    ctx: Arc<EagerRuntime>,
418    op: StdTensorOp,
419    inputs: &[&EagerTensor],
420    outputs: Vec<Tensor>,
421    session: &mut EagerSession<'_>,
422) -> Result<Vec<EagerTensor>> {
423    if outputs.len() != op.output_count() {
424        return Err(Error::Internal(format!(
425            "expected {} eager outputs for {:?}, got {}",
426            op.output_count(),
427            op,
428            outputs.len()
429        )));
430    }
431
432    if !eager_grad_recording_enabled()
433        || (!eager_capture_active() && !inputs.iter().any(|input| input.requires_grad))
434    {
435        return outputs
436            .into_iter()
437            .map(|output| EagerTensor::new_untracked_result(Arc::clone(&ctx), output))
438            .collect();
439    }
440
441    let output_refs: Vec<&Tensor> = outputs.iter().collect();
442    let recorded = session.record_outputs(&op, &output_refs, inputs)?;
443    if recorded.traces.len() != outputs.len() {
444        return Err(Error::Internal(format!(
445            "expected {} eager traces for {:?}, got {}",
446            outputs.len(),
447            op,
448            recorded.traces.len()
449        )));
450    }
451    let results = recorded
452        .traces
453        .into_iter()
454        .zip(recorded.semantic_traces)
455        .zip(outputs)
456        .map(|((trace, semantic_trace), output)| {
457            if trace.requires_grad {
458                EagerTensor::new_result_with_semantic_trace(
459                    Arc::clone(&ctx),
460                    trace.key,
461                    output,
462                    trace.requires_grad,
463                    trace.trace,
464                    semantic_trace,
465                )
466            } else {
467                EagerTensor::new_unregistered_result_with_semantic_trace(
468                    Arc::clone(&ctx),
469                    trace.key,
470                    output,
471                    trace.requires_grad,
472                    trace.trace,
473                    semantic_trace,
474                )
475            }
476        })
477        .collect::<Result<Vec<_>>>()?;
478    crate::eager::finish_residuals(&op, inputs, &results.iter().collect::<Vec<_>>())?;
479    Ok(results)
480}
tenferro_ad/extension.rs

tenferro_ad/
extension.rs
