1use std::cell::{Cell, RefCell};
2use std::cmp::Reverse;
3use std::collections::HashMap;
4use std::env;
5use std::fmt;
6use std::mem::{size_of, size_of_val};
7#[cfg(test)]
8use std::sync::atomic::{AtomicUsize, Ordering};
9use std::sync::{Arc, Mutex, MutexGuard, OnceLock, Weak};
10use std::time::{Duration, Instant};
11
12use lru::LruCache;
13
14use crate::extension_cache::{ExtensionCacheLimits, ExtensionCacheSelector, ExtensionCacheStore};
15#[cfg(test)]
16use computegraph::graph::Graph;
17use computegraph::ValueKey;
18#[cfg(test)]
19use computegraph::ValueRef;
20use tenferro_cpu::{CpuBackend, CpuBackendError, CpuPlacement};
21#[cfg(feature = "cuda")]
22use tenferro_gpu::CudaBackend;
23#[cfg(feature = "webgpu")]
24use tenferro_gpu::WebGpuBackend;
25#[cfg(test)]
26use tenferro_ops::input_key::TensorInputKey;
27use tenferro_ops::{std_tensor_op::StdTensorOp, SymDim, TensorMeta};
28use tenferro_runtime::ad_support::{compile_ad_source, ones_tensor};
29use tenferro_runtime::program::{ProgramValueMetadata, SemanticFingerprint, SemanticProgram};
30use tenferro_runtime::{
31 CompiledGraph, CoreCapabilityBundle, EngineId, ErrorPhase, ExecutionContextIdentity,
32 ExtensionModule, GraphCompiler, HardwareClassId, PreparedCompiledGraph, RegistrationIdentity,
33 Runtime, RuntimeConfigError, RuntimeConfigSnapshot, RuntimeEpoch, TracedTensor,
34};
35#[cfg(test)]
36use tenferro_tensor::TypedTensor;
37use tenferro_tensor::{BackendSession, BackendSessionHost};
38use tenferro_tensor::{
39 CacheStats, DType, IntoShapeVec, Tensor, TensorBackend, TensorElementwise, TensorRead,
40 TensorScalar, TensorValue,
41};
42
43use crate::eager_backend::{
44 cpu_runtime_engine_id, cpu_runtime_hardware_class, eager_runtime_for_backend, EagerBackend,
45};
46#[cfg(test)]
47use crate::eager_exec::exec_standard_op_on_tensor_reads_in_session;
48use crate::eager_exec::{exec_op_on_tensor_reads_with_runtime, exec_op_on_tensors_with_runtime};
49use crate::error::{ContextId, Error, Result};
50#[cfg(test)]
51use crate::metadata::push_metadata_scope;
52use crate::metadata::{
53 metadata_scopes_for_scope, register_scoped_metadata_batch, register_scoped_value_metadata,
54 tensor_meta_from_tensor, GlobalMetadataScope,
55};
56use crate::semantic_extension::SemanticExtensionRuleSet;
57use crate::traced::{derivative_trace_from_frozen_program, next_input_key};
58use crate::transform_cache::{AdTransformCache, AdTransformCacheLimits};
59
60use crate::AdContext;
61
62pub(crate) type GradSlot = Arc<Mutex<Option<Arc<Tensor>>>>;
63pub(crate) type WeakGradSlot = Weak<Mutex<Option<Arc<Tensor>>>>;
64
65#[derive(Clone, Debug)]
66pub(crate) struct EagerTrace;
67
68#[cfg(test)]
69pub(crate) static CPU_RUNTIME_SELECTION_REFRESHES: AtomicUsize = AtomicUsize::new(0);
70
71struct CpuRuntimeSelection {
72 snapshot: Arc<RuntimeConfigSnapshot>,
73 epoch: RuntimeEpoch,
74 engine_id: EngineId,
75 registration_identity: RegistrationIdentity,
76 capabilities: CoreCapabilityBundle,
77}
78
79#[derive(Debug, Default, Clone)]
80struct EagerOpProfileEntry {
81 calls: usize,
82 total_time: Duration,
83}
84
85thread_local! {
86 static EAGER_OP_PROFILE_STATE: RefCell<HashMap<&'static str, EagerOpProfileEntry>> =
87 RefCell::new(HashMap::new());
88 static EAGER_NO_GRAD_DEPTH: Cell<usize> = const { Cell::new(0) };
89 #[cfg(test)]
90 static EAGER_OP_PROFILE_ENABLED_OVERRIDE: RefCell<Option<bool>> = const { RefCell::new(None) };
91 #[cfg(test)]
92 static EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE: RefCell<Option<Option<usize>>> = const { RefCell::new(None) };
93 #[cfg(test)]
94 static EAGER_SEMANTIC_VJP_ENABLED_OVERRIDE: RefCell<Option<bool>> = const { RefCell::new(None) };
95}
96
97#[cfg(test)]
98pub(crate) static EAGER_SEMANTIC_VJP_EXECUTIONS: AtomicUsize = AtomicUsize::new(0);
99
100pub(crate) fn eager_grad_recording_enabled() -> bool {
101 EAGER_NO_GRAD_DEPTH.with(|depth| depth.get() == 0)
102}
103
104fn eager_semantic_vjp_enabled() -> bool {
105 #[cfg(test)]
106 if let Some(value) = EAGER_SEMANTIC_VJP_ENABLED_OVERRIDE.with(|state| *state.borrow()) {
107 return value;
108 }
109
110 static ENABLED: OnceLock<bool> = OnceLock::new();
113 *ENABLED.get_or_init(|| env::var("TENFERRO_EAGER_SEMANTIC_VJP").map_or(true, |v| v != "0"))
114}
115
116#[derive(Debug)]
140pub struct EagerNoGradGuard {
141 active: bool,
142}
143
144impl Drop for EagerNoGradGuard {
145 fn drop(&mut self) {
146 if !self.active {
147 return;
148 }
149 EAGER_NO_GRAD_DEPTH.with(|depth| {
150 depth.set(depth.get().saturating_sub(1));
151 });
152 self.active = false;
153 }
154}
155
156pub(crate) fn eager_op_profile_enabled() -> bool {
157 #[cfg(test)]
158 if let Some(value) = EAGER_OP_PROFILE_ENABLED_OVERRIDE.with(|state| *state.borrow()) {
159 return value;
160 }
161
162 static ENABLED: OnceLock<bool> = OnceLock::new();
163 *ENABLED.get_or_init(|| env::var("TENFERRO_PROFILE_EAGER_OP_AGG").is_ok())
164}
165
166pub(crate) fn eager_op_profile_start() -> Option<Instant> {
167 eager_op_profile_enabled().then(Instant::now)
168}
169
170pub(crate) fn record_eager_op_profile(section: &'static str, elapsed: Duration) {
171 if !eager_op_profile_enabled() {
172 return;
173 }
174 EAGER_OP_PROFILE_STATE.with(|state| {
175 let mut state = state.borrow_mut();
176 let entry = state.entry(section).or_default();
177 entry.calls += 1;
178 entry.total_time += elapsed;
179 });
180}
181
182pub(crate) fn profile_eager_op_section<T>(section: &'static str, f: impl FnOnce() -> T) -> T {
183 if !eager_op_profile_enabled() {
184 return f();
185 }
186 let started = Instant::now();
187 let result = f();
188 record_eager_op_profile(section, started.elapsed());
189 result
190}
191
192pub(crate) fn maybe_print_eager_op_profile() {
193 if !eager_op_profile_enabled() {
194 return;
195 }
196 let Some(print_every) = eager_op_profile_print_every() else {
197 return;
198 };
199 if print_every == 0 {
200 return;
201 }
202
203 let should_print = EAGER_OP_PROFILE_STATE.with(|state| {
204 state
205 .borrow()
206 .get("nary_op.total")
207 .is_some_and(|entry| entry.calls % print_every == 0)
208 });
209 if should_print {
210 print_and_reset_eager_op_profile();
211 }
212}
213
214fn eager_op_profile_print_every() -> Option<usize> {
215 #[cfg(test)]
216 if let Some(value) = EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE.with(|state| *state.borrow()) {
217 return value;
218 }
219
220 env::var("TENFERRO_PROFILE_EAGER_OP_PRINT_EVERY")
221 .ok()?
222 .parse()
223 .ok()
224}
225
226pub(crate) fn print_and_reset_eager_op_profile() {
227 EAGER_OP_PROFILE_STATE.with(|state| {
228 let mut entries: Vec<_> = state
229 .borrow()
230 .iter()
231 .map(|(section, entry)| (*section, entry.clone()))
232 .collect();
233 state.borrow_mut().clear();
234 entries.sort_by_key(|(_, entry)| Reverse(entry.total_time));
235
236 eprintln!("=== tenferro eager op profile ===");
237 for (section, entry) in entries {
238 let Some(per_call_us) = eager_op_profile_per_call_us(&entry) else {
239 continue;
240 };
241 eprintln!(
242 "{section}: calls={} total={:.6}ms per_call={:.3}us",
243 entry.calls,
244 entry.total_time.as_secs_f64() * 1.0e3,
245 per_call_us,
246 );
247 }
248 });
249}
250
251fn eager_op_profile_per_call_us(entry: &EagerOpProfileEntry) -> Option<f64> {
252 (entry.calls != 0).then(|| entry.total_time.as_secs_f64() * 1.0e6 / entry.calls as f64)
253}
254
255fn runtime_config_error(op: &'static str, source: RuntimeConfigError) -> Error {
256 Error::runtime_state_source(op, ErrorPhase::Execution, source)
257}
258
259fn runtime_state_source<E>(op: &'static str, source: E) -> Error
260where
261 E: std::error::Error + Send + Sync + 'static,
262{
263 Error::runtime_state_source(op, ErrorPhase::Execution, source)
264}
265
266fn cpu_runtime_bridge_unsupported(message: impl Into<String>) -> Error {
267 Error::unsupported(
268 "CpuPlacementBoundEager::refresh_runtime_selection",
269 ErrorPhase::Execution,
270 message,
271 )
272}
273
274fn select_cpu_runtime(runtime: &Runtime) -> Result<CpuRuntimeSelection> {
275 let snapshot = runtime
276 .snapshot()
277 .map_err(|source| runtime_state_source("EagerRuntime::runtime_snapshot", source))?;
278 let engine_id = cpu_runtime_engine_id()
279 .map_err(|source| runtime_config_error("EagerRuntime::cpu_runtime_engine_id", source))?;
280 let expected_hardware = cpu_runtime_hardware_class().map_err(|source| {
281 runtime_config_error("EagerRuntime::cpu_runtime_hardware_class", source)
282 })?;
283 let engine = snapshot
284 .engine(&engine_id)
285 .ok_or_else(|| cpu_runtime_bridge_unsupported("missing CPU runtime engine"))?;
286 validate_cpu_runtime_engine(
287 engine.context_identity(),
288 engine.hardware_class(),
289 engine.capabilities(),
290 &expected_hardware,
291 )?;
292 let epoch = snapshot.epoch();
293 let registration_identity = engine.registration_identity();
294 let capabilities = engine.capabilities().clone();
295 Ok(CpuRuntimeSelection {
296 snapshot,
297 epoch,
298 engine_id,
299 registration_identity,
300 capabilities,
301 })
302}
303
304fn validate_cpu_runtime_engine(
305 context_identity: ExecutionContextIdentity,
306 hardware_class: &HardwareClassId,
307 capabilities: &CoreCapabilityBundle,
308 expected_hardware: &HardwareClassId,
309) -> Result<()> {
310 if context_identity != ExecutionContextIdentity::of::<CpuBackend>() {
311 return Err(cpu_runtime_bridge_unsupported(
312 "CPU runtime context mismatch",
313 ));
314 }
315 if hardware_class != expected_hardware {
316 return Err(cpu_runtime_bridge_unsupported(
317 "CPU runtime hardware mismatch",
318 ));
319 }
320 if capabilities.elementwise().is_none() {
321 return Err(cpu_runtime_bridge_unsupported(
322 "missing CPU runtime capability: elementwise",
323 ));
324 }
325 if capabilities.reduction().is_none() {
326 return Err(cpu_runtime_bridge_unsupported(
327 "missing CPU runtime capability: reduction",
328 ));
329 }
330 if capabilities.indexing().is_none() {
331 return Err(cpu_runtime_bridge_unsupported(
332 "missing CPU runtime capability: indexing",
333 ));
334 }
335 if capabilities.dot_general().is_none() {
336 return Err(cpu_runtime_bridge_unsupported(
337 "missing CPU runtime capability: dot_general",
338 ));
339 }
340 if capabilities.layout().is_none() {
341 return Err(cpu_runtime_bridge_unsupported(
342 "missing CPU runtime capability: layout",
343 ));
344 }
345 Ok(())
346}
347
348#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
352pub struct EagerRuntimeCacheStats {
353 pub extensions: CacheStats,
355 pub ad_transforms: CacheStats,
357 pub prepared_derivatives: CacheStats,
359}
360
361#[cfg(test)]
362pub(crate) struct EagerGraphExecution {
363 pub(crate) outputs: Vec<Arc<Tensor>>,
364}
365
366pub struct CpuPlacementBoundEager {
391 runtime: Arc<EagerRuntime>,
392 backend: CpuBackend,
393 snapshot: Arc<RuntimeConfigSnapshot>,
394 epoch: RuntimeEpoch,
395 engine_id: EngineId,
396 registration_identity: RegistrationIdentity,
397 capabilities: CoreCapabilityBundle,
398}
399
400impl fmt::Debug for CpuPlacementBoundEager {
401 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
402 f.debug_struct("CpuPlacementBoundEager")
403 .field("runtime_id", &self.runtime.id())
404 .field("placement", &self.backend.placement())
405 .field("runtime_epoch", &self.epoch)
406 .field("engine_id", &self.engine_id)
407 .field("registration_identity", &self.registration_identity)
408 .finish_non_exhaustive()
409 }
410}
411
412impl CpuPlacementBoundEager {
413 fn refresh_runtime_selection(&mut self) -> Result<()> {
414 let current_epoch = self.runtime.runtime.epoch().map_err(|source| {
415 runtime_state_source("CpuPlacementBoundEager::refresh_runtime_selection", source)
416 })?;
417 if current_epoch == self.epoch {
418 return Ok(());
419 }
420
421 #[cfg(test)]
422 CPU_RUNTIME_SELECTION_REFRESHES.fetch_add(1, Ordering::SeqCst);
423
424 let selection = select_cpu_runtime(&self.runtime.runtime)?;
425 self.snapshot = selection.snapshot;
426 self.epoch = selection.epoch;
427 self.engine_id = selection.engine_id;
428 self.registration_identity = selection.registration_identity;
429 self.capabilities = selection.capabilities;
430 Ok(())
431 }
432
433 pub fn runtime_id(&self) -> ContextId {
447 self.runtime.id()
448 }
449
450 pub fn placement(&self) -> CpuPlacement {
464 self.backend.placement()
465 }
466
467 pub fn with_eager_session<R: Send>(
508 &mut self,
509 f: impl FnOnce(&mut dyn BackendSession) -> Result<R> + Send,
510 ) -> Result<R> {
511 self.refresh_runtime_selection()?;
512 self.backend.with_backend_session(f)
513 }
514}
515
516pub struct EagerRuntime {
536 id: ContextId,
537 runtime: Runtime,
538 backend: Mutex<EagerBackend>,
542 extension_install_lock: Mutex<()>,
543 pub(crate) extension_caches: Mutex<ExtensionCacheStore>,
544 semantic_extension_rules: SemanticExtensionRuleSet,
545 grad_slots: Mutex<HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>,
546 value_records: Mutex<HashMap<ValueKey<StdTensorOp>, Weak<EagerTensorRecord>>>,
547 value_ptr_records: Mutex<HashMap<usize, Weak<EagerTensorRecord>>>,
548 ad_transform_cache: Arc<AdTransformCache>,
549 prepared_derivative_cache: Mutex<PreparedDerivativeCache>,
553}
554
555impl fmt::Debug for EagerRuntime {
556 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
557 let mut debug = f.debug_struct("EagerRuntime");
558 debug.field("id", &self.id);
559 debug.field("runtime_id", &self.runtime.id());
560 debug.field("runtime_epoch", &self.runtime.epoch().ok());
561 match self.backend.try_lock() {
562 Ok(backend) => {
563 debug.field("backend", &*backend);
564 }
565 Err(_) => {
566 debug.field("backend", &"<locked>");
567 }
568 }
569 match self.extension_caches.try_lock() {
570 Ok(caches) => {
571 debug.field(
572 "extension_cache_stats",
573 &caches.stats(ExtensionCacheSelector::All),
574 );
575 }
576 Err(_) => {
577 debug.field("extension_cache_stats", &"<locked>");
578 }
579 }
580 match self.extension_install_lock.try_lock() {
581 Ok(_) => {
582 debug.field("extension_install_lock", &"<unlocked>");
583 }
584 Err(_) => {
585 debug.field("extension_install_lock", &"<locked>");
586 }
587 }
588 debug.field("semantic_extension_rules", &self.semantic_extension_rules);
589 match self.grad_slots.try_lock() {
590 Ok(slots) => {
591 debug.field("grad_slots_len", &slots.len());
592 }
593 Err(_) => {
594 debug.field("grad_slots_len", &"<locked>");
595 }
596 }
597 match self.value_records.try_lock() {
598 Ok(records) => {
599 debug.field("value_records_len", &records.len());
600 }
601 Err(_) => {
602 debug.field("value_records_len", &"<locked>");
603 }
604 }
605 match self.value_ptr_records.try_lock() {
606 Ok(records) => {
607 debug.field("value_ptr_records_len", &records.len());
608 }
609 Err(_) => {
610 debug.field("value_ptr_records_len", &"<locked>");
611 }
612 }
613 match self.ad_transform_cache.stats() {
614 Ok(stats) => {
615 debug.field("ad_transform_cache_stats", &stats);
616 }
617 Err(err) => {
618 debug.field("ad_transform_cache_stats", &format_args!("{err}"));
619 }
620 }
621 match self.prepared_derivative_cache.try_lock() {
622 Ok(cache) => {
623 debug.field("prepared_derivative_cache_stats", &cache.stats());
624 }
625 Err(_) => {
626 debug.field("prepared_derivative_cache_stats", &"<locked>");
627 }
628 }
629 debug.finish_non_exhaustive()
630 }
631}
632
633impl EagerRuntime {
634 pub(crate) fn lock_backend(&self) -> Result<MutexGuard<'_, EagerBackend>> {
635 self.backend.lock().map_err(|_| {
636 Error::runtime_state("eager_backend", ErrorPhase::Execution, "lock poisoned")
637 })
638 }
639
640 fn lock_extension_caches(&self) -> Result<MutexGuard<'_, ExtensionCacheStore>> {
641 self.extension_caches.lock().map_err(|_| {
642 Error::runtime_state(
643 "eager_extension_caches",
644 ErrorPhase::Execution,
645 "lock poisoned",
646 )
647 })
648 }
649
650 fn lock_extension_install(&self) -> Result<MutexGuard<'_, ()>> {
651 self.extension_install_lock.lock().map_err(|_| {
652 Error::runtime_state(
653 "eager_extension_install",
654 ErrorPhase::Execution,
655 "lock poisoned",
656 )
657 })
658 }
659
660 fn lock_prepared_derivative_cache(&self) -> Result<MutexGuard<'_, PreparedDerivativeCache>> {
661 self.prepared_derivative_cache.lock().map_err(|_| {
662 Error::runtime_state(
663 "prepared_derivative_cache",
664 ErrorPhase::Execution,
665 "lock poisoned",
666 )
667 })
668 }
669
670 fn lock_grad_slots(
671 &self,
672 ) -> Result<MutexGuard<'_, HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>> {
673 self.grad_slots.lock().map_err(|_| {
674 Error::runtime_state(
675 "eager_gradient_slots",
676 ErrorPhase::Execution,
677 "lock poisoned",
678 )
679 })
680 }
681
682 fn lock_value_records(
683 &self,
684 ) -> Result<MutexGuard<'_, HashMap<ValueKey<StdTensorOp>, Weak<EagerTensorRecord>>>> {
685 self.value_records.lock().map_err(|_| {
686 Error::runtime_state(
687 "eager_value_registry",
688 ErrorPhase::Execution,
689 "lock poisoned",
690 )
691 })
692 }
693
694 fn lock_value_ptr_records(
695 &self,
696 ) -> Result<MutexGuard<'_, HashMap<usize, Weak<EagerTensorRecord>>>> {
697 self.value_ptr_records.lock().map_err(|_| {
698 Error::runtime_state(
699 "eager_value_pointer_registry",
700 ErrorPhase::Execution,
701 "lock poisoned",
702 )
703 })
704 }
705
706 fn from_backend(backend: EagerBackend) -> Result<Self> {
707 Self::from_backend_with_rules_and_cache(
708 backend,
709 SemanticExtensionRuleSet::default(),
710 Arc::new(AdTransformCache::new()),
711 )
712 }
713
714 fn from_backend_with_rules_and_cache(
715 backend: EagerBackend,
716 semantic_extension_rules: SemanticExtensionRuleSet,
717 ad_transform_cache: Arc<AdTransformCache>,
718 ) -> Result<Self> {
719 let runtime = eager_runtime_for_backend(&backend)
720 .map_err(|source| runtime_config_error("EagerRuntime::from_backend", source))?;
721 Ok(Self {
722 id: ContextId::fresh(),
723 runtime,
724 backend: Mutex::new(backend),
725 extension_install_lock: Mutex::new(()),
726 extension_caches: Mutex::new(ExtensionCacheStore::new()),
727 semantic_extension_rules,
728 grad_slots: Mutex::new(HashMap::new()),
729 value_records: Mutex::new(HashMap::new()),
730 value_ptr_records: Mutex::new(HashMap::new()),
731 ad_transform_cache,
732 prepared_derivative_cache: Mutex::new(PreparedDerivativeCache::default()),
733 })
734 }
735
736 pub fn new() -> Result<Arc<Self>> {
754 Self::with_cpu_backend(CpuBackend::new())
755 }
756
757 pub fn with_cpu_backend(backend: CpuBackend) -> Result<Arc<Self>> {
776 Ok(Arc::new(Self::from_backend(EagerBackend::cpu(backend))?))
777 }
778
779 pub fn on_cpu(self: &Arc<Self>, placement: CpuPlacement) -> Result<CpuPlacementBoundEager> {
805 let backend = {
806 let backend = self.lock_backend()?;
807 backend.cpu_snapshot().ok_or_else(|| {
808 Error::unsupported(
809 "EagerRuntime::on_cpu",
810 ErrorPhase::Execution,
811 "the eager runtime is not CPU-backed",
812 )
813 })?
814 };
815 let selection = select_cpu_runtime(&self.runtime)?;
816 let backend = backend.for_placement(placement).map_err(|source| {
817 let error: tenferro_tensor::Error = CpuBackendError::Placement {
818 op: "EagerRuntime::on_cpu",
819 source,
820 }
821 .into();
822 Error::from(error)
823 })?;
824 Ok(CpuPlacementBoundEager {
825 runtime: Arc::clone(self),
826 backend,
827 snapshot: selection.snapshot,
828 epoch: selection.epoch,
829 engine_id: selection.engine_id,
830 registration_identity: selection.registration_identity,
831 capabilities: selection.capabilities,
832 })
833 }
834
835 pub fn with_cpu_backend_and_ad_context(
855 backend: CpuBackend,
856 ad: &AdContext,
857 ) -> Result<Arc<Self>> {
858 Ok(Arc::new(Self::from_backend_with_rules_and_cache(
859 EagerBackend::cpu(backend),
860 ad.semantic_extension_rules().clone(),
861 ad.ad_transform_cache(),
862 )?))
863 }
864
865 #[cfg(feature = "cuda")]
877 pub fn with_cuda_backend(backend: CudaBackend) -> Result<Arc<Self>> {
884 Ok(Arc::new(Self::from_backend(EagerBackend::cuda(backend))?))
885 }
886
887 #[cfg(feature = "cuda")]
899 pub fn with_cuda_backend_and_ad_context(
906 backend: CudaBackend,
907 ad: &AdContext,
908 ) -> Result<Arc<Self>> {
909 Ok(Arc::new(Self::from_backend_with_rules_and_cache(
910 EagerBackend::cuda(backend),
911 ad.semantic_extension_rules().clone(),
912 ad.ad_transform_cache(),
913 )?))
914 }
915
916 #[cfg(feature = "webgpu")]
928 pub fn with_webgpu_backend(backend: WebGpuBackend) -> Result<Arc<Self>> {
935 Ok(Arc::new(Self::from_backend(EagerBackend::webgpu(backend))?))
936 }
937
938 #[cfg(feature = "webgpu")]
950 pub fn with_webgpu_backend_and_ad_context(
957 backend: WebGpuBackend,
958 ad: &AdContext,
959 ) -> Result<Arc<Self>> {
960 Ok(Arc::new(Self::from_backend_with_rules_and_cache(
961 EagerBackend::webgpu(backend),
962 ad.semantic_extension_rules().clone(),
963 ad.ad_transform_cache(),
964 )?))
965 }
966
967 pub fn id(&self) -> ContextId {
980 self.id
981 }
982
983 pub fn no_grad(&self) -> EagerNoGradGuard {
1007 EAGER_NO_GRAD_DEPTH.with(|depth| {
1008 depth.set(depth.get().saturating_add(1));
1009 });
1010 EagerNoGradGuard { active: true }
1011 }
1012
1013 pub fn install_extension_module(
1024 &self,
1025 module: Arc<dyn ExtensionModule>,
1026 ) -> Result<RuntimeEpoch> {
1027 let _install_guard = self.lock_extension_install()?;
1028 self.runtime
1029 .reconfigure(|edit| {
1030 edit.replace_extension_module(module)?;
1031 Ok(())
1032 })
1033 .map_err(|source| {
1034 runtime_state_source("EagerRuntime::install_extension_module", source)
1035 })
1036 }
1037
1038 pub fn clear_extension_caches(&self) -> Result<()> {
1057 self.lock_extension_caches()?.clear();
1058 Ok(())
1059 }
1060
1061 pub fn clear_caches(&self) -> Result<()> {
1082 self.clear_extension_caches()?;
1083 self.clear_ad_transform_caches()?;
1084 self.clear_prepared_derivative_cache()?;
1085 Ok(())
1086 }
1087
1088 pub fn clear_prepared_derivative_cache(&self) -> Result<()> {
1107 self.lock_prepared_derivative_cache()?.clear();
1108 Ok(())
1109 }
1110
1111 pub fn cache_stats(&self) -> Result<EagerRuntimeCacheStats> {
1132 Ok(EagerRuntimeCacheStats {
1133 extensions: self
1134 .lock_extension_caches()?
1135 .stats(ExtensionCacheSelector::All),
1136 ad_transforms: self.ad_transform_cache.stats()?,
1137 prepared_derivatives: self.lock_prepared_derivative_cache()?.stats(),
1138 })
1139 }
1140
1141 pub fn ad_transform_cache_limits(&self) -> Result<AdTransformCacheLimits> {
1159 self.ad_transform_cache.limits()
1160 }
1161
1162 pub fn set_ad_transform_cache_limits(&self, limits: AdTransformCacheLimits) -> Result<()> {
1183 self.ad_transform_cache.set_limits(limits)
1184 }
1185
1186 pub fn clear_ad_transform_caches(&self) -> Result<()> {
1205 self.ad_transform_cache.clear()
1206 }
1207
1208 pub fn prepared_derivative_cache_limits(&self) -> Result<AdTransformCacheLimits> {
1226 Ok(self.lock_prepared_derivative_cache()?.limits())
1227 }
1228
1229 pub fn set_prepared_derivative_cache_limits(
1250 &self,
1251 limits: AdTransformCacheLimits,
1252 ) -> Result<()> {
1253 self.lock_prepared_derivative_cache()?.set_limits(limits);
1254 Ok(())
1255 }
1256
1257 pub fn extension_cache_limits(&self) -> Result<ExtensionCacheLimits> {
1264 Ok(self.lock_extension_caches()?.limits())
1265 }
1266
1267 pub fn set_extension_cache_limits(&self, limits: ExtensionCacheLimits) -> Result<()> {
1274 self.lock_extension_caches()?.set_limits(limits);
1275 Ok(())
1276 }
1277
1278 pub fn with_execution_session<R: Send>(
1308 &self,
1309 f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
1310 ) -> Result<R> {
1311 let mut backend = self.lock_backend()?;
1312 Ok(backend.with_backend_session(f))
1313 }
1314
1315 pub fn with_extension_execution_context<R: Send>(
1351 &self,
1352 f: impl FnOnce(
1353 &mut tenferro_runtime::ExtensionExecutionContext<'_, dyn BackendSession + '_>,
1354 ) -> R
1355 + Send,
1356 ) -> Result<R> {
1357 let mut backend = self.lock_backend()?;
1358 let mut extension_cache_guard = self.lock_extension_caches()?;
1359 let extension_caches: &mut ExtensionCacheStore = &mut extension_cache_guard;
1360 Ok(backend.with_backend_session(move |session| {
1361 let mut extension_ctx =
1362 tenferro_runtime::ExtensionExecutionContext::new(session, extension_caches);
1363 f(&mut extension_ctx)
1364 }))
1365 }
1366
1367 pub(crate) fn materialize_value(&self, value: &TensorValue) -> Result<Tensor> {
1368 if let Some(tensor) = value.as_tensor_arc() {
1369 return Ok(tensor.as_ref().clone());
1370 }
1371
1372 let mut backend = self.lock_backend()?;
1373 backend
1374 .with_backend_session(|exec| exec.to_contiguous_read(value.tensor_read()))
1375 .map_err(Error::from)
1376 }
1377
1378 pub fn synchronize(&self) -> Result<()> {
1399 self.lock_backend()?.synchronize().map_err(Error::from)
1400 }
1401
1402 fn exec_outputs_with_runtime<R>(
1403 &self,
1404 lock_backend_section: &'static str,
1405 exec_section: &'static str,
1406 op: &StdTensorOp,
1407 execute: impl FnOnce(&mut EagerBackend, Option<&Runtime>) -> Result<R>,
1408 ) -> Result<R> {
1409 let mut backend = profile_eager_op_section(lock_backend_section, || self.lock_backend())?;
1413 let runtime = matches!(op, StdTensorOp::Extension(_)).then_some(&self.runtime);
1414 profile_eager_op_section(exec_section, || execute(&mut backend, runtime))
1415 }
1416
1417 pub(crate) fn exec_outputs(&self, op: &StdTensorOp, inputs: &[&Tensor]) -> Result<Vec<Tensor>> {
1418 self.exec_outputs_with_runtime(
1419 "exec_outputs.lock_backend",
1420 "exec_outputs.exec_op",
1421 op,
1422 |backend, runtime| exec_op_on_tensors_with_runtime(op, inputs, backend, runtime),
1423 )
1424 }
1425
1426 pub(crate) fn exec_outputs_read(
1427 &self,
1428 op: &StdTensorOp,
1429 inputs: &[TensorRead<'_>],
1430 ) -> Result<Vec<Tensor>> {
1431 self.exec_outputs_with_runtime(
1432 "exec_outputs_read.lock_backend",
1433 "exec_outputs_read.exec_op",
1434 op,
1435 |backend, runtime| exec_op_on_tensor_reads_with_runtime(op, inputs, backend, runtime),
1436 )
1437 }
1438
1439 #[cfg(test)]
1440 pub(crate) fn exec_standard_graph_outputs(
1441 &self,
1442 graph: &Graph<StdTensorOp>,
1443 initial_data: &HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>,
1444 ) -> Result<EagerGraphExecution> {
1445 let mut backend =
1446 profile_eager_op_section("exec_graph.lock_backend", || self.lock_backend())?;
1447 let mut all_values = initial_data.clone();
1448
1449 profile_eager_op_section("exec_graph.with_backend_session", || {
1450 backend.with_backend_session(|exec| -> Result<()> {
1451 for op_node in graph.operations() {
1452 let outputs = {
1453 let input_values = op_node
1454 .inputs
1455 .iter()
1456 .map(|input| {
1457 let key = match input {
1458 ValueRef::Local(local_id) => &graph.values()[*local_id].key,
1459 ValueRef::External(key) => key,
1460 };
1461 all_values.get(key).cloned().ok_or_else(|| {
1462 Error::Internal(format!(
1463 "standard graph eager execution missing value for {key:?}"
1464 ))
1465 })
1466 })
1467 .collect::<Result<Vec<_>>>()?;
1468 let input_reads = input_values
1469 .iter()
1470 .map(|value| TensorRead::from_tensor(value.as_ref()))
1471 .collect::<Vec<_>>();
1472 exec_standard_op_on_tensor_reads_in_session(
1473 &op_node.operation,
1474 &input_reads,
1475 exec,
1476 )?
1477 };
1478
1479 if outputs.len() != op_node.outputs.len() {
1480 return Err(Error::Internal(format!(
1481 "standard graph eager execution expected {} outputs for {:?}, got {}",
1482 op_node.outputs.len(),
1483 op_node.operation,
1484 outputs.len()
1485 )));
1486 }
1487
1488 for (output_id, output) in op_node.outputs.iter().zip(outputs) {
1489 let key = graph.values()[*output_id].key.clone();
1490 all_values.insert(key, Arc::new(output));
1491 }
1492 }
1493 Ok(())
1494 })
1495 })?;
1496
1497 let outputs = graph
1498 .outputs()
1499 .iter()
1500 .map(|&output_id| {
1501 let key = &graph.values()[output_id].key;
1502 all_values.get(key).cloned().ok_or_else(|| {
1503 Error::Internal(format!(
1504 "standard graph eager execution missing graph output {key:?}"
1505 ))
1506 })
1507 })
1508 .collect::<Result<Vec<_>>>()?;
1509
1510 Ok(EagerGraphExecution { outputs })
1511 }
1512
1513 pub(crate) fn try_register_grad_slot(
1514 &self,
1515 key: &ValueKey<StdTensorOp>,
1516 slot: &GradSlot,
1517 ) -> Result<()> {
1518 self.lock_grad_slots()?
1519 .insert(key.clone(), Arc::downgrade(slot));
1520 Ok(())
1521 }
1522
1523 pub(crate) fn try_register_value_record(
1524 &self,
1525 key: &ValueKey<StdTensorOp>,
1526 record: &Arc<EagerTensorRecord>,
1527 ) -> Result<()> {
1528 self.lock_value_records()?
1529 .insert(key.clone(), Arc::downgrade(record));
1530 self.try_register_value_record_ptr(record)?;
1531 Ok(())
1532 }
1533
1534 pub(crate) fn try_register_value_record_ptr(
1535 &self,
1536 record: &Arc<EagerTensorRecord>,
1537 ) -> Result<()> {
1538 let tensor = match record.value.as_tensor_arc() {
1539 Some(tensor) => Some(Arc::clone(tensor)),
1540 None => record.materialized_cache.get().cloned(),
1541 };
1542 let Some(tensor) = tensor else {
1543 return Ok(());
1544 };
1545 self.lock_value_ptr_records()?
1546 .insert(tensor_ptr(&tensor), Arc::downgrade(record));
1547 Ok(())
1548 }
1549
1550 pub(crate) fn value_record(
1551 &self,
1552 key: &ValueKey<StdTensorOp>,
1553 ) -> Result<Option<Arc<EagerTensorRecord>>> {
1554 let mut records = self.lock_value_records()?;
1555 let Some(record) = records.get(key).cloned() else {
1556 return Ok(None);
1557 };
1558 match record.upgrade() {
1559 Some(record) => Ok(Some(record)),
1560 None => {
1561 records.remove(key);
1562 Ok(None)
1563 }
1564 }
1565 }
1566
1567 pub fn clear_grads(&self) -> Result<()> {
1596 let live_slots = {
1597 let mut live_slots = Vec::new();
1598 self.lock_grad_slots()?.retain(|_, slot| {
1599 if let Some(slot) = slot.upgrade() {
1600 live_slots.push(slot);
1601 true
1602 } else {
1603 false
1604 }
1605 });
1606 live_slots
1607 };
1608
1609 let mut poisoned_slot = false;
1610 for slot in live_slots {
1611 match slot.lock() {
1612 Ok(mut current) => {
1613 *current = None;
1614 }
1615 Err(_) => {
1616 poisoned_slot = true;
1617 }
1618 }
1619 }
1620 if poisoned_slot {
1621 return Err(Error::runtime_state(
1622 "eager_gradient_slot",
1623 ErrorPhase::Execution,
1624 "lock poisoned",
1625 ));
1626 }
1627 Ok(())
1628 }
1629
1630 pub fn constant_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
1656 EagerTensor::new_leaf(Arc::clone(self), tensor, false)
1657 }
1658
1659 pub fn variable_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
1685 EagerTensor::new_leaf(Arc::clone(self), tensor, true)
1686 }
1687
1688 pub fn grad(self: &Arc<Self>, output: &EagerTensor, wrt: &EagerTensor) -> Result<EagerTensor> {
1718 self.grad_optional(output, wrt)?
1719 .ok_or_else(|| Error::Internal(format!("grad output is inactive for {:?}", wrt.key)))
1720 }
1721
1722 pub fn grad_optional(
1750 self: &Arc<Self>,
1751 output: &EagerTensor,
1752 wrt: &EagerTensor,
1753 ) -> Result<Option<EagerTensor>> {
1754 if !output.shape().is_empty() {
1755 return Err(Error::NonScalarGrad {
1756 shape: output.shape().to_vec(),
1757 });
1758 }
1759
1760 let value = output.materialized_arc()?;
1761 let seed = {
1762 let mut backend = self.lock_backend()?;
1763 one_like_tensor(value.as_ref(), &mut *backend)?
1764 };
1765 let seed = EagerTensor::new_result_arc(
1766 Arc::clone(self),
1767 eager_val_key(),
1768 Arc::new(seed),
1769 false,
1770 None,
1771 Vec::new(),
1772 )?;
1773 self.vjp_optional(output, wrt, &seed)
1774 }
1775
1776 pub fn vjp(
1806 self: &Arc<Self>,
1807 output: &EagerTensor,
1808 wrt: &EagerTensor,
1809 cotangent: &EagerTensor,
1810 ) -> Result<EagerTensor> {
1811 self.vjp_optional(output, wrt, cotangent)?
1812 .ok_or_else(|| Error::Internal(format!("vjp output is inactive for {:?}", wrt.key)))
1813 }
1814
1815 pub fn vjp_optional(
1848 self: &Arc<Self>,
1849 output: &EagerTensor,
1850 wrt: &EagerTensor,
1851 cotangent: &EagerTensor,
1852 ) -> Result<Option<EagerTensor>> {
1853 validate_same_runtime(self, output, "vjp output")?;
1854 validate_same_runtime(self, wrt, "vjp wrt")?;
1855 validate_same_runtime(self, cotangent, "vjp cotangent")?;
1856 validate_seed_tensor("vjp", output, cotangent)?;
1857 match semantic_eager_vjp_optional(self, output, wrt, cotangent)? {
1859 Some(result) => Ok(result),
1860 None => Ok(None),
1861 }
1862 }
1863
1864 pub fn jvp(
1894 self: &Arc<Self>,
1895 output: &EagerTensor,
1896 wrt: &EagerTensor,
1897 tangent: &EagerTensor,
1898 ) -> Result<EagerTensor> {
1899 self.jvp_optional(output, wrt, tangent)?
1900 .ok_or_else(|| Error::Internal(format!("jvp output is inactive for {:?}", wrt.key)))
1901 }
1902
1903 pub fn jvp_optional(
1936 self: &Arc<Self>,
1937 output: &EagerTensor,
1938 wrt: &EagerTensor,
1939 tangent: &EagerTensor,
1940 ) -> Result<Option<EagerTensor>> {
1941 validate_same_runtime(self, output, "jvp output")?;
1942 validate_same_runtime(self, wrt, "jvp wrt")?;
1943 validate_same_runtime(self, tangent, "jvp tangent")?;
1944 validate_seed_tensor("jvp", wrt, tangent)?;
1945 match semantic_eager_jvp_optional(self, output, wrt, tangent)? {
1947 Some(result) => Ok(result),
1948 None => Ok(None),
1949 }
1950 }
1951
1952 fn store_grads(
1953 &self,
1954 cotangents: &HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>,
1955 backend: &mut EagerBackend,
1956 ) -> Result<()> {
1957 let mut updates = Vec::new();
1958
1959 {
1960 let mut slots = self.lock_grad_slots()?;
1961 slots.retain(|key, slot| {
1962 let Some(slot) = slot.upgrade() else {
1963 return false;
1964 };
1965
1966 if let Some(incoming) = cotangents.get(key) {
1967 updates.push((slot, Arc::clone(incoming)));
1968 }
1969
1970 true
1971 });
1972 }
1973
1974 for (slot, incoming) in updates {
1975 let mut current = slot.lock().map_err(|_| {
1976 Error::runtime_state(
1977 "eager_gradient_slot",
1978 ErrorPhase::Execution,
1979 "lock poisoned",
1980 )
1981 })?;
1982 let next = match current.as_ref() {
1983 Some(existing) => Arc::new(backend.add(existing.as_ref(), incoming.as_ref())?),
1984 None => incoming,
1985 };
1986 *current = Some(next);
1987 }
1988
1989 Ok(())
1990 }
1991}
1992
1993#[derive(Clone, Debug, PartialEq, Eq, Hash)]
1994struct PreparedDerivativeCacheKey {
1995 semantic_fingerprint: SemanticFingerprint,
1996 runtime_epoch: RuntimeEpoch,
1997 wrt_input_index: usize,
1998 input_metadata: Box<[ProgramValueMetadata]>,
1999}
2000
2001#[derive(Debug)]
2003struct PreparedDerivative {
2004 program: Arc<CompiledGraph>,
2005 prepared: Arc<PreparedCompiledGraph>,
2006 seed_input_index: usize,
2007 derivative_output_index: usize,
2008}
2009
2010#[derive(Debug)]
2011struct PreparedDerivativeCache {
2012 limits: AdTransformCacheLimits,
2013 entries: LruCache<PreparedDerivativeCacheKey, PreparedDerivativeCacheEntry>,
2014 stats: CacheStats,
2015}
2016
2017impl PreparedDerivativeCache {
2018 fn limits(&self) -> AdTransformCacheLimits {
2019 self.limits
2020 }
2021
2022 fn set_limits(&mut self, limits: AdTransformCacheLimits) {
2023 self.limits = limits;
2024 self.evict_to_limits();
2025 }
2026
2027 fn clear(&mut self) {
2028 let clears = self.stats.clears.saturating_add(1);
2029 self.entries.clear();
2030 self.stats = CacheStats {
2031 clears,
2032 ..CacheStats::empty()
2033 };
2034 }
2035
2036 fn stats(&self) -> CacheStats {
2037 self.stats
2038 }
2039
2040 fn get(&mut self, key: &PreparedDerivativeCacheKey) -> Option<Arc<PreparedDerivative>> {
2041 match self.entries.get(key) {
2042 Some(entry) => {
2043 self.stats.hits = self.stats.hits.saturating_add(1);
2044 Some(Arc::clone(&entry.value))
2045 }
2046 None => {
2047 self.stats.misses = self.stats.misses.saturating_add(1);
2048 None
2049 }
2050 }
2051 }
2052
2053 fn insert(&mut self, key: PreparedDerivativeCacheKey, value: Arc<PreparedDerivative>) {
2054 let retained_bytes = prepared_derivative_cache_entry_retained_bytes(&key, value.as_ref());
2055 let entry = PreparedDerivativeCacheEntry {
2056 value,
2057 retained_bytes,
2058 };
2059 self.stats.retained_bytes = self.stats.retained_bytes.saturating_add(retained_bytes);
2060 if let Some((_old_key, old_entry)) = self.entries.push(key, entry) {
2061 self.stats.retained_bytes = self
2062 .stats
2063 .retained_bytes
2064 .saturating_sub(old_entry.retained_bytes);
2065 }
2066 self.stats.entries = self.entries.len();
2067 self.evict_to_limits();
2068 }
2069
2070 fn evict_to_limits(&mut self) {
2071 while self.entries.len() > self.limits.max_entries().get()
2072 || self
2073 .limits
2074 .max_retained_bytes()
2075 .is_some_and(|limit| self.stats.retained_bytes > limit.get())
2076 {
2077 let Some((_key, entry)) = self.entries.pop_lru() else {
2078 break;
2079 };
2080 self.stats.retained_bytes = self
2081 .stats
2082 .retained_bytes
2083 .saturating_sub(entry.retained_bytes);
2084 self.stats.evictions = self.stats.evictions.saturating_add(1);
2085 }
2086 self.stats.entries = self.entries.len();
2087 }
2088}
2089
2090impl Default for PreparedDerivativeCache {
2091 fn default() -> Self {
2092 Self {
2093 limits: AdTransformCacheLimits::default(),
2094 entries: LruCache::unbounded(),
2095 stats: CacheStats::empty(),
2096 }
2097 }
2098}
2099
2100#[derive(Debug)]
2101struct PreparedDerivativeCacheEntry {
2102 value: Arc<PreparedDerivative>,
2103 retained_bytes: usize,
2104}
2105
2106fn prepared_derivative_cache_entry_retained_bytes(
2107 key: &PreparedDerivativeCacheKey,
2108 value: &PreparedDerivative,
2109) -> usize {
2110 size_of::<PreparedDerivativeCacheKey>()
2111 .saturating_add(
2112 key.input_metadata
2113 .len()
2114 .saturating_mul(size_of::<ProgramValueMetadata>()),
2115 )
2116 .saturating_add(size_of::<PreparedDerivative>())
2117 .saturating_add(compiled_graph_retained_bytes(value.program.as_ref()))
2118 .saturating_add(prepared_compiled_graph_retained_bytes(
2119 value.prepared.as_ref(),
2120 value.program.as_ref(),
2121 ))
2122}
2123
2124fn prepared_compiled_graph_retained_bytes(
2125 prepared: &PreparedCompiledGraph,
2126 derivative_program: &CompiledGraph,
2127) -> usize {
2128 size_of_val(prepared).saturating_add(compiled_graph_retained_bytes(derivative_program))
2129}
2130
2131fn compiled_graph_retained_bytes(program: &CompiledGraph) -> usize {
2132 size_of::<CompiledGraph>()
2133 .saturating_add(size_of_val(program.input_keys()))
2134 .saturating_add(program.bindings().len().saturating_mul(size_of::<usize>()))
2135 .saturating_add(semantic_program_retained_bytes(program.program()))
2136}
2137
2138fn semantic_program_retained_bytes(program: &SemanticProgram) -> usize {
2139 size_of::<SemanticProgram>()
2140 .saturating_add(size_of_val(program.inputs()))
2141 .saturating_add(size_of_val(program.outputs()))
2142 .saturating_add(
2143 program
2144 .operations()
2145 .len()
2146 .saturating_mul(size_of::<usize>()),
2147 )
2148 .saturating_add(
2149 program
2150 .shape_guards()
2151 .len()
2152 .saturating_mul(size_of::<usize>()),
2153 )
2154}
2155
2156fn semantic_eager_vjp_optional(
2157 ctx: &Arc<EagerRuntime>,
2158 output: &EagerTensor,
2159 wrt: &EagerTensor,
2160 cotangent: &EagerTensor,
2161) -> Result<Option<Option<EagerTensor>>> {
2162 if !eager_semantic_vjp_enabled() {
2163 return Ok(None);
2164 }
2165 let (Some(output_trace), Some(wrt_trace)) =
2166 (output.semantic_trace.as_ref(), wrt.semantic_trace.as_ref())
2167 else {
2168 return Ok(None);
2169 };
2170 let Some(wrt_key) = wrt_trace.input_key() else {
2171 return Ok(None);
2172 };
2173 if !output_trace.has_attached_input_key(&wrt_key) {
2174 return Ok(None);
2175 }
2176
2177 let mut compiler = GraphCompiler::new();
2180 let source = compile_ad_source(&mut compiler, output_trace)?;
2181 if source.output_count() != 1
2182 || source.input_keys().len() != source.input_count()
2183 || source.bindings().len() != source.input_count()
2184 {
2185 return Ok(None);
2186 }
2187 let Some(wrt_input_index) = source.input_key_index(&wrt_key) else {
2188 return Ok(None);
2189 };
2190
2191 let cache_key = PreparedDerivativeCacheKey {
2193 semantic_fingerprint: source.program().semantic_fingerprint(),
2194 runtime_epoch: ctx.runtime.epoch().map_err(|source| {
2195 Error::runtime_state_source("semantic_eager_vjp", ErrorPhase::Execution, source)
2196 })?,
2197 wrt_input_index,
2198 input_metadata: source.frozen_program().input_metadata_with_bound_shapes(),
2199 };
2200 let prepared = { ctx.lock_prepared_derivative_cache()?.get(&cache_key) };
2201 let (seed_input_index, derivative_output_index, derivative_program, prepared_runtime) =
2202 if let Some(prepared) = prepared {
2203 (
2204 prepared.seed_input_index,
2205 prepared.derivative_output_index,
2206 Arc::clone(&prepared.program),
2207 Some(Arc::clone(&prepared.prepared)),
2208 )
2209 } else {
2210 let mut active_inputs = vec![false; source.input_count()];
2211 if let Some(active) = active_inputs.get_mut(wrt_input_index) {
2212 *active = true;
2213 } else {
2214 return Ok(None);
2215 }
2216 let active_outputs = vec![true; source.output_count()];
2217 let ad = AdContext::with_rules_and_transform_cache(
2218 ctx.semantic_extension_rules.clone(),
2219 Arc::clone(&ctx.ad_transform_cache),
2220 );
2221 let derivative = ad
2222 .vjp_program(source.frozen_program(), &active_inputs, &active_outputs)
2223 .map_err(|source| {
2224 Error::runtime_state_source(
2225 "semantic_eager_vjp",
2226 ErrorPhase::GraphBuild,
2227 source,
2228 )
2229 })?;
2230 let seed_input_index = derivative
2231 .derivative_input_indices()
2232 .first()
2233 .copied()
2234 .flatten();
2235 let derivative_output_index = derivative
2236 .derivative_output_indices()
2237 .get(wrt_input_index)
2238 .copied()
2239 .flatten();
2240 let (Some(seed_input_index), Some(derivative_output_index)) =
2241 (seed_input_index, derivative_output_index)
2242 else {
2243 return Ok(Some(None));
2244 };
2245 let program = Arc::new(compiler.compile_frozen_program(derivative.frozen())?);
2246 (seed_input_index, derivative_output_index, program, None)
2247 };
2248
2249 let cotangent_tensor = cotangent.materialized_arc()?;
2250 let input_count = derivative_program.input_count();
2251 let mut owned_inputs = vec![None; input_count];
2252 for (source_input_index, (_, tensor)) in source.bindings().iter().enumerate() {
2253 let Some(slot) = owned_inputs.get_mut(source_input_index) else {
2254 return Err(Error::Internal(format!(
2255 "semantic eager VJP derivative program has no primal input slot {source_input_index}"
2256 )));
2257 };
2258 *slot = Some(tensor.clone());
2259 }
2260 let Some(slot) = owned_inputs.get_mut(seed_input_index) else {
2261 return Err(Error::Internal(format!(
2262 "semantic eager VJP seed input index {seed_input_index} is outside {} inputs",
2263 owned_inputs.len()
2264 )));
2265 };
2266 *slot = Some(cotangent_tensor.as_ref().clone());
2267 let input_refs = owned_inputs
2268 .iter()
2269 .enumerate()
2270 .map(|(index, tensor)| {
2271 tensor.as_ref().ok_or_else(|| {
2272 Error::Internal(format!(
2273 "semantic eager VJP derivative input {index} was not populated"
2274 ))
2275 })
2276 })
2277 .collect::<Result<Vec<_>>>()?;
2278 let prepared_runtime = if let Some(prepared_runtime) = prepared_runtime {
2279 prepared_runtime
2280 } else {
2281 let prepared_runtime = Arc::new(
2282 ctx.runtime
2283 .prepare_compiled(&derivative_program, &input_refs)?,
2284 );
2285 let entry = Arc::new(PreparedDerivative {
2286 program: Arc::clone(&derivative_program),
2287 prepared: Arc::clone(&prepared_runtime),
2288 seed_input_index,
2289 derivative_output_index,
2290 });
2291 ctx.lock_prepared_derivative_cache()?
2292 .insert(cache_key, entry);
2293 prepared_runtime
2294 };
2295 let outputs = ctx.runtime.run_prepared(&prepared_runtime, &input_refs)?;
2296 let Some(result) = outputs.get(derivative_output_index).cloned() else {
2297 return Err(Error::Internal(format!(
2298 "semantic eager VJP derivative output index {derivative_output_index} is outside {} outputs",
2299 outputs.len()
2300 )));
2301 };
2302 let cotangent_trace =
2303 TracedTensor::from_tensor_arc_symbolic_shape(Arc::clone(&cotangent_tensor))?;
2304 let semantic_trace = derivative_trace_from_frozen_program(
2305 &source,
2306 derivative_program.frozen_program(),
2307 derivative_output_index,
2308 &[(seed_input_index, Arc::clone(&cotangent_tensor))],
2309 &[output_trace, wrt_trace, &cotangent_trace],
2310 None,
2311 "semantic_eager_vjp",
2312 )?;
2313
2314 #[cfg(test)]
2315 EAGER_SEMANTIC_VJP_EXECUTIONS.fetch_add(1, Ordering::Relaxed);
2316
2317 Ok(Some(Some(EagerTensor::new_result_arc_with_semantic_trace(
2318 Arc::clone(ctx),
2319 eager_val_key(),
2320 Arc::new(result),
2321 true,
2322 None,
2323 Some(semantic_trace),
2324 Vec::new(),
2325 )?)))
2326}
2327
2328fn semantic_eager_jvp_optional(
2329 ctx: &Arc<EagerRuntime>,
2330 output: &EagerTensor,
2331 wrt: &EagerTensor,
2332 tangent: &EagerTensor,
2333) -> Result<Option<Option<EagerTensor>>> {
2334 if !eager_semantic_vjp_enabled() {
2335 return Ok(None);
2336 }
2337 let (Some(output_trace), Some(wrt_trace)) =
2338 (output.semantic_trace.as_ref(), wrt.semantic_trace.as_ref())
2339 else {
2340 return Ok(None);
2341 };
2342 let Some(wrt_key) = wrt_trace.input_key() else {
2343 return Ok(None);
2344 };
2345 if !output_trace.has_attached_input_key(&wrt_key) {
2346 return Ok(None);
2347 }
2348
2349 let mut compiler = GraphCompiler::new();
2350 let source = compile_ad_source(&mut compiler, output_trace)?;
2351 if source.output_count() != 1
2352 || source.input_keys().len() != source.input_count()
2353 || source.bindings().len() != source.input_count()
2354 {
2355 return Ok(None);
2356 }
2357 let Some(wrt_input_index) = source.input_key_index(&wrt_key) else {
2358 return Ok(None);
2359 };
2360
2361 let mut active_inputs = vec![false; source.input_count()];
2362 if let Some(active) = active_inputs.get_mut(wrt_input_index) {
2363 *active = true;
2364 } else {
2365 return Ok(None);
2366 }
2367 let ad = AdContext::with_rules_and_transform_cache(
2368 ctx.semantic_extension_rules.clone(),
2369 Arc::clone(&ctx.ad_transform_cache),
2370 );
2371 let derivative = ad
2372 .jvp_program(source.frozen_program(), &active_inputs)
2373 .map_err(|source| {
2374 Error::runtime_state_source("semantic_eager_jvp", ErrorPhase::GraphBuild, source)
2375 })?;
2376 let Some(seed_input_index) = derivative
2378 .derivative_input_indices()
2379 .get(wrt_input_index)
2380 .copied()
2381 .flatten()
2382 else {
2383 return Ok(Some(None));
2384 };
2385 let Some(derivative_output_index) = derivative
2388 .derivative_output_indices()
2389 .first()
2390 .copied()
2391 .flatten()
2392 else {
2393 return Ok(Some(None));
2394 };
2395
2396 let derivative_program = compiler.compile_frozen_program(derivative.frozen())?;
2397 let tangent_tensor = tangent.materialized_arc()?;
2398 let input_count = derivative_program.input_count();
2399 let mut owned_inputs = vec![None; input_count];
2400 for (source_input_index, (_, tensor)) in source.bindings().iter().enumerate() {
2401 let Some(slot) = owned_inputs.get_mut(source_input_index) else {
2402 return Err(Error::Internal(format!(
2403 "semantic eager JVP derivative program has no primal input slot {source_input_index}"
2404 )));
2405 };
2406 *slot = Some(tensor.clone());
2407 }
2408 let Some(slot) = owned_inputs.get_mut(seed_input_index) else {
2409 return Err(Error::Internal(format!(
2410 "semantic eager JVP seed input index {seed_input_index} is outside {} inputs",
2411 owned_inputs.len()
2412 )));
2413 };
2414 *slot = Some(tangent_tensor.as_ref().clone());
2415 let input_refs = owned_inputs
2416 .iter()
2417 .enumerate()
2418 .map(|(index, tensor)| {
2419 tensor.as_ref().ok_or_else(|| {
2420 Error::Internal(format!(
2421 "semantic eager JVP derivative input {index} was not populated"
2422 ))
2423 })
2424 })
2425 .collect::<Result<Vec<_>>>()?;
2426 let outputs = ctx.runtime.run_compiled(&derivative_program, &input_refs)?;
2427 let Some(result) = outputs.get(derivative_output_index).cloned() else {
2428 return Err(Error::Internal(format!(
2429 "semantic eager JVP derivative output index {derivative_output_index} is outside {} outputs",
2430 outputs.len()
2431 )));
2432 };
2433 let tangent_trace = TracedTensor::from_tensor_arc_symbolic_shape(Arc::clone(&tangent_tensor))?;
2434 let semantic_trace = derivative_trace_from_frozen_program(
2435 &source,
2436 derivative.frozen(),
2437 derivative_output_index,
2438 &[(seed_input_index, Arc::clone(&tangent_tensor))],
2439 &[output_trace, wrt_trace, &tangent_trace],
2440 None,
2441 "semantic_eager_jvp",
2442 )?;
2443
2444 Ok(Some(Some(EagerTensor::new_result_arc_with_semantic_trace(
2445 Arc::clone(ctx),
2446 eager_val_key(),
2447 Arc::new(result),
2448 true,
2449 None,
2450 Some(semantic_trace),
2451 Vec::new(),
2452 )?)))
2453}
2454
2455fn validate_same_runtime(
2456 runtime: &Arc<EagerRuntime>,
2457 tensor: &EagerTensor,
2458 role: &'static str,
2459) -> Result<()> {
2460 if tensor.ctx_id() != runtime.id() {
2461 return Err(Error::ContextMismatch {
2462 lhs: runtime.id(),
2463 rhs: tensor.ctx_id(),
2464 });
2465 }
2466 let _ = role;
2467 Ok(())
2468}
2469
2470pub(crate) fn tensor_ptr(tensor: &Arc<Tensor>) -> usize {
2471 Arc::as_ptr(tensor) as usize
2472}
2473
2474fn validate_seed_tensor(op: &'static str, primal: &EagerTensor, seed: &EagerTensor) -> Result<()> {
2475 if primal.dtype() != seed.dtype() {
2476 return Err(
2477 tenferro_tensor::Error::dtype_mismatch(op, primal.dtype(), seed.dtype()).into(),
2478 );
2479 }
2480 if primal.shape() != seed.shape() {
2481 return Err(
2482 tenferro_tensor::Error::shape_mismatch(op, primal.shape(), seed.shape()).into(),
2483 );
2484 }
2485 Ok(())
2486}
2487
2488#[derive(Clone)]
2514pub struct EagerTensor {
2515 pub(crate) value: Arc<TensorValue>,
2516 materialized_cache: Arc<OnceLock<Arc<Tensor>>>,
2517 pub(crate) key: ValueKey<StdTensorOp>,
2518 pub(crate) trace: Option<EagerTrace>,
2519 pub(crate) semantic_trace: Option<TracedTensor>,
2520 pub(crate) requires_grad: bool,
2521 grad_slot: GradSlot,
2522 pub(crate) metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
2523 pub(crate) ctx: Arc<EagerRuntime>,
2524 _record: Arc<EagerTensorRecord>,
2525}
2526
2527pub(crate) struct EagerTensorRecord {
2528 value: Arc<TensorValue>,
2529 materialized_cache: Arc<OnceLock<Arc<Tensor>>>,
2530 key: ValueKey<StdTensorOp>,
2531 trace: Option<EagerTrace>,
2532 semantic_trace: Option<TracedTensor>,
2533 requires_grad: bool,
2534 grad_slot: GradSlot,
2535 metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
2536 ctx: Arc<EagerRuntime>,
2537}
2538
2539struct EagerTensorParts {
2540 ctx: Arc<EagerRuntime>,
2541 key: ValueKey<StdTensorOp>,
2542 requires_grad: bool,
2543 trace: Option<EagerTrace>,
2544 semantic_trace: Option<TracedTensor>,
2545 value: Arc<TensorValue>,
2546 metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
2547 register_value: bool,
2548}
2549
2550impl fmt::Debug for EagerTensor {
2551 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2552 f.debug_struct("EagerTensor")
2553 .field("dtype", &self.dtype())
2554 .field("shape", &self.shape())
2555 .field("key", &self.key)
2556 .field("requires_grad", &self.requires_grad)
2557 .field("has_trace", &self.trace.is_some())
2558 .field("has_semantic_trace", &self.semantic_trace.is_some())
2559 .field("ctx_id", &self.ctx_id())
2560 .finish_non_exhaustive()
2561 }
2562}
2563
2564impl EagerTensor {
2565 pub fn from_tensor_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
2586 Self::new_leaf(ctx, tensor, false)
2587 }
2588
2589 pub fn from_vec_col_major_in<T: TensorScalar>(
2601 shape: impl IntoShapeVec,
2602 data: Vec<T>,
2603 ctx: Arc<EagerRuntime>,
2604 ) -> Result<Self> {
2605 Self::from_tensor_in(Tensor::from_vec_col_major(shape, data)?, ctx)
2606 }
2607
2608 pub fn requires_grad_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
2629 Self::new_leaf(ctx, tensor, true)
2630 }
2631
2632 pub(crate) fn new_leaf(
2633 ctx: Arc<EagerRuntime>,
2634 tensor: Tensor,
2635 requires_grad: bool,
2636 ) -> Result<Self> {
2637 let key = eager_val_key();
2638 let tensor = Arc::new(tensor);
2639 let semantic_trace = TracedTensor::from_tensor_arc_symbolic_shape(Arc::clone(&tensor))?;
2640 let metadata_scope =
2641 register_scoped_value_metadata(key.clone(), tensor_meta_from_tensor(tensor.as_ref()))
2642 .map_err(|err| {
2643 Error::runtime_state_source("eager leaf metadata", ErrorPhase::GraphBuild, err)
2644 })?;
2645 Self::from_parts(EagerTensorParts {
2646 ctx,
2647 key,
2648 requires_grad,
2649 trace: None,
2650 semantic_trace: Some(semantic_trace),
2651 value: Arc::new(TensorValue::from_tensor_arc(tensor)),
2652 metadata_scopes: metadata_scopes_for_scope(metadata_scope),
2653 register_value: true,
2654 })
2655 }
2656
2657 pub(crate) fn new_result_arc(
2658 ctx: Arc<EagerRuntime>,
2659 key: ValueKey<StdTensorOp>,
2660 tensor: Arc<Tensor>,
2661 requires_grad: bool,
2662 trace: Option<EagerTrace>,
2663 metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
2664 ) -> Result<Self> {
2665 Self::new_result_arc_with_semantic_trace(
2666 ctx,
2667 key,
2668 tensor,
2669 requires_grad,
2670 trace,
2671 None,
2672 metadata_scopes,
2673 )
2674 }
2675
2676 pub(crate) fn new_result_arc_with_semantic_trace(
2677 ctx: Arc<EagerRuntime>,
2678 key: ValueKey<StdTensorOp>,
2679 tensor: Arc<Tensor>,
2680 requires_grad: bool,
2681 trace: Option<EagerTrace>,
2682 semantic_trace: Option<TracedTensor>,
2683 metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
2684 ) -> Result<Self> {
2685 Self::from_parts(EagerTensorParts {
2686 ctx,
2687 key,
2688 requires_grad,
2689 trace,
2690 semantic_trace,
2691 value: Arc::new(TensorValue::from_tensor_arc(tensor)),
2692 metadata_scopes,
2693 register_value: true,
2694 })
2695 }
2696
2697 pub(crate) fn new_unregistered_result_arc_with_semantic_trace(
2698 ctx: Arc<EagerRuntime>,
2699 key: ValueKey<StdTensorOp>,
2700 tensor: Arc<Tensor>,
2701 requires_grad: bool,
2702 trace: Option<EagerTrace>,
2703 semantic_trace: Option<TracedTensor>,
2704 metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
2705 ) -> Result<Self> {
2706 Self::from_parts(EagerTensorParts {
2707 ctx,
2708 key,
2709 requires_grad,
2710 trace,
2711 semantic_trace,
2712 value: Arc::new(TensorValue::from_tensor_arc(tensor)),
2713 metadata_scopes,
2714 register_value: false,
2715 })
2716 }
2717
2718 pub(crate) fn new_result_value(
2719 ctx: Arc<EagerRuntime>,
2720 key: ValueKey<StdTensorOp>,
2721 value: TensorValue,
2722 requires_grad: bool,
2723 trace: Option<EagerTrace>,
2724 semantic_trace: Option<TracedTensor>,
2725 metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
2726 ) -> Result<Self> {
2727 Self::from_parts(EagerTensorParts {
2728 ctx,
2729 key,
2730 requires_grad,
2731 trace,
2732 semantic_trace,
2733 value: Arc::new(value),
2734 metadata_scopes,
2735 register_value: true,
2736 })
2737 }
2738
2739 fn from_parts(parts: EagerTensorParts) -> Result<Self> {
2740 let EagerTensorParts {
2741 ctx,
2742 key,
2743 requires_grad,
2744 trace,
2745 semantic_trace,
2746 value,
2747 metadata_scopes,
2748 register_value,
2749 } = parts;
2750 let grad_slot = Arc::new(Mutex::new(None));
2751 if requires_grad {
2752 ctx.try_register_grad_slot(&key, &grad_slot)?;
2753 }
2754 let materialized_cache = Arc::new(OnceLock::new());
2755 let record = Arc::new(EagerTensorRecord {
2756 value: Arc::clone(&value),
2757 materialized_cache: Arc::clone(&materialized_cache),
2758 key: key.clone(),
2759 trace: trace.clone(),
2760 semantic_trace: semantic_trace.clone(),
2761 requires_grad,
2762 grad_slot: Arc::clone(&grad_slot),
2763 metadata_scopes: metadata_scopes.clone(),
2764 ctx: Arc::clone(&ctx),
2765 });
2766 if register_value {
2767 ctx.try_register_value_record(&key, &record)?;
2768 }
2769
2770 Ok(Self {
2771 value,
2772 materialized_cache,
2773 key,
2774 trace,
2775 semantic_trace,
2776 requires_grad,
2777 grad_slot,
2778 metadata_scopes,
2779 ctx,
2780 _record: record,
2781 })
2782 }
2783
2784 pub(crate) fn new_untracked_result(ctx: Arc<EagerRuntime>, tensor: Tensor) -> Result<Self> {
2785 Ok(Self::new_untracked_value_result(
2786 ctx,
2787 TensorValue::from_tensor(tensor),
2788 ))
2789 }
2790
2791 pub(crate) fn new_untracked_value_result(ctx: Arc<EagerRuntime>, value: TensorValue) -> Self {
2792 Self::new_untracked_value_result_with_semantic_trace(ctx, value, None)
2793 }
2794
2795 pub(crate) fn new_untracked_value_result_with_semantic_trace(
2796 ctx: Arc<EagerRuntime>,
2797 value: TensorValue,
2798 semantic_trace: Option<TracedTensor>,
2799 ) -> Self {
2800 let value = Arc::new(value);
2801 let materialized_cache = Arc::new(OnceLock::new());
2802 let key = eager_val_key();
2803 let grad_slot = Arc::new(Mutex::new(None));
2804 let record = Arc::new(EagerTensorRecord {
2805 value: Arc::clone(&value),
2806 materialized_cache: Arc::clone(&materialized_cache),
2807 key: key.clone(),
2808 trace: None,
2809 semantic_trace: semantic_trace.clone(),
2810 requires_grad: false,
2811 grad_slot: Arc::clone(&grad_slot),
2812 metadata_scopes: Vec::new(),
2813 ctx: Arc::clone(&ctx),
2814 });
2815 Self {
2816 value,
2817 materialized_cache,
2818 key,
2819 trace: None,
2820 semantic_trace,
2821 requires_grad: false,
2822 grad_slot,
2823 metadata_scopes: Vec::new(),
2824 ctx,
2825 _record: record,
2826 }
2827 }
2828
2829 pub(crate) fn from_record(record: Arc<EagerTensorRecord>) -> Self {
2830 Self {
2831 value: Arc::clone(&record.value),
2832 materialized_cache: Arc::clone(&record.materialized_cache),
2833 key: record.key.clone(),
2834 trace: record.trace.clone(),
2835 semantic_trace: record.semantic_trace.clone(),
2836 requires_grad: record.requires_grad,
2837 grad_slot: Arc::clone(&record.grad_slot),
2838 metadata_scopes: record.metadata_scopes.clone(),
2839 ctx: Arc::clone(&record.ctx),
2840 _record: record,
2841 }
2842 }
2843
2844 pub fn detach(&self) -> Self {
2864 let semantic_trace = self.value.as_tensor_arc().and_then(|tensor| {
2865 TracedTensor::from_tensor_arc_symbolic_shape(Arc::clone(tensor)).ok()
2866 });
2867 Self::new_untracked_value_result_with_semantic_trace(
2868 self.ctx.clone(),
2869 self.value.as_ref().clone(),
2870 semantic_trace,
2871 )
2872 }
2873
2874 pub fn detach_into(&self, ctx: &Arc<EagerRuntime>) -> Result<Self> {
2898 Self::from_tensor_in(self.to_tensor()?, Arc::clone(ctx))
2899 }
2900
2901 pub fn materialized(&self) -> Result<Arc<Tensor>> {
2920 self.materialized_arc()
2921 }
2922
2923 pub fn dtype(&self) -> DType {
2926 self.value.dtype()
2927 }
2928
2929 pub fn shape(&self) -> &[usize] {
2932 self.value.shape()
2933 }
2934
2935 pub fn tensor_read(&self) -> TensorRead<'_> {
2941 self.value.tensor_read()
2942 }
2943
2944 pub fn to_tensor(&self) -> Result<Tensor> {
2955 self.ctx.materialize_value(self.value.as_ref())
2956 }
2957
2958 pub(crate) fn materialized_arc(&self) -> Result<Arc<Tensor>> {
2959 if let Some(tensor) = self.value.as_tensor_arc() {
2960 self.ctx.try_register_value_record_ptr(&self._record)?;
2961 return Ok(Arc::clone(tensor));
2962 }
2963 if let Some(tensor) = self.materialized_cache.get() {
2964 self.ctx.try_register_value_record_ptr(&self._record)?;
2965 return Ok(Arc::clone(tensor));
2966 }
2967
2968 let materialized = Arc::new(self.ctx.materialize_value(self.value.as_ref())?);
2969 let _ = self.materialized_cache.set(Arc::clone(&materialized));
2970 self.ctx.try_register_value_record_ptr(&self._record)?;
2971 Ok(self
2972 .materialized_cache
2973 .get()
2974 .map(Arc::clone)
2975 .unwrap_or(materialized))
2976 }
2977
2978 #[cfg(test)]
2979 pub(crate) fn materialized_cache_is_initialized(&self) -> bool {
2980 self.materialized_cache.get().is_some()
2981 }
2982
2983 pub fn grad(&self) -> Result<Option<Arc<Tensor>>> {
3013 self.grad_slot
3014 .lock()
3015 .map_err(|_| {
3016 Error::runtime_state(
3017 "eager_gradient_slot",
3018 ErrorPhase::Execution,
3019 "lock poisoned",
3020 )
3021 })
3022 .map(|slot| slot.clone())
3023 }
3024
3025 pub fn clear_grad(&self) -> Result<()> {
3054 *self.grad_slot.lock().map_err(|_| {
3055 Error::runtime_state(
3056 "eager_gradient_slot",
3057 ErrorPhase::Execution,
3058 "lock poisoned",
3059 )
3060 })? = None;
3061 Ok(())
3062 }
3063
3064 pub fn tracks_grad(&self) -> bool {
3086 self.requires_grad
3087 }
3088
3089 #[cfg(test)]
3090 fn debug_trace_saved_value_count(&self) -> Option<usize> {
3091 None
3092 }
3093
3094 pub fn ctx_id(&self) -> ContextId {
3109 self.ctx.id()
3110 }
3111
3112 pub fn runtime(&self) -> &Arc<EagerRuntime> {
3114 &self.ctx
3115 }
3116
3117 pub fn same_context(&self, other: &Self) -> bool {
3133 self.ctx_id() == other.ctx_id()
3134 }
3135
3136 #[cfg(test)]
3137 pub(crate) fn standard_graph_op(
3138 inputs: &[&Self],
3139 build_graph: impl FnOnce(&[TensorInputKey]) -> Result<Arc<Graph<StdTensorOp>>>,
3140 ) -> Result<Vec<Self>> {
3141 let Some(first) = inputs.first() else {
3142 return Err(Error::Internal(
3143 "standard eager graph op requires at least one input tensor".to_string(),
3144 ));
3145 };
3146 let ctx = Arc::clone(&first.ctx);
3147 for tensor in inputs.iter().skip(1) {
3148 if !first.same_context(tensor) {
3149 return Err(Error::ContextMismatch {
3150 lhs: first.ctx_id(),
3151 rhs: tensor.ctx_id(),
3152 });
3153 }
3154 }
3155
3156 let graph_input_keys = (0..inputs.len())
3157 .map(|_| next_input_key())
3158 .collect::<Vec<_>>();
3159 let graph = build_graph(&graph_input_keys)?;
3160 let initial_data = graph_input_keys
3161 .iter()
3162 .zip(inputs.iter())
3163 .map(|(key, tensor)| Ok((ValueKey::Input(key.clone()), tensor.materialized_arc()?)))
3164 .collect::<Result<HashMap<_, _>>>()?;
3165 let execution = ctx.exec_standard_graph_outputs(graph.as_ref(), &initial_data)?;
3166 if execution.outputs.len() != graph.outputs().len() {
3167 return Err(Error::Internal(format!(
3168 "standard eager graph op expected {} graph outputs, got {}",
3169 graph.outputs().len(),
3170 execution.outputs.len()
3171 )));
3172 }
3173
3174 if !eager_grad_recording_enabled() || !inputs.iter().any(|input| input.requires_grad) {
3175 return execution
3176 .outputs
3177 .into_iter()
3178 .map(|output| {
3179 Self::new_unregistered_result_arc_with_semantic_trace(
3180 Arc::clone(&ctx),
3181 eager_val_key(),
3182 output,
3183 false,
3184 None,
3185 None,
3186 Vec::new(),
3187 )
3188 })
3189 .collect();
3190 }
3191
3192 let recorded = record_eager_graph_outputs(
3193 graph.as_ref(),
3194 &graph_input_keys,
3195 &execution.outputs,
3196 inputs,
3197 )?;
3198 if recorded.traces.len() != execution.outputs.len() {
3199 return Err(Error::Internal(format!(
3200 "standard eager graph op expected {} eager traces, got {}",
3201 execution.outputs.len(),
3202 recorded.traces.len()
3203 )));
3204 }
3205
3206 let mut metadata_scopes = vec![Arc::clone(&recorded.metadata_scope)];
3207 for input in inputs {
3208 for scope in &input.metadata_scopes {
3209 push_metadata_scope(&mut metadata_scopes, Arc::clone(scope));
3210 }
3211 }
3212
3213 recorded
3214 .traces
3215 .into_iter()
3216 .zip(recorded.semantic_traces)
3217 .zip(execution.outputs)
3218 .map(|((trace, semantic_trace), output)| {
3219 Self::new_result_arc_with_semantic_trace(
3220 Arc::clone(&ctx),
3221 trace.key,
3222 output,
3223 trace.requires_grad,
3224 trace.trace,
3225 semantic_trace,
3226 metadata_scopes.clone(),
3227 )
3228 })
3229 .collect()
3230 }
3231
3232 pub fn backward(&self) -> Result<HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>> {
3265 if !self.shape().is_empty() {
3266 return Err(Error::NonScalarGrad {
3267 shape: self.shape().to_vec(),
3268 });
3269 }
3270
3271 let value = self.materialized_arc()?;
3272 let seed = {
3273 let mut backend = self.ctx.lock_backend()?;
3274 Arc::new(one_like_tensor(value.as_ref(), &mut *backend)?)
3275 };
3276 self.backward_from_seed(seed)
3277 }
3278
3279 pub fn backward_with(
3314 &self,
3315 cotangent: &EagerTensor,
3316 ) -> Result<HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>> {
3317 if !self.same_context(cotangent) {
3318 return Err(Error::ContextMismatch {
3319 lhs: self.ctx_id(),
3320 rhs: cotangent.ctx_id(),
3321 });
3322 }
3323 validate_seed_tensor("backward", self, cotangent)?;
3324 self.backward_from_seed(cotangent.materialized_arc()?)
3325 }
3326
3327 fn backward_from_seed(
3328 &self,
3329 seed: Arc<Tensor>,
3330 ) -> Result<HashMap<ValueKey<StdTensorOp>, Arc<Tensor>>> {
3331 let cotangent = EagerTensor::new_result_arc(
3332 Arc::clone(&self.ctx),
3333 eager_val_key(),
3334 seed,
3335 false,
3336 None,
3337 Vec::new(),
3338 )?;
3339 let candidate_keys = {
3340 let mut slots = self.ctx.lock_grad_slots()?;
3341 let mut keys = Vec::new();
3342 slots.retain(|key, slot| {
3343 if slot.upgrade().is_some() {
3344 keys.push(key.clone());
3345 true
3346 } else {
3347 false
3348 }
3349 });
3350 keys
3351 };
3352
3353 let mut cotangents = HashMap::new();
3354 for key in candidate_keys {
3355 let Some(record) = self.ctx.value_record(&key)? else {
3356 continue;
3357 };
3358 if !record.requires_grad {
3359 continue;
3360 }
3361 let wrt = EagerTensor::from_record(record);
3362 let Some(grad) = self.ctx.vjp_optional(self, &wrt, &cotangent)? else {
3363 continue;
3364 };
3365 cotangents.insert(key, grad.materialized_arc()?);
3366 }
3367 let mut backend = self.ctx.lock_backend()?;
3368 self.ctx.store_grads(&cotangents, &mut backend)?;
3369 Ok(cotangents)
3370 }
3371}
3372
3373pub(crate) fn eager_val_key() -> ValueKey<StdTensorOp> {
3374 ValueKey::Input(next_input_key())
3375}
3376
3377pub(crate) struct RecordedEagerTrace {
3378 pub(crate) key: ValueKey<StdTensorOp>,
3379 pub(crate) trace: Option<EagerTrace>,
3380 pub(crate) requires_grad: bool,
3381}
3382
3383pub(crate) struct RecordedEagerOutputs {
3384 pub(crate) traces: Vec<RecordedEagerTrace>,
3385 pub(crate) semantic_traces: Vec<Option<TracedTensor>>,
3386 pub(crate) metadata_scope: Arc<GlobalMetadataScope>,
3387}
3388
3389pub(crate) fn record_eager_outputs(
3390 op: &StdTensorOp,
3391 outputs: &[Arc<Tensor>],
3392 inputs: &[&EagerTensor],
3393) -> Result<RecordedEagerOutputs> {
3394 let semantic_traces = record_semantic_eager_outputs(op, outputs.len(), inputs)?;
3395 let output_metadata = outputs
3396 .iter()
3397 .map(|output| tensor_meta_from_tensor(output.as_ref()));
3398 record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
3399}
3400
3401pub(crate) fn record_eager_value_outputs(
3402 op: &StdTensorOp,
3403 outputs: &[&TensorValue],
3404 inputs: &[&EagerTensor],
3405) -> Result<RecordedEagerOutputs> {
3406 let semantic_traces = record_semantic_eager_outputs(op, outputs.len(), inputs)?;
3407 let output_metadata = outputs.iter().map(|output| tensor_meta_from_value(output));
3408 record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
3409}
3410
3411fn record_semantic_eager_outputs(
3412 op: &StdTensorOp,
3413 output_count: usize,
3414 inputs: &[&EagerTensor],
3415) -> Result<Vec<Option<TracedTensor>>> {
3416 let Some(semantic_inputs) = inputs
3417 .iter()
3418 .map(|input| input.semantic_trace.as_ref())
3419 .collect::<Option<Vec<_>>>()
3420 else {
3421 return Ok(vec![None; output_count]);
3422 };
3423 let semantic_outputs = match op {
3424 StdTensorOp::Extension(ext) => {
3425 tenferro_runtime::extension::apply(Arc::clone(ext), &semantic_inputs)?
3426 }
3427 _ => tenferro_runtime::extension::apply_standard_op(op.clone(), &semantic_inputs)?,
3428 };
3429 if semantic_outputs.len() != output_count {
3430 return Err(Error::Internal(format!(
3431 "semantic eager recording expected {output_count} outputs for {op:?}, got {}",
3432 semantic_outputs.len()
3433 )));
3434 }
3435 Ok(semantic_outputs.into_iter().map(Some).collect())
3436}
3437
3438#[cfg(test)]
3439fn record_eager_graph_outputs(
3440 graph: &Graph<StdTensorOp>,
3441 graph_input_keys: &[TensorInputKey],
3442 outputs: &[Arc<Tensor>],
3443 inputs: &[&EagerTensor],
3444) -> Result<RecordedEagerOutputs> {
3445 let semantic_traces = record_semantic_eager_graph_outputs(graph, graph_input_keys, inputs)?;
3446 let output_metadata = outputs
3447 .iter()
3448 .map(|output| tensor_meta_from_tensor(output.as_ref()));
3449 record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
3450}
3451
3452#[cfg(test)]
3453fn record_semantic_eager_graph_outputs(
3454 graph: &Graph<StdTensorOp>,
3455 graph_input_keys: &[TensorInputKey],
3456 inputs: &[&EagerTensor],
3457) -> Result<Vec<Option<TracedTensor>>> {
3458 let Some(semantic_inputs) = inputs
3459 .iter()
3460 .map(|input| input.semantic_trace.as_ref())
3461 .collect::<Option<Vec<_>>>()
3462 else {
3463 return Ok(vec![None; graph.outputs().len()]);
3464 };
3465 if graph_input_keys.len() != semantic_inputs.len() {
3466 return Err(Error::Internal(format!(
3467 "semantic graph recording expected {} input keys, got {}",
3468 semantic_inputs.len(),
3469 graph_input_keys.len()
3470 )));
3471 }
3472
3473 let mut values = HashMap::new();
3474 for (key, tensor) in graph_input_keys.iter().zip(semantic_inputs) {
3475 values.insert(ValueKey::Input(key.clone()), tensor.clone());
3476 }
3477
3478 for op_node in graph.operations() {
3479 let input_values = op_node
3480 .inputs
3481 .iter()
3482 .map(|input| {
3483 let key = match input {
3484 ValueRef::Local(local_id) => &graph.values()[*local_id].key,
3485 ValueRef::External(key) => key,
3486 };
3487 values.get(key).cloned().ok_or_else(|| {
3488 Error::Internal(format!(
3489 "semantic graph recording missing value for {key:?}"
3490 ))
3491 })
3492 })
3493 .collect::<Result<Vec<_>>>()?;
3494 let input_refs = input_values.iter().collect::<Vec<_>>();
3495 let semantic_outputs = match &op_node.operation {
3496 StdTensorOp::Extension(ext) => {
3497 tenferro_runtime::extension::apply(Arc::clone(ext), &input_refs)?
3498 }
3499 op => tenferro_runtime::extension::apply_standard_op(op.clone(), &input_refs)?,
3500 };
3501 if semantic_outputs.len() != op_node.outputs.len() {
3502 return Err(Error::Internal(format!(
3503 "semantic graph recording expected {} outputs for {:?}, got {}",
3504 op_node.outputs.len(),
3505 op_node.operation,
3506 semantic_outputs.len()
3507 )));
3508 }
3509 for (output_id, output) in op_node.outputs.iter().copied().zip(semantic_outputs) {
3510 values.insert(graph.values()[output_id].key.clone(), output);
3511 }
3512 }
3513
3514 graph
3515 .outputs()
3516 .iter()
3517 .map(|&output_id| {
3518 let key = &graph.values()[output_id].key;
3519 values.get(key).cloned().map(Some).ok_or_else(|| {
3520 Error::Internal(format!(
3521 "semantic graph recording missing output for {key:?}"
3522 ))
3523 })
3524 })
3525 .collect()
3526}
3527
3528fn record_eager_outputs_from_metadata(
3529 output_metadata: impl IntoIterator<Item = TensorMeta>,
3530 semantic_traces: Vec<Option<TracedTensor>>,
3531 inputs: &[&EagerTensor],
3532) -> Result<RecordedEagerOutputs> {
3533 let output_metadata = output_metadata.into_iter().collect::<Vec<_>>();
3534 if semantic_traces.len() != output_metadata.len() {
3535 return Err(Error::Internal(format!(
3536 "eager recording expected {} semantic traces, got {}",
3537 output_metadata.len(),
3538 semantic_traces.len()
3539 )));
3540 }
3541 let requires_grad =
3542 eager_grad_recording_enabled() && inputs.iter().any(|input| input.requires_grad);
3543 let mut registrations = Vec::with_capacity(output_metadata.len());
3544 let traces = output_metadata
3545 .into_iter()
3546 .map(|metadata| {
3547 let key = eager_val_key();
3548 registrations.push((key.clone(), metadata));
3549 RecordedEagerTrace {
3550 key,
3551 trace: None,
3552 requires_grad,
3553 }
3554 })
3555 .collect();
3556
3557 Ok(RecordedEagerOutputs {
3558 traces,
3559 semantic_traces,
3560 metadata_scope: Arc::new(register_scoped_metadata_batch(registrations)?),
3561 })
3562}
3563
3564fn tensor_meta_from_value(value: &TensorValue) -> TensorMeta {
3565 TensorMeta::exact(
3566 value.dtype(),
3567 value.shape().iter().copied().map(SymDim::from).collect(),
3568 )
3569}
3570
3571pub(crate) fn exec_single_output(
3572 op: &StdTensorOp,
3573 inputs: &[&Tensor],
3574 ctx: &EagerRuntime,
3575) -> Result<Tensor> {
3576 let mut outputs = ctx.exec_outputs(op, inputs)?;
3577 if outputs.len() != 1 {
3578 return Err(Error::Internal(format!(
3579 "expected one eager output for {:?}, got {}",
3580 op,
3581 outputs.len()
3582 )));
3583 }
3584 Ok(profile_eager_op_section(
3585 "exec_single_output.remove_output",
3586 || outputs.remove(0),
3587 ))
3588}
3589
3590pub(crate) fn exec_single_output_read(
3591 op: &StdTensorOp,
3592 inputs: &[TensorRead<'_>],
3593 ctx: &EagerRuntime,
3594) -> Result<Tensor> {
3595 let mut outputs = ctx.exec_outputs_read(op, inputs)?;
3596 if outputs.len() != 1 {
3597 return Err(Error::Internal(format!(
3598 "expected one eager output for {:?}, got {}",
3599 op,
3600 outputs.len()
3601 )));
3602 }
3603 Ok(profile_eager_op_section(
3604 "exec_single_output_read.remove_output",
3605 || outputs.remove(0),
3606 ))
3607}
3608
3609#[cfg(test)]
3610pub(crate) fn zero_like_tensor<B: TensorBackend>(
3611 input: &Tensor,
3612 backend: &mut B,
3613) -> Result<Tensor> {
3614 let host = match input {
3615 Tensor::F32(tensor) => Tensor::F32(TypedTensor::zeros(tensor.shape().to_vec())?),
3616 Tensor::F64(tensor) => Tensor::F64(TypedTensor::zeros(tensor.shape().to_vec())?),
3617 Tensor::I32(tensor) => Tensor::I32(TypedTensor::zeros(tensor.shape().to_vec())?),
3618 Tensor::I64(tensor) => Tensor::I64(TypedTensor::zeros(tensor.shape().to_vec())?),
3619 Tensor::Bool(tensor) => Tensor::Bool(TypedTensor::from_vec_col_major(
3620 tensor.shape().to_vec(),
3621 vec![false; tensor.n_elements()],
3622 )?),
3623 Tensor::C32(tensor) => Tensor::C32(TypedTensor::zeros(tensor.shape().to_vec())?),
3624 Tensor::C64(tensor) => Tensor::C64(TypedTensor::zeros(tensor.shape().to_vec())?),
3625 };
3626 backend.upload_host_tensor(&host).map_err(Error::from)
3627}
3628
3629pub(crate) fn one_like_tensor<B: TensorBackend>(input: &Tensor, backend: &mut B) -> Result<Tensor> {
3630 let host = ones_tensor(input.dtype(), input.shape().to_vec())?;
3631 backend.upload_host_tensor(&host).map_err(Error::from)
3632}
3633
3634#[cfg(test)]
3635mod tests;