use std::cell::{Cell, RefCell};
use std::cmp::Reverse;
use std::collections::HashMap;
use std::env;
use std::fmt;
use std::mem::{size_of, size_of_val};
#[cfg(test)]
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, MutexGuard, OnceLock, Weak};
use std::time::{Duration, Instant};
use lru::LruCache;
use crate::extension::{
validate_eager_extension_target, EagerExtensionBackendKind, EagerExtensionTarget,
};
use crate::extension_cache::{ExtensionCacheLimits, ExtensionCacheSelector, ExtensionCacheStore};
#[cfg(test)]
use computegraph::graph::Graph;
use computegraph::ValueKey;
#[cfg(test)]
use computegraph::ValueRef;
use tenferro_cpu::{CpuBackend, CpuBackendError, CpuPlacement};
#[cfg(feature = "cuda")]
use tenferro_gpu::cuda::CudaBackend;
#[cfg(feature = "webgpu")]
use tenferro_gpu::webgpu::WebGpuBackend;
#[cfg(test)]
use tenferro_ops::input_key::TensorInputKey;
use tenferro_ops::{std_tensor_op::StdTensorOp, SymDim, TensorMeta};
use tenferro_runtime::ad_support::{compile_ad_source, ones_tensor, RetainedValue};
use tenferro_runtime::program::{ProgramValueMetadata, SemanticFingerprint, SemanticProgram};
use tenferro_runtime::{
CompiledGraph, CoreCapabilityBundle, EngineId, ErrorPhase, ExecutionContextIdentity,
ExtensionModule, GraphCompiler, HardwareClassId, PreparedCompiledGraph, RegistrationIdentity,
Runtime, RuntimeConfigError, RuntimeConfigSnapshot, RuntimeEpoch, TracedTensor,
};
#[cfg(test)]
use tenferro_tensor::TypedTensor;
use tenferro_tensor::{
AllocationGroup, CacheStats, DType, DescriptorSlot, GroupError, IntoShapeVec, Tensor,
TensorBackend, TensorRead, TensorScalar, TensorValue, TensorView,
};
use tenferro_tensor::{BackendSession, BackendSessionHost};
#[cfg(feature = "cuda")]
use crate::eager_backend::cuda_runtime_engine_id;
use crate::eager_backend::{
cpu_runtime_engine_id, cpu_runtime_hardware_class, eager_runtime_for_backend, EagerBackend,
};
#[cfg(test)]
use crate::eager_exec::exec_standard_op_on_tensor_reads_in_session;
use crate::eager_exec::{exec_op_on_tensor_reads_with_runtime, exec_op_on_tensors_with_runtime};
use crate::error::{ContextId, Error, Result};
#[cfg(test)]
use crate::metadata::push_metadata_scope;
use crate::metadata::{
metadata_scopes_for_scope, register_scoped_metadata_batch, register_scoped_value_metadata,
tensor_meta_from_tensor, GlobalMetadataScope,
};
use crate::semantic_extension::SemanticExtensionRuleSet;
use crate::traced::{derivative_trace_from_frozen_program, next_input_key};
use crate::transform_cache::{AdTransformCache, AdTransformCacheLimits};
use crate::AdContext;
pub(crate) type GradSlot = Arc<Mutex<Option<Arc<AdValueRecord>>>>;
pub(crate) type WeakGradSlot = Weak<Mutex<Option<Arc<AdValueRecord>>>>;
#[derive(Clone, Debug)]
pub(crate) struct EagerTrace;
#[cfg(test)]
pub(crate) static CPU_RUNTIME_SELECTION_REFRESHES: AtomicUsize = AtomicUsize::new(0);
struct CpuRuntimeSelection {
snapshot: Arc<RuntimeConfigSnapshot>,
epoch: RuntimeEpoch,
engine_id: EngineId,
registration_identity: RegistrationIdentity,
capabilities: CoreCapabilityBundle,
}
#[derive(Debug, Default, Clone)]
struct EagerOpProfileEntry {
calls: usize,
total_time: Duration,
}
thread_local! {
static EAGER_OP_PROFILE_STATE: RefCell<HashMap<&'static str, EagerOpProfileEntry>> =
RefCell::new(HashMap::new());
static EAGER_NO_GRAD_DEPTH: Cell<usize> = const { Cell::new(0) };
#[cfg(test)]
static EAGER_OP_PROFILE_ENABLED_OVERRIDE: RefCell<Option<bool>> = const { RefCell::new(None) };
#[cfg(test)]
static EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE: RefCell<Option<Option<usize>>> = const { RefCell::new(None) };
#[cfg(test)]
static EAGER_SEMANTIC_VJP_ENABLED_OVERRIDE: RefCell<Option<bool>> = const { RefCell::new(None) };
}
#[cfg(test)]
pub(crate) static EAGER_SEMANTIC_VJP_EXECUTIONS: AtomicUsize = AtomicUsize::new(0);
pub(crate) fn eager_grad_recording_enabled() -> bool {
EAGER_NO_GRAD_DEPTH.with(|depth| depth.get() == 0)
}
fn eager_semantic_vjp_enabled() -> bool {
#[cfg(test)]
if let Some(value) = EAGER_SEMANTIC_VJP_ENABLED_OVERRIDE.with(|state| *state.borrow()) {
return value;
}
static ENABLED: OnceLock<bool> = OnceLock::new();
*ENABLED.get_or_init(|| env::var("TENFERRO_EAGER_SEMANTIC_VJP").map_or(true, |v| v != "0"))
}
#[derive(Debug)]
pub struct EagerNoGradGuard {
active: bool,
}
impl Drop for EagerNoGradGuard {
fn drop(&mut self) {
if !self.active {
return;
}
EAGER_NO_GRAD_DEPTH.with(|depth| {
depth.set(depth.get().saturating_sub(1));
});
self.active = false;
}
}
pub(crate) fn eager_op_profile_enabled() -> bool {
#[cfg(test)]
if let Some(value) = EAGER_OP_PROFILE_ENABLED_OVERRIDE.with(|state| *state.borrow()) {
return value;
}
static ENABLED: OnceLock<bool> = OnceLock::new();
*ENABLED.get_or_init(|| env::var("TENFERRO_PROFILE_EAGER_OP_AGG").is_ok())
}
pub(crate) fn eager_op_profile_start() -> Option<Instant> {
eager_op_profile_enabled().then(Instant::now)
}
pub(crate) fn record_eager_op_profile(section: &'static str, elapsed: Duration) {
if !eager_op_profile_enabled() {
return;
}
EAGER_OP_PROFILE_STATE.with(|state| {
let mut state = state.borrow_mut();
let entry = state.entry(section).or_default();
entry.calls += 1;
entry.total_time += elapsed;
});
}
pub(crate) fn profile_eager_op_section<T>(section: &'static str, f: impl FnOnce() -> T) -> T {
if !eager_op_profile_enabled() {
return f();
}
let started = Instant::now();
let result = f();
record_eager_op_profile(section, started.elapsed());
result
}
pub(crate) fn maybe_print_eager_op_profile() {
if !eager_op_profile_enabled() {
return;
}
let Some(print_every) = eager_op_profile_print_every() else {
return;
};
if print_every == 0 {
return;
}
let should_print = EAGER_OP_PROFILE_STATE.with(|state| {
state
.borrow()
.get("nary_op.total")
.is_some_and(|entry| entry.calls % print_every == 0)
});
if should_print {
print_and_reset_eager_op_profile();
}
}
fn eager_op_profile_print_every() -> Option<usize> {
#[cfg(test)]
if let Some(value) = EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE.with(|state| *state.borrow()) {
return value;
}
env::var("TENFERRO_PROFILE_EAGER_OP_PRINT_EVERY")
.ok()?
.parse()
.ok()
}
pub(crate) fn print_and_reset_eager_op_profile() {
EAGER_OP_PROFILE_STATE.with(|state| {
let mut entries: Vec<_> = state
.borrow()
.iter()
.map(|(section, entry)| (*section, entry.clone()))
.collect();
state.borrow_mut().clear();
entries.sort_by_key(|(_, entry)| Reverse(entry.total_time));
eprintln!("=== tenferro eager op profile ===");
for (section, entry) in entries {
let Some(per_call_us) = eager_op_profile_per_call_us(&entry) else {
continue;
};
eprintln!(
"{section}: calls={} total={:.6}ms per_call={:.3}us",
entry.calls,
entry.total_time.as_secs_f64() * 1.0e3,
per_call_us,
);
}
});
}
fn eager_op_profile_per_call_us(entry: &EagerOpProfileEntry) -> Option<f64> {
(entry.calls != 0).then(|| entry.total_time.as_secs_f64() * 1.0e6 / entry.calls as f64)
}
fn runtime_config_error(op: &'static str, source: RuntimeConfigError) -> Error {
Error::runtime_state_source(op, ErrorPhase::Execution, source)
}
fn runtime_state_source<E>(op: &'static str, source: E) -> Error
where
E: std::error::Error + Send + Sync + 'static,
{
Error::runtime_state_source(op, ErrorPhase::Execution, source)
}
fn cpu_runtime_bridge_unsupported(message: impl Into<String>) -> Error {
Error::unsupported(
"CpuPlacementBoundEager::refresh_runtime_selection",
ErrorPhase::Execution,
message,
)
}
fn select_cpu_runtime(runtime: &Runtime) -> Result<CpuRuntimeSelection> {
let snapshot = runtime
.snapshot()
.map_err(|source| runtime_state_source("EagerRuntime::runtime_snapshot", source))?;
let engine_id = cpu_runtime_engine_id()
.map_err(|source| runtime_config_error("EagerRuntime::cpu_runtime_engine_id", source))?;
let expected_hardware = cpu_runtime_hardware_class().map_err(|source| {
runtime_config_error("EagerRuntime::cpu_runtime_hardware_class", source)
})?;
let engine = snapshot
.engine(&engine_id)
.ok_or_else(|| cpu_runtime_bridge_unsupported("missing CPU runtime engine"))?;
validate_cpu_runtime_engine(
engine.context_identity(),
engine.hardware_class(),
engine.capabilities(),
&expected_hardware,
)?;
let epoch = snapshot.epoch();
let registration_identity = engine.registration_identity();
let capabilities = engine.capabilities().clone();
Ok(CpuRuntimeSelection {
snapshot,
epoch,
engine_id,
registration_identity,
capabilities,
})
}
fn validate_cpu_runtime_engine(
context_identity: ExecutionContextIdentity,
hardware_class: &HardwareClassId,
capabilities: &CoreCapabilityBundle,
expected_hardware: &HardwareClassId,
) -> Result<()> {
if context_identity != ExecutionContextIdentity::of::<CpuBackend>() {
return Err(cpu_runtime_bridge_unsupported(
"CPU runtime context mismatch",
));
}
if hardware_class != expected_hardware {
return Err(cpu_runtime_bridge_unsupported(
"CPU runtime hardware mismatch",
));
}
if capabilities.elementwise().is_none() {
return Err(cpu_runtime_bridge_unsupported(
"missing CPU runtime capability: elementwise",
));
}
if capabilities.reduction().is_none() {
return Err(cpu_runtime_bridge_unsupported(
"missing CPU runtime capability: reduction",
));
}
if capabilities.indexing().is_none() {
return Err(cpu_runtime_bridge_unsupported(
"missing CPU runtime capability: indexing",
));
}
if capabilities.dot_general().is_none() {
return Err(cpu_runtime_bridge_unsupported(
"missing CPU runtime capability: dot_general",
));
}
if capabilities.layout().is_none() {
return Err(cpu_runtime_bridge_unsupported(
"missing CPU runtime capability: layout",
));
}
Ok(())
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct EagerRuntimeCacheStats {
pub extensions: CacheStats,
pub ad_transforms: CacheStats,
pub prepared_derivatives: CacheStats,
}
#[cfg(test)]
pub(crate) struct EagerGraphExecution {
pub(crate) outputs: Vec<Tensor>,
}
#[derive(Debug)]
pub struct ValueGuard<'a> {
view: TensorView<'a>,
}
impl<'a> ValueGuard<'a> {
pub fn dtype(&self) -> DType {
self.view.dtype()
}
pub fn shape(&self) -> &[usize] {
self.view.shape()
}
pub fn as_tensor_view(&self) -> &TensorView<'_> {
&self.view
}
pub fn as_slice<T: TensorScalar>(&self) -> tenferro_tensor::Result<&'a [T]> {
self.view.as_slice()
}
fn duplicate_host_tensor(&self) -> tenferro_tensor::Result<Tensor> {
match &self.view {
TensorView::F32(view) => {
<f32 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
}
TensorView::F64(view) => {
<f64 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
}
TensorView::I32(view) => {
<i32 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
}
TensorView::I64(view) => {
<i64 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
}
TensorView::Bool(view) => <bool as TensorScalar>::into_tensor(
view.shape().to_vec(),
view.as_slice()?.to_vec(),
),
TensorView::C32(view) => <num_complex::Complex32 as TensorScalar>::into_tensor(
view.shape().to_vec(),
view.as_slice()?.to_vec(),
),
TensorView::C64(view) => <num_complex::Complex64 as TensorScalar>::into_tensor(
view.shape().to_vec(),
view.as_slice()?.to_vec(),
),
}
}
}
#[derive(Clone, Debug)]
pub struct GradientValue {
record: Arc<AdValueRecord>,
ctx: Arc<EagerRuntime>,
}
impl GradientValue {
pub fn dtype(&self) -> DType {
self.record.dtype()
}
pub fn shape(&self) -> &[usize] {
self.record.shape()
}
pub fn value(&self) -> Result<ValueGuard<'_>> {
self.record.value("GradientValue::value")
}
pub fn tensor_read(&self) -> Result<TensorRead<'_>> {
self.record.tensor_read("GradientValue::tensor_read")
}
pub fn as_slice<T: TensorScalar>(&self) -> tenferro_tensor::Result<&[T]> {
self.record
.value("GradientValue::as_slice")
.map_err(|error| {
tenferro_tensor::Error::runtime_state_source("GradientValue::as_slice", error)
})?
.as_slice()
}
pub fn to_tensor(&self) -> Result<Tensor> {
let value = self
.record
.value("GradientValue::to_tensor")
.map_err(|error| {
Error::runtime_state_source(
"GradientValue::to_tensor",
ErrorPhase::Execution,
error,
)
})?;
match value.duplicate_host_tensor() {
Ok(tensor) => Ok(tensor),
Err(_) => {
let read = self.record.tensor_read("GradientValue::to_tensor")?;
self.ctx
.with_execution_session(|session| session.to_contiguous_read(read))?
.map_err(Error::from)
}
}
}
}
#[derive(Debug)]
pub struct Gradients {
group: AllocationGroup,
slots: HashMap<ValueKey<StdTensorOp>, DescriptorSlot>,
}
impl Gradients {
fn from_tensors(tensors: HashMap<ValueKey<StdTensorOp>, Tensor>) -> Result<Self> {
let (keys, values): (Vec<_>, Vec<_>) = tensors.into_iter().unzip();
let (group, bindings) = AllocationGroup::from_tensors(values).map_err(|error| {
Error::runtime_state_source("Gradients::from_tensors", ErrorPhase::Execution, error)
})?;
let slots = keys.into_iter().zip(bindings).collect();
Ok(Self { group, slots })
}
pub fn len(&self) -> usize {
self.slots.len()
}
pub fn is_empty(&self) -> bool {
self.slots.is_empty()
}
pub fn grad(&self, key: &ValueKey<StdTensorOp>) -> Option<TensorView<'_>> {
let slot = self.slots.get(key).copied()?;
let mut reads = self.group.read_views(std::slice::from_ref(&slot)).ok()?;
match reads.pop()? {
TensorRead::View(view) => Some(view),
TensorRead::Tensor(_) => None,
}
}
pub fn take_grad(
&mut self,
key: &ValueKey<StdTensorOp>,
) -> tenferro_tensor::Result<Option<Tensor>> {
let Some(&slot) = self.slots.get(key) else {
return Ok(None);
};
let tensor = self.group.take_tensor(slot).map_err(|error| {
tenferro_tensor::Error::runtime_state_source("Gradients::take_grad", error)
})?;
self.slots.remove(key);
Ok(Some(tensor))
}
}
#[derive(Debug)]
pub enum IntoValueError<H> {
NotUnique(H),
Extract { value: H, error: GroupError },
}
impl<H> std::fmt::Display for IntoValueError<H> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::NotUnique(_) => formatter.write_str("eager value is retained by another handle"),
Self::Extract { error, .. } => {
write!(formatter, "eager value extraction failed: {error}")
}
}
}
}
impl<H: std::fmt::Debug + Send + Sync + 'static> std::error::Error for IntoValueError<H> {}
#[derive(Debug)]
struct RetentionContainer {
group: AllocationGroup,
}
#[derive(Debug)]
pub(crate) struct AdValueRecord {
container: Arc<RetentionContainer>,
slot: DescriptorSlot,
dtype: DType,
shape: Box<[usize]>,
}
impl AdValueRecord {
fn from_group(
group: AllocationGroup,
slot: DescriptorSlot,
dtype: DType,
shape: Vec<usize>,
) -> Arc<Self> {
Arc::new(Self {
container: Arc::new(RetentionContainer { group }),
slot,
dtype,
shape: shape.into_boxed_slice(),
})
}
fn from_tensor(tensor: Tensor, op: &'static str) -> Result<Arc<Self>> {
let dtype = tensor.dtype();
let shape = tensor.shape().to_vec();
let (group, bindings) = AllocationGroup::from_tensors(vec![tensor])
.map_err(|error| Error::runtime_state_source(op, ErrorPhase::Execution, error))?;
let slot = bindings.first().copied().ok_or_else(|| {
Error::runtime_state(op, ErrorPhase::Execution, "empty allocation-group binding")
})?;
Ok(Self::from_group(group, slot, dtype, shape))
}
fn tensor_read(&self, op: &'static str) -> Result<TensorRead<'_>> {
let mut reads = self
.container
.group
.read_views(std::slice::from_ref(&self.slot))
.map_err(|error| Error::runtime_state_source(op, ErrorPhase::Execution, error))?;
reads.pop().ok_or_else(|| {
Error::runtime_state(op, ErrorPhase::Execution, "empty allocation-group binding")
})
}
fn value(&self, op: &'static str) -> Result<ValueGuard<'_>> {
match self.tensor_read(op)? {
TensorRead::View(view) => Ok(ValueGuard { view }),
TensorRead::Tensor(_) => Err(Error::runtime_state(
op,
ErrorPhase::Execution,
"allocation-group value did not produce a borrowed descriptor view",
)),
}
}
fn dtype(&self) -> DType {
self.dtype
}
fn shape(&self) -> &[usize] {
&self.shape
}
}
pub struct CpuPlacementBoundEager {
runtime: Arc<EagerRuntime>,
backend: CpuBackend,
snapshot: Arc<RuntimeConfigSnapshot>,
epoch: RuntimeEpoch,
engine_id: EngineId,
registration_identity: RegistrationIdentity,
capabilities: CoreCapabilityBundle,
}
impl fmt::Debug for CpuPlacementBoundEager {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("CpuPlacementBoundEager")
.field("runtime_id", &self.runtime.id())
.field("placement", &self.backend.placement())
.field("runtime_epoch", &self.epoch)
.field("engine_id", &self.engine_id)
.field("registration_identity", &self.registration_identity)
.finish_non_exhaustive()
}
}
impl CpuPlacementBoundEager {
fn refresh_runtime_selection(&mut self) -> Result<()> {
let current_epoch = self.runtime.runtime.epoch().map_err(|source| {
runtime_state_source("CpuPlacementBoundEager::refresh_runtime_selection", source)
})?;
if current_epoch == self.epoch {
return Ok(());
}
#[cfg(test)]
CPU_RUNTIME_SELECTION_REFRESHES.fetch_add(1, Ordering::SeqCst);
let selection = select_cpu_runtime(&self.runtime.runtime)?;
self.snapshot = selection.snapshot;
self.epoch = selection.epoch;
self.engine_id = selection.engine_id;
self.registration_identity = selection.registration_identity;
self.capabilities = selection.capabilities;
Ok(())
}
pub fn runtime_id(&self) -> ContextId {
self.runtime.id()
}
pub fn placement(&self) -> CpuPlacement {
self.backend.placement()
}
pub fn with_eager_session<R: Send>(
&mut self,
f: impl FnOnce(&mut dyn BackendSession) -> Result<R> + Send,
) -> Result<R> {
self.refresh_runtime_selection()?;
self.backend.with_backend_session(f)
}
}
pub struct EagerRuntime {
id: ContextId,
runtime: Runtime,
backend: Mutex<EagerBackend>,
extension_install_lock: Mutex<()>,
pub(crate) extension_caches: Mutex<ExtensionCacheStore>,
semantic_extension_rules: SemanticExtensionRuleSet,
grad_slots: Mutex<HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>,
value_records: Mutex<HashMap<ValueKey<StdTensorOp>, Weak<EagerTensorRecord>>>,
ad_transform_cache: Arc<AdTransformCache>,
prepared_derivative_cache: Mutex<PreparedDerivativeCache>,
}
impl fmt::Debug for EagerRuntime {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut debug = f.debug_struct("EagerRuntime");
debug.field("id", &self.id);
debug.field("runtime_id", &self.runtime.id());
debug.field("runtime_epoch", &self.runtime.epoch().ok());
match self.backend.try_lock() {
Ok(backend) => {
debug.field("backend", &*backend);
}
Err(_) => {
debug.field("backend", &"<locked>");
}
}
match self.extension_caches.try_lock() {
Ok(caches) => {
debug.field(
"extension_cache_stats",
&caches.stats(ExtensionCacheSelector::All),
);
}
Err(_) => {
debug.field("extension_cache_stats", &"<locked>");
}
}
match self.extension_install_lock.try_lock() {
Ok(_) => {
debug.field("extension_install_lock", &"<unlocked>");
}
Err(_) => {
debug.field("extension_install_lock", &"<locked>");
}
}
debug.field("semantic_extension_rules", &self.semantic_extension_rules);
match self.grad_slots.try_lock() {
Ok(slots) => {
debug.field("grad_slots_len", &slots.len());
}
Err(_) => {
debug.field("grad_slots_len", &"<locked>");
}
}
match self.value_records.try_lock() {
Ok(records) => {
debug.field("value_records_len", &records.len());
}
Err(_) => {
debug.field("value_records_len", &"<locked>");
}
}
match self.ad_transform_cache.stats() {
Ok(stats) => {
debug.field("ad_transform_cache_stats", &stats);
}
Err(err) => {
debug.field("ad_transform_cache_stats", &format_args!("{err}"));
}
}
match self.prepared_derivative_cache.try_lock() {
Ok(cache) => {
debug.field("prepared_derivative_cache_stats", &cache.stats());
}
Err(_) => {
debug.field("prepared_derivative_cache_stats", &"<locked>");
}
}
debug.finish_non_exhaustive()
}
}
impl EagerRuntime {
pub(crate) fn lock_backend(&self) -> Result<MutexGuard<'_, EagerBackend>> {
self.backend.lock().map_err(|_| {
Error::runtime_state("eager_backend", ErrorPhase::Execution, "lock poisoned")
})
}
fn lock_extension_caches(&self) -> Result<MutexGuard<'_, ExtensionCacheStore>> {
self.extension_caches.lock().map_err(|_| {
Error::runtime_state(
"eager_extension_caches",
ErrorPhase::Execution,
"lock poisoned",
)
})
}
fn lock_extension_install(&self) -> Result<MutexGuard<'_, ()>> {
self.extension_install_lock.lock().map_err(|_| {
Error::runtime_state(
"eager_extension_install",
ErrorPhase::Execution,
"lock poisoned",
)
})
}
fn lock_prepared_derivative_cache(&self) -> Result<MutexGuard<'_, PreparedDerivativeCache>> {
self.prepared_derivative_cache.lock().map_err(|_| {
Error::runtime_state(
"prepared_derivative_cache",
ErrorPhase::Execution,
"lock poisoned",
)
})
}
fn lock_grad_slots(
&self,
) -> Result<MutexGuard<'_, HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>> {
self.grad_slots.lock().map_err(|_| {
Error::runtime_state(
"eager_gradient_slots",
ErrorPhase::Execution,
"lock poisoned",
)
})
}
fn lock_value_records(
&self,
) -> Result<MutexGuard<'_, HashMap<ValueKey<StdTensorOp>, Weak<EagerTensorRecord>>>> {
self.value_records.lock().map_err(|_| {
Error::runtime_state(
"eager_value_registry",
ErrorPhase::Execution,
"lock poisoned",
)
})
}
fn from_backend(backend: EagerBackend) -> Result<Self> {
Self::from_backend_with_rules_and_cache(
backend,
SemanticExtensionRuleSet::default(),
Arc::new(AdTransformCache::new()),
)
}
fn from_backend_with_rules_and_cache(
backend: EagerBackend,
semantic_extension_rules: SemanticExtensionRuleSet,
ad_transform_cache: Arc<AdTransformCache>,
) -> Result<Self> {
let runtime = eager_runtime_for_backend(&backend)
.map_err(|source| runtime_config_error("EagerRuntime::from_backend", source))?;
Ok(Self {
id: ContextId::fresh(),
runtime,
backend: Mutex::new(backend),
extension_install_lock: Mutex::new(()),
extension_caches: Mutex::new(ExtensionCacheStore::new()),
semantic_extension_rules,
grad_slots: Mutex::new(HashMap::new()),
value_records: Mutex::new(HashMap::new()),
ad_transform_cache,
prepared_derivative_cache: Mutex::new(PreparedDerivativeCache::default()),
})
}
pub fn new() -> Result<Arc<Self>> {
Self::with_cpu_backend(CpuBackend::new())
}
pub fn with_cpu_backend(backend: CpuBackend) -> Result<Arc<Self>> {
Ok(Arc::new(Self::from_backend(EagerBackend::cpu(backend))?))
}
pub fn on_cpu(self: &Arc<Self>, placement: CpuPlacement) -> Result<CpuPlacementBoundEager> {
let backend = {
let backend = self.lock_backend()?;
backend.cpu_snapshot().ok_or_else(|| {
Error::unsupported(
"EagerRuntime::on_cpu",
ErrorPhase::Execution,
"the eager runtime is not CPU-backed",
)
})?
};
let selection = select_cpu_runtime(&self.runtime)?;
let backend = backend.for_placement(placement).map_err(|source| {
let error: tenferro_tensor::Error = CpuBackendError::Placement {
op: "EagerRuntime::on_cpu",
source,
}
.into();
Error::from(error)
})?;
Ok(CpuPlacementBoundEager {
runtime: Arc::clone(self),
backend,
snapshot: selection.snapshot,
epoch: selection.epoch,
engine_id: selection.engine_id,
registration_identity: selection.registration_identity,
capabilities: selection.capabilities,
})
}
pub fn with_cpu_backend_and_ad_context(
backend: CpuBackend,
ad: &AdContext,
) -> Result<Arc<Self>> {
Ok(Arc::new(Self::from_backend_with_rules_and_cache(
EagerBackend::cpu(backend),
ad.semantic_extension_rules().clone(),
ad.ad_transform_cache(),
)?))
}
#[cfg(feature = "cuda")]
pub fn with_cuda_backend(backend: CudaBackend) -> Result<Arc<Self>> {
Ok(Arc::new(Self::from_backend(EagerBackend::cuda(backend))?))
}
#[cfg(feature = "cuda")]
pub fn with_cuda_backend_and_ad_context(
backend: CudaBackend,
ad: &AdContext,
) -> Result<Arc<Self>> {
Ok(Arc::new(Self::from_backend_with_rules_and_cache(
EagerBackend::cuda(backend),
ad.semantic_extension_rules().clone(),
ad.ad_transform_cache(),
)?))
}
#[cfg(feature = "webgpu")]
pub fn with_webgpu_backend(backend: WebGpuBackend) -> Result<Arc<Self>> {
Ok(Arc::new(Self::from_backend(EagerBackend::webgpu(backend))?))
}
#[cfg(feature = "webgpu")]
pub fn with_webgpu_backend_and_ad_context(
backend: WebGpuBackend,
ad: &AdContext,
) -> Result<Arc<Self>> {
Ok(Arc::new(Self::from_backend_with_rules_and_cache(
EagerBackend::webgpu(backend),
ad.semantic_extension_rules().clone(),
ad.ad_transform_cache(),
)?))
}
pub fn id(&self) -> ContextId {
self.id
}
pub fn no_grad(&self) -> EagerNoGradGuard {
EAGER_NO_GRAD_DEPTH.with(|depth| {
depth.set(depth.get().saturating_add(1));
});
EagerNoGradGuard { active: true }
}
pub fn install_extension_module(
&self,
module: Arc<dyn ExtensionModule>,
) -> Result<RuntimeEpoch> {
let _install_guard = self.lock_extension_install()?;
self.runtime
.reconfigure(|edit| {
edit.replace_extension_module(module)?;
Ok(())
})
.map_err(|source| {
runtime_state_source("EagerRuntime::install_extension_module", source)
})
}
pub(crate) fn ensure_extension_module_for_engine(
&self,
module: Arc<dyn ExtensionModule>,
family_id: &'static str,
engine_id: &EngineId,
) -> Result<RuntimeEpoch> {
let _install_guard = self.lock_extension_install()?;
self.runtime
.reconfigure(|edit| {
edit.ensure_extension_module_for_engine(module, family_id, engine_id)?;
Ok(())
})
.map_err(|source| {
runtime_state_source("EagerRuntime::ensure_extension_module_for_engine", source)
})
}
pub(crate) fn runtime(&self) -> &Runtime {
&self.runtime
}
pub(crate) fn eager_extension_target(&self) -> Result<EagerExtensionTarget> {
let (engine_id, backend_kind) = {
let backend = self.lock_backend()?;
match &*backend {
EagerBackend::Cpu(_) => (
cpu_runtime_engine_id().map_err(|source| {
runtime_config_error("EagerRuntime::eager_extension_target", source)
})?,
EagerExtensionBackendKind::Cpu,
),
#[cfg(test)]
EagerBackend::Recording(_) => {
return Err(Error::unsupported(
"EagerRuntime::eager_extension_target",
ErrorPhase::Execution,
"the recording backend has no registered eager extension engine",
));
}
#[cfg(feature = "cuda")]
EagerBackend::Cuda(_) => (
cuda_runtime_engine_id().map_err(|source| {
runtime_config_error("EagerRuntime::eager_extension_target", source)
})?,
EagerExtensionBackendKind::Cuda,
),
#[cfg(feature = "webgpu")]
EagerBackend::WebGpu(_) => (
tenferro_gpu::webgpu::webgpu_runtime_engine_id().map_err(|source| {
runtime_config_error("EagerRuntime::eager_extension_target", source)
})?,
EagerExtensionBackendKind::WebGpu,
),
}
};
let target = EagerExtensionTarget {
engine_id,
backend_kind,
};
validate_eager_extension_target(&self.runtime, &target)?;
Ok(target)
}
pub fn clear_extension_caches(&self) -> Result<()> {
self.lock_extension_caches()?.clear();
Ok(())
}
pub fn clear_caches(&self) -> Result<()> {
self.clear_extension_caches()?;
self.clear_ad_transform_caches()?;
self.clear_prepared_derivative_cache()?;
Ok(())
}
pub fn clear_prepared_derivative_cache(&self) -> Result<()> {
self.lock_prepared_derivative_cache()?.clear();
Ok(())
}
pub fn cache_stats(&self) -> Result<EagerRuntimeCacheStats> {
Ok(EagerRuntimeCacheStats {
extensions: self
.lock_extension_caches()?
.stats(ExtensionCacheSelector::All),
ad_transforms: self.ad_transform_cache.stats()?,
prepared_derivatives: self.lock_prepared_derivative_cache()?.stats(),
})
}
pub fn ad_transform_cache_limits(&self) -> Result<AdTransformCacheLimits> {
self.ad_transform_cache.limits()
}
pub fn set_ad_transform_cache_limits(&self, limits: AdTransformCacheLimits) -> Result<()> {
self.ad_transform_cache.set_limits(limits)
}
pub fn clear_ad_transform_caches(&self) -> Result<()> {
self.ad_transform_cache.clear()
}
pub fn prepared_derivative_cache_limits(&self) -> Result<AdTransformCacheLimits> {
Ok(self.lock_prepared_derivative_cache()?.limits())
}
pub fn set_prepared_derivative_cache_limits(
&self,
limits: AdTransformCacheLimits,
) -> Result<()> {
self.lock_prepared_derivative_cache()?.set_limits(limits);
Ok(())
}
pub fn extension_cache_limits(&self) -> Result<ExtensionCacheLimits> {
Ok(self.lock_extension_caches()?.limits())
}
pub fn set_extension_cache_limits(&self, limits: ExtensionCacheLimits) -> Result<()> {
self.lock_extension_caches()?.set_limits(limits);
Ok(())
}
pub fn with_execution_session<R: Send>(
&self,
f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
) -> Result<R> {
let mut backend = self.lock_backend()?;
Ok(backend.with_backend_session(f))
}
pub fn with_extension_execution_context<R: Send>(
&self,
f: impl FnOnce(
&mut tenferro_runtime::ExtensionExecutionContext<'_, dyn BackendSession + '_>,
) -> R
+ Send,
) -> Result<R> {
let mut backend = self.lock_backend()?;
let mut extension_cache_guard = self.lock_extension_caches()?;
let extension_caches: &mut ExtensionCacheStore = &mut extension_cache_guard;
Ok(backend.with_backend_session(move |session| {
let mut extension_ctx =
tenferro_runtime::ExtensionExecutionContext::new(session, extension_caches);
f(&mut extension_ctx)
}))
}
pub fn synchronize(&self) -> Result<()> {
self.lock_backend()?.synchronize().map_err(Error::from)
}
fn exec_outputs_with_runtime<R>(
&self,
lock_backend_section: &'static str,
exec_section: &'static str,
op: &StdTensorOp,
execute: impl FnOnce(&mut EagerBackend, Option<&Runtime>) -> Result<R>,
) -> Result<R> {
let mut backend = profile_eager_op_section(lock_backend_section, || self.lock_backend())?;
let runtime = matches!(op, StdTensorOp::Extension(_)).then_some(&self.runtime);
profile_eager_op_section(exec_section, || execute(&mut backend, runtime))
}
pub(crate) fn exec_outputs(&self, op: &StdTensorOp, inputs: &[&Tensor]) -> Result<Vec<Tensor>> {
self.exec_outputs_with_runtime(
"exec_outputs.lock_backend",
"exec_outputs.exec_op",
op,
|backend, runtime| exec_op_on_tensors_with_runtime(op, inputs, backend, runtime),
)
}
pub(crate) fn exec_outputs_read(
&self,
op: &StdTensorOp,
inputs: &[TensorRead<'_>],
) -> Result<Vec<Tensor>> {
self.exec_outputs_with_runtime(
"exec_outputs_read.lock_backend",
"exec_outputs_read.exec_op",
op,
|backend, runtime| exec_op_on_tensor_reads_with_runtime(op, inputs, backend, runtime),
)
}
#[cfg(test)]
pub(crate) fn exec_standard_graph_outputs(
&self,
graph: &Graph<StdTensorOp>,
initial_data: HashMap<ValueKey<StdTensorOp>, Tensor>,
) -> Result<EagerGraphExecution> {
let mut backend =
profile_eager_op_section("exec_graph.lock_backend", || self.lock_backend())?;
let mut all_values = initial_data;
profile_eager_op_section("exec_graph.with_backend_session", || {
backend.with_backend_session(|exec| -> Result<()> {
for op_node in graph.operations() {
let outputs = {
let input_values = op_node
.inputs
.iter()
.map(|input| {
let key = match input {
ValueRef::Local(local_id) => &graph.values()[*local_id].key,
ValueRef::External(key) => key,
};
all_values.get(key).ok_or_else(|| {
Error::Internal(format!(
"standard graph eager execution missing value for {key:?}"
))
})
})
.collect::<Result<Vec<_>>>()?;
let input_reads = input_values
.iter()
.map(|value| TensorRead::from_tensor(value))
.collect::<Vec<_>>();
exec_standard_op_on_tensor_reads_in_session(
&op_node.operation,
&input_reads,
exec,
)?
};
if outputs.len() != op_node.outputs.len() {
return Err(Error::Internal(format!(
"standard graph eager execution expected {} outputs for {:?}, got {}",
op_node.outputs.len(),
op_node.operation,
outputs.len()
)));
}
for (output_id, output) in op_node.outputs.iter().zip(outputs) {
let key = graph.values()[*output_id].key.clone();
all_values.insert(key, output);
}
}
Ok(())
})
})?;
let outputs = graph
.outputs()
.iter()
.map(|&output_id| {
let key = &graph.values()[output_id].key;
all_values
.get(key)
.ok_or_else(|| {
Error::Internal(format!(
"standard graph eager execution missing graph output {key:?}"
))
})?
.duplicate()
.map_err(Error::from)
})
.collect::<Result<Vec<_>>>()?;
Ok(EagerGraphExecution { outputs })
}
pub(crate) fn try_register_grad_slot(
&self,
key: &ValueKey<StdTensorOp>,
slot: &GradSlot,
) -> Result<()> {
self.lock_grad_slots()?
.insert(key.clone(), Arc::downgrade(slot));
Ok(())
}
pub(crate) fn try_register_value_record(
&self,
key: &ValueKey<StdTensorOp>,
record: &Arc<EagerTensorRecord>,
) -> Result<()> {
self.lock_value_records()?
.insert(key.clone(), Arc::downgrade(record));
Ok(())
}
pub(crate) fn value_record(
&self,
key: &ValueKey<StdTensorOp>,
) -> Result<Option<Arc<EagerTensorRecord>>> {
let mut records = self.lock_value_records()?;
let Some(record) = records.get(key).cloned() else {
return Ok(None);
};
match record.upgrade() {
Some(record) => Ok(Some(record)),
None => {
records.remove(key);
Ok(None)
}
}
}
pub fn clear_grads(&self) -> Result<()> {
let live_slots = {
let mut live_slots = Vec::new();
self.lock_grad_slots()?.retain(|_, slot| {
if let Some(slot) = slot.upgrade() {
live_slots.push(slot);
true
} else {
false
}
});
live_slots
};
let mut poisoned_slot = false;
for slot in live_slots {
match slot.lock() {
Ok(mut current) => {
*current = None;
}
Err(_) => {
poisoned_slot = true;
}
}
}
if poisoned_slot {
return Err(Error::runtime_state(
"eager_gradient_slot",
ErrorPhase::Execution,
"lock poisoned",
));
}
Ok(())
}
pub fn constant_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
EagerTensor::new_leaf(Arc::clone(self), tensor, false)
}
pub fn variable_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
EagerTensor::new_leaf(Arc::clone(self), tensor, true)
}
pub fn grad(self: &Arc<Self>, output: &EagerTensor, wrt: &EagerTensor) -> Result<EagerTensor> {
self.grad_optional(output, wrt)?
.ok_or_else(|| Error::Internal(format!("grad output is inactive for {:?}", wrt.key)))
}
pub fn grad_optional(
self: &Arc<Self>,
output: &EagerTensor,
wrt: &EagerTensor,
) -> Result<Option<EagerTensor>> {
if !output.shape().is_empty() {
return Err(Error::NonScalarGrad {
shape: output.shape().to_vec(),
});
}
let value = output.to_tensor()?;
let seed = {
let mut backend = self.lock_backend()?;
one_like_tensor(&value, &mut *backend)?
};
let seed = EagerTensor::new_result(
Arc::clone(self),
eager_val_key(),
seed,
false,
None,
Vec::new(),
)?;
self.vjp_optional(output, wrt, &seed)
}
pub fn vjp(
self: &Arc<Self>,
output: &EagerTensor,
wrt: &EagerTensor,
cotangent: &EagerTensor,
) -> Result<EagerTensor> {
self.vjp_optional(output, wrt, cotangent)?
.ok_or_else(|| Error::Internal(format!("vjp output is inactive for {:?}", wrt.key)))
}
pub fn vjp_optional(
self: &Arc<Self>,
output: &EagerTensor,
wrt: &EagerTensor,
cotangent: &EagerTensor,
) -> Result<Option<EagerTensor>> {
validate_same_runtime(self, output, "vjp output")?;
validate_same_runtime(self, wrt, "vjp wrt")?;
validate_same_runtime(self, cotangent, "vjp cotangent")?;
validate_seed_tensor("vjp", output, cotangent)?;
match semantic_eager_vjp_optional(self, output, wrt, cotangent)? {
Some(result) => Ok(result),
None => Ok(None),
}
}
pub fn jvp(
self: &Arc<Self>,
output: &EagerTensor,
wrt: &EagerTensor,
tangent: &EagerTensor,
) -> Result<EagerTensor> {
self.jvp_optional(output, wrt, tangent)?
.ok_or_else(|| Error::Internal(format!("jvp output is inactive for {:?}", wrt.key)))
}
pub fn jvp_optional(
self: &Arc<Self>,
output: &EagerTensor,
wrt: &EagerTensor,
tangent: &EagerTensor,
) -> Result<Option<EagerTensor>> {
validate_same_runtime(self, output, "jvp output")?;
validate_same_runtime(self, wrt, "jvp wrt")?;
validate_same_runtime(self, tangent, "jvp tangent")?;
validate_seed_tensor("jvp", wrt, tangent)?;
match semantic_eager_jvp_optional(self, output, wrt, tangent)? {
Some(result) => Ok(result),
None => Ok(None),
}
}
fn store_grads(
&self,
cotangents: &HashMap<ValueKey<StdTensorOp>, Tensor>,
backend: &mut EagerBackend,
) -> Result<()> {
let mut updates = Vec::new();
{
let mut slots = self.lock_grad_slots()?;
slots.retain(|key, slot| {
let Some(slot) = slot.upgrade() else {
return false;
};
if let Some(incoming) = cotangents.get(key) {
updates.push((slot, incoming));
}
true
});
}
for (slot, incoming) in updates {
let mut current = slot.lock().map_err(|_| {
Error::runtime_state(
"eager_gradient_slot",
ErrorPhase::Execution,
"lock poisoned",
)
})?;
let next = match current.as_ref() {
Some(existing) => {
let existing_read = existing.tensor_read("EagerRuntime::store_grads")?;
let incoming_read = TensorRead::from_tensor(incoming);
let tensor = backend
.with_backend_session(|session| {
session.add_read(existing_read, incoming_read)
})
.map_err(Error::from)?;
AdValueRecord::from_tensor(tensor, "EagerRuntime::store_grads")?
}
None => {
let duplicate = backend
.with_backend_session(|session| {
session.to_contiguous_read(TensorRead::from_tensor(incoming))
})
.map_err(Error::from)?;
AdValueRecord::from_tensor(duplicate, "EagerRuntime::store_grads")?
}
};
*current = Some(next);
}
Ok(())
}
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct PreparedDerivativeCacheKey {
semantic_fingerprint: SemanticFingerprint,
runtime_epoch: RuntimeEpoch,
wrt_input_index: usize,
input_metadata: Box<[ProgramValueMetadata]>,
}
#[derive(Debug)]
struct PreparedDerivative {
program: Arc<CompiledGraph>,
prepared: Arc<PreparedCompiledGraph>,
seed_input_index: usize,
derivative_output_index: usize,
}
#[derive(Debug)]
struct PreparedDerivativeCache {
limits: AdTransformCacheLimits,
entries: LruCache<PreparedDerivativeCacheKey, PreparedDerivativeCacheEntry>,
stats: CacheStats,
}
impl PreparedDerivativeCache {
fn limits(&self) -> AdTransformCacheLimits {
self.limits
}
fn set_limits(&mut self, limits: AdTransformCacheLimits) {
self.limits = limits;
self.evict_to_limits();
}
fn clear(&mut self) {
let clears = self.stats.clears.saturating_add(1);
self.entries.clear();
self.stats = CacheStats {
clears,
..CacheStats::empty()
};
}
fn stats(&self) -> CacheStats {
self.stats
}
fn get(&mut self, key: &PreparedDerivativeCacheKey) -> Option<Arc<PreparedDerivative>> {
match self.entries.get(key) {
Some(entry) => {
self.stats.hits = self.stats.hits.saturating_add(1);
Some(Arc::clone(&entry.value))
}
None => {
self.stats.misses = self.stats.misses.saturating_add(1);
None
}
}
}
fn insert(&mut self, key: PreparedDerivativeCacheKey, value: Arc<PreparedDerivative>) {
let retained_bytes = prepared_derivative_cache_entry_retained_bytes(&key, value.as_ref());
let entry = PreparedDerivativeCacheEntry {
value,
retained_bytes,
};
self.stats.retained_bytes = self.stats.retained_bytes.saturating_add(retained_bytes);
if let Some((_old_key, old_entry)) = self.entries.push(key, entry) {
self.stats.retained_bytes = self
.stats
.retained_bytes
.saturating_sub(old_entry.retained_bytes);
}
self.stats.entries = self.entries.len();
self.evict_to_limits();
}
fn evict_to_limits(&mut self) {
while self.entries.len() > self.limits.max_entries().get()
|| self
.limits
.max_retained_bytes()
.is_some_and(|limit| self.stats.retained_bytes > limit.get())
{
let Some((_key, entry)) = self.entries.pop_lru() else {
break;
};
self.stats.retained_bytes = self
.stats
.retained_bytes
.saturating_sub(entry.retained_bytes);
self.stats.evictions = self.stats.evictions.saturating_add(1);
}
self.stats.entries = self.entries.len();
}
}
impl Default for PreparedDerivativeCache {
fn default() -> Self {
Self {
limits: AdTransformCacheLimits::default(),
entries: LruCache::unbounded(),
stats: CacheStats::empty(),
}
}
}
#[derive(Debug)]
struct PreparedDerivativeCacheEntry {
value: Arc<PreparedDerivative>,
retained_bytes: usize,
}
fn prepared_derivative_cache_entry_retained_bytes(
key: &PreparedDerivativeCacheKey,
value: &PreparedDerivative,
) -> usize {
size_of::<PreparedDerivativeCacheKey>()
.saturating_add(
key.input_metadata
.len()
.saturating_mul(size_of::<ProgramValueMetadata>()),
)
.saturating_add(size_of::<PreparedDerivative>())
.saturating_add(compiled_graph_retained_bytes(value.program.as_ref()))
.saturating_add(prepared_compiled_graph_retained_bytes(
value.prepared.as_ref(),
value.program.as_ref(),
))
}
fn prepared_compiled_graph_retained_bytes(
prepared: &PreparedCompiledGraph,
derivative_program: &CompiledGraph,
) -> usize {
size_of_val(prepared).saturating_add(compiled_graph_retained_bytes(derivative_program))
}
fn compiled_graph_retained_bytes(program: &CompiledGraph) -> usize {
size_of::<CompiledGraph>()
.saturating_add(size_of_val(program.input_keys()))
.saturating_add(program.bindings().len().saturating_mul(size_of::<usize>()))
.saturating_add(semantic_program_retained_bytes(program.program()))
}
fn semantic_program_retained_bytes(program: &SemanticProgram) -> usize {
size_of::<SemanticProgram>()
.saturating_add(size_of_val(program.inputs()))
.saturating_add(size_of_val(program.outputs()))
.saturating_add(
program
.operations()
.len()
.saturating_mul(size_of::<usize>()),
)
.saturating_add(
program
.shape_guards()
.len()
.saturating_mul(size_of::<usize>()),
)
}
fn semantic_eager_vjp_optional(
ctx: &Arc<EagerRuntime>,
output: &EagerTensor,
wrt: &EagerTensor,
cotangent: &EagerTensor,
) -> Result<Option<Option<EagerTensor>>> {
if !eager_semantic_vjp_enabled() {
return Ok(None);
}
let (Some(output_trace), Some(wrt_trace)) =
(output.semantic_trace.as_ref(), wrt.semantic_trace.as_ref())
else {
return Ok(None);
};
let Some(wrt_key) = wrt_trace.input_key() else {
return Ok(None);
};
if !output_trace.has_attached_input_key(&wrt_key) {
return Ok(None);
}
let mut compiler = GraphCompiler::new();
let source = compile_ad_source(&mut compiler, output_trace)?;
if source.output_count() != 1
|| source.input_keys().len() != source.input_count()
|| source.bindings().len() != source.input_count()
{
return Ok(None);
}
let Some(wrt_input_index) = source.input_key_index(&wrt_key) else {
return Ok(None);
};
let cache_key = PreparedDerivativeCacheKey {
semantic_fingerprint: source.program().semantic_fingerprint(),
runtime_epoch: ctx.runtime.epoch().map_err(|source| {
Error::runtime_state_source("semantic_eager_vjp", ErrorPhase::Execution, source)
})?,
wrt_input_index,
input_metadata: source.frozen_program().input_metadata_with_bound_shapes(),
};
let prepared = { ctx.lock_prepared_derivative_cache()?.get(&cache_key) };
let (seed_input_index, derivative_output_index, derivative_program, prepared_runtime) =
if let Some(prepared) = prepared {
(
prepared.seed_input_index,
prepared.derivative_output_index,
Arc::clone(&prepared.program),
Some(Arc::clone(&prepared.prepared)),
)
} else {
let mut active_inputs = vec![false; source.input_count()];
if let Some(active) = active_inputs.get_mut(wrt_input_index) {
*active = true;
} else {
return Ok(None);
}
let active_outputs = vec![true; source.output_count()];
let ad = AdContext::with_rules_and_transform_cache(
ctx.semantic_extension_rules.clone(),
Arc::clone(&ctx.ad_transform_cache),
);
let derivative = ad
.vjp_program(source.frozen_program(), &active_inputs, &active_outputs)
.map_err(|source| {
Error::runtime_state_source(
"semantic_eager_vjp",
ErrorPhase::GraphBuild,
source,
)
})?;
let seed_input_index = derivative
.derivative_input_indices()
.first()
.copied()
.flatten();
let derivative_output_index = derivative
.derivative_output_indices()
.get(wrt_input_index)
.copied()
.flatten();
let (Some(seed_input_index), Some(derivative_output_index)) =
(seed_input_index, derivative_output_index)
else {
return Ok(Some(None));
};
let program = Arc::new(compiler.compile_frozen_program(derivative.frozen())?);
(seed_input_index, derivative_output_index, program, None)
};
let cotangent_tensor = Arc::new(RetainedValue::from_tensor(cotangent.to_tensor()?));
let input_count = derivative_program.input_count();
let mut owned_inputs: Vec<Option<Tensor>> = (0..input_count).map(|_| None).collect();
for (source_input_index, (_, tensor)) in source.bindings().iter().enumerate() {
let Some(slot) = owned_inputs.get_mut(source_input_index) else {
return Err(Error::Internal(format!(
"semantic eager VJP derivative program has no primal input slot {source_input_index}"
)));
};
*slot = Some(copy_value_for_runtime(ctx, tensor)?);
}
let Some(slot) = owned_inputs.get_mut(seed_input_index) else {
return Err(Error::Internal(format!(
"semantic eager VJP seed input index {seed_input_index} is outside {} inputs",
owned_inputs.len()
)));
};
*slot = Some(copy_value_for_runtime(ctx, cotangent_tensor.as_ref())?);
let input_refs = owned_inputs
.iter()
.enumerate()
.map(|(index, tensor)| {
tensor.as_ref().ok_or_else(|| {
Error::Internal(format!(
"semantic eager VJP derivative input {index} was not populated"
))
})
})
.collect::<Result<Vec<_>>>()?;
let prepared_runtime = if let Some(prepared_runtime) = prepared_runtime {
prepared_runtime
} else {
let prepared_runtime = Arc::new(
ctx.runtime
.prepare_compiled(&derivative_program, &input_refs)?,
);
let entry = Arc::new(PreparedDerivative {
program: Arc::clone(&derivative_program),
prepared: Arc::clone(&prepared_runtime),
seed_input_index,
derivative_output_index,
});
ctx.lock_prepared_derivative_cache()?
.insert(cache_key, entry);
prepared_runtime
};
let outputs = ctx.runtime.run_prepared(&prepared_runtime, &input_refs)?;
let output_count = outputs.len();
let Some(result) = outputs.into_iter().nth(derivative_output_index) else {
return Err(Error::Internal(format!(
"semantic eager VJP derivative output index {derivative_output_index} is outside {} outputs",
output_count
)));
};
let cotangent_trace =
TracedTensor::from_shared_tensor_value_symbolic_shape(Arc::clone(&cotangent_tensor))?;
let semantic_trace = derivative_trace_from_frozen_program(
&source,
derivative_program.frozen_program(),
derivative_output_index,
&[(seed_input_index, Arc::clone(&cotangent_tensor))],
&[output_trace, wrt_trace, &cotangent_trace],
None,
"semantic_eager_vjp",
)?;
#[cfg(test)]
EAGER_SEMANTIC_VJP_EXECUTIONS.fetch_add(1, Ordering::Relaxed);
Ok(Some(Some(EagerTensor::new_result_with_semantic_trace(
Arc::clone(ctx),
eager_val_key(),
result,
true,
None,
Some(semantic_trace),
Vec::new(),
)?)))
}
fn semantic_eager_jvp_optional(
ctx: &Arc<EagerRuntime>,
output: &EagerTensor,
wrt: &EagerTensor,
tangent: &EagerTensor,
) -> Result<Option<Option<EagerTensor>>> {
if !eager_semantic_vjp_enabled() {
return Ok(None);
}
let (Some(output_trace), Some(wrt_trace)) =
(output.semantic_trace.as_ref(), wrt.semantic_trace.as_ref())
else {
return Ok(None);
};
let Some(wrt_key) = wrt_trace.input_key() else {
return Ok(None);
};
if !output_trace.has_attached_input_key(&wrt_key) {
return Ok(None);
}
let mut compiler = GraphCompiler::new();
let source = compile_ad_source(&mut compiler, output_trace)?;
if source.output_count() != 1
|| source.input_keys().len() != source.input_count()
|| source.bindings().len() != source.input_count()
{
return Ok(None);
}
let Some(wrt_input_index) = source.input_key_index(&wrt_key) else {
return Ok(None);
};
let mut active_inputs = vec![false; source.input_count()];
if let Some(active) = active_inputs.get_mut(wrt_input_index) {
*active = true;
} else {
return Ok(None);
}
let ad = AdContext::with_rules_and_transform_cache(
ctx.semantic_extension_rules.clone(),
Arc::clone(&ctx.ad_transform_cache),
);
let derivative = ad
.jvp_program(source.frozen_program(), &active_inputs)
.map_err(|source| {
Error::runtime_state_source("semantic_eager_jvp", ErrorPhase::GraphBuild, source)
})?;
let Some(seed_input_index) = derivative
.derivative_input_indices()
.get(wrt_input_index)
.copied()
.flatten()
else {
return Ok(Some(None));
};
let Some(derivative_output_index) = derivative
.derivative_output_indices()
.first()
.copied()
.flatten()
else {
return Ok(Some(None));
};
let derivative_program = compiler.compile_frozen_program(derivative.frozen())?;
let tangent_tensor = Arc::new(RetainedValue::from_tensor(tangent.to_tensor()?));
let input_count = derivative_program.input_count();
let mut owned_inputs: Vec<Option<Tensor>> = (0..input_count).map(|_| None).collect();
for (source_input_index, (_, tensor)) in source.bindings().iter().enumerate() {
let Some(slot) = owned_inputs.get_mut(source_input_index) else {
return Err(Error::Internal(format!(
"semantic eager JVP derivative program has no primal input slot {source_input_index}"
)));
};
*slot = Some(copy_value_for_runtime(ctx, tensor)?);
}
let Some(slot) = owned_inputs.get_mut(seed_input_index) else {
return Err(Error::Internal(format!(
"semantic eager JVP seed input index {seed_input_index} is outside {} inputs",
owned_inputs.len()
)));
};
*slot = Some(copy_value_for_runtime(ctx, tangent_tensor.as_ref())?);
let input_refs = owned_inputs
.iter()
.enumerate()
.map(|(index, tensor)| {
tensor.as_ref().ok_or_else(|| {
Error::Internal(format!(
"semantic eager JVP derivative input {index} was not populated"
))
})
})
.collect::<Result<Vec<_>>>()?;
let outputs = ctx.runtime.run_compiled(&derivative_program, &input_refs)?;
let output_count = outputs.len();
let Some(result) = outputs.into_iter().nth(derivative_output_index) else {
return Err(Error::Internal(format!(
"semantic eager JVP derivative output index {derivative_output_index} is outside {} outputs",
output_count
)));
};
let tangent_trace =
TracedTensor::from_shared_tensor_value_symbolic_shape(Arc::clone(&tangent_tensor))?;
let semantic_trace = derivative_trace_from_frozen_program(
&source,
derivative.frozen(),
derivative_output_index,
&[(seed_input_index, Arc::clone(&tangent_tensor))],
&[output_trace, wrt_trace, &tangent_trace],
None,
"semantic_eager_jvp",
)?;
Ok(Some(Some(EagerTensor::new_result_with_semantic_trace(
Arc::clone(ctx),
eager_val_key(),
result,
true,
None,
Some(semantic_trace),
Vec::new(),
)?)))
}
fn validate_same_runtime(
runtime: &Arc<EagerRuntime>,
tensor: &EagerTensor,
role: &'static str,
) -> Result<()> {
if tensor.ctx_id() != runtime.id() {
return Err(Error::ContextMismatch {
lhs: runtime.id(),
rhs: tensor.ctx_id(),
});
}
let _ = role;
Ok(())
}
fn copy_value_for_runtime(ctx: &EagerRuntime, value: &RetainedValue) -> Result<Tensor> {
let read = value.tensor_read().map_err(|error| {
Error::runtime_state_source("copy_value_for_runtime", ErrorPhase::Execution, error)
})?;
ctx.with_execution_session(|session| session.to_contiguous_read(read))?
.map_err(Error::from)
}
fn validate_seed_tensor(op: &'static str, primal: &EagerTensor, seed: &EagerTensor) -> Result<()> {
if primal.dtype() != seed.dtype() {
return Err(
tenferro_tensor::Error::dtype_mismatch(op, primal.dtype(), seed.dtype()).into(),
);
}
if primal.shape() != seed.shape() {
return Err(
tenferro_tensor::Error::shape_mismatch(op, primal.shape(), seed.shape()).into(),
);
}
Ok(())
}
#[derive(Clone)]
pub struct EagerTensor {
pub(crate) key: ValueKey<StdTensorOp>,
pub(crate) trace: Option<EagerTrace>,
pub(crate) semantic_trace: Option<TracedTensor>,
pub(crate) requires_grad: bool,
grad_slot: GradSlot,
pub(crate) metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
pub(crate) ctx: Arc<EagerRuntime>,
_record: Arc<EagerTensorRecord>,
}
pub(crate) struct EagerTensorRecord {
value: Arc<AdValueRecord>,
key: ValueKey<StdTensorOp>,
trace: Option<EagerTrace>,
semantic_trace: Option<TracedTensor>,
requires_grad: bool,
grad_slot: GradSlot,
metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
ctx: Arc<EagerRuntime>,
}
struct EagerTensorParts {
ctx: Arc<EagerRuntime>,
key: ValueKey<StdTensorOp>,
requires_grad: bool,
trace: Option<EagerTrace>,
semantic_trace: Option<TracedTensor>,
value: Arc<AdValueRecord>,
metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
register_value: bool,
}
impl fmt::Debug for EagerTensor {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("EagerTensor")
.field("dtype", &self.dtype())
.field("shape", &self.shape())
.field("key", &self.key)
.field("requires_grad", &self.requires_grad)
.field("has_trace", &self.trace.is_some())
.field("has_semantic_trace", &self.semantic_trace.is_some())
.field("ctx_id", &self.ctx_id())
.finish_non_exhaustive()
}
}
impl EagerTensor {
pub fn from_tensor_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
Self::new_leaf(ctx, tensor, false)
}
pub fn from_vec_col_major_in<T: TensorScalar>(
shape: impl IntoShapeVec,
data: Vec<T>,
ctx: Arc<EagerRuntime>,
) -> Result<Self> {
Self::from_tensor_in(Tensor::from_vec_col_major(shape, data)?, ctx)
}
pub fn requires_grad_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
Self::new_leaf(ctx, tensor, true)
}
pub(crate) fn new_leaf(
ctx: Arc<EagerRuntime>,
tensor: Tensor,
requires_grad: bool,
) -> Result<Self> {
let key = eager_val_key();
let semantic_tensor = ctx
.with_execution_session(|session| {
session.to_contiguous_read(TensorRead::from_tensor(&tensor))
})?
.map_err(Error::from)?;
let semantic_value = Arc::new(RetainedValue::from_tensor(semantic_tensor));
let semantic_trace = TracedTensor::from_shared_tensor_value_symbolic_shape(semantic_value)?;
let metadata_scope =
register_scoped_value_metadata(key.clone(), tensor_meta_from_tensor(&tensor)).map_err(
|err| {
Error::runtime_state_source("eager leaf metadata", ErrorPhase::GraphBuild, err)
},
)?;
let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_leaf")?;
Self::from_parts(EagerTensorParts {
ctx,
key,
requires_grad,
trace: None,
semantic_trace: Some(semantic_trace),
value,
metadata_scopes: metadata_scopes_for_scope(metadata_scope),
register_value: true,
})
}
pub(crate) fn new_result(
ctx: Arc<EagerRuntime>,
key: ValueKey<StdTensorOp>,
tensor: Tensor,
requires_grad: bool,
trace: Option<EagerTrace>,
metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
) -> Result<Self> {
Self::new_result_with_semantic_trace(
ctx,
key,
tensor,
requires_grad,
trace,
None,
metadata_scopes,
)
}
pub(crate) fn new_result_with_semantic_trace(
ctx: Arc<EagerRuntime>,
key: ValueKey<StdTensorOp>,
tensor: Tensor,
requires_grad: bool,
trace: Option<EagerTrace>,
semantic_trace: Option<TracedTensor>,
metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
) -> Result<Self> {
let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_result")?;
Self::from_parts(EagerTensorParts {
ctx,
key,
requires_grad,
trace,
semantic_trace,
value,
metadata_scopes,
register_value: true,
})
}
pub(crate) fn new_unregistered_result_with_semantic_trace(
ctx: Arc<EagerRuntime>,
key: ValueKey<StdTensorOp>,
tensor: Tensor,
requires_grad: bool,
trace: Option<EagerTrace>,
semantic_trace: Option<TracedTensor>,
metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
) -> Result<Self> {
let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_unregistered_result")?;
Self::from_parts(EagerTensorParts {
ctx,
key,
requires_grad,
trace,
semantic_trace,
value,
metadata_scopes,
register_value: false,
})
}
pub(crate) fn new_result_value(
ctx: Arc<EagerRuntime>,
key: ValueKey<StdTensorOp>,
value: TensorValue,
requires_grad: bool,
trace: Option<EagerTrace>,
semantic_trace: Option<TracedTensor>,
metadata_scopes: Vec<Arc<GlobalMetadataScope>>,
) -> Result<Self> {
let (group, slot, dtype, shape) = value.try_into_group_parts().map_err(|_| {
Error::runtime_state(
"EagerTensor::new_result_value",
ErrorPhase::Execution,
"a TensorValue could not be transferred into its allocation group",
)
})?;
let value = AdValueRecord::from_group(group, slot, dtype, shape);
Self::from_parts(EagerTensorParts {
ctx,
key,
requires_grad,
trace,
semantic_trace,
value,
metadata_scopes,
register_value: true,
})
}
fn from_parts(parts: EagerTensorParts) -> Result<Self> {
let EagerTensorParts {
ctx,
key,
requires_grad,
trace,
semantic_trace,
value,
metadata_scopes,
register_value,
} = parts;
let grad_slot = Arc::new(Mutex::new(None));
if requires_grad {
ctx.try_register_grad_slot(&key, &grad_slot)?;
}
let record = Arc::new(EagerTensorRecord {
value: Arc::clone(&value),
key: key.clone(),
trace: trace.clone(),
semantic_trace: semantic_trace.clone(),
requires_grad,
grad_slot: Arc::clone(&grad_slot),
metadata_scopes: metadata_scopes.clone(),
ctx: Arc::clone(&ctx),
});
if register_value {
ctx.try_register_value_record(&key, &record)?;
}
Ok(Self {
key,
trace,
semantic_trace,
requires_grad,
grad_slot,
metadata_scopes,
ctx,
_record: record,
})
}
pub(crate) fn new_untracked_result(ctx: Arc<EagerRuntime>, tensor: Tensor) -> Result<Self> {
let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_untracked_result")?;
Ok(Self::new_untracked_value_record(ctx, value, None))
}
pub(crate) fn new_untracked_value_result(
ctx: Arc<EagerRuntime>,
value: TensorValue,
) -> Result<Self> {
Self::new_untracked_value_result_with_semantic_trace(ctx, value, None)
}
pub(crate) fn new_untracked_value_result_with_semantic_trace(
ctx: Arc<EagerRuntime>,
value: TensorValue,
semantic_trace: Option<TracedTensor>,
) -> Result<Self> {
let (group, slot, dtype, shape) = value.try_into_group_parts().map_err(|_| {
Error::runtime_state(
"EagerTensor::new_untracked_value_result",
ErrorPhase::Execution,
"a TensorValue could not be transferred into its allocation group",
)
})?;
let value = AdValueRecord::from_group(group, slot, dtype, shape);
Ok(Self::new_untracked_value_record(ctx, value, semantic_trace))
}
fn new_untracked_value_record(
ctx: Arc<EagerRuntime>,
value: Arc<AdValueRecord>,
semantic_trace: Option<TracedTensor>,
) -> Self {
let key = eager_val_key();
let grad_slot = Arc::new(Mutex::new(None));
let record = Arc::new(EagerTensorRecord {
value,
key: key.clone(),
trace: None,
semantic_trace: semantic_trace.clone(),
requires_grad: false,
grad_slot: Arc::clone(&grad_slot),
metadata_scopes: Vec::new(),
ctx: Arc::clone(&ctx),
});
Self {
key,
trace: None,
semantic_trace,
requires_grad: false,
grad_slot,
metadata_scopes: Vec::new(),
ctx,
_record: record,
}
}
pub(crate) fn from_record(record: Arc<EagerTensorRecord>) -> Self {
Self {
key: record.key.clone(),
trace: record.trace.clone(),
semantic_trace: record.semantic_trace.clone(),
requires_grad: record.requires_grad,
grad_slot: Arc::clone(&record.grad_slot),
metadata_scopes: record.metadata_scopes.clone(),
ctx: Arc::clone(&record.ctx),
_record: record,
}
}
pub fn detach(&self) -> Self {
let semantic_trace = self
.duplicate_value()
.ok()
.and_then(|tensor| TracedTensor::from_tensor_symbolic_shape(tensor).ok());
Self::new_untracked_value_record(
self.ctx.clone(),
Arc::clone(&self._record.value),
semantic_trace,
)
}
pub fn detach_into(&self, ctx: &Arc<EagerRuntime>) -> Result<Self> {
Self::from_tensor_in(self.to_tensor()?, Arc::clone(ctx))
}
pub fn value(&self) -> Result<ValueGuard<'_>> {
self._record.value.value("EagerTensor::value")
}
pub fn duplicate_value(&self) -> Result<Tensor> {
let value = self.value()?;
match value.duplicate_host_tensor() {
Ok(tensor) => Ok(tensor),
Err(_) => {
let read = self
._record
.value
.tensor_read("EagerTensor::duplicate_value")?;
self.ctx
.with_execution_session(|session| session.to_contiguous_read(read))?
.map_err(Error::from)
}
}
}
#[allow(clippy::result_large_err)]
pub fn into_value(self) -> std::result::Result<Tensor, IntoValueError<Self>> {
if Arc::strong_count(&self._record) != 1 {
return Err(IntoValueError::NotUnique(self));
}
let Self { _record, .. } = self;
let record = match Arc::try_unwrap(_record) {
Ok(record) => record,
Err(record) => return Err(IntoValueError::NotUnique(Self::from_record(record))),
};
let EagerTensorRecord {
value,
key,
trace,
semantic_trace,
requires_grad,
grad_slot,
metadata_scopes,
ctx,
} = record;
let value = match Arc::try_unwrap(value) {
Ok(value) => value,
Err(value) => {
let record = Arc::new(EagerTensorRecord {
value,
key,
trace,
semantic_trace,
requires_grad,
grad_slot,
metadata_scopes,
ctx,
});
return Err(IntoValueError::NotUnique(Self::from_record(record)));
}
};
let AdValueRecord {
container,
slot,
dtype,
shape,
} = value;
let container = match Arc::try_unwrap(container) {
Ok(container) => container,
Err(container) => {
let record = Arc::new(EagerTensorRecord {
value: Arc::new(AdValueRecord {
container,
slot,
dtype,
shape,
}),
key,
trace,
semantic_trace,
requires_grad,
grad_slot,
metadata_scopes,
ctx,
});
return Err(IntoValueError::NotUnique(Self::from_record(record)));
}
};
match container.group.into_tensor(slot) {
Ok(tensor) => Ok(tensor),
Err((group, error)) => {
let record = Arc::new(EagerTensorRecord {
value: Arc::new(AdValueRecord {
container: Arc::new(RetentionContainer { group }),
slot,
dtype,
shape,
}),
key,
trace,
semantic_trace,
requires_grad,
grad_slot,
metadata_scopes,
ctx,
});
Err(IntoValueError::Extract {
value: Self::from_record(record),
error,
})
}
}
}
pub fn dtype(&self) -> DType {
self._record.value.dtype()
}
pub fn shape(&self) -> &[usize] {
self._record.value.shape()
}
pub fn tensor_read(&self) -> TensorRead<'_> {
self._record
.value
.tensor_read("EagerTensor::tensor_read")
.expect("validated eager value record")
}
pub fn to_tensor(&self) -> Result<Tensor> {
self.duplicate_value()
}
pub fn grad(&self) -> Result<Option<GradientValue>> {
self.grad_slot
.lock()
.map_err(|_| {
Error::runtime_state(
"eager_gradient_slot",
ErrorPhase::Execution,
"lock poisoned",
)
})
.map(|slot| {
slot.as_ref().map(|record| GradientValue {
record: Arc::clone(record),
ctx: Arc::clone(&self.ctx),
})
})
}
pub fn clear_grad(&self) -> Result<()> {
*self.grad_slot.lock().map_err(|_| {
Error::runtime_state(
"eager_gradient_slot",
ErrorPhase::Execution,
"lock poisoned",
)
})? = None;
Ok(())
}
pub fn tracks_grad(&self) -> bool {
self.requires_grad
}
#[cfg(test)]
fn debug_trace_saved_value_count(&self) -> Option<usize> {
None
}
pub fn ctx_id(&self) -> ContextId {
self.ctx.id()
}
pub fn runtime(&self) -> &Arc<EagerRuntime> {
&self.ctx
}
pub fn same_context(&self, other: &Self) -> bool {
self.ctx_id() == other.ctx_id()
}
#[cfg(test)]
pub(crate) fn standard_graph_op(
inputs: &[&Self],
build_graph: impl FnOnce(&[TensorInputKey]) -> Result<Arc<Graph<StdTensorOp>>>,
) -> Result<Vec<Self>> {
let Some(first) = inputs.first() else {
return Err(Error::Internal(
"standard eager graph op requires at least one input tensor".to_string(),
));
};
let ctx = Arc::clone(&first.ctx);
for tensor in inputs.iter().skip(1) {
if !first.same_context(tensor) {
return Err(Error::ContextMismatch {
lhs: first.ctx_id(),
rhs: tensor.ctx_id(),
});
}
}
let graph_input_keys = (0..inputs.len())
.map(|_| next_input_key())
.collect::<Vec<_>>();
let graph = build_graph(&graph_input_keys)?;
let initial_data = graph_input_keys
.iter()
.zip(inputs.iter())
.map(|(key, tensor)| Ok((ValueKey::Input(key.clone()), tensor.to_tensor()?)))
.collect::<Result<HashMap<_, _>>>()?;
let execution = ctx.exec_standard_graph_outputs(graph.as_ref(), initial_data)?;
if execution.outputs.len() != graph.outputs().len() {
return Err(Error::Internal(format!(
"standard eager graph op expected {} graph outputs, got {}",
graph.outputs().len(),
execution.outputs.len()
)));
}
if !eager_grad_recording_enabled() || !inputs.iter().any(|input| input.requires_grad) {
return execution
.outputs
.into_iter()
.map(|output| {
Self::new_unregistered_result_with_semantic_trace(
Arc::clone(&ctx),
eager_val_key(),
output,
false,
None,
None,
Vec::new(),
)
})
.collect();
}
let recorded = record_eager_graph_outputs(
graph.as_ref(),
&graph_input_keys,
&execution.outputs,
inputs,
)?;
if recorded.traces.len() != execution.outputs.len() {
return Err(Error::Internal(format!(
"standard eager graph op expected {} eager traces, got {}",
execution.outputs.len(),
recorded.traces.len()
)));
}
let mut metadata_scopes = vec![Arc::clone(&recorded.metadata_scope)];
for input in inputs {
for scope in &input.metadata_scopes {
push_metadata_scope(&mut metadata_scopes, Arc::clone(scope));
}
}
recorded
.traces
.into_iter()
.zip(recorded.semantic_traces)
.zip(execution.outputs)
.map(|((trace, semantic_trace), output)| {
Self::new_result_with_semantic_trace(
Arc::clone(&ctx),
trace.key,
output,
trace.requires_grad,
trace.trace,
semantic_trace,
metadata_scopes.clone(),
)
})
.collect()
}
pub fn backward(&self) -> Result<Gradients> {
if !self.shape().is_empty() {
return Err(Error::NonScalarGrad {
shape: self.shape().to_vec(),
});
}
let value = self.to_tensor()?;
let seed = {
let mut backend = self.ctx.lock_backend()?;
one_like_tensor(&value, &mut *backend)?
};
self.backward_from_seed(seed)
}
pub fn backward_with(&self, cotangent: &EagerTensor) -> Result<Gradients> {
if !self.same_context(cotangent) {
return Err(Error::ContextMismatch {
lhs: self.ctx_id(),
rhs: cotangent.ctx_id(),
});
}
validate_seed_tensor("backward", self, cotangent)?;
self.backward_from_seed(cotangent.to_tensor()?)
}
fn backward_from_seed(&self, seed: Tensor) -> Result<Gradients> {
let cotangent = EagerTensor::new_result(
Arc::clone(&self.ctx),
eager_val_key(),
seed,
false,
None,
Vec::new(),
)?;
let candidate_keys = {
let mut slots = self.ctx.lock_grad_slots()?;
let mut keys = Vec::new();
slots.retain(|key, slot| {
if slot.upgrade().is_some() {
keys.push(key.clone());
true
} else {
false
}
});
keys
};
let mut cotangents = HashMap::new();
for key in candidate_keys {
let Some(record) = self.ctx.value_record(&key)? else {
continue;
};
if !record.requires_grad {
continue;
}
let wrt = EagerTensor::from_record(record);
let Some(grad) = self.ctx.vjp_optional(self, &wrt, &cotangent)? else {
continue;
};
let tensor = match grad.into_value() {
Ok(tensor) => tensor,
Err(IntoValueError::NotUnique(handle)) => handle.duplicate_value()?,
Err(IntoValueError::Extract { error, .. }) => {
return Err(Error::runtime_state_source(
"EagerTensor::backward",
ErrorPhase::Execution,
error,
));
}
};
cotangents.insert(key, tensor);
}
let mut backend = self.ctx.lock_backend()?;
self.ctx.store_grads(&cotangents, &mut backend)?;
Gradients::from_tensors(cotangents)
}
}
pub(crate) fn eager_val_key() -> ValueKey<StdTensorOp> {
ValueKey::Input(next_input_key())
}
pub(crate) struct RecordedEagerTrace {
pub(crate) key: ValueKey<StdTensorOp>,
pub(crate) trace: Option<EagerTrace>,
pub(crate) requires_grad: bool,
}
pub(crate) struct RecordedEagerOutputs {
pub(crate) traces: Vec<RecordedEagerTrace>,
pub(crate) semantic_traces: Vec<Option<TracedTensor>>,
pub(crate) metadata_scope: Arc<GlobalMetadataScope>,
}
pub(crate) fn record_eager_outputs(
op: &StdTensorOp,
outputs: &[&Tensor],
inputs: &[&EagerTensor],
) -> Result<RecordedEagerOutputs> {
let semantic_traces = record_semantic_eager_outputs(op, outputs.len(), inputs)?;
let output_metadata = outputs.iter().map(|output| tensor_meta_from_tensor(output));
record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
}
pub(crate) fn record_eager_value_outputs(
op: &StdTensorOp,
outputs: &[&TensorValue],
inputs: &[&EagerTensor],
) -> Result<RecordedEagerOutputs> {
let semantic_traces = record_semantic_eager_outputs(op, outputs.len(), inputs)?;
let output_metadata = outputs.iter().map(|output| tensor_meta_from_value(output));
record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
}
fn record_semantic_eager_outputs(
op: &StdTensorOp,
output_count: usize,
inputs: &[&EagerTensor],
) -> Result<Vec<Option<TracedTensor>>> {
let Some(semantic_inputs) = inputs
.iter()
.map(|input| input.semantic_trace.as_ref())
.collect::<Option<Vec<_>>>()
else {
return Ok(vec![None; output_count]);
};
let semantic_outputs = match op {
StdTensorOp::Extension(ext) => {
tenferro_runtime::extension::apply(Arc::clone(ext), &semantic_inputs)?
}
_ => tenferro_runtime::extension::apply_standard_op(op.clone(), &semantic_inputs)?,
};
if semantic_outputs.len() != output_count {
return Err(Error::Internal(format!(
"semantic eager recording expected {output_count} outputs for {op:?}, got {}",
semantic_outputs.len()
)));
}
Ok(semantic_outputs.into_iter().map(Some).collect())
}
#[cfg(test)]
fn record_eager_graph_outputs(
graph: &Graph<StdTensorOp>,
graph_input_keys: &[TensorInputKey],
outputs: &[Tensor],
inputs: &[&EagerTensor],
) -> Result<RecordedEagerOutputs> {
let semantic_traces = record_semantic_eager_graph_outputs(graph, graph_input_keys, inputs)?;
let output_metadata = outputs.iter().map(tensor_meta_from_tensor);
record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
}
#[cfg(test)]
fn record_semantic_eager_graph_outputs(
graph: &Graph<StdTensorOp>,
graph_input_keys: &[TensorInputKey],
inputs: &[&EagerTensor],
) -> Result<Vec<Option<TracedTensor>>> {
let Some(semantic_inputs) = inputs
.iter()
.map(|input| input.semantic_trace.as_ref())
.collect::<Option<Vec<_>>>()
else {
return Ok(vec![None; graph.outputs().len()]);
};
if graph_input_keys.len() != semantic_inputs.len() {
return Err(Error::Internal(format!(
"semantic graph recording expected {} input keys, got {}",
semantic_inputs.len(),
graph_input_keys.len()
)));
}
let mut values = HashMap::new();
for (key, tensor) in graph_input_keys.iter().zip(semantic_inputs) {
values.insert(ValueKey::Input(key.clone()), tensor.clone());
}
for op_node in graph.operations() {
let input_values = op_node
.inputs
.iter()
.map(|input| {
let key = match input {
ValueRef::Local(local_id) => &graph.values()[*local_id].key,
ValueRef::External(key) => key,
};
values.get(key).cloned().ok_or_else(|| {
Error::Internal(format!(
"semantic graph recording missing value for {key:?}"
))
})
})
.collect::<Result<Vec<_>>>()?;
let input_refs = input_values.iter().collect::<Vec<_>>();
let semantic_outputs = match &op_node.operation {
StdTensorOp::Extension(ext) => {
tenferro_runtime::extension::apply(Arc::clone(ext), &input_refs)?
}
op => tenferro_runtime::extension::apply_standard_op(op.clone(), &input_refs)?,
};
if semantic_outputs.len() != op_node.outputs.len() {
return Err(Error::Internal(format!(
"semantic graph recording expected {} outputs for {:?}, got {}",
op_node.outputs.len(),
op_node.operation,
semantic_outputs.len()
)));
}
for (output_id, output) in op_node.outputs.iter().copied().zip(semantic_outputs) {
values.insert(graph.values()[output_id].key.clone(), output);
}
}
graph
.outputs()
.iter()
.map(|&output_id| {
let key = &graph.values()[output_id].key;
values.get(key).cloned().map(Some).ok_or_else(|| {
Error::Internal(format!(
"semantic graph recording missing output for {key:?}"
))
})
})
.collect()
}
fn record_eager_outputs_from_metadata(
output_metadata: impl IntoIterator<Item = TensorMeta>,
semantic_traces: Vec<Option<TracedTensor>>,
inputs: &[&EagerTensor],
) -> Result<RecordedEagerOutputs> {
let output_metadata = output_metadata.into_iter().collect::<Vec<_>>();
if semantic_traces.len() != output_metadata.len() {
return Err(Error::Internal(format!(
"eager recording expected {} semantic traces, got {}",
output_metadata.len(),
semantic_traces.len()
)));
}
let requires_grad =
eager_grad_recording_enabled() && inputs.iter().any(|input| input.requires_grad);
let mut registrations = Vec::with_capacity(output_metadata.len());
let traces = output_metadata
.into_iter()
.map(|metadata| {
let key = eager_val_key();
registrations.push((key.clone(), metadata));
RecordedEagerTrace {
key,
trace: None,
requires_grad,
}
})
.collect();
Ok(RecordedEagerOutputs {
traces,
semantic_traces,
metadata_scope: Arc::new(register_scoped_metadata_batch(registrations)?),
})
}
fn tensor_meta_from_value(value: &TensorValue) -> TensorMeta {
TensorMeta::exact(
value.dtype(),
value.shape().iter().copied().map(SymDim::from).collect(),
)
}
pub(crate) fn exec_single_output(
op: &StdTensorOp,
inputs: &[&Tensor],
ctx: &EagerRuntime,
) -> Result<Tensor> {
let mut outputs = ctx.exec_outputs(op, inputs)?;
if outputs.len() != 1 {
return Err(Error::Internal(format!(
"expected one eager output for {:?}, got {}",
op,
outputs.len()
)));
}
Ok(profile_eager_op_section(
"exec_single_output.remove_output",
|| outputs.remove(0),
))
}
pub(crate) fn exec_single_output_read(
op: &StdTensorOp,
inputs: &[TensorRead<'_>],
ctx: &EagerRuntime,
) -> Result<Tensor> {
let mut outputs = ctx.exec_outputs_read(op, inputs)?;
if outputs.len() != 1 {
return Err(Error::Internal(format!(
"expected one eager output for {:?}, got {}",
op,
outputs.len()
)));
}
Ok(profile_eager_op_section(
"exec_single_output_read.remove_output",
|| outputs.remove(0),
))
}
#[cfg(test)]
pub(crate) fn zero_like_tensor<B: TensorBackend>(
input: &Tensor,
backend: &mut B,
) -> Result<Tensor> {
let host = match input {
Tensor::F32(tensor) => Tensor::F32(TypedTensor::zeros(tensor.shape().to_vec())?),
Tensor::F64(tensor) => Tensor::F64(TypedTensor::zeros(tensor.shape().to_vec())?),
Tensor::I32(tensor) => Tensor::I32(TypedTensor::zeros(tensor.shape().to_vec())?),
Tensor::I64(tensor) => Tensor::I64(TypedTensor::zeros(tensor.shape().to_vec())?),
Tensor::Bool(tensor) => Tensor::Bool(TypedTensor::from_vec_col_major(
tensor.shape().to_vec(),
vec![false; tensor.n_elements()],
)?),
Tensor::C32(tensor) => Tensor::C32(TypedTensor::zeros(tensor.shape().to_vec())?),
Tensor::C64(tensor) => Tensor::C64(TypedTensor::zeros(tensor.shape().to_vec())?),
};
backend
.upload_host_tensor(TensorRead::from_tensor(&host))
.map_err(Error::from)
}
pub(crate) fn one_like_tensor<B: TensorBackend>(input: &Tensor, backend: &mut B) -> Result<Tensor> {
let host = ones_tensor(input.dtype(), input.shape().to_vec())?;
backend
.upload_host_tensor(TensorRead::from_tensor(&host))
.map_err(Error::from)
}
#[cfg(test)]
mod tests;