1use std::borrow::Cow;
2use std::cell::{Cell, RefCell};
3use std::cmp::Reverse;
4use std::collections::HashMap;
5use std::env;
6use std::fmt;
7use std::marker::PhantomData;
8use std::mem::{size_of, size_of_val};
9use std::rc::Rc;
10#[cfg(test)]
11use std::sync::atomic::{AtomicUsize, Ordering};
12use std::sync::{Arc, Mutex, MutexGuard, OnceLock, Weak};
13use std::time::{Duration, Instant};
14
15use lru::LruCache;
16
17use crate::extension::{
18 validate_eager_extension_target, EagerExtensionBackendKind, EagerExtensionTarget,
19};
20use crate::extension_cache::{ExtensionCacheLimits, ExtensionCacheSelector, ExtensionCacheStore};
21#[cfg(test)]
22use computegraph::graph::Graph;
23use computegraph::ValueKey;
24#[cfg(test)]
25use computegraph::ValueRef;
26use tenferro_cpu::{CpuBackend, CpuBackendError, CpuPlacement};
27#[cfg(feature = "cuda")]
28use tenferro_gpu::cuda::CudaBackend;
29#[cfg(feature = "webgpu")]
30use tenferro_gpu::webgpu::WebGpuBackend;
31#[cfg(test)]
32use tenferro_ops::input_key::TensorInputKey;
33use tenferro_ops::{std_tensor_op::StdTensorOp, SymDim, TensorMeta};
34use tenferro_runtime::ad_support::{
35 analyze_deferred_semantic_trace, compile_ad_source, ones_tensor, RetainedValue,
36};
37use tenferro_runtime::program::{ProgramValueMetadata, SemanticFingerprint, SemanticProgram};
38use tenferro_runtime::{
39 CompiledGraph, CoreCapabilityBundle, EngineId, ErrorPhase, ExecutionContextIdentity,
40 ExtensionModule, GraphCompiler, HardwareClassId, PreparedCompiledGraph, RegistrationIdentity,
41 Runtime, RuntimeConfigError, RuntimeConfigSnapshot, RuntimeEpoch, TracedTensor,
42};
43#[cfg(test)]
44use tenferro_tensor::TypedTensor;
45use tenferro_tensor::{
46 AllocationGroup, CacheStats, DType, DescriptorSlot, GroupError, IntoShapeVec, Tensor,
47 TensorBackend, TensorRead, TensorScalar, TensorValue, TensorView,
48};
49use tenferro_tensor::{BackendSession, BackendSessionHost};
50
51#[cfg(feature = "cuda")]
52use crate::eager_backend::cuda_runtime_engine_id;
53use crate::eager_backend::{
54 cpu_runtime_engine_id, cpu_runtime_hardware_class, eager_runtime_for_backend, EagerBackend,
55};
56#[cfg(test)]
57use crate::eager_exec::exec_standard_op_on_tensor_reads_in_session;
58use crate::eager_exec::{
59 eager_input_promotion_plan, exec_op_on_tensor_reads_with_runtime,
60 exec_op_on_tensors_with_runtime,
61};
62use crate::error::{ContextId, Error, Result};
63use crate::metadata::tensor_meta_from_tensor;
64use crate::semantic_extension::SemanticExtensionRuleSet;
65use crate::traced::{derivative_trace_from_frozen_program, next_input_key};
66use crate::transform_cache::{AdTransformCache, AdTransformCacheLimits};
67
68use crate::AdContext;
69
70pub(crate) type GradSlot = Arc<Mutex<Option<Arc<AdValueRecord>>>>;
71pub(crate) type WeakGradSlot = Weak<Mutex<Option<Arc<AdValueRecord>>>>;
72
73#[derive(Clone, Debug)]
74pub(crate) struct EagerTrace;
75
76#[cfg(test)]
77pub(crate) static CPU_RUNTIME_SELECTION_REFRESHES: AtomicUsize = AtomicUsize::new(0);
78
79struct CpuRuntimeSelection {
80 snapshot: Arc<RuntimeConfigSnapshot>,
81 epoch: RuntimeEpoch,
82 engine_id: EngineId,
83 registration_identity: RegistrationIdentity,
84 capabilities: CoreCapabilityBundle,
85}
86
87#[derive(Debug, Default, Clone)]
88struct EagerOpProfileEntry {
89 calls: usize,
90 total_time: Duration,
91}
92
93thread_local! {
94 static EAGER_OP_PROFILE_STATE: RefCell<HashMap<&'static str, EagerOpProfileEntry>> =
95 RefCell::new(HashMap::new());
96 static EAGER_NO_GRAD_DEPTH: Cell<usize> = const { Cell::new(0) };
97 static EAGER_CAPTURE_DEPTH: Cell<usize> = const { Cell::new(0) };
98 #[cfg(test)]
99 static EAGER_OP_PROFILE_ENABLED_OVERRIDE: RefCell<Option<bool>> = const { RefCell::new(None) };
100 #[cfg(test)]
101 static EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE: RefCell<Option<Option<usize>>> = const { RefCell::new(None) };
102 #[cfg(test)]
103 static EAGER_SEMANTIC_VJP_ENABLED_OVERRIDE: RefCell<Option<bool>> = const { RefCell::new(None) };
104}
105
106#[cfg(test)]
107pub(crate) static EAGER_SEMANTIC_VJP_EXECUTIONS: AtomicUsize = AtomicUsize::new(0);
108
109pub(crate) fn eager_grad_recording_enabled() -> bool {
110 EAGER_NO_GRAD_DEPTH.with(|depth| depth.get() == 0)
111}
112
113pub(crate) fn eager_capture_active() -> bool {
114 EAGER_CAPTURE_DEPTH.with(|depth| depth.get() > 0)
115}
116
117fn eager_semantic_vjp_enabled() -> bool {
118 #[cfg(test)]
119 if let Some(value) = EAGER_SEMANTIC_VJP_ENABLED_OVERRIDE.with(|state| *state.borrow()) {
120 return value;
121 }
122
123 static ENABLED: OnceLock<bool> = OnceLock::new();
126 *ENABLED.get_or_init(|| env::var("TENFERRO_EAGER_SEMANTIC_VJP").map_or(true, |v| v != "0"))
127}
128
129#[derive(Debug)]
153pub struct EagerNoGradGuard {
154 active: bool,
155 _not_send: PhantomData<Rc<()>>,
158}
159
160impl Drop for EagerNoGradGuard {
161 fn drop(&mut self) {
162 if !self.active {
163 return;
164 }
165 EAGER_NO_GRAD_DEPTH.with(|depth| {
166 depth.set(depth.get().saturating_sub(1));
167 });
168 self.active = false;
169 }
170}
171
172#[derive(Debug)]
206pub struct EagerTraceCaptureGuard {
207 active: bool,
208 _not_send: PhantomData<Rc<()>>,
211}
212
213impl Drop for EagerTraceCaptureGuard {
214 fn drop(&mut self) {
215 if !self.active {
216 return;
217 }
218 EAGER_CAPTURE_DEPTH.with(|depth| {
219 depth.set(depth.get().saturating_sub(1));
220 });
221 self.active = false;
222 }
223}
224
225pub(crate) fn eager_op_profile_enabled() -> bool {
226 #[cfg(test)]
227 if let Some(value) = EAGER_OP_PROFILE_ENABLED_OVERRIDE.with(|state| *state.borrow()) {
228 return value;
229 }
230
231 static ENABLED: OnceLock<bool> = OnceLock::new();
232 *ENABLED.get_or_init(|| env::var("TENFERRO_PROFILE_EAGER_OP_AGG").is_ok())
233}
234
235pub(crate) fn eager_op_profile_start() -> Option<Instant> {
236 eager_op_profile_enabled().then(Instant::now)
237}
238
239pub(crate) fn record_eager_op_profile(section: &'static str, elapsed: Duration) {
240 if !eager_op_profile_enabled() {
241 return;
242 }
243 EAGER_OP_PROFILE_STATE.with(|state| {
244 let mut state = state.borrow_mut();
245 let entry = state.entry(section).or_default();
246 entry.calls += 1;
247 entry.total_time += elapsed;
248 });
249}
250
251pub(crate) fn profile_eager_op_section<T>(section: &'static str, f: impl FnOnce() -> T) -> T {
252 if !eager_op_profile_enabled() {
253 return f();
254 }
255 let started = Instant::now();
256 let result = f();
257 record_eager_op_profile(section, started.elapsed());
258 result
259}
260
261pub(crate) fn maybe_print_eager_op_profile() {
262 if !eager_op_profile_enabled() {
263 return;
264 }
265 let Some(print_every) = eager_op_profile_print_every() else {
266 return;
267 };
268 if print_every == 0 {
269 return;
270 }
271
272 let should_print = EAGER_OP_PROFILE_STATE.with(|state| {
273 state
274 .borrow()
275 .get("nary_op.total")
276 .is_some_and(|entry| entry.calls % print_every == 0)
277 });
278 if should_print {
279 print_and_reset_eager_op_profile();
280 }
281}
282
283fn eager_op_profile_print_every() -> Option<usize> {
284 #[cfg(test)]
285 if let Some(value) = EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE.with(|state| *state.borrow()) {
286 return value;
287 }
288
289 env::var("TENFERRO_PROFILE_EAGER_OP_PRINT_EVERY")
290 .ok()?
291 .parse()
292 .ok()
293}
294
295pub(crate) fn print_and_reset_eager_op_profile() {
296 EAGER_OP_PROFILE_STATE.with(|state| {
297 let mut entries: Vec<_> = state
298 .borrow()
299 .iter()
300 .map(|(section, entry)| (*section, entry.clone()))
301 .collect();
302 state.borrow_mut().clear();
303 entries.sort_by_key(|(_, entry)| Reverse(entry.total_time));
304
305 eprintln!("=== tenferro eager op profile ===");
306 for (section, entry) in entries {
307 let Some(per_call_us) = eager_op_profile_per_call_us(&entry) else {
308 continue;
309 };
310 eprintln!(
311 "{section}: calls={} total={:.6}ms per_call={:.3}us",
312 entry.calls,
313 entry.total_time.as_secs_f64() * 1.0e3,
314 per_call_us,
315 );
316 }
317 });
318}
319
320fn eager_op_profile_per_call_us(entry: &EagerOpProfileEntry) -> Option<f64> {
321 (entry.calls != 0).then(|| entry.total_time.as_secs_f64() * 1.0e6 / entry.calls as f64)
322}
323
324fn runtime_config_error(op: &'static str, source: RuntimeConfigError) -> Error {
325 Error::runtime_state_source(op, ErrorPhase::Execution, source)
326}
327
328fn runtime_state_source<E>(op: &'static str, source: E) -> Error
329where
330 E: std::error::Error + Send + Sync + 'static,
331{
332 Error::runtime_state_source(op, ErrorPhase::Execution, source)
333}
334
335fn cpu_runtime_bridge_unsupported(message: impl Into<String>) -> Error {
336 Error::unsupported(
337 "CpuPlacementBoundEager::refresh_runtime_selection",
338 ErrorPhase::Execution,
339 message,
340 )
341}
342
343fn select_cpu_runtime(runtime: &Runtime) -> Result<CpuRuntimeSelection> {
344 let snapshot = runtime
345 .snapshot()
346 .map_err(|source| runtime_state_source("EagerRuntime::runtime_snapshot", source))?;
347 let engine_id = cpu_runtime_engine_id()
348 .map_err(|source| runtime_config_error("EagerRuntime::cpu_runtime_engine_id", source))?;
349 let expected_hardware = cpu_runtime_hardware_class().map_err(|source| {
350 runtime_config_error("EagerRuntime::cpu_runtime_hardware_class", source)
351 })?;
352 let engine = snapshot
353 .engine(&engine_id)
354 .ok_or_else(|| cpu_runtime_bridge_unsupported("missing CPU runtime engine"))?;
355 validate_cpu_runtime_engine(
356 engine.context_identity(),
357 engine.hardware_class(),
358 engine.capabilities(),
359 &expected_hardware,
360 )?;
361 let epoch = snapshot.epoch();
362 let registration_identity = engine.registration_identity();
363 let capabilities = engine.capabilities().clone();
364 Ok(CpuRuntimeSelection {
365 snapshot,
366 epoch,
367 engine_id,
368 registration_identity,
369 capabilities,
370 })
371}
372
373fn validate_cpu_runtime_engine(
374 context_identity: ExecutionContextIdentity,
375 hardware_class: &HardwareClassId,
376 capabilities: &CoreCapabilityBundle,
377 expected_hardware: &HardwareClassId,
378) -> Result<()> {
379 if context_identity != ExecutionContextIdentity::of::<CpuBackend>() {
380 return Err(cpu_runtime_bridge_unsupported(
381 "CPU runtime context mismatch",
382 ));
383 }
384 if hardware_class != expected_hardware {
385 return Err(cpu_runtime_bridge_unsupported(
386 "CPU runtime hardware mismatch",
387 ));
388 }
389 if capabilities.elementwise().is_none() {
390 return Err(cpu_runtime_bridge_unsupported(
391 "missing CPU runtime capability: elementwise",
392 ));
393 }
394 if capabilities.reduction().is_none() {
395 return Err(cpu_runtime_bridge_unsupported(
396 "missing CPU runtime capability: reduction",
397 ));
398 }
399 if capabilities.indexing().is_none() {
400 return Err(cpu_runtime_bridge_unsupported(
401 "missing CPU runtime capability: indexing",
402 ));
403 }
404 if capabilities.dot_general().is_none() {
405 return Err(cpu_runtime_bridge_unsupported(
406 "missing CPU runtime capability: dot_general",
407 ));
408 }
409 if capabilities.layout().is_none() {
410 return Err(cpu_runtime_bridge_unsupported(
411 "missing CPU runtime capability: layout",
412 ));
413 }
414 Ok(())
415}
416
417#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
421pub struct EagerRuntimeCacheStats {
422 pub extensions: CacheStats,
424 pub ad_transforms: CacheStats,
426 pub prepared_derivatives: CacheStats,
428}
429
430#[cfg(test)]
431pub(crate) struct EagerGraphExecution {
432 pub(crate) outputs: Vec<Tensor>,
433}
434
435#[derive(Debug)]
456pub struct ValueGuard<'a> {
457 view: TensorView<'a>,
458}
459
460impl<'a> ValueGuard<'a> {
461 pub fn dtype(&self) -> DType {
463 self.view.dtype()
464 }
465
466 pub fn shape(&self) -> &[usize] {
468 self.view.shape()
469 }
470
471 pub fn as_tensor_view(&self) -> &TensorView<'_> {
473 &self.view
474 }
475
476 pub fn as_slice<T: TensorScalar>(&self) -> tenferro_tensor::Result<&'a [T]> {
488 self.view.as_slice()
489 }
490
491 fn duplicate_host_tensor(&self) -> tenferro_tensor::Result<Tensor> {
492 match &self.view {
493 TensorView::F32(view) => {
494 <f32 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
495 }
496 TensorView::F64(view) => {
497 <f64 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
498 }
499 TensorView::I32(view) => {
500 <i32 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
501 }
502 TensorView::I64(view) => {
503 <i64 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
504 }
505 TensorView::Bool(view) => <bool as TensorScalar>::into_tensor(
506 view.shape().to_vec(),
507 view.as_slice()?.to_vec(),
508 ),
509 TensorView::C32(view) => <num_complex::Complex32 as TensorScalar>::into_tensor(
510 view.shape().to_vec(),
511 view.as_slice()?.to_vec(),
512 ),
513 TensorView::C64(view) => <num_complex::Complex64 as TensorScalar>::into_tensor(
514 view.shape().to_vec(),
515 view.as_slice()?.to_vec(),
516 ),
517 }
518 }
519}
520
521#[derive(Clone, Debug)]
541pub struct GradientValue {
542 record: Arc<AdValueRecord>,
543 ctx: Arc<EagerRuntime>,
544}
545
546impl GradientValue {
547 pub fn dtype(&self) -> DType {
549 self.record.dtype()
550 }
551
552 pub fn shape(&self) -> &[usize] {
554 self.record.shape()
555 }
556
557 pub fn value(&self) -> Result<ValueGuard<'_>> {
564 self.record.value("GradientValue::value")
565 }
566
567 pub fn tensor_read(&self) -> Result<TensorRead<'_>> {
574 self.record.tensor_read("GradientValue::tensor_read")
575 }
576
577 pub fn as_slice<T: TensorScalar>(&self) -> tenferro_tensor::Result<&[T]> {
586 self.record
587 .value("GradientValue::as_slice")
588 .map_err(|error| {
589 tenferro_tensor::Error::runtime_state_source("GradientValue::as_slice", error)
590 })?
591 .as_slice()
592 }
593
594 pub fn to_tensor(&self) -> Result<Tensor> {
602 let value = self
603 .record
604 .value("GradientValue::to_tensor")
605 .map_err(|error| {
606 Error::runtime_state_source(
607 "GradientValue::to_tensor",
608 ErrorPhase::Execution,
609 error,
610 )
611 })?;
612 match value.duplicate_host_tensor() {
613 Ok(tensor) => Ok(tensor),
614 Err(_) => {
615 let read = self.record.tensor_read("GradientValue::to_tensor")?;
616 self.ctx
617 .with_execution_session(|session| session.to_contiguous_read(read))?
618 .map_err(Error::from)
619 }
620 }
621 }
622}
623
624#[derive(Debug)]
643pub struct Gradients {
644 group: AllocationGroup,
645 slots: HashMap<ValueKey<StdTensorOp>, DescriptorSlot>,
646}
647
648impl Gradients {
649 fn from_tensors(tensors: HashMap<ValueKey<StdTensorOp>, Tensor>) -> Result<Self> {
650 let (keys, values): (Vec<_>, Vec<_>) = tensors.into_iter().unzip();
651 let (group, bindings) = AllocationGroup::from_tensors(values).map_err(|error| {
652 Error::runtime_state_source("Gradients::from_tensors", ErrorPhase::Execution, error)
653 })?;
654 let slots = keys.into_iter().zip(bindings).collect();
655 Ok(Self { group, slots })
656 }
657
658 pub fn len(&self) -> usize {
660 self.slots.len()
661 }
662
663 pub fn is_empty(&self) -> bool {
665 self.slots.is_empty()
666 }
667
668 pub fn grad(&self, key: &ValueKey<StdTensorOp>) -> Option<TensorView<'_>> {
670 let slot = self.slots.get(key).copied()?;
671 let mut reads = self.group.read_views(std::slice::from_ref(&slot)).ok()?;
672 match reads.pop()? {
673 TensorRead::View(view) => Some(view),
674 TensorRead::Tensor(_) => None,
675 }
676 }
677
678 pub fn take_grad(
686 &mut self,
687 key: &ValueKey<StdTensorOp>,
688 ) -> tenferro_tensor::Result<Option<Tensor>> {
689 let Some(&slot) = self.slots.get(key) else {
690 return Ok(None);
691 };
692 let tensor = self.group.take_tensor(slot).map_err(|error| {
693 tenferro_tensor::Error::runtime_state_source("Gradients::take_grad", error)
694 })?;
695 self.slots.remove(key);
696 Ok(Some(tensor))
697 }
698}
699
700#[derive(Debug)]
721pub enum IntoValueError<H> {
722 NotUnique(H),
724 Extract { value: H, error: GroupError },
726}
727
728impl<H> std::fmt::Display for IntoValueError<H> {
729 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
730 match self {
731 Self::NotUnique(_) => formatter.write_str("eager value is retained by another handle"),
732 Self::Extract { error, .. } => {
733 write!(formatter, "eager value extraction failed: {error}")
734 }
735 }
736 }
737}
738
739impl<H: std::fmt::Debug + Send + Sync + 'static> std::error::Error for IntoValueError<H> {}
740
741#[derive(Debug)]
743struct RetentionContainer {
744 group: AllocationGroup,
745}
746
747#[derive(Debug)]
749pub(crate) struct AdValueRecord {
750 container: Arc<RetentionContainer>,
751 slot: DescriptorSlot,
752 dtype: DType,
753 shape: Box<[usize]>,
754}
755
756impl AdValueRecord {
757 fn from_group(
758 group: AllocationGroup,
759 slot: DescriptorSlot,
760 dtype: DType,
761 shape: Vec<usize>,
762 ) -> Arc<Self> {
763 Arc::new(Self {
764 container: Arc::new(RetentionContainer { group }),
765 slot,
766 dtype,
767 shape: shape.into_boxed_slice(),
768 })
769 }
770
771 fn from_tensor(tensor: Tensor, op: &'static str) -> Result<Arc<Self>> {
772 let dtype = tensor.dtype();
773 let shape = tensor.shape().to_vec();
774 let (group, bindings) = AllocationGroup::from_tensors(vec![tensor])
775 .map_err(|error| Error::runtime_state_source(op, ErrorPhase::Execution, error))?;
776 let slot = bindings.first().copied().ok_or_else(|| {
777 Error::runtime_state(op, ErrorPhase::Execution, "empty allocation-group binding")
778 })?;
779 Ok(Self::from_group(group, slot, dtype, shape))
780 }
781
782 fn tensor_read(&self, op: &'static str) -> Result<TensorRead<'_>> {
783 let mut reads = self
784 .container
785 .group
786 .read_views(std::slice::from_ref(&self.slot))
787 .map_err(|error| Error::runtime_state_source(op, ErrorPhase::Execution, error))?;
788 reads.pop().ok_or_else(|| {
789 Error::runtime_state(op, ErrorPhase::Execution, "empty allocation-group binding")
790 })
791 }
792
793 fn value(&self, op: &'static str) -> Result<ValueGuard<'_>> {
794 match self.tensor_read(op)? {
795 TensorRead::View(view) => Ok(ValueGuard { view }),
796 TensorRead::Tensor(_) => Err(Error::runtime_state(
797 op,
798 ErrorPhase::Execution,
799 "allocation-group value did not produce a borrowed descriptor view",
800 )),
801 }
802 }
803
804 fn dtype(&self) -> DType {
805 self.dtype
806 }
807
808 fn shape(&self) -> &[usize] {
809 &self.shape
810 }
811}
812
813pub struct CpuPlacementBoundEager {
838 runtime: Arc<EagerRuntime>,
839 backend: CpuBackend,
840 snapshot: Arc<RuntimeConfigSnapshot>,
841 epoch: RuntimeEpoch,
842 engine_id: EngineId,
843 registration_identity: RegistrationIdentity,
844 capabilities: CoreCapabilityBundle,
845}
846
847impl fmt::Debug for CpuPlacementBoundEager {
848 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
849 f.debug_struct("CpuPlacementBoundEager")
850 .field("runtime_id", &self.runtime.id())
851 .field("placement", &self.backend.placement())
852 .field("runtime_epoch", &self.epoch)
853 .field("engine_id", &self.engine_id)
854 .field("registration_identity", &self.registration_identity)
855 .finish_non_exhaustive()
856 }
857}
858
859impl CpuPlacementBoundEager {
860 fn refresh_runtime_selection(&mut self) -> Result<()> {
861 let current_epoch = self.runtime.runtime.epoch().map_err(|source| {
862 runtime_state_source("CpuPlacementBoundEager::refresh_runtime_selection", source)
863 })?;
864 if current_epoch == self.epoch {
865 return Ok(());
866 }
867
868 #[cfg(test)]
869 CPU_RUNTIME_SELECTION_REFRESHES.fetch_add(1, Ordering::SeqCst);
870
871 let selection = select_cpu_runtime(&self.runtime.runtime)?;
872 self.snapshot = selection.snapshot;
873 self.epoch = selection.epoch;
874 self.engine_id = selection.engine_id;
875 self.registration_identity = selection.registration_identity;
876 self.capabilities = selection.capabilities;
877 Ok(())
878 }
879
880 pub fn runtime_id(&self) -> ContextId {
894 self.runtime.id()
895 }
896
897 pub fn placement(&self) -> CpuPlacement {
911 self.backend.placement()
912 }
913
914 pub fn with_eager_session<R: Send>(
955 &mut self,
956 f: impl FnOnce(&mut dyn BackendSession) -> Result<R> + Send,
957 ) -> Result<R> {
958 self.refresh_runtime_selection()?;
959 self.backend.with_backend_session(f)
960 }
961}
962
963pub struct EagerRuntime {
983 id: ContextId,
984 runtime: Runtime,
985 backend: Mutex<EagerBackend>,
989 extension_install_lock: Mutex<()>,
990 pub(crate) extension_caches: Mutex<ExtensionCacheStore>,
991 semantic_extension_rules: SemanticExtensionRuleSet,
992 grad_slots: Mutex<HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>,
993 value_records: Mutex<HashMap<ValueKey<StdTensorOp>, Weak<EagerTensorRecord>>>,
994 ad_transform_cache: Arc<AdTransformCache>,
995 prepared_derivative_cache: Mutex<PreparedDerivativeCache>,
999}
1000
1001impl fmt::Debug for EagerRuntime {
1002 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1003 let mut debug = f.debug_struct("EagerRuntime");
1004 debug.field("id", &self.id);
1005 debug.field("runtime_id", &self.runtime.id());
1006 debug.field("runtime_epoch", &self.runtime.epoch().ok());
1007 match self.backend.try_lock() {
1008 Ok(backend) => {
1009 debug.field("backend", &*backend);
1010 }
1011 Err(_) => {
1012 debug.field("backend", &"<locked>");
1013 }
1014 }
1015 match self.extension_caches.try_lock() {
1016 Ok(caches) => {
1017 debug.field(
1018 "extension_cache_stats",
1019 &caches.stats(ExtensionCacheSelector::All),
1020 );
1021 }
1022 Err(_) => {
1023 debug.field("extension_cache_stats", &"<locked>");
1024 }
1025 }
1026 match self.extension_install_lock.try_lock() {
1027 Ok(_) => {
1028 debug.field("extension_install_lock", &"<unlocked>");
1029 }
1030 Err(_) => {
1031 debug.field("extension_install_lock", &"<locked>");
1032 }
1033 }
1034 debug.field("semantic_extension_rules", &self.semantic_extension_rules);
1035 match self.grad_slots.try_lock() {
1036 Ok(slots) => {
1037 debug.field("grad_slots_len", &slots.len());
1038 }
1039 Err(_) => {
1040 debug.field("grad_slots_len", &"<locked>");
1041 }
1042 }
1043 match self.value_records.try_lock() {
1044 Ok(records) => {
1045 debug.field("value_records_len", &records.len());
1046 }
1047 Err(_) => {
1048 debug.field("value_records_len", &"<locked>");
1049 }
1050 }
1051 match self.ad_transform_cache.stats() {
1052 Ok(stats) => {
1053 debug.field("ad_transform_cache_stats", &stats);
1054 }
1055 Err(err) => {
1056 debug.field("ad_transform_cache_stats", &format_args!("{err}"));
1057 }
1058 }
1059 match self.prepared_derivative_cache.try_lock() {
1060 Ok(cache) => {
1061 debug.field("prepared_derivative_cache_stats", &cache.stats());
1062 }
1063 Err(_) => {
1064 debug.field("prepared_derivative_cache_stats", &"<locked>");
1065 }
1066 }
1067 debug.finish_non_exhaustive()
1068 }
1069}
1070
1071impl EagerRuntime {
1072 pub(crate) fn lock_backend(&self) -> Result<MutexGuard<'_, EagerBackend>> {
1073 self.backend.lock().map_err(|_| {
1074 Error::runtime_state("eager_backend", ErrorPhase::Execution, "lock poisoned")
1075 })
1076 }
1077
1078 fn lock_extension_caches(&self) -> Result<MutexGuard<'_, ExtensionCacheStore>> {
1079 self.extension_caches.lock().map_err(|_| {
1080 Error::runtime_state(
1081 "eager_extension_caches",
1082 ErrorPhase::Execution,
1083 "lock poisoned",
1084 )
1085 })
1086 }
1087
1088 fn lock_extension_install(&self) -> Result<MutexGuard<'_, ()>> {
1089 self.extension_install_lock.lock().map_err(|_| {
1090 Error::runtime_state(
1091 "eager_extension_install",
1092 ErrorPhase::Execution,
1093 "lock poisoned",
1094 )
1095 })
1096 }
1097
1098 fn lock_prepared_derivative_cache(&self) -> Result<MutexGuard<'_, PreparedDerivativeCache>> {
1099 self.prepared_derivative_cache.lock().map_err(|_| {
1100 Error::runtime_state(
1101 "prepared_derivative_cache",
1102 ErrorPhase::Execution,
1103 "lock poisoned",
1104 )
1105 })
1106 }
1107
1108 fn lock_grad_slots(
1109 &self,
1110 ) -> Result<MutexGuard<'_, HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>> {
1111 self.grad_slots.lock().map_err(|_| {
1112 Error::runtime_state(
1113 "eager_gradient_slots",
1114 ErrorPhase::Execution,
1115 "lock poisoned",
1116 )
1117 })
1118 }
1119
1120 fn lock_value_records(
1121 &self,
1122 ) -> Result<MutexGuard<'_, HashMap<ValueKey<StdTensorOp>, Weak<EagerTensorRecord>>>> {
1123 self.value_records.lock().map_err(|_| {
1124 Error::runtime_state(
1125 "eager_value_registry",
1126 ErrorPhase::Execution,
1127 "lock poisoned",
1128 )
1129 })
1130 }
1131
1132 fn from_backend(backend: EagerBackend) -> Result<Self> {
1133 Self::from_backend_with_rules_and_cache(
1134 backend,
1135 SemanticExtensionRuleSet::default(),
1136 Arc::new(AdTransformCache::new()),
1137 )
1138 }
1139
1140 fn from_backend_with_rules_and_cache(
1141 backend: EagerBackend,
1142 semantic_extension_rules: SemanticExtensionRuleSet,
1143 ad_transform_cache: Arc<AdTransformCache>,
1144 ) -> Result<Self> {
1145 let runtime = eager_runtime_for_backend(&backend)
1146 .map_err(|source| runtime_config_error("EagerRuntime::from_backend", source))?;
1147 Ok(Self {
1148 id: ContextId::fresh(),
1149 runtime,
1150 backend: Mutex::new(backend),
1151 extension_install_lock: Mutex::new(()),
1152 extension_caches: Mutex::new(ExtensionCacheStore::new()),
1153 semantic_extension_rules,
1154 grad_slots: Mutex::new(HashMap::new()),
1155 value_records: Mutex::new(HashMap::new()),
1156 ad_transform_cache,
1157 prepared_derivative_cache: Mutex::new(PreparedDerivativeCache::default()),
1158 })
1159 }
1160
1161 pub fn new() -> Result<Arc<Self>> {
1179 Self::with_cpu_backend(CpuBackend::new())
1180 }
1181
1182 pub fn with_cpu_backend(backend: CpuBackend) -> Result<Arc<Self>> {
1201 Ok(Arc::new(Self::from_backend(EagerBackend::cpu(backend))?))
1202 }
1203
1204 pub fn on_cpu(self: &Arc<Self>, placement: CpuPlacement) -> Result<CpuPlacementBoundEager> {
1230 let backend = {
1231 let backend = self.lock_backend()?;
1232 backend.cpu_snapshot().ok_or_else(|| {
1233 Error::unsupported(
1234 "EagerRuntime::on_cpu",
1235 ErrorPhase::Execution,
1236 "the eager runtime is not CPU-backed",
1237 )
1238 })?
1239 };
1240 let selection = select_cpu_runtime(&self.runtime)?;
1241 let backend = backend.for_placement(placement).map_err(|source| {
1242 let error: tenferro_tensor::Error = CpuBackendError::Placement {
1243 op: "EagerRuntime::on_cpu",
1244 source,
1245 }
1246 .into();
1247 Error::from(error)
1248 })?;
1249 Ok(CpuPlacementBoundEager {
1250 runtime: Arc::clone(self),
1251 backend,
1252 snapshot: selection.snapshot,
1253 epoch: selection.epoch,
1254 engine_id: selection.engine_id,
1255 registration_identity: selection.registration_identity,
1256 capabilities: selection.capabilities,
1257 })
1258 }
1259
1260 pub fn with_cpu_backend_and_ad_context(
1280 backend: CpuBackend,
1281 ad: &AdContext,
1282 ) -> Result<Arc<Self>> {
1283 Ok(Arc::new(Self::from_backend_with_rules_and_cache(
1284 EagerBackend::cpu(backend),
1285 ad.semantic_extension_rules().clone(),
1286 ad.ad_transform_cache(),
1287 )?))
1288 }
1289
1290 #[cfg(feature = "cuda")]
1302 pub fn with_cuda_backend(backend: CudaBackend) -> Result<Arc<Self>> {
1309 Ok(Arc::new(Self::from_backend(EagerBackend::cuda(backend))?))
1310 }
1311
1312 #[cfg(feature = "cuda")]
1324 pub fn with_cuda_backend_and_ad_context(
1331 backend: CudaBackend,
1332 ad: &AdContext,
1333 ) -> Result<Arc<Self>> {
1334 Ok(Arc::new(Self::from_backend_with_rules_and_cache(
1335 EagerBackend::cuda(backend),
1336 ad.semantic_extension_rules().clone(),
1337 ad.ad_transform_cache(),
1338 )?))
1339 }
1340
1341 #[cfg(feature = "webgpu")]
1353 pub fn with_webgpu_backend(backend: WebGpuBackend) -> Result<Arc<Self>> {
1360 Ok(Arc::new(Self::from_backend(EagerBackend::webgpu(backend))?))
1361 }
1362
1363 #[cfg(feature = "webgpu")]
1375 pub fn with_webgpu_backend_and_ad_context(
1382 backend: WebGpuBackend,
1383 ad: &AdContext,
1384 ) -> Result<Arc<Self>> {
1385 Ok(Arc::new(Self::from_backend_with_rules_and_cache(
1386 EagerBackend::webgpu(backend),
1387 ad.semantic_extension_rules().clone(),
1388 ad.ad_transform_cache(),
1389 )?))
1390 }
1391
1392 pub fn id(&self) -> ContextId {
1405 self.id
1406 }
1407
1408 pub fn no_grad(&self) -> EagerNoGradGuard {
1432 EAGER_NO_GRAD_DEPTH.with(|depth| {
1433 depth.set(depth.get().saturating_add(1));
1434 });
1435 EagerNoGradGuard {
1436 active: true,
1437 _not_send: PhantomData,
1438 }
1439 }
1440
1441 pub fn capture_trace(&self) -> EagerTraceCaptureGuard {
1445 EAGER_CAPTURE_DEPTH.with(|depth| {
1446 depth.set(depth.get().saturating_add(1));
1447 });
1448 EagerTraceCaptureGuard {
1449 active: true,
1450 _not_send: PhantomData,
1451 }
1452 }
1453
1454 pub fn install_extension_module(
1469 &self,
1470 module: Arc<dyn ExtensionModule>,
1471 ) -> Result<RuntimeEpoch> {
1472 let snapshot = self.runtime.snapshot().map_err(|source| {
1473 runtime_state_source("EagerRuntime::install_extension_module", source)
1474 })?;
1475 if snapshot.has_extension_module_identical(&module) {
1476 return Ok(snapshot.epoch());
1477 }
1478 let _install_guard = self.lock_extension_install()?;
1479 self.runtime
1480 .reconfigure(|edit| {
1481 edit.replace_extension_module(module)?;
1482 Ok(())
1483 })
1484 .map_err(|source| {
1485 runtime_state_source("EagerRuntime::install_extension_module", source)
1486 })
1487 }
1488
1489 pub(crate) fn ensure_extension_module_for_engine(
1490 &self,
1491 module: Arc<dyn ExtensionModule>,
1492 family_id: &'static str,
1493 engine_id: &EngineId,
1494 ) -> Result<RuntimeEpoch> {
1495 let snapshot = self.runtime.snapshot().map_err(|source| {
1496 runtime_state_source("EagerRuntime::ensure_extension_module_for_engine", source)
1497 })?;
1498 if snapshot.has_extension_module_engine(module.module_id(), family_id, engine_id) {
1499 return Ok(snapshot.epoch());
1500 }
1501 let _install_guard = self.lock_extension_install()?;
1502 self.runtime
1503 .reconfigure(|edit| {
1504 edit.ensure_extension_module_for_engine(module, family_id, engine_id)?;
1505 Ok(())
1506 })
1507 .map_err(|source| {
1508 runtime_state_source("EagerRuntime::ensure_extension_module_for_engine", source)
1509 })
1510 }
1511
1512 pub(crate) fn runtime(&self) -> &Runtime {
1513 &self.runtime
1514 }
1515
1516 pub(crate) fn eager_extension_target(&self) -> Result<EagerExtensionTarget> {
1517 let (engine_id, backend_kind) = {
1518 let backend = self.lock_backend()?;
1519 match &*backend {
1520 EagerBackend::Cpu(_) => (
1521 cpu_runtime_engine_id().map_err(|source| {
1522 runtime_config_error("EagerRuntime::eager_extension_target", source)
1523 })?,
1524 EagerExtensionBackendKind::Cpu,
1525 ),
1526 #[cfg(test)]
1527 EagerBackend::Recording(_) => {
1528 return Err(Error::unsupported(
1529 "EagerRuntime::eager_extension_target",
1530 ErrorPhase::Execution,
1531 "the recording backend has no registered eager extension engine",
1532 ));
1533 }
1534 #[cfg(feature = "cuda")]
1535 EagerBackend::Cuda(_) => (
1536 cuda_runtime_engine_id().map_err(|source| {
1537 runtime_config_error("EagerRuntime::eager_extension_target", source)
1538 })?,
1539 EagerExtensionBackendKind::Cuda,
1540 ),
1541 #[cfg(feature = "webgpu")]
1542 EagerBackend::WebGpu(_) => (
1543 tenferro_gpu::webgpu::webgpu_runtime_engine_id().map_err(|source| {
1544 runtime_config_error("EagerRuntime::eager_extension_target", source)
1545 })?,
1546 EagerExtensionBackendKind::WebGpu,
1547 ),
1548 }
1549 };
1550 let target = EagerExtensionTarget {
1551 engine_id,
1552 backend_kind,
1553 };
1554 validate_eager_extension_target(&self.runtime, &target)?;
1555 Ok(target)
1556 }
1557
1558 pub fn clear_extension_caches(&self) -> Result<()> {
1577 self.lock_extension_caches()?.clear();
1578 Ok(())
1579 }
1580
1581 pub fn clear_caches(&self) -> Result<()> {
1602 self.clear_extension_caches()?;
1603 self.clear_ad_transform_caches()?;
1604 self.clear_prepared_derivative_cache()?;
1605 Ok(())
1606 }
1607
1608 pub fn clear_prepared_derivative_cache(&self) -> Result<()> {
1627 self.lock_prepared_derivative_cache()?.clear();
1628 Ok(())
1629 }
1630
1631 pub fn cache_stats(&self) -> Result<EagerRuntimeCacheStats> {
1652 Ok(EagerRuntimeCacheStats {
1653 extensions: self
1654 .lock_extension_caches()?
1655 .stats(ExtensionCacheSelector::All),
1656 ad_transforms: self.ad_transform_cache.stats()?,
1657 prepared_derivatives: self.lock_prepared_derivative_cache()?.stats(),
1658 })
1659 }
1660
1661 pub fn ad_transform_cache_limits(&self) -> Result<AdTransformCacheLimits> {
1679 self.ad_transform_cache.limits()
1680 }
1681
1682 pub fn set_ad_transform_cache_limits(&self, limits: AdTransformCacheLimits) -> Result<()> {
1703 self.ad_transform_cache.set_limits(limits)
1704 }
1705
1706 pub fn clear_ad_transform_caches(&self) -> Result<()> {
1725 self.ad_transform_cache.clear()
1726 }
1727
1728 pub fn prepared_derivative_cache_limits(&self) -> Result<AdTransformCacheLimits> {
1746 Ok(self.lock_prepared_derivative_cache()?.limits())
1747 }
1748
1749 pub fn set_prepared_derivative_cache_limits(
1770 &self,
1771 limits: AdTransformCacheLimits,
1772 ) -> Result<()> {
1773 self.lock_prepared_derivative_cache()?.set_limits(limits);
1774 Ok(())
1775 }
1776
1777 pub fn extension_cache_limits(&self) -> Result<ExtensionCacheLimits> {
1784 Ok(self.lock_extension_caches()?.limits())
1785 }
1786
1787 pub fn set_extension_cache_limits(&self, limits: ExtensionCacheLimits) -> Result<()> {
1794 self.lock_extension_caches()?.set_limits(limits);
1795 Ok(())
1796 }
1797
1798 pub fn with_execution_session<R: Send>(
1828 &self,
1829 f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
1830 ) -> Result<R> {
1831 let mut backend = self.lock_backend()?;
1832 Ok(backend.with_backend_session(f))
1833 }
1834
1835 pub fn with_extension_execution_context<R: Send>(
1871 &self,
1872 f: impl FnOnce(
1873 &mut tenferro_runtime::ExtensionExecutionContext<'_, dyn BackendSession + '_>,
1874 ) -> R
1875 + Send,
1876 ) -> Result<R> {
1877 let mut backend = self.lock_backend()?;
1878 let mut extension_cache_guard = self.lock_extension_caches()?;
1879 let extension_caches: &mut ExtensionCacheStore = &mut extension_cache_guard;
1880 Ok(backend.with_backend_session(move |session| {
1881 let mut extension_ctx =
1882 tenferro_runtime::ExtensionExecutionContext::new(session, extension_caches);
1883 f(&mut extension_ctx)
1884 }))
1885 }
1886
1887 pub(crate) fn with_extension_erased_context<R: Send>(
1896 &self,
1897 f: impl FnOnce(&mut tenferro_runtime::ErasedExecutionContext<'_>, &mut ExtensionCacheStore) -> R
1898 + Send,
1899 ) -> Result<R> {
1900 let mut backend = self.lock_backend()?;
1901 let mut extension_cache_guard = self.lock_extension_caches()?;
1902 let extension_caches: &mut ExtensionCacheStore = &mut extension_cache_guard;
1903 let mut erased = backend.erased_context();
1904 Ok(f(&mut erased, extension_caches))
1905 }
1906
1907 pub fn synchronize(&self) -> Result<()> {
1928 self.lock_backend()?.synchronize().map_err(Error::from)
1929 }
1930
1931 fn exec_outputs_with_runtime<R>(
1932 &self,
1933 lock_backend_section: &'static str,
1934 exec_section: &'static str,
1935 op: &StdTensorOp,
1936 execute: impl FnOnce(&mut EagerBackend, Option<&Runtime>) -> Result<R>,
1937 ) -> Result<R> {
1938 let mut backend = profile_eager_op_section(lock_backend_section, || self.lock_backend())?;
1942 let runtime = matches!(op, StdTensorOp::Extension(_)).then_some(&self.runtime);
1943 profile_eager_op_section(exec_section, || execute(&mut backend, runtime))
1944 }
1945
1946 pub(crate) fn exec_outputs(&self, op: &StdTensorOp, inputs: &[&Tensor]) -> Result<Vec<Tensor>> {
1947 self.exec_outputs_with_runtime(
1948 "exec_outputs.lock_backend",
1949 "exec_outputs.exec_op",
1950 op,
1951 |backend, runtime| exec_op_on_tensors_with_runtime(op, inputs, backend, runtime),
1952 )
1953 }
1954
1955 pub(crate) fn exec_outputs_read(
1956 &self,
1957 op: &StdTensorOp,
1958 inputs: &[TensorRead<'_>],
1959 ) -> Result<Vec<Tensor>> {
1960 self.exec_outputs_with_runtime(
1961 "exec_outputs_read.lock_backend",
1962 "exec_outputs_read.exec_op",
1963 op,
1964 |backend, runtime| exec_op_on_tensor_reads_with_runtime(op, inputs, backend, runtime),
1965 )
1966 }
1967
1968 #[cfg(test)]
1969 pub(crate) fn exec_standard_graph_outputs(
1970 &self,
1971 graph: &Graph<StdTensorOp>,
1972 initial_data: HashMap<ValueKey<StdTensorOp>, Tensor>,
1973 ) -> Result<EagerGraphExecution> {
1974 let mut backend =
1975 profile_eager_op_section("exec_graph.lock_backend", || self.lock_backend())?;
1976 let mut all_values = initial_data;
1977
1978 profile_eager_op_section("exec_graph.with_backend_session", || {
1979 backend.with_backend_session(|exec| -> Result<()> {
1980 for op_node in graph.operations() {
1981 let outputs = {
1982 let input_values = op_node
1983 .inputs
1984 .iter()
1985 .map(|input| {
1986 let key = match input {
1987 ValueRef::Local(local_id) => &graph.values()[*local_id].key,
1988 ValueRef::External(key) => key,
1989 };
1990 all_values.get(key).ok_or_else(|| {
1991 Error::Internal(format!(
1992 "standard graph eager execution missing value for {key:?}"
1993 ))
1994 })
1995 })
1996 .collect::<Result<Vec<_>>>()?;
1997 let input_reads = input_values
1998 .iter()
1999 .map(|value| TensorRead::from_tensor(value))
2000 .collect::<Vec<_>>();
2001 exec_standard_op_on_tensor_reads_in_session(
2002 &op_node.operation,
2003 &input_reads,
2004 exec,
2005 )?
2006 };
2007
2008 if outputs.len() != op_node.outputs.len() {
2009 return Err(Error::Internal(format!(
2010 "standard graph eager execution expected {} outputs for {:?}, got {}",
2011 op_node.outputs.len(),
2012 op_node.operation,
2013 outputs.len()
2014 )));
2015 }
2016
2017 for (output_id, output) in op_node.outputs.iter().zip(outputs) {
2018 let key = graph.values()[*output_id].key.clone();
2019 all_values.insert(key, output);
2020 }
2021 }
2022 Ok(())
2023 })
2024 })?;
2025
2026 let outputs = graph
2027 .outputs()
2028 .iter()
2029 .map(|&output_id| {
2030 let key = &graph.values()[output_id].key;
2031 all_values
2032 .get(key)
2033 .ok_or_else(|| {
2034 Error::Internal(format!(
2035 "standard graph eager execution missing graph output {key:?}"
2036 ))
2037 })?
2038 .duplicate()
2039 .map_err(Error::from)
2040 })
2041 .collect::<Result<Vec<_>>>()?;
2042
2043 Ok(EagerGraphExecution { outputs })
2044 }
2045
2046 pub(crate) fn try_register_grad_slot(
2047 &self,
2048 key: &ValueKey<StdTensorOp>,
2049 slot: &GradSlot,
2050 ) -> Result<()> {
2051 self.lock_grad_slots()?
2052 .insert(key.clone(), Arc::downgrade(slot));
2053 Ok(())
2054 }
2055
2056 pub(crate) fn try_register_value_record(
2057 &self,
2058 key: &ValueKey<StdTensorOp>,
2059 record: &Arc<EagerTensorRecord>,
2060 ) -> Result<()> {
2061 self.lock_value_records()?
2062 .insert(key.clone(), Arc::downgrade(record));
2063 Ok(())
2064 }
2065
2066 pub(crate) fn value_record(
2067 &self,
2068 key: &ValueKey<StdTensorOp>,
2069 ) -> Result<Option<Arc<EagerTensorRecord>>> {
2070 let mut records = self.lock_value_records()?;
2071 let Some(record) = records.get(key).cloned() else {
2072 return Ok(None);
2073 };
2074 match record.upgrade() {
2075 Some(record) => Ok(Some(record)),
2076 None => {
2077 records.remove(key);
2078 Ok(None)
2079 }
2080 }
2081 }
2082
2083 pub fn clear_grads(&self) -> Result<()> {
2112 let live_slots = {
2113 let mut live_slots = Vec::new();
2114 self.lock_grad_slots()?.retain(|_, slot| {
2115 if let Some(slot) = slot.upgrade() {
2116 live_slots.push(slot);
2117 true
2118 } else {
2119 false
2120 }
2121 });
2122 live_slots
2123 };
2124
2125 let mut poisoned_slot = false;
2126 for slot in live_slots {
2127 match slot.lock() {
2128 Ok(mut current) => {
2129 *current = None;
2130 }
2131 Err(_) => {
2132 poisoned_slot = true;
2133 }
2134 }
2135 }
2136 if poisoned_slot {
2137 return Err(Error::runtime_state(
2138 "eager_gradient_slot",
2139 ErrorPhase::Execution,
2140 "lock poisoned",
2141 ));
2142 }
2143 Ok(())
2144 }
2145
2146 pub fn constant_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
2172 EagerTensor::new_leaf(Arc::clone(self), tensor, false)
2173 }
2174
2175 pub fn variable_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
2201 EagerTensor::new_leaf(Arc::clone(self), tensor, true)
2202 }
2203
2204 pub fn grad(self: &Arc<Self>, output: &EagerTensor, wrt: &EagerTensor) -> Result<EagerTensor> {
2234 self.grad_optional(output, wrt)?
2235 .ok_or_else(|| Error::Internal(format!("grad output is inactive for {:?}", wrt.key)))
2236 }
2237
2238 pub fn grad_optional(
2266 self: &Arc<Self>,
2267 output: &EagerTensor,
2268 wrt: &EagerTensor,
2269 ) -> Result<Option<EagerTensor>> {
2270 if !output.shape().is_empty() {
2271 return Err(Error::NonScalarGrad {
2272 shape: output.shape().to_vec(),
2273 });
2274 }
2275
2276 let value = output.to_tensor()?;
2277 let seed = {
2278 let mut backend = self.lock_backend()?;
2279 one_like_tensor(&value, &mut *backend)?
2280 };
2281 let seed = EagerTensor::new_result(Arc::clone(self), eager_val_key(), seed, false, None)?;
2282 self.vjp_optional(output, wrt, &seed)
2283 }
2284
2285 pub fn vjp(
2315 self: &Arc<Self>,
2316 output: &EagerTensor,
2317 wrt: &EagerTensor,
2318 cotangent: &EagerTensor,
2319 ) -> Result<EagerTensor> {
2320 self.vjp_optional(output, wrt, cotangent)?
2321 .ok_or_else(|| Error::Internal(format!("vjp output is inactive for {:?}", wrt.key)))
2322 }
2323
2324 pub fn vjp_optional(
2357 self: &Arc<Self>,
2358 output: &EagerTensor,
2359 wrt: &EagerTensor,
2360 cotangent: &EagerTensor,
2361 ) -> Result<Option<EagerTensor>> {
2362 validate_same_runtime(self, output, "vjp output")?;
2363 validate_same_runtime(self, wrt, "vjp wrt")?;
2364 validate_same_runtime(self, cotangent, "vjp cotangent")?;
2365 validate_seed_tensor("vjp", output, cotangent)?;
2366 match semantic_eager_vjp_optional(self, output, wrt, cotangent)? {
2368 Some(result) => Ok(result),
2369 None => Ok(None),
2370 }
2371 }
2372
2373 pub fn jvp(
2403 self: &Arc<Self>,
2404 output: &EagerTensor,
2405 wrt: &EagerTensor,
2406 tangent: &EagerTensor,
2407 ) -> Result<EagerTensor> {
2408 self.jvp_optional(output, wrt, tangent)?
2409 .ok_or_else(|| Error::Internal(format!("jvp output is inactive for {:?}", wrt.key)))
2410 }
2411
2412 pub fn jvp_optional(
2445 self: &Arc<Self>,
2446 output: &EagerTensor,
2447 wrt: &EagerTensor,
2448 tangent: &EagerTensor,
2449 ) -> Result<Option<EagerTensor>> {
2450 validate_same_runtime(self, output, "jvp output")?;
2451 validate_same_runtime(self, wrt, "jvp wrt")?;
2452 validate_same_runtime(self, tangent, "jvp tangent")?;
2453 validate_seed_tensor("jvp", wrt, tangent)?;
2454 match semantic_eager_jvp_optional(self, output, wrt, tangent)? {
2456 Some(result) => Ok(result),
2457 None => Ok(None),
2458 }
2459 }
2460
2461 fn store_grads(
2462 &self,
2463 cotangents: &HashMap<ValueKey<StdTensorOp>, Tensor>,
2464 backend: &mut EagerBackend,
2465 ) -> Result<()> {
2466 let mut updates = Vec::new();
2467
2468 {
2469 let mut slots = self.lock_grad_slots()?;
2470 slots.retain(|key, slot| {
2471 let Some(slot) = slot.upgrade() else {
2472 return false;
2473 };
2474
2475 if let Some(incoming) = cotangents.get(key) {
2476 updates.push((slot, incoming));
2477 }
2478
2479 true
2480 });
2481 }
2482
2483 for (slot, incoming) in updates {
2484 let mut current = slot.lock().map_err(|_| {
2485 Error::runtime_state(
2486 "eager_gradient_slot",
2487 ErrorPhase::Execution,
2488 "lock poisoned",
2489 )
2490 })?;
2491 let next = match current.as_ref() {
2492 Some(existing) => {
2493 let existing_read = existing.tensor_read("EagerRuntime::store_grads")?;
2494 let incoming_read = TensorRead::from_tensor(incoming);
2495 let tensor = backend
2496 .with_backend_session(|session| {
2497 session.add_read(existing_read, incoming_read)
2498 })
2499 .map_err(Error::from)?;
2500 AdValueRecord::from_tensor(tensor, "EagerRuntime::store_grads")?
2501 }
2502 None => {
2503 let duplicate = backend
2504 .with_backend_session(|session| {
2505 session.to_contiguous_read(TensorRead::from_tensor(incoming))
2506 })
2507 .map_err(Error::from)?;
2508 AdValueRecord::from_tensor(duplicate, "EagerRuntime::store_grads")?
2509 }
2510 };
2511 *current = Some(next);
2512 }
2513
2514 Ok(())
2515 }
2516}
2517
2518#[derive(Clone, Debug, PartialEq, Eq, Hash)]
2519struct PreparedDerivativeCacheKey {
2520 semantic_fingerprint: SemanticFingerprint,
2521 runtime_epoch: RuntimeEpoch,
2522 wrt_input_index: usize,
2523 input_metadata: Box<[ProgramValueMetadata]>,
2524}
2525
2526#[derive(Debug)]
2528struct PreparedDerivative {
2529 program: Arc<CompiledGraph>,
2530 prepared: Arc<PreparedCompiledGraph>,
2531 seed_input_index: usize,
2532 derivative_output_index: usize,
2533}
2534
2535#[derive(Debug)]
2536struct PreparedDerivativeCache {
2537 limits: AdTransformCacheLimits,
2538 entries: LruCache<PreparedDerivativeCacheKey, PreparedDerivativeCacheEntry>,
2539 stats: CacheStats,
2540}
2541
2542impl PreparedDerivativeCache {
2543 fn limits(&self) -> AdTransformCacheLimits {
2544 self.limits
2545 }
2546
2547 fn set_limits(&mut self, limits: AdTransformCacheLimits) {
2548 self.limits = limits;
2549 self.evict_to_limits();
2550 }
2551
2552 fn clear(&mut self) {
2553 let clears = self.stats.clears.saturating_add(1);
2554 self.entries.clear();
2555 self.stats = CacheStats {
2556 clears,
2557 ..CacheStats::empty()
2558 };
2559 }
2560
2561 fn stats(&self) -> CacheStats {
2562 self.stats
2563 }
2564
2565 fn get(&mut self, key: &PreparedDerivativeCacheKey) -> Option<Arc<PreparedDerivative>> {
2566 match self.entries.get(key) {
2567 Some(entry) => {
2568 self.stats.hits = self.stats.hits.saturating_add(1);
2569 Some(Arc::clone(&entry.value))
2570 }
2571 None => {
2572 self.stats.misses = self.stats.misses.saturating_add(1);
2573 None
2574 }
2575 }
2576 }
2577
2578 fn insert(&mut self, key: PreparedDerivativeCacheKey, value: Arc<PreparedDerivative>) {
2579 let retained_bytes = prepared_derivative_cache_entry_retained_bytes(&key, value.as_ref());
2580 let entry = PreparedDerivativeCacheEntry {
2581 value,
2582 retained_bytes,
2583 };
2584 self.stats.retained_bytes = self.stats.retained_bytes.saturating_add(retained_bytes);
2585 if let Some((_old_key, old_entry)) = self.entries.push(key, entry) {
2586 self.stats.retained_bytes = self
2587 .stats
2588 .retained_bytes
2589 .saturating_sub(old_entry.retained_bytes);
2590 }
2591 self.stats.entries = self.entries.len();
2592 self.evict_to_limits();
2593 }
2594
2595 fn evict_to_limits(&mut self) {
2596 while self.entries.len() > self.limits.max_entries().get()
2597 || self
2598 .limits
2599 .max_retained_bytes()
2600 .is_some_and(|limit| self.stats.retained_bytes > limit.get())
2601 {
2602 let Some((_key, entry)) = self.entries.pop_lru() else {
2603 break;
2604 };
2605 self.stats.retained_bytes = self
2606 .stats
2607 .retained_bytes
2608 .saturating_sub(entry.retained_bytes);
2609 self.stats.evictions = self.stats.evictions.saturating_add(1);
2610 }
2611 self.stats.entries = self.entries.len();
2612 }
2613}
2614
2615impl Default for PreparedDerivativeCache {
2616 fn default() -> Self {
2617 Self {
2618 limits: AdTransformCacheLimits::default(),
2619 entries: LruCache::unbounded(),
2620 stats: CacheStats::empty(),
2621 }
2622 }
2623}
2624
2625#[derive(Debug)]
2626struct PreparedDerivativeCacheEntry {
2627 value: Arc<PreparedDerivative>,
2628 retained_bytes: usize,
2629}
2630
2631fn prepared_derivative_cache_entry_retained_bytes(
2632 key: &PreparedDerivativeCacheKey,
2633 value: &PreparedDerivative,
2634) -> usize {
2635 size_of::<PreparedDerivativeCacheKey>()
2636 .saturating_add(
2637 key.input_metadata
2638 .len()
2639 .saturating_mul(size_of::<ProgramValueMetadata>()),
2640 )
2641 .saturating_add(size_of::<PreparedDerivative>())
2642 .saturating_add(compiled_graph_retained_bytes(value.program.as_ref()))
2643 .saturating_add(prepared_compiled_graph_retained_bytes(
2644 value.prepared.as_ref(),
2645 value.program.as_ref(),
2646 ))
2647}
2648
2649fn prepared_compiled_graph_retained_bytes(
2650 prepared: &PreparedCompiledGraph,
2651 derivative_program: &CompiledGraph,
2652) -> usize {
2653 size_of_val(prepared).saturating_add(compiled_graph_retained_bytes(derivative_program))
2654}
2655
2656fn compiled_graph_retained_bytes(program: &CompiledGraph) -> usize {
2657 size_of::<CompiledGraph>()
2658 .saturating_add(size_of_val(program.input_keys()))
2659 .saturating_add(program.bindings().len().saturating_mul(size_of::<usize>()))
2660 .saturating_add(semantic_program_retained_bytes(program.program()))
2661}
2662
2663fn semantic_program_retained_bytes(program: &SemanticProgram) -> usize {
2664 size_of::<SemanticProgram>()
2665 .saturating_add(size_of_val(program.inputs()))
2666 .saturating_add(size_of_val(program.outputs()))
2667 .saturating_add(
2668 program
2669 .operations()
2670 .len()
2671 .saturating_mul(size_of::<usize>()),
2672 )
2673 .saturating_add(
2674 program
2675 .shape_guards()
2676 .len()
2677 .saturating_mul(size_of::<usize>()),
2678 )
2679}
2680
2681fn semantic_eager_vjp_optional(
2682 ctx: &Arc<EagerRuntime>,
2683 output: &EagerTensor,
2684 wrt: &EagerTensor,
2685 cotangent: &EagerTensor,
2686) -> Result<Option<Option<EagerTensor>>> {
2687 if !eager_semantic_vjp_enabled() {
2688 return Ok(None);
2689 }
2690 let (Some(raw_output_trace), Some(wrt_trace)) =
2691 (output.semantic_trace.as_ref(), wrt.semantic_trace.as_ref())
2692 else {
2693 return Ok(None);
2694 };
2695 let Some(wrt_key) = wrt_trace.input_key() else {
2696 return Ok(None);
2697 };
2698 if !raw_output_trace.has_attached_input_key(&wrt_key) {
2699 return Ok(None);
2700 }
2701
2702 let output_trace = analyze_deferred_semantic_trace(raw_output_trace)?;
2707
2708 let mut compiler = GraphCompiler::new();
2711 let source = compile_ad_source(&mut compiler, &output_trace)?;
2712 if source.output_count() != 1
2713 || source.input_keys().len() != source.input_count()
2714 || source.bindings().len() != source.input_count()
2715 {
2716 return Ok(None);
2717 }
2718 let Some(wrt_input_index) = source.input_key_index(&wrt_key) else {
2719 return Ok(None);
2720 };
2721
2722 let cache_key = PreparedDerivativeCacheKey {
2724 semantic_fingerprint: source.program().semantic_fingerprint(),
2725 runtime_epoch: ctx.runtime.epoch().map_err(|source| {
2726 Error::runtime_state_source("semantic_eager_vjp", ErrorPhase::Execution, source)
2727 })?,
2728 wrt_input_index,
2729 input_metadata: source.frozen_program().input_metadata_with_bound_shapes(),
2730 };
2731 let prepared = { ctx.lock_prepared_derivative_cache()?.get(&cache_key) };
2732 let (seed_input_index, derivative_output_index, derivative_program, prepared_runtime) =
2733 if let Some(prepared) = prepared {
2734 (
2735 prepared.seed_input_index,
2736 prepared.derivative_output_index,
2737 Arc::clone(&prepared.program),
2738 Some(Arc::clone(&prepared.prepared)),
2739 )
2740 } else {
2741 let mut active_inputs = vec![false; source.input_count()];
2742 if let Some(active) = active_inputs.get_mut(wrt_input_index) {
2743 *active = true;
2744 } else {
2745 return Ok(None);
2746 }
2747 let active_outputs = vec![true; source.output_count()];
2748 let ad = AdContext::with_rules_and_transform_cache(
2749 ctx.semantic_extension_rules.clone(),
2750 Arc::clone(&ctx.ad_transform_cache),
2751 );
2752 let derivative = ad
2753 .vjp_program(source.frozen_program(), &active_inputs, &active_outputs)
2754 .map_err(|source| {
2755 Error::runtime_state_source(
2756 "semantic_eager_vjp",
2757 ErrorPhase::GraphBuild,
2758 source,
2759 )
2760 })?;
2761 let seed_input_index = derivative
2762 .derivative_input_indices()
2763 .first()
2764 .copied()
2765 .flatten();
2766 let derivative_output_index = derivative
2767 .derivative_output_indices()
2768 .get(wrt_input_index)
2769 .copied()
2770 .flatten();
2771 let (Some(seed_input_index), Some(derivative_output_index)) =
2772 (seed_input_index, derivative_output_index)
2773 else {
2774 return Ok(Some(None));
2775 };
2776 let program = Arc::new(compiler.compile_frozen_program(derivative.frozen())?);
2777 (seed_input_index, derivative_output_index, program, None)
2778 };
2779
2780 let cotangent_tensor = Arc::new(RetainedValue::from_tensor(cotangent.to_tensor()?));
2781 let input_count = derivative_program.input_count();
2782 let mut owned_inputs: Vec<Option<Tensor>> = (0..input_count).map(|_| None).collect();
2783 for (source_input_index, (_, tensor)) in source.bindings().iter().enumerate() {
2784 let Some(slot) = owned_inputs.get_mut(source_input_index) else {
2785 return Err(Error::Internal(format!(
2786 "semantic eager VJP derivative program has no primal input slot {source_input_index}"
2787 )));
2788 };
2789 *slot = Some(copy_value_for_runtime(ctx, tensor)?);
2790 }
2791 let Some(slot) = owned_inputs.get_mut(seed_input_index) else {
2792 return Err(Error::Internal(format!(
2793 "semantic eager VJP seed input index {seed_input_index} is outside {} inputs",
2794 owned_inputs.len()
2795 )));
2796 };
2797 *slot = Some(copy_value_for_runtime(ctx, cotangent_tensor.as_ref())?);
2798 let input_refs = owned_inputs
2799 .iter()
2800 .enumerate()
2801 .map(|(index, tensor)| {
2802 tensor.as_ref().ok_or_else(|| {
2803 Error::Internal(format!(
2804 "semantic eager VJP derivative input {index} was not populated"
2805 ))
2806 })
2807 })
2808 .collect::<Result<Vec<_>>>()?;
2809 let prepared_runtime = if let Some(prepared_runtime) = prepared_runtime {
2810 prepared_runtime
2811 } else {
2812 let prepared_runtime = Arc::new(
2813 ctx.runtime
2814 .prepare_compiled(&derivative_program, &input_refs)?,
2815 );
2816 let entry = Arc::new(PreparedDerivative {
2817 program: Arc::clone(&derivative_program),
2818 prepared: Arc::clone(&prepared_runtime),
2819 seed_input_index,
2820 derivative_output_index,
2821 });
2822 ctx.lock_prepared_derivative_cache()?
2823 .insert(cache_key, entry);
2824 prepared_runtime
2825 };
2826 let outputs = ctx.runtime.run_prepared(&prepared_runtime, &input_refs)?;
2827 let output_count = outputs.len();
2828 let Some(result) = outputs.into_iter().nth(derivative_output_index) else {
2829 return Err(Error::Internal(format!(
2830 "semantic eager VJP derivative output index {derivative_output_index} is outside {} outputs",
2831 output_count
2832 )));
2833 };
2834 let cotangent_trace =
2835 TracedTensor::from_shared_tensor_value_symbolic_shape(Arc::clone(&cotangent_tensor))?;
2836 let semantic_trace = derivative_trace_from_frozen_program(
2837 &source,
2838 derivative_program.frozen_program(),
2839 derivative_output_index,
2840 &[(seed_input_index, Arc::clone(&cotangent_tensor))],
2841 &[&output_trace, wrt_trace, &cotangent_trace],
2842 None,
2843 "semantic_eager_vjp",
2844 )?;
2845
2846 #[cfg(test)]
2847 EAGER_SEMANTIC_VJP_EXECUTIONS.fetch_add(1, Ordering::Relaxed);
2848
2849 Ok(Some(Some(EagerTensor::new_result_with_semantic_trace(
2850 Arc::clone(ctx),
2851 eager_val_key(),
2852 result,
2853 true,
2854 None,
2855 Some(semantic_trace),
2856 )?)))
2857}
2858
2859fn semantic_eager_jvp_optional(
2860 ctx: &Arc<EagerRuntime>,
2861 output: &EagerTensor,
2862 wrt: &EagerTensor,
2863 tangent: &EagerTensor,
2864) -> Result<Option<Option<EagerTensor>>> {
2865 if !eager_semantic_vjp_enabled() {
2866 return Ok(None);
2867 }
2868 let (Some(raw_output_trace), Some(wrt_trace)) =
2869 (output.semantic_trace.as_ref(), wrt.semantic_trace.as_ref())
2870 else {
2871 return Ok(None);
2872 };
2873 let Some(wrt_key) = wrt_trace.input_key() else {
2874 return Ok(None);
2875 };
2876 if !raw_output_trace.has_attached_input_key(&wrt_key) {
2877 return Ok(None);
2878 }
2879
2880 let output_trace = analyze_deferred_semantic_trace(raw_output_trace)?;
2883
2884 let mut compiler = GraphCompiler::new();
2885 let source = compile_ad_source(&mut compiler, &output_trace)?;
2886 if source.output_count() != 1
2887 || source.input_keys().len() != source.input_count()
2888 || source.bindings().len() != source.input_count()
2889 {
2890 return Ok(None);
2891 }
2892 let Some(wrt_input_index) = source.input_key_index(&wrt_key) else {
2893 return Ok(None);
2894 };
2895
2896 let mut active_inputs = vec![false; source.input_count()];
2897 if let Some(active) = active_inputs.get_mut(wrt_input_index) {
2898 *active = true;
2899 } else {
2900 return Ok(None);
2901 }
2902 let ad = AdContext::with_rules_and_transform_cache(
2903 ctx.semantic_extension_rules.clone(),
2904 Arc::clone(&ctx.ad_transform_cache),
2905 );
2906 let derivative = ad
2907 .jvp_program(source.frozen_program(), &active_inputs)
2908 .map_err(|source| {
2909 Error::runtime_state_source("semantic_eager_jvp", ErrorPhase::GraphBuild, source)
2910 })?;
2911 let Some(seed_input_index) = derivative
2913 .derivative_input_indices()
2914 .get(wrt_input_index)
2915 .copied()
2916 .flatten()
2917 else {
2918 return Ok(Some(None));
2919 };
2920 let Some(derivative_output_index) = derivative
2923 .derivative_output_indices()
2924 .first()
2925 .copied()
2926 .flatten()
2927 else {
2928 return Ok(Some(None));
2929 };
2930
2931 let derivative_program = compiler.compile_frozen_program(derivative.frozen())?;
2932 let tangent_tensor = Arc::new(RetainedValue::from_tensor(tangent.to_tensor()?));
2933 let input_count = derivative_program.input_count();
2934 let mut owned_inputs: Vec<Option<Tensor>> = (0..input_count).map(|_| None).collect();
2935 for (source_input_index, (_, tensor)) in source.bindings().iter().enumerate() {
2936 let Some(slot) = owned_inputs.get_mut(source_input_index) else {
2937 return Err(Error::Internal(format!(
2938 "semantic eager JVP derivative program has no primal input slot {source_input_index}"
2939 )));
2940 };
2941 *slot = Some(copy_value_for_runtime(ctx, tensor)?);
2942 }
2943 let Some(slot) = owned_inputs.get_mut(seed_input_index) else {
2944 return Err(Error::Internal(format!(
2945 "semantic eager JVP seed input index {seed_input_index} is outside {} inputs",
2946 owned_inputs.len()
2947 )));
2948 };
2949 *slot = Some(copy_value_for_runtime(ctx, tangent_tensor.as_ref())?);
2950 let input_refs = owned_inputs
2951 .iter()
2952 .enumerate()
2953 .map(|(index, tensor)| {
2954 tensor.as_ref().ok_or_else(|| {
2955 Error::Internal(format!(
2956 "semantic eager JVP derivative input {index} was not populated"
2957 ))
2958 })
2959 })
2960 .collect::<Result<Vec<_>>>()?;
2961 let outputs = ctx.runtime.run_compiled(&derivative_program, &input_refs)?;
2962 let output_count = outputs.len();
2963 let Some(result) = outputs.into_iter().nth(derivative_output_index) else {
2964 return Err(Error::Internal(format!(
2965 "semantic eager JVP derivative output index {derivative_output_index} is outside {} outputs",
2966 output_count
2967 )));
2968 };
2969 let tangent_trace =
2970 TracedTensor::from_shared_tensor_value_symbolic_shape(Arc::clone(&tangent_tensor))?;
2971 let semantic_trace = derivative_trace_from_frozen_program(
2972 &source,
2973 derivative.frozen(),
2974 derivative_output_index,
2975 &[(seed_input_index, Arc::clone(&tangent_tensor))],
2976 &[&output_trace, wrt_trace, &tangent_trace],
2977 None,
2978 "semantic_eager_jvp",
2979 )?;
2980
2981 Ok(Some(Some(EagerTensor::new_result_with_semantic_trace(
2982 Arc::clone(ctx),
2983 eager_val_key(),
2984 result,
2985 true,
2986 None,
2987 Some(semantic_trace),
2988 )?)))
2989}
2990
2991fn validate_same_runtime(
2992 runtime: &Arc<EagerRuntime>,
2993 tensor: &EagerTensor,
2994 role: &'static str,
2995) -> Result<()> {
2996 if tensor.ctx_id() != runtime.id() {
2997 return Err(Error::ContextMismatch {
2998 lhs: runtime.id(),
2999 rhs: tensor.ctx_id(),
3000 });
3001 }
3002 let _ = role;
3003 Ok(())
3004}
3005
3006fn copy_value_for_runtime(ctx: &EagerRuntime, value: &RetainedValue) -> Result<Tensor> {
3007 let read = value.tensor_read().map_err(|error| {
3008 Error::runtime_state_source("copy_value_for_runtime", ErrorPhase::Execution, error)
3009 })?;
3010 ctx.with_execution_session(|session| session.to_contiguous_read(read))?
3011 .map_err(Error::from)
3012}
3013
3014fn validate_seed_tensor(op: &'static str, primal: &EagerTensor, seed: &EagerTensor) -> Result<()> {
3015 if primal.dtype() != seed.dtype() {
3016 return Err(
3017 tenferro_tensor::Error::dtype_mismatch(op, primal.dtype(), seed.dtype()).into(),
3018 );
3019 }
3020 if primal.shape() != seed.shape() {
3021 return Err(
3022 tenferro_tensor::Error::shape_mismatch(op, primal.shape(), seed.shape()).into(),
3023 );
3024 }
3025 Ok(())
3026}
3027
3028#[derive(Clone)]
3054pub struct EagerTensor {
3055 pub(crate) key: ValueKey<StdTensorOp>,
3056 pub(crate) trace: Option<EagerTrace>,
3057 pub(crate) semantic_trace: Option<TracedTensor>,
3058 pub(crate) requires_grad: bool,
3059 grad_slot: GradSlot,
3060 pub(crate) ctx: Arc<EagerRuntime>,
3061 _record: Arc<EagerTensorRecord>,
3062}
3063
3064pub(crate) struct EagerTensorRecord {
3065 value: Arc<AdValueRecord>,
3066 key: ValueKey<StdTensorOp>,
3067 trace: Option<EagerTrace>,
3068 semantic_trace: Option<TracedTensor>,
3069 requires_grad: bool,
3070 grad_slot: GradSlot,
3071 ctx: Arc<EagerRuntime>,
3072}
3073
3074struct EagerTensorParts {
3075 ctx: Arc<EagerRuntime>,
3076 key: ValueKey<StdTensorOp>,
3077 requires_grad: bool,
3078 trace: Option<EagerTrace>,
3079 semantic_trace: Option<TracedTensor>,
3080 value: Arc<AdValueRecord>,
3081 register_value: bool,
3082}
3083
3084impl fmt::Debug for EagerTensor {
3085 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
3086 f.debug_struct("EagerTensor")
3087 .field("dtype", &self.dtype())
3088 .field("shape", &self.shape())
3089 .field("key", &self.key)
3090 .field("requires_grad", &self.requires_grad)
3091 .field("has_trace", &self.trace.is_some())
3092 .field("has_semantic_trace", &self.semantic_trace.is_some())
3093 .field("ctx_id", &self.ctx_id())
3094 .finish_non_exhaustive()
3095 }
3096}
3097
3098impl EagerTensor {
3099 pub fn from_tensor_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
3120 Self::new_leaf(ctx, tensor, false)
3121 }
3122
3123 pub fn from_vec_col_major_in<T: TensorScalar>(
3135 shape: impl IntoShapeVec,
3136 data: Vec<T>,
3137 ctx: Arc<EagerRuntime>,
3138 ) -> Result<Self> {
3139 Self::from_tensor_in(Tensor::from_vec_col_major(shape, data)?, ctx)
3140 }
3141
3142 pub fn requires_grad_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
3163 Self::new_leaf(ctx, tensor, true)
3164 }
3165
3166 pub(crate) fn new_leaf(
3167 ctx: Arc<EagerRuntime>,
3168 tensor: Tensor,
3169 requires_grad: bool,
3170 ) -> Result<Self> {
3171 let key = eager_val_key();
3172 let semantic_tensor = ctx
3173 .with_execution_session(|session| {
3174 session.to_contiguous_read(TensorRead::from_tensor(&tensor))
3175 })?
3176 .map_err(Error::from)?;
3177 let semantic_value = Arc::new(RetainedValue::from_tensor(semantic_tensor));
3178 let semantic_trace = TracedTensor::from_shared_tensor_value_symbolic_shape(semantic_value)?;
3179 let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_leaf")?;
3186 Self::from_parts(EagerTensorParts {
3187 ctx,
3188 key,
3189 requires_grad,
3190 trace: None,
3191 semantic_trace: Some(semantic_trace),
3192 value,
3193 register_value: true,
3194 })
3195 }
3196
3197 pub(crate) fn new_result(
3198 ctx: Arc<EagerRuntime>,
3199 key: ValueKey<StdTensorOp>,
3200 tensor: Tensor,
3201 requires_grad: bool,
3202 trace: Option<EagerTrace>,
3203 ) -> Result<Self> {
3204 Self::new_result_with_semantic_trace(ctx, key, tensor, requires_grad, trace, None)
3205 }
3206
3207 pub(crate) fn new_result_with_semantic_trace(
3208 ctx: Arc<EagerRuntime>,
3209 key: ValueKey<StdTensorOp>,
3210 tensor: Tensor,
3211 requires_grad: bool,
3212 trace: Option<EagerTrace>,
3213 semantic_trace: Option<TracedTensor>,
3214 ) -> Result<Self> {
3215 let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_result")?;
3216 Self::from_parts(EagerTensorParts {
3217 ctx,
3218 key,
3219 requires_grad,
3220 trace,
3221 semantic_trace,
3222 value,
3223 register_value: true,
3224 })
3225 }
3226
3227 pub(crate) fn new_unregistered_result_with_semantic_trace(
3228 ctx: Arc<EagerRuntime>,
3229 key: ValueKey<StdTensorOp>,
3230 tensor: Tensor,
3231 requires_grad: bool,
3232 trace: Option<EagerTrace>,
3233 semantic_trace: Option<TracedTensor>,
3234 ) -> Result<Self> {
3235 let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_unregistered_result")?;
3236 Self::from_parts(EagerTensorParts {
3237 ctx,
3238 key,
3239 requires_grad,
3240 trace,
3241 semantic_trace,
3242 value,
3243 register_value: false,
3244 })
3245 }
3246
3247 pub(crate) fn new_result_value(
3248 ctx: Arc<EagerRuntime>,
3249 key: ValueKey<StdTensorOp>,
3250 value: TensorValue,
3251 requires_grad: bool,
3252 trace: Option<EagerTrace>,
3253 semantic_trace: Option<TracedTensor>,
3254 ) -> Result<Self> {
3255 let (group, slot, dtype, shape) = value.try_into_group_parts().map_err(|_| {
3256 Error::runtime_state(
3257 "EagerTensor::new_result_value",
3258 ErrorPhase::Execution,
3259 "a TensorValue could not be transferred into its allocation group",
3260 )
3261 })?;
3262 let value = AdValueRecord::from_group(group, slot, dtype, shape);
3263 Self::from_parts(EagerTensorParts {
3264 ctx,
3265 key,
3266 requires_grad,
3267 trace,
3268 semantic_trace,
3269 value,
3270 register_value: true,
3271 })
3272 }
3273
3274 fn from_parts(parts: EagerTensorParts) -> Result<Self> {
3275 let EagerTensorParts {
3276 ctx,
3277 key,
3278 requires_grad,
3279 trace,
3280 semantic_trace,
3281 value,
3282 register_value,
3283 } = parts;
3284 let grad_slot = Arc::new(Mutex::new(None));
3285 if requires_grad {
3286 ctx.try_register_grad_slot(&key, &grad_slot)?;
3287 }
3288 let record = Arc::new(EagerTensorRecord {
3289 value: Arc::clone(&value),
3290 key: key.clone(),
3291 trace: trace.clone(),
3292 semantic_trace: semantic_trace.clone(),
3293 requires_grad,
3294 grad_slot: Arc::clone(&grad_slot),
3295 ctx: Arc::clone(&ctx),
3296 });
3297 if register_value {
3298 ctx.try_register_value_record(&key, &record)?;
3299 }
3300
3301 Ok(Self {
3302 key,
3303 trace,
3304 semantic_trace,
3305 requires_grad,
3306 grad_slot,
3307 ctx,
3308 _record: record,
3309 })
3310 }
3311
3312 pub(crate) fn new_untracked_result(ctx: Arc<EagerRuntime>, tensor: Tensor) -> Result<Self> {
3313 let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_untracked_result")?;
3314 Ok(Self::new_untracked_value_record(ctx, value, None))
3315 }
3316
3317 pub(crate) fn new_untracked_value_result(
3318 ctx: Arc<EagerRuntime>,
3319 value: TensorValue,
3320 ) -> Result<Self> {
3321 Self::new_untracked_value_result_with_semantic_trace(ctx, value, None)
3322 }
3323
3324 pub(crate) fn new_untracked_value_result_with_semantic_trace(
3325 ctx: Arc<EagerRuntime>,
3326 value: TensorValue,
3327 semantic_trace: Option<TracedTensor>,
3328 ) -> Result<Self> {
3329 let (group, slot, dtype, shape) = value.try_into_group_parts().map_err(|_| {
3330 Error::runtime_state(
3331 "EagerTensor::new_untracked_value_result",
3332 ErrorPhase::Execution,
3333 "a TensorValue could not be transferred into its allocation group",
3334 )
3335 })?;
3336 let value = AdValueRecord::from_group(group, slot, dtype, shape);
3337 Ok(Self::new_untracked_value_record(ctx, value, semantic_trace))
3338 }
3339
3340 fn new_untracked_value_record(
3341 ctx: Arc<EagerRuntime>,
3342 value: Arc<AdValueRecord>,
3343 semantic_trace: Option<TracedTensor>,
3344 ) -> Self {
3345 let key = eager_val_key();
3346 let grad_slot = Arc::new(Mutex::new(None));
3347 let record = Arc::new(EagerTensorRecord {
3348 value,
3349 key: key.clone(),
3350 trace: None,
3351 semantic_trace: semantic_trace.clone(),
3352 requires_grad: false,
3353 grad_slot: Arc::clone(&grad_slot),
3354 ctx: Arc::clone(&ctx),
3355 });
3356 Self {
3357 key,
3358 trace: None,
3359 semantic_trace,
3360 requires_grad: false,
3361 grad_slot,
3362 ctx,
3363 _record: record,
3364 }
3365 }
3366
3367 pub(crate) fn from_record(record: Arc<EagerTensorRecord>) -> Self {
3368 Self {
3369 key: record.key.clone(),
3370 trace: record.trace.clone(),
3371 semantic_trace: record.semantic_trace.clone(),
3372 requires_grad: record.requires_grad,
3373 grad_slot: Arc::clone(&record.grad_slot),
3374 ctx: Arc::clone(&record.ctx),
3375 _record: record,
3376 }
3377 }
3378
3379 pub fn detach(&self) -> Self {
3399 let semantic_trace = self
3400 .duplicate_value()
3401 .ok()
3402 .and_then(|tensor| TracedTensor::from_tensor_symbolic_shape(tensor).ok());
3403 Self::new_untracked_value_record(
3404 self.ctx.clone(),
3405 Arc::clone(&self._record.value),
3406 semantic_trace,
3407 )
3408 }
3409
3410 pub fn detach_into(&self, ctx: &Arc<EagerRuntime>) -> Result<Self> {
3434 Self::from_tensor_in(self.to_tensor()?, Arc::clone(ctx))
3435 }
3436
3437 pub fn value(&self) -> Result<ValueGuard<'_>> {
3444 self._record.value.value("EagerTensor::value")
3445 }
3446
3447 pub fn duplicate_value(&self) -> Result<Tensor> {
3471 let value = self.value()?;
3472 match value.duplicate_host_tensor() {
3473 Ok(tensor) => Ok(tensor),
3474 Err(_) => {
3475 let read = self
3476 ._record
3477 .value
3478 .tensor_read("EagerTensor::duplicate_value")?;
3479 self.ctx
3480 .with_execution_session(|session| session.to_contiguous_read(read))?
3481 .map_err(Error::from)
3482 }
3483 }
3484 }
3485
3486 #[allow(clippy::result_large_err)]
3489 pub fn into_value(self) -> std::result::Result<Tensor, IntoValueError<Self>> {
3519 if Arc::strong_count(&self._record) != 1 {
3520 return Err(IntoValueError::NotUnique(self));
3521 }
3522 let Self { _record, .. } = self;
3523 let record = match Arc::try_unwrap(_record) {
3524 Ok(record) => record,
3525 Err(record) => return Err(IntoValueError::NotUnique(Self::from_record(record))),
3526 };
3527 let EagerTensorRecord {
3528 value,
3529 key,
3530 trace,
3531 semantic_trace,
3532 requires_grad,
3533 grad_slot,
3534 ctx,
3535 } = record;
3536 let value = match Arc::try_unwrap(value) {
3537 Ok(value) => value,
3538 Err(value) => {
3539 let record = Arc::new(EagerTensorRecord {
3540 value,
3541 key,
3542 trace,
3543 semantic_trace,
3544 requires_grad,
3545 grad_slot,
3546 ctx,
3547 });
3548 return Err(IntoValueError::NotUnique(Self::from_record(record)));
3549 }
3550 };
3551 let AdValueRecord {
3552 container,
3553 slot,
3554 dtype,
3555 shape,
3556 } = value;
3557 let container = match Arc::try_unwrap(container) {
3558 Ok(container) => container,
3559 Err(container) => {
3560 let record = Arc::new(EagerTensorRecord {
3561 value: Arc::new(AdValueRecord {
3562 container,
3563 slot,
3564 dtype,
3565 shape,
3566 }),
3567 key,
3568 trace,
3569 semantic_trace,
3570 requires_grad,
3571 grad_slot,
3572 ctx,
3573 });
3574 return Err(IntoValueError::NotUnique(Self::from_record(record)));
3575 }
3576 };
3577 match container.group.into_tensor(slot) {
3578 Ok(tensor) => Ok(tensor),
3579 Err((group, error)) => {
3580 let record = Arc::new(EagerTensorRecord {
3581 value: Arc::new(AdValueRecord {
3582 container: Arc::new(RetentionContainer { group }),
3583 slot,
3584 dtype,
3585 shape,
3586 }),
3587 key,
3588 trace,
3589 semantic_trace,
3590 requires_grad,
3591 grad_slot,
3592 ctx,
3593 });
3594 Err(IntoValueError::Extract {
3595 value: Self::from_record(record),
3596 error,
3597 })
3598 }
3599 }
3600 }
3601
3602 pub fn dtype(&self) -> DType {
3605 self._record.value.dtype()
3606 }
3607
3608 pub fn shape(&self) -> &[usize] {
3611 self._record.value.shape()
3612 }
3613
3614 pub fn tensor_read(&self) -> TensorRead<'_> {
3625 self._record
3626 .value
3627 .tensor_read("EagerTensor::tensor_read")
3628 .expect("validated eager value record")
3629 }
3630
3631 pub fn to_tensor(&self) -> Result<Tensor> {
3642 self.duplicate_value()
3643 }
3644
3645 pub fn grad(&self) -> Result<Option<GradientValue>> {
3675 self.grad_slot
3676 .lock()
3677 .map_err(|_| {
3678 Error::runtime_state(
3679 "eager_gradient_slot",
3680 ErrorPhase::Execution,
3681 "lock poisoned",
3682 )
3683 })
3684 .map(|slot| {
3685 slot.as_ref().map(|record| GradientValue {
3686 record: Arc::clone(record),
3687 ctx: Arc::clone(&self.ctx),
3688 })
3689 })
3690 }
3691
3692 pub fn clear_grad(&self) -> Result<()> {
3721 *self.grad_slot.lock().map_err(|_| {
3722 Error::runtime_state(
3723 "eager_gradient_slot",
3724 ErrorPhase::Execution,
3725 "lock poisoned",
3726 )
3727 })? = None;
3728 Ok(())
3729 }
3730
3731 pub fn tracks_grad(&self) -> bool {
3753 self.requires_grad
3754 }
3755
3756 #[cfg(test)]
3757 fn debug_trace_saved_value_count(&self) -> Option<usize> {
3758 None
3759 }
3760
3761 pub fn ctx_id(&self) -> ContextId {
3776 self.ctx.id()
3777 }
3778
3779 pub fn runtime(&self) -> &Arc<EagerRuntime> {
3781 &self.ctx
3782 }
3783
3784 pub fn same_context(&self, other: &Self) -> bool {
3800 self.ctx_id() == other.ctx_id()
3801 }
3802
3803 #[cfg(test)]
3804 pub(crate) fn standard_graph_op(
3805 inputs: &[&Self],
3806 build_graph: impl FnOnce(&[TensorInputKey]) -> Result<Arc<Graph<StdTensorOp>>>,
3807 ) -> Result<Vec<Self>> {
3808 let Some(first) = inputs.first() else {
3809 return Err(Error::Internal(
3810 "standard eager graph op requires at least one input tensor".to_string(),
3811 ));
3812 };
3813 let ctx = Arc::clone(&first.ctx);
3814 for tensor in inputs.iter().skip(1) {
3815 if !first.same_context(tensor) {
3816 return Err(Error::ContextMismatch {
3817 lhs: first.ctx_id(),
3818 rhs: tensor.ctx_id(),
3819 });
3820 }
3821 }
3822
3823 let graph_input_keys = (0..inputs.len())
3824 .map(|_| next_input_key())
3825 .collect::<Vec<_>>();
3826 let graph = build_graph(&graph_input_keys)?;
3827 let initial_data = graph_input_keys
3828 .iter()
3829 .zip(inputs.iter())
3830 .map(|(key, tensor)| Ok((ValueKey::Input(key.clone()), tensor.to_tensor()?)))
3831 .collect::<Result<HashMap<_, _>>>()?;
3832 let execution = ctx.exec_standard_graph_outputs(graph.as_ref(), initial_data)?;
3833 if execution.outputs.len() != graph.outputs().len() {
3834 return Err(Error::Internal(format!(
3835 "standard eager graph op expected {} graph outputs, got {}",
3836 graph.outputs().len(),
3837 execution.outputs.len()
3838 )));
3839 }
3840
3841 if !eager_grad_recording_enabled() || !inputs.iter().any(|input| input.requires_grad) {
3842 return execution
3843 .outputs
3844 .into_iter()
3845 .map(|output| {
3846 Self::new_unregistered_result_with_semantic_trace(
3847 Arc::clone(&ctx),
3848 eager_val_key(),
3849 output,
3850 false,
3851 None,
3852 None,
3853 )
3854 })
3855 .collect();
3856 }
3857
3858 let recorded = record_eager_graph_outputs(
3859 graph.as_ref(),
3860 &graph_input_keys,
3861 &execution.outputs,
3862 inputs,
3863 )?;
3864 if recorded.traces.len() != execution.outputs.len() {
3865 return Err(Error::Internal(format!(
3866 "standard eager graph op expected {} eager traces, got {}",
3867 execution.outputs.len(),
3868 recorded.traces.len()
3869 )));
3870 }
3871
3872 recorded
3873 .traces
3874 .into_iter()
3875 .zip(recorded.semantic_traces)
3876 .zip(execution.outputs)
3877 .map(|((trace, semantic_trace), output)| {
3878 Self::new_result_with_semantic_trace(
3879 Arc::clone(&ctx),
3880 trace.key,
3881 output,
3882 trace.requires_grad,
3883 trace.trace,
3884 semantic_trace,
3885 )
3886 })
3887 .collect()
3888 }
3889
3890 pub fn backward(&self) -> Result<Gradients> {
3923 if !self.shape().is_empty() {
3924 return Err(Error::NonScalarGrad {
3925 shape: self.shape().to_vec(),
3926 });
3927 }
3928
3929 let value = self.to_tensor()?;
3930 let seed = {
3931 let mut backend = self.ctx.lock_backend()?;
3932 one_like_tensor(&value, &mut *backend)?
3933 };
3934 self.backward_from_seed(seed)
3935 }
3936
3937 pub fn backward_with(&self, cotangent: &EagerTensor) -> Result<Gradients> {
3972 if !self.same_context(cotangent) {
3973 return Err(Error::ContextMismatch {
3974 lhs: self.ctx_id(),
3975 rhs: cotangent.ctx_id(),
3976 });
3977 }
3978 validate_seed_tensor("backward", self, cotangent)?;
3979 self.backward_from_seed(cotangent.to_tensor()?)
3980 }
3981
3982 fn backward_from_seed(&self, seed: Tensor) -> Result<Gradients> {
3983 let cotangent =
3984 EagerTensor::new_result(Arc::clone(&self.ctx), eager_val_key(), seed, false, None)?;
3985 let candidate_keys = {
3986 let mut slots = self.ctx.lock_grad_slots()?;
3987 let mut keys = Vec::new();
3988 slots.retain(|key, slot| {
3989 if slot.upgrade().is_some() {
3990 keys.push(key.clone());
3991 true
3992 } else {
3993 false
3994 }
3995 });
3996 keys
3997 };
3998
3999 let mut cotangents = HashMap::new();
4000 for key in candidate_keys {
4001 let Some(record) = self.ctx.value_record(&key)? else {
4002 continue;
4003 };
4004 if !record.requires_grad {
4005 continue;
4006 }
4007 let wrt = EagerTensor::from_record(record);
4008 let Some(grad) = self.ctx.vjp_optional(self, &wrt, &cotangent)? else {
4009 continue;
4010 };
4011 let tensor = match grad.into_value() {
4012 Ok(tensor) => tensor,
4013 Err(IntoValueError::NotUnique(handle)) => handle.duplicate_value()?,
4014 Err(IntoValueError::Extract { error, .. }) => {
4015 return Err(Error::runtime_state_source(
4016 "EagerTensor::backward",
4017 ErrorPhase::Execution,
4018 error,
4019 ));
4020 }
4021 };
4022 cotangents.insert(key, tensor);
4023 }
4024 let mut backend = self.ctx.lock_backend()?;
4025 self.ctx.store_grads(&cotangents, &mut backend)?;
4026 Gradients::from_tensors(cotangents)
4027 }
4028}
4029
4030pub(crate) fn eager_val_key() -> ValueKey<StdTensorOp> {
4031 ValueKey::Input(next_input_key())
4032}
4033
4034pub(crate) struct RecordedEagerTrace {
4035 pub(crate) key: ValueKey<StdTensorOp>,
4036 pub(crate) trace: Option<EagerTrace>,
4037 pub(crate) requires_grad: bool,
4038}
4039
4040pub(crate) struct RecordedEagerOutputs {
4041 pub(crate) traces: Vec<RecordedEagerTrace>,
4042 pub(crate) semantic_traces: Vec<Option<TracedTensor>>,
4043}
4044
4045pub(crate) fn record_eager_outputs(
4046 op: &StdTensorOp,
4047 outputs: &[&Tensor],
4048 inputs: &[&EagerTensor],
4049) -> Result<RecordedEagerOutputs> {
4050 let output_metadata = outputs
4051 .iter()
4052 .map(|output| tensor_meta_from_tensor(output))
4053 .collect::<Vec<_>>();
4054 record_eager_outputs_inner(op, output_metadata, inputs)
4055}
4056
4057pub(crate) fn record_eager_value_outputs(
4058 op: &StdTensorOp,
4059 outputs: &[&TensorValue],
4060 inputs: &[&EagerTensor],
4061) -> Result<RecordedEagerOutputs> {
4062 let output_metadata = outputs
4063 .iter()
4064 .map(|output| tensor_meta_from_value(output))
4065 .collect::<Vec<_>>();
4066 record_eager_outputs_inner(op, output_metadata, inputs)
4067}
4068
4069fn record_eager_outputs_inner(
4070 op: &StdTensorOp,
4071 output_metadata: Vec<TensorMeta>,
4072 inputs: &[&EagerTensor],
4073) -> Result<RecordedEagerOutputs> {
4074 let semantic_traces = record_semantic_eager_outputs(op, &output_metadata, inputs)?;
4075 record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
4076}
4077
4078fn record_semantic_eager_outputs(
4079 op: &StdTensorOp,
4080 output_metadata: &[TensorMeta],
4081 inputs: &[&EagerTensor],
4082) -> Result<Vec<Option<TracedTensor>>> {
4083 let mut owned_constants = Vec::<TracedTensor>::new();
4089 for input in inputs {
4090 if input.semantic_trace.is_none() {
4091 owned_constants.push(TracedTensor::from_tensor_symbolic_shape(
4092 input.to_tensor()?,
4093 )?);
4094 }
4095 }
4096 let mut constants = owned_constants.iter();
4097 let mut semantic_inputs: Vec<&TracedTensor> = inputs
4098 .iter()
4099 .map(|input| {
4100 input
4101 .semantic_trace
4102 .as_ref()
4103 .unwrap_or_else(|| constants.next().expect("materialized constant"))
4104 })
4105 .collect();
4106 let promotion_plan =
4107 eager_input_promotion_plan(op, inputs.len(), |index| inputs[index].dtype());
4108 let promoted_semantic_inputs = if semantic_inputs.iter().enumerate().any(|(index, semantic)| {
4112 semantic.dtype != promotion_plan.target_dtype(index, semantic.dtype)
4113 }) {
4114 Some(
4115 semantic_inputs
4116 .iter()
4117 .enumerate()
4118 .map(|(index, &semantic)| {
4119 let target = promotion_plan.target_dtype(index, semantic.dtype);
4120 if semantic.dtype == target {
4121 Ok(Cow::Borrowed(semantic))
4122 } else {
4123 semantic.cast(target).map(Cow::Owned)
4124 }
4125 })
4126 .collect::<Result<Vec<Cow<'_, TracedTensor>>>>()?,
4127 )
4128 } else {
4129 None
4130 };
4131 if let Some(promoted_semantic_inputs) = &promoted_semantic_inputs {
4132 semantic_inputs = promoted_semantic_inputs.iter().map(Cow::as_ref).collect();
4133 }
4134 let exact_semantic_inputs = if matches!(op, StdTensorOp::Concatenate { .. }) {
4135 Some(
4136 semantic_inputs
4137 .iter()
4138 .zip(inputs)
4139 .map(|(&semantic, input)| {
4140 if semantic.is_concrete_shape() {
4141 Ok(semantic.clone())
4142 } else {
4143 semantic.reshape(input.shape())
4144 }
4145 })
4146 .collect::<Result<Vec<_>>>()?,
4147 )
4148 } else {
4149 None
4150 };
4151 if let Some(exact_semantic_inputs) = &exact_semantic_inputs {
4152 semantic_inputs = exact_semantic_inputs.iter().collect();
4153 }
4154 let outputs = tenferro_runtime::extension::append_raw_eager_outputs(
4158 op.clone(),
4159 &semantic_inputs,
4160 output_metadata,
4161 )?;
4162 Ok(outputs.into_iter().map(Some).collect())
4163}
4164
4165#[cfg(test)]
4166fn record_eager_graph_outputs(
4167 graph: &Graph<StdTensorOp>,
4168 graph_input_keys: &[TensorInputKey],
4169 outputs: &[Tensor],
4170 inputs: &[&EagerTensor],
4171) -> Result<RecordedEagerOutputs> {
4172 let semantic_traces = record_semantic_eager_graph_outputs(graph, graph_input_keys, inputs)?;
4173 let output_metadata = outputs.iter().map(tensor_meta_from_tensor);
4174 record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
4175}
4176
4177#[cfg(test)]
4178fn record_semantic_eager_graph_outputs(
4179 graph: &Graph<StdTensorOp>,
4180 graph_input_keys: &[TensorInputKey],
4181 inputs: &[&EagerTensor],
4182) -> Result<Vec<Option<TracedTensor>>> {
4183 let Some(semantic_inputs) = inputs
4184 .iter()
4185 .map(|input| input.semantic_trace.as_ref())
4186 .collect::<Option<Vec<_>>>()
4187 else {
4188 return Ok(vec![None; graph.outputs().len()]);
4189 };
4190 if graph_input_keys.len() != semantic_inputs.len() {
4191 return Err(Error::Internal(format!(
4192 "semantic graph recording expected {} input keys, got {}",
4193 semantic_inputs.len(),
4194 graph_input_keys.len()
4195 )));
4196 }
4197
4198 let mut values = HashMap::new();
4199 for (key, tensor) in graph_input_keys.iter().zip(semantic_inputs) {
4200 values.insert(ValueKey::Input(key.clone()), tensor.clone());
4201 }
4202
4203 for op_node in graph.operations() {
4204 let input_values = op_node
4205 .inputs
4206 .iter()
4207 .map(|input| {
4208 let key = match input {
4209 ValueRef::Local(local_id) => &graph.values()[*local_id].key,
4210 ValueRef::External(key) => key,
4211 };
4212 values.get(key).cloned().ok_or_else(|| {
4213 Error::Internal(format!(
4214 "semantic graph recording missing value for {key:?}"
4215 ))
4216 })
4217 })
4218 .collect::<Result<Vec<_>>>()?;
4219 let input_refs = input_values.iter().collect::<Vec<_>>();
4220 let semantic_outputs = match &op_node.operation {
4221 StdTensorOp::Extension(ext) => {
4222 tenferro_runtime::extension::apply(Arc::clone(ext), &input_refs)?
4223 }
4224 op => tenferro_runtime::extension::apply_standard_op(op.clone(), &input_refs)?,
4225 };
4226 if semantic_outputs.len() != op_node.outputs.len() {
4227 return Err(Error::Internal(format!(
4228 "semantic graph recording expected {} outputs for {:?}, got {}",
4229 op_node.outputs.len(),
4230 op_node.operation,
4231 semantic_outputs.len()
4232 )));
4233 }
4234 for (output_id, output) in op_node.outputs.iter().copied().zip(semantic_outputs) {
4235 values.insert(graph.values()[output_id].key.clone(), output);
4236 }
4237 }
4238
4239 graph
4240 .outputs()
4241 .iter()
4242 .map(|&output_id| {
4243 let key = &graph.values()[output_id].key;
4244 values.get(key).cloned().map(Some).ok_or_else(|| {
4245 Error::Internal(format!(
4246 "semantic graph recording missing output for {key:?}"
4247 ))
4248 })
4249 })
4250 .collect()
4251}
4252
4253fn record_eager_outputs_from_metadata(
4254 output_metadata: impl IntoIterator<Item = TensorMeta>,
4255 semantic_traces: Vec<Option<TracedTensor>>,
4256 inputs: &[&EagerTensor],
4257) -> Result<RecordedEagerOutputs> {
4258 let output_metadata = output_metadata.into_iter().collect::<Vec<_>>();
4259 if semantic_traces.len() != output_metadata.len() {
4260 return Err(Error::Internal(format!(
4261 "eager recording expected {} semantic traces, got {}",
4262 output_metadata.len(),
4263 semantic_traces.len()
4264 )));
4265 }
4266 let requires_grad =
4267 eager_grad_recording_enabled() && inputs.iter().any(|input| input.requires_grad);
4268 let trace_count = output_metadata.len();
4269 let traces = (0..trace_count)
4270 .map(|_| RecordedEagerTrace {
4271 key: eager_val_key(),
4272 trace: None,
4273 requires_grad,
4274 })
4275 .collect();
4276
4277 Ok(RecordedEagerOutputs {
4278 traces,
4279 semantic_traces,
4280 })
4281}
4282
4283fn tensor_meta_from_value(value: &TensorValue) -> TensorMeta {
4284 TensorMeta::exact(
4285 value.dtype(),
4286 value.shape().iter().copied().map(SymDim::from).collect(),
4287 )
4288}
4289
4290pub(crate) fn exec_single_output(
4291 op: &StdTensorOp,
4292 inputs: &[&Tensor],
4293 ctx: &EagerRuntime,
4294) -> Result<Tensor> {
4295 let mut outputs = ctx.exec_outputs(op, inputs)?;
4296 if outputs.len() != 1 {
4297 return Err(Error::Internal(format!(
4298 "expected one eager output for {:?}, got {}",
4299 op,
4300 outputs.len()
4301 )));
4302 }
4303 Ok(profile_eager_op_section(
4304 "exec_single_output.remove_output",
4305 || outputs.remove(0),
4306 ))
4307}
4308
4309pub(crate) fn exec_single_output_read(
4310 op: &StdTensorOp,
4311 inputs: &[TensorRead<'_>],
4312 ctx: &EagerRuntime,
4313) -> Result<Tensor> {
4314 let mut outputs = ctx.exec_outputs_read(op, inputs)?;
4315 if outputs.len() != 1 {
4316 return Err(Error::Internal(format!(
4317 "expected one eager output for {:?}, got {}",
4318 op,
4319 outputs.len()
4320 )));
4321 }
4322 Ok(profile_eager_op_section(
4323 "exec_single_output_read.remove_output",
4324 || outputs.remove(0),
4325 ))
4326}
4327
4328#[cfg(test)]
4329pub(crate) fn zero_like_tensor<B: TensorBackend>(
4330 input: &Tensor,
4331 backend: &mut B,
4332) -> Result<Tensor> {
4333 let host = match input {
4334 Tensor::F32(tensor) => Tensor::F32(TypedTensor::zeros(tensor.shape().to_vec())?),
4335 Tensor::F64(tensor) => Tensor::F64(TypedTensor::zeros(tensor.shape().to_vec())?),
4336 Tensor::I32(tensor) => Tensor::I32(TypedTensor::zeros(tensor.shape().to_vec())?),
4337 Tensor::I64(tensor) => Tensor::I64(TypedTensor::zeros(tensor.shape().to_vec())?),
4338 Tensor::Bool(tensor) => Tensor::Bool(TypedTensor::from_vec_col_major(
4339 tensor.shape().to_vec(),
4340 vec![false; tensor.n_elements()],
4341 )?),
4342 Tensor::C32(tensor) => Tensor::C32(TypedTensor::zeros(tensor.shape().to_vec())?),
4343 Tensor::C64(tensor) => Tensor::C64(TypedTensor::zeros(tensor.shape().to_vec())?),
4344 };
4345 backend
4346 .upload_host_tensor(TensorRead::from_tensor(&host))
4347 .map_err(Error::from)
4348}
4349
4350pub(crate) fn one_like_tensor<B: TensorBackend>(input: &Tensor, backend: &mut B) -> Result<Tensor> {
4351 let host = ones_tensor(input.dtype(), input.shape().to_vec())?;
4352 backend
4353 .upload_host_tensor(TensorRead::from_tensor(&host))
4354 .map_err(Error::from)
4355}
4356
4357#[cfg(test)]
4358mod tests;