diff --git a/Cargo.toml b/Cargo.toml index 50a04b728..a8aa360ad 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,7 @@ members = [ "tenferro-device", "tenferro-tensor", "tenferro-cubecl", + "tenferro-extension-macros", "tenferro-ops", "tenferro-einsum", "tenferro", @@ -35,9 +36,12 @@ criterion = "0.5" serde = "1" serde_json = "1" lru = "0.12" -chainrules-core = { git = "https://github.com/tensor4all/chainrules-rs.git", rev = "cba354d", package = "chainrules" } -chainrules = { git = "https://github.com/tensor4all/chainrules-rs.git", rev = "cba354d" } -tidu = { git = "https://github.com/tensor4all/tidu-rs.git", rev = "97e4ae1" } +proc-macro2 = "1" +quote = "1" +syn = "2" +chainrules-core = { git = "https://github.com/tensor4all/chainrules-rs.git", rev = "da3f2b50ba2e6a0738737e4e4b2ec8f5c8243469", package = "chainrules" } +chainrules = { git = "https://github.com/tensor4all/chainrules-rs.git", rev = "da3f2b50ba2e6a0738737e4e4b2ec8f5c8243469" } +tidu = { git = "https://github.com/tensor4all/tidu-rs.git", rev = "1d4a4c1d1c3923b79a509467ea4af4be2a148aaf" } strided-view = { git = "https://github.com/tensor4all/strided-rs", rev = "ea37986f4d9cb99cfc1f62a1d4aea561cb3a9722" } strided-traits = { git = "https://github.com/tensor4all/strided-rs", rev = "ea37986f4d9cb99cfc1f62a1d4aea561cb3a9722" } strided-perm = { git = "https://github.com/tensor4all/strided-rs", rev = "ea37986f4d9cb99cfc1f62a1d4aea561cb3a9722", features = ["parallel"] } diff --git a/docs/spec/extension-op.md b/docs/spec/extension-op.md index 4b9b7b71e..36bf3ffeb 100644 --- a/docs/spec/extension-op.md +++ b/docs/spec/extension-op.md @@ -188,16 +188,13 @@ pub trait ExtensionOp: std::fmt::Debug + Send + Sync + 'static { inputs: &[&tenferro_tensor::Tensor], ) -> tenferro_tensor::Result>; - // ----- AD: linearize and transpose_rule (Section 10) ----- + // ----- Backwards-compatible inline AD hooks (Section 10) ----- /// Emit the linear (JVP) rule. /// - /// MUST only emit ops in the core `StdTensorOp` vocabulary (the - /// `PrimitiveOp`-implementing set). MUST NOT emit nested - /// `StdTensorOp::Extension` variants (no extension-to-extension - /// lowering in AD). MUST respect `OpMode::Linear { active_mask }` - /// on every emitted op, following the same convention as core - /// primitives (see `tenferro-ops/src/ad/semiring.rs`). + /// Legacy source-compatible inline hook. AD dispatch uses registered + /// `ExtensionAdRule` providers; new extension crates SHOULD register a + /// rule instead of relying on this method. fn linearize( &self, builder: &mut computegraph::fragment::FragmentBuilder, @@ -209,12 +206,9 @@ pub trait ExtensionOp: std::fmt::Debug + Send + Sync + 'static { /// Emit the transpose (VJP) rule. /// - /// MUST only emit ops in the core `StdTensorOp` vocabulary. The - /// returned vector MUST have length `self.n_inputs()`: one entry - /// per primal input (with `None` for inactive tangent slots). The - /// `inputs` slice provides `ValRef` handles that AD rules resolve - /// through `ShapeGuardContext::shape_of` / `dtype_of` / - /// `metadata_of`. + /// Legacy source-compatible inline hook. AD dispatch uses registered + /// `ExtensionAdRule` providers; new extension crates SHOULD register a + /// rule instead of relying on this method. fn transpose_rule( &self, emitter: &mut dyn computegraph::OpEmitter, @@ -292,9 +286,9 @@ op interner, AD rule caching, and structural graph comparison. - An implementer whose `payload_hash` disagrees with `payload_eq` breaks `HashMap`-keyed caches. Symptom: AD caches return wrong cotangents or miss. -- An implementer whose `linearize` emits a nested `Extension` variant - breaks AD closure. Symptom: downstream panic in `todo_linearize` / cache - corruption. +- An implementer whose registered AD rule emits an `Extension` whose family + has no registered AD rule gets `ADRuleError::Unsupported` on the next AD + pass. --- @@ -324,6 +318,20 @@ Example: "tenferro-ext-tropical.fused_dot_general.v1" ``` +Extension crates MAY use the `ExtensionFamilyId` derive macro re-exported by +`tenferro::extension` / `tenferro_ops` to generate this string as an inherent +`FAMILY_ID` constant: + +```rust +use tenferro_ops::ExtensionFamilyId; + +#[derive(ExtensionFamilyId)] +#[tenferro_extension(namespace = "my-crate", name = "fft", version = 1)] +struct FftOp; + +assert_eq!(FftOp::FAMILY_ID, "my-crate.fft.v1"); +``` + ### Uniqueness `family_id` uniqueness is enforced at registration (see Section 9). The @@ -690,53 +698,47 @@ required failure mode. ### Method signatures -See Section 4 for the canonical `linearize` and `transpose_rule` -signatures. They mirror `PrimitiveOp::linearize` and -`PrimitiveOp::transpose_rule` in `ad-contract.md`, with two differences: - -1. They are methods on `ExtensionOp` rather than on `PrimitiveOp`. The - core enum's `PrimitiveOp` impl routes the `Extension` arm to these - methods via a dispatcher in `tenferro-ops/src/ad/mod.rs`: - - ```rust - // Conceptual: - pub fn linearize( - op: &StdTensorOp, - builder: &mut FragmentBuilder, - primal_in: &[GlobalValKey], - primal_out: &[GlobalValKey], - tangent_in: &[Option], - ctx: &mut context::ShapeGuardContext, - ) -> Vec> { - match op { - // ... existing arms ... - StdTensorOp::Extension(ext) => { - ext.linearize(builder, primal_in, primal_out, tangent_in, ctx) - } - _ => todo_linearize(op), - } - } - ``` - -2. They accept `&mut dyn OpEmitter` rather than `&mut impl - OpEmitter` for object safety. The dispatcher boxes the - generic emitter as `&mut dyn` at the call site so trait objects can - work through the vtable. +Extension AD is registered independently from the primal factory through +`register_extension_rule(Arc)`. Rule signatures mirror +`PrimitiveOp::try_linearize` and `PrimitiveOp::try_transpose_rule` and return +`ADRuleResult<_>` so missing rules can propagate without panic: -### AD closure +```rust +pub trait ExtensionAdRule: Debug + Send + Sync + 'static { + fn family_id(&self) -> &'static str; -`linearize` and `transpose_rule` MUST only emit values in the core -`StdTensorOp` vocabulary. They MUST NOT emit `StdTensorOp::Extension` -nodes, either their own or any other extension's. This preserves the -`ad-contract.md` closure rule: every op in a cotangent graph -implements `PrimitiveOp`, because every such op is a core `StdTensorOp` -variant. + fn linearize( + &self, + op: &dyn ExtensionOp, + builder: &mut FragmentBuilder, + primal_in: &[GlobalValKey], + primal_out: &[GlobalValKey], + tangent_in: &[Option], + ctx: &mut ShapeGuardContext, + ) -> ADRuleResult>>; + + fn transpose_rule( + &self, + op: &dyn ExtensionOp, + emitter: &mut dyn OpEmitter, + cotangent_out: &[Option], + inputs: &[ValRef], + mode: &OpMode, + ctx: &mut ShapeGuardContext, + ) -> ADRuleResult>>; +} +``` + +The `op` argument is the concrete extension payload as a trait object. +Rules that need payload parameters should downcast via `op.as_any()`. + +### AD closure -Rationale: nesting extensions in AD would make AD closure contingent -on the registry state at AD-time rather than at graph-build time, -breaking determinism. The restriction to core ops is the AD closure -invariant: every arrow emits values only in the core op vocabulary, and -extension mechanisms lower to this same vocabulary for the backward pass. +`linearize` and `transpose_rule` may emit core `StdTensorOp` values and +`StdTensorOp::Extension` values. Emitted extension families MUST have their +own registered `ExtensionAdRule` before a subsequent AD pass reaches them. +This keeps out-of-tree operations in the same compute graph while preserving +the `PrimitiveOp` closure invariant at the `StdTensorOp` carrier level. ### `ShapeGuardContext` interaction @@ -761,14 +763,9 @@ inside the extension's AD rules. ### Failure signature -- Emitting a nested `Extension` from `linearize` is caught at the next - AD pass; depending on cache state it may surface as a `todo!` panic - in `todo_linearize`, or as a wrong-cotangent on the next gradient - call. Extensions MUST NOT do this. - Dispatcher reaching a `StdTensorOp::Extension` variant for an - `ExtensionOp` whose `linearize` / `transpose_rule` is missing (for - example because the trait was impl'd with `todo!()`) is treated as - a programming error and is permitted to panic — see Section 12. + `family_id` with no registered `ExtensionAdRule` returns + `ADRuleError::Unsupported` with the family ID and rule kind. --- @@ -835,7 +832,7 @@ these error types / behaviours in the listed scenarios. | Backend lacks a capability the extension needs | The extension's `eager_execute` SHOULD return `Error::BackendFailure` with a descriptive message that includes `family_id` and the missing capability name. The core pipeline MUST NOT fall back to a different backend. | | Graph references an unregistered `family_id` at eager-execute time | Return `Error::Unsupported { op: "extension", message: ": not registered" }`. | | Graph references an unregistered `family_id` at compile time | Return `Error::Unsupported` from `compile_std_to_exec`. | -| AD rules (`linearize` / `transpose_rule`) encounter an `Extension` whose implementation is missing the AD method (e.g. `todo!()`) | **Panic** with `family_id` in the panic message. This is a programming error in the extension crate, not a recoverable runtime condition. | +| AD rules (`linearize` / `transpose_rule`) encounter an `Extension` with no registered `ExtensionAdRule` | Return `ADRuleError::Unsupported` with `family_id` and rule kind; traced `grad` / eager `backward` propagate it through `tenferro::Error`. | | Hash collision on `family_id` (second registration attempt) | Registry MUST reject with `RegistrationError::Duplicate`. | | Arity mismatch: `n_inputs()` disagrees with the `primal_in.len()` the dispatcher passed | `Error::InvalidConfig { op: "extension", message: "family_id=: expected N inputs, got M" }`. | | Output shape disagrees with `infer_output_meta` result length | `Error::InvalidConfig` with `family_id` and the mismatched counts. | diff --git a/tenferro-extension-macros/Cargo.toml b/tenferro-extension-macros/Cargo.toml new file mode 100644 index 000000000..1336fb58f --- /dev/null +++ b/tenferro-extension-macros/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "tenferro-extension-macros" +version.workspace = true +edition.workspace = true +license.workspace = true +authors.workspace = true +publish.workspace = true + +[lib] +proc-macro = true + +[dependencies] +proc-macro2.workspace = true +quote.workspace = true +syn = { workspace = true, features = ["derive"] } diff --git a/tenferro-extension-macros/src/lib.rs b/tenferro-extension-macros/src/lib.rs new file mode 100644 index 000000000..88c59b409 --- /dev/null +++ b/tenferro-extension-macros/src/lib.rs @@ -0,0 +1,214 @@ +//! Procedural macros for tenferro extension crates. +//! +//! # Examples +//! +//! ``` +//! use tenferro_extension_macros::ExtensionFamilyId; +//! +//! #[derive(ExtensionFamilyId)] +//! #[tenferro_extension(namespace = "my-crate", name = "fft", version = 1)] +//! struct FftOp; +//! +//! assert_eq!(FftOp::FAMILY_ID, "my-crate.fft.v1"); +//! ``` + +use proc_macro::TokenStream; +use quote::quote; +use syn::parse::{Parse, ParseStream}; +use syn::{parse_macro_input, DeriveInput, Expr, ExprLit, Lit, Token}; + +#[derive(Debug, Default)] +struct ExtensionArgs { + namespace: Option, + name: Option, + version: Option, +} + +impl Parse for ExtensionArgs { + fn parse(input: ParseStream<'_>) -> syn::Result { + let mut args = Self::default(); + while !input.is_empty() { + let key: syn::Ident = input.parse()?; + input.parse::()?; + let value: Expr = input.parse()?; + match key.to_string().as_str() { + "namespace" => args.namespace = Some(expect_string(value, "namespace")?), + "name" => args.name = Some(expect_string(value, "name")?), + "version" => args.version = Some(expect_u64(value, "version")?), + other => { + return Err(syn::Error::new( + key.span(), + format!("unsupported tenferro_extension argument {other:?}"), + )); + } + } + if input.is_empty() { + break; + } + input.parse::()?; + } + Ok(args) + } +} + +/// Derive an inherent `FAMILY_ID` constant for an extension payload type. +/// +/// The required attribute is: +/// `#[tenferro_extension(namespace = "...", version = N)]`. +/// `name = "..."` is optional; when omitted, the Rust type name is converted +/// to snake_case. +#[proc_macro_derive(ExtensionFamilyId, attributes(tenferro_extension))] +pub fn derive_extension_family_id(input: TokenStream) -> TokenStream { + let input = parse_macro_input!(input as DeriveInput); + match expand_extension_family_id(input) { + Ok(tokens) => tokens.into(), + Err(err) => err.to_compile_error().into(), + } +} + +fn expand_extension_family_id(input: DeriveInput) -> syn::Result { + let mut parsed = None; + for attr in &input.attrs { + if attr.path().is_ident("tenferro_extension") { + let args = attr.parse_args::()?; + parsed = Some(args); + } + } + let args = parsed.ok_or_else(|| { + syn::Error::new_spanned( + &input.ident, + "missing #[tenferro_extension(namespace = \"...\", version = N)]", + ) + })?; + let namespace = args.namespace.ok_or_else(|| { + syn::Error::new_spanned(&input.ident, "missing tenferro_extension namespace") + })?; + let version = args.version.ok_or_else(|| { + syn::Error::new_spanned(&input.ident, "missing tenferro_extension version") + })?; + let name = args + .name + .unwrap_or_else(|| to_snake_case(&input.ident.to_string())); + let family_id = format!("{namespace}.{name}.v{version}"); + let ident = input.ident; + + Ok(quote! { + impl #ident { + /// Stable extension family identifier generated by `ExtensionFamilyId`. + pub const FAMILY_ID: &'static str = #family_id; + } + }) +} + +fn expect_string(value: Expr, field: &str) -> syn::Result { + match value { + Expr::Lit(ExprLit { + lit: Lit::Str(value), + .. + }) => Ok(value.value()), + other => Err(syn::Error::new_spanned( + other, + format!("{field} must be a string literal"), + )), + } +} + +fn expect_u64(value: Expr, field: &str) -> syn::Result { + match value { + Expr::Lit(ExprLit { + lit: Lit::Int(value), + .. + }) => value.base10_parse(), + other => Err(syn::Error::new_spanned( + other, + format!("{field} must be an integer literal"), + )), + } +} + +fn to_snake_case(input: &str) -> String { + let mut out = String::new(); + let mut prev_lower_or_digit = false; + for ch in input.chars() { + if ch.is_ascii_uppercase() { + if prev_lower_or_digit { + out.push('_'); + } + out.push(ch.to_ascii_lowercase()); + prev_lower_or_digit = false; + } else { + prev_lower_or_digit = ch.is_ascii_lowercase() || ch.is_ascii_digit(); + out.push(ch); + } + } + out +} + +#[cfg(test)] +mod tests { + use super::{expand_extension_family_id, to_snake_case, ExtensionArgs}; + use syn::DeriveInput; + + #[test] + fn snake_case_type_name() { + assert_eq!(to_snake_case("FftOp"), "fft_op"); + assert_eq!(to_snake_case("ScaleBy2"), "scale_by2"); + } + + #[test] + fn derive_uses_snake_case_type_name_when_name_is_omitted() { + let input: DeriveInput = syn::parse_quote! { + #[tenferro_extension(namespace = "my-crate", version = 2)] + struct FftPlanOp; + }; + + let tokens = expand_extension_family_id(input).expect("derive should expand"); + + assert!(tokens.to_string().contains("\"my-crate.fft_plan_op.v2\"")); + } + + #[test] + fn derive_reports_missing_attribute_fields() { + let missing_attr: DeriveInput = syn::parse_quote! { + struct MissingAttr; + }; + assert!(expand_extension_family_id(missing_attr) + .expect_err("missing attribute should fail") + .to_string() + .contains("missing #[tenferro_extension")); + + let missing_namespace: DeriveInput = syn::parse_quote! { + #[tenferro_extension(version = 1)] + struct MissingNamespace; + }; + assert!(expand_extension_family_id(missing_namespace) + .expect_err("missing namespace should fail") + .to_string() + .contains("missing tenferro_extension namespace")); + + let missing_version: DeriveInput = syn::parse_quote! { + #[tenferro_extension(namespace = "my-crate")] + struct MissingVersion; + }; + assert!(expand_extension_family_id(missing_version) + .expect_err("missing version should fail") + .to_string() + .contains("missing tenferro_extension version")); + } + + #[test] + fn derive_attribute_parser_rejects_unknown_and_wrong_typed_values() { + assert!(syn::parse_str::(r#"unknown = "x""#) + .expect_err("unknown argument should fail") + .to_string() + .contains("unsupported tenferro_extension argument")); + assert!(syn::parse_str::("namespace = 1") + .expect_err("namespace must be string") + .to_string() + .contains("namespace must be a string literal")); + assert!(syn::parse_str::(r#"version = "1""#) + .expect_err("version must be integer") + .to_string() + .contains("version must be an integer literal")); + } +} diff --git a/tenferro-ops/Cargo.toml b/tenferro-ops/Cargo.toml index ed74dbfb2..e7a70db00 100644 --- a/tenferro-ops/Cargo.toml +++ b/tenferro-ops/Cargo.toml @@ -14,6 +14,7 @@ cubecl = ["tenferro-tensor/cubecl"] [dependencies] tenferro-tensor = { path = "../tenferro-tensor", default-features = false } +tenferro-extension-macros = { path = "../tenferro-extension-macros" } computegraph.workspace = true chainrules-core.workspace = true num-complex.workspace = true diff --git a/tenferro-ops/src/ad/mod.rs b/tenferro-ops/src/ad/mod.rs index 459ecab5f..4907451fd 100644 --- a/tenferro-ops/src/ad/mod.rs +++ b/tenferro-ops/src/ad/mod.rs @@ -25,6 +25,9 @@ use computegraph::fragment::FragmentBuilder; use computegraph::types::{GlobalValKey, LocalValId, OpMode, ValRef}; use computegraph::OpEmitter; +use chainrules_core::ADRuleResult; + +use crate::ext_op::{linearize_extension_rule, transpose_extension_rule}; use crate::std_tensor_op::StdTensorOp; fn todo_transpose_rule(op: &StdTensorOp) -> ! { @@ -55,7 +58,22 @@ pub fn linearize( tangent_in: &[Option], ctx: &mut context::ShapeGuardContext, ) -> Vec> { - match op { + match try_linearize(op, builder, primal_in, primal_out, tangent_in, ctx) { + Ok(tangents) => tangents, + Err(err) => panic!("{err}"), + } +} + +/// Fallible forward-mode AD (JVP) for `StdTensorOp`. +pub fn try_linearize( + op: &StdTensorOp, + builder: &mut FragmentBuilder, + primal_in: &[GlobalValKey], + primal_out: &[GlobalValKey], + tangent_in: &[Option], + ctx: &mut context::ShapeGuardContext, +) -> ADRuleResult>> { + let tangents = match op { // Semiring-arithmetic family (Add/Mul/Neg/Conj form a commutative // semiring over the supported scalar dtypes). StdTensorOp::Add => semiring::linearize_add(builder, tangent_in), @@ -225,9 +243,17 @@ pub fn linearize( // Extension substrate. StdTensorOp::Extension(ext) => { - ext.linearize(builder, primal_in, primal_out, tangent_in, ctx) - } - } + return linearize_extension_rule( + ext.as_ref(), + builder, + primal_in, + primal_out, + tangent_in, + ctx, + ); + } + }; + Ok(tangents) } /// Reverse-mode AD (VJP) for `StdTensorOp`: given the primal op, its @@ -244,7 +270,22 @@ pub fn transpose_rule( mode: &OpMode, ctx: &mut context::ShapeGuardContext, ) -> Vec> { - match op { + match try_transpose_rule(op, emitter, cotangent_out, inputs, mode, ctx) { + Ok(cotangents) => cotangents, + Err(err) => panic!("{err}"), + } +} + +/// Fallible reverse-mode AD (VJP) for `StdTensorOp`. +pub fn try_transpose_rule( + op: &StdTensorOp, + emitter: &mut impl OpEmitter, + cotangent_out: &[Option], + inputs: &[ValRef], + mode: &OpMode, + ctx: &mut context::ShapeGuardContext, +) -> ADRuleResult>> { + let cotangents = match op { // Semiring-arithmetic family. StdTensorOp::Add => semiring::transpose_add(cotangent_out), StdTensorOp::Mul => semiring::transpose_mul(emitter, cotangent_out, inputs, mode), @@ -407,11 +448,19 @@ pub fn transpose_rule( // Extension substrate. StdTensorOp::Extension(ext) => { let emitter_dyn: &mut dyn OpEmitter = emitter; - ext.transpose_rule(emitter_dyn, cotangent_out, inputs, mode, ctx) + return transpose_extension_rule( + ext.as_ref(), + emitter_dyn, + cotangent_out, + inputs, + mode, + ctx, + ); } _ => todo_transpose_rule(op), - } + }; + Ok(cotangents) } #[cfg(test)] diff --git a/tenferro-ops/src/ext_op.rs b/tenferro-ops/src/ext_op.rs index c0a6fa358..ba97cb9fb 100644 --- a/tenferro-ops/src/ext_op.rs +++ b/tenferro-ops/src/ext_op.rs @@ -12,9 +12,10 @@ //! type-erased `Arc` carrier can satisfy //! `Clone + Hash + Eq + Send + Sync + 'static` (computegraph's //! `GraphOp` requirements). -//! - AD rules ([`ExtensionOp::linearize`] and [`ExtensionOp::transpose_rule`]) -//! MUST emit only core [`StdTensorOp`] values — never another `Extension` -//! variant — preserving ad-contract.md's closure invariant. +//! - AD rules are registered separately through [`ExtensionAdRule`] and +//! [`register_extension_rule`]. A rule may emit core [`StdTensorOp`] values +//! and registered `Extension` values so out-of-tree operations remain in the +//! same graph. //! - [`ExtensionFactory`] is registered at program start via //! [`register_extension`]; the registry is an //! `OnceLock>>>` @@ -41,6 +42,7 @@ use std::fmt::Debug; use std::hash::{Hash, Hasher}; use std::sync::{Arc, OnceLock, RwLock}; +use chainrules_core::{ADRuleError, ADRuleKind, ADRuleResult}; use computegraph::fragment::FragmentBuilder; use computegraph::types::{GlobalValKey, LocalValId, OpMode, ValRef}; use computegraph::OpEmitter; @@ -62,9 +64,7 @@ use crate::sym_dim::SymDim; /// - shape / dtype inference via [`infer_output_meta`][Self::infer_output_meta]; /// - forward dispatch via [`eager_execute`][Self::eager_execute] (used by both /// the eager and compiled paths); -/// - AD via [`linearize`][Self::linearize] and -/// [`transpose_rule`][Self::transpose_rule], which MUST emit only core -/// [`StdTensorOp`] values. +/// - AD via a separately registered [`ExtensionAdRule`]. /// /// # Downcast convention /// @@ -188,32 +188,77 @@ pub trait ExtensionOp: Debug + Send + Sync + 'static { /// Emit the linear (JVP) rule. /// - /// MUST only emit ops in the core [`StdTensorOp`] vocabulary. MUST NOT - /// emit nested `StdTensorOp::Extension` variants. MUST respect - /// `OpMode::Linear { active_mask }` on every emitted op. + /// This legacy inline hook is retained so existing impl blocks remain + /// source-compatible. AD dispatch uses registered [`ExtensionAdRule`] + /// providers; new extension crates should register a rule instead of + /// relying on this method. fn linearize( &self, + _builder: &mut FragmentBuilder, + _primal_in: &[GlobalValKey], + _primal_out: &[GlobalValKey], + _tangent_in: &[Option], + _ctx: &mut ShapeGuardContext, + ) -> Vec> { + panic!( + "extension family {:?} has no inline linearize rule; register an ExtensionAdRule", + self.family_id() + ) + } + + /// Emit the transpose (VJP) rule. + /// + /// This legacy inline hook is retained so existing impl blocks remain + /// source-compatible. AD dispatch uses registered [`ExtensionAdRule`] + /// providers; new extension crates should register a rule instead of + /// relying on this method. + fn transpose_rule( + &self, + _emitter: &mut dyn OpEmitter, + _cotangent_out: &[Option], + _inputs: &[ValRef], + _mode: &OpMode, + _ctx: &mut ShapeGuardContext, + ) -> Vec> { + panic!( + "extension family {:?} has no inline transpose rule; register an ExtensionAdRule", + self.family_id() + ) + } +} + +/// AD rule provider for an extension family. +/// +/// Rules are registered independently from [`ExtensionFactory`] so an +/// out-of-tree crate can provide a primal operation and AD behavior as separate +/// components. Rule methods receive the concrete [`ExtensionOp`] payload as a +/// trait object; implementations should downcast through [`ExtensionOp::as_any`] +/// when they need payload-specific parameters. +pub trait ExtensionAdRule: Debug + Send + Sync + 'static { + /// The extension family this rule handles. + fn family_id(&self) -> &'static str; + + /// Emit the linear (JVP) rule. + fn linearize( + &self, + op: &dyn ExtensionOp, builder: &mut FragmentBuilder, primal_in: &[GlobalValKey], primal_out: &[GlobalValKey], tangent_in: &[Option], ctx: &mut ShapeGuardContext, - ) -> Vec>; + ) -> ADRuleResult>>; /// Emit the transpose (VJP) rule. - /// - /// MUST only emit ops in the core [`StdTensorOp`] vocabulary. The returned - /// vector MUST have length `self.n_inputs()` (one entry per primal input, - /// with `None` for inactive tangent slots). The `inputs` slice provides - /// [`ValRef`] handles that AD rules resolve through `ShapeGuardContext`. fn transpose_rule( &self, + op: &dyn ExtensionOp, emitter: &mut dyn OpEmitter, cotangent_out: &[Option], inputs: &[ValRef], mode: &OpMode, ctx: &mut ShapeGuardContext, - ) -> Vec>; + ) -> ADRuleResult>>; } /// Factory trait used at registration time. @@ -257,6 +302,9 @@ pub enum ExtensionRegistryError { /// A factory with the same `family_id` was already registered. #[error("family_id {family_id:?} already registered")] Duplicate { family_id: &'static str }, + /// An AD rule with the same `family_id` was already registered. + #[error("AD rule for family_id {family_id:?} already registered")] + DuplicateRule { family_id: &'static str }, /// The `family_id` does not match the namespaced format /// `"..v"`. #[error("family_id {family_id:?} does not match the namespaced format")] @@ -264,12 +312,18 @@ pub enum ExtensionRegistryError { } type FactoryMap = HashMap<&'static str, Arc>; +type RuleMap = HashMap<&'static str, Arc>; fn registry() -> &'static RwLock { static REG: OnceLock> = OnceLock::new(); REG.get_or_init(|| RwLock::new(HashMap::new())) } +fn rule_registry() -> &'static RwLock { + static REG: OnceLock> = OnceLock::new(); + REG.get_or_init(|| RwLock::new(HashMap::new())) +} + fn is_valid_family_id(family_id: &str) -> bool { // Required shape: `..v` with at least one non-empty // `` chunk, at least one non-empty `` chunk (which may itself @@ -341,6 +395,29 @@ pub fn register_extension( Ok(()) } +/// Register a new extension AD rule. +/// +/// The rule's `family_id` uses the same validation as +/// [`register_extension`]. Registering a rule does not require registering a +/// factory first; this lets crates split primal construction and AD support +/// across modules or optional features. +pub fn register_extension_rule( + rule: Arc, +) -> Result<(), ExtensionRegistryError> { + let family_id = rule.family_id(); + if !is_valid_family_id(family_id) { + return Err(ExtensionRegistryError::MalformedFamilyId { family_id }); + } + let mut guard = rule_registry() + .write() + .expect("extension rule registry RwLock poisoned"); + if guard.contains_key(family_id) { + return Err(ExtensionRegistryError::DuplicateRule { family_id }); + } + guard.insert(family_id, rule); + Ok(()) +} + /// Look up a factory by `family_id`. /// /// Returns `None` if no factory is registered for the given identifier. @@ -361,6 +438,51 @@ pub fn lookup_extension_factory(family_id: &str) -> Option Option> { + rule_registry() + .read() + .expect("extension rule registry RwLock poisoned") + .get(family_id) + .cloned() +} + +/// Emit a registered extension linearization rule. +pub fn linearize_extension_rule( + op: &dyn ExtensionOp, + builder: &mut FragmentBuilder, + primal_in: &[GlobalValKey], + primal_out: &[GlobalValKey], + tangent_in: &[Option], + ctx: &mut ShapeGuardContext, +) -> ADRuleResult>> { + match lookup_extension_rule(op.family_id()) { + Some(rule) => rule.linearize(op, builder, primal_in, primal_out, tangent_in, ctx), + None => Err(ADRuleError::unsupported( + op.family_id(), + ADRuleKind::Linearize, + )), + } +} + +/// Emit a registered extension transpose rule. +pub fn transpose_extension_rule( + op: &dyn ExtensionOp, + emitter: &mut dyn OpEmitter, + cotangent_out: &[Option], + inputs: &[ValRef], + mode: &OpMode, + ctx: &mut ShapeGuardContext, +) -> ADRuleResult>> { + match lookup_extension_rule(op.family_id()) { + Some(rule) => rule.transpose_rule(op, emitter, cotangent_out, inputs, mode, ctx), + None => Err(ADRuleError::unsupported( + op.family_id(), + ADRuleKind::Transpose, + )), + } +} + /// Returns `true` when a factory with `family_id` is currently registered. /// /// # Examples @@ -374,6 +496,11 @@ pub fn is_extension_registered(family_id: &str) -> bool { lookup_extension_factory(family_id).is_some() } +/// Returns `true` when an AD rule with `family_id` is currently registered. +pub fn is_extension_rule_registered(family_id: &str) -> bool { + lookup_extension_rule(family_id).is_some() +} + /// Thin adapter that lets a generic `H: Hasher` satisfy the object-safe /// `&mut dyn Hasher` signature required by [`ExtensionOp::payload_hash`]. /// diff --git a/tenferro-ops/src/lib.rs b/tenferro-ops/src/lib.rs index 4cda0fa31..57350f758 100644 --- a/tenferro-ops/src/lib.rs +++ b/tenferro-ops/src/lib.rs @@ -8,11 +8,14 @@ pub mod sym_dim; pub use ad::context::{ShapeGuard, ShapeGuardContext, TensorMeta}; pub use ext_op::{ - is_extension_registered, lookup_extension_factory, register_extension, ExtensionFactory, - ExtensionOp, ExtensionRegistryError, + is_extension_registered, is_extension_rule_registered, linearize_extension_rule, + lookup_extension_factory, lookup_extension_rule, register_extension, register_extension_rule, + transpose_extension_rule, ExtensionAdRule, ExtensionFactory, ExtensionOp, + ExtensionRegistryError, }; pub use shape_extent::ShapeExtent; pub use sym_dim::SymDim; +pub use tenferro_extension_macros::ExtensionFamilyId; pub use tenferro_tensor::config; #[cfg(test)] diff --git a/tenferro-ops/src/std_tensor_op.rs b/tenferro-ops/src/std_tensor_op.rs index 4510ffece..c5eb3b4b4 100644 --- a/tenferro-ops/src/std_tensor_op.rs +++ b/tenferro-ops/src/std_tensor_op.rs @@ -1,7 +1,7 @@ use std::hash::{Hash, Hasher}; use std::sync::Arc; -use chainrules_core::PrimitiveOp; +use chainrules_core::{ADRuleResult, PrimitiveOp}; use computegraph::fragment::FragmentBuilder; use computegraph::types::{GlobalValKey, LocalValId, OpMode, ValRef}; use computegraph::{GraphOp, OpEmitter}; @@ -703,6 +703,17 @@ impl PrimitiveOp for StdTensorOp { crate::ad::linearize(self, builder, primal_in, primal_out, tangent_in, ctx) } + fn try_linearize( + &self, + builder: &mut FragmentBuilder, + primal_in: &[GlobalValKey], + primal_out: &[GlobalValKey], + tangent_in: &[Option], + ctx: &mut Self::ADContext, + ) -> ADRuleResult>> { + crate::ad::try_linearize(self, builder, primal_in, primal_out, tangent_in, ctx) + } + fn transpose_rule( &self, emitter: &mut impl OpEmitter, @@ -713,4 +724,15 @@ impl PrimitiveOp for StdTensorOp { ) -> Vec> { crate::ad::transpose_rule(self, emitter, cotangent_out, inputs, mode, ctx) } + + fn try_transpose_rule( + &self, + emitter: &mut impl OpEmitter, + cotangent_out: &[Option], + inputs: &[ValRef], + mode: &OpMode, + ctx: &mut Self::ADContext, + ) -> ADRuleResult>> { + crate::ad::try_transpose_rule(self, emitter, cotangent_out, inputs, mode, ctx) + } } diff --git a/tenferro-ops/src/tests/ext_op_tests.rs b/tenferro-ops/src/tests/ext_op_tests.rs index cf6193aac..180420817 100644 --- a/tenferro-ops/src/tests/ext_op_tests.rs +++ b/tenferro-ops/src/tests/ext_op_tests.rs @@ -1,11 +1,25 @@ //! Coverage tests for `ext_op` registry + validation. +use std::any::Any; +use std::hash::Hasher; +use std::panic::{catch_unwind, AssertUnwindSafe}; use std::sync::Arc; +use chainrules_core::{ADRuleKind, ADRuleResult}; +use computegraph::fragment::FragmentBuilder; +use computegraph::types::{GlobalValKey, LocalValId, OpMode, ValRef}; +use computegraph::OpEmitter; + +use crate::ad::context::ShapeGuardContext; use crate::ext_op::{ - is_extension_registered, lookup_extension_factory, register_extension, ExtensionFactory, + is_extension_registered, is_extension_rule_registered, linearize_extension_rule, + lookup_extension_factory, lookup_extension_rule, register_extension, register_extension_rule, + transpose_extension_rule, ExtensionAdRule, ExtensionFactory, ExtensionOp, ExtensionRegistryError, }; +use crate::std_tensor_op::StdTensorOp; +use crate::{ExtensionFamilyId, SymDim}; +use tenferro_tensor::{DType, Tensor}; #[derive(Debug)] struct CoverageFamily { @@ -23,6 +37,88 @@ impl ExtensionFactory for CoverageFamily { // exercise the default-impl body. } +#[derive(Debug)] +struct CoverageRule { + family: &'static str, +} + +#[derive(Clone, Debug)] +struct NoInlineRuleOp; + +#[derive(ExtensionFamilyId)] +#[tenferro_extension(namespace = "covtest", name = "macro_rule", version = 1)] +struct MacroRuleFamily; + +impl ExtensionOp for NoInlineRuleOp { + fn family_id(&self) -> &'static str { + "covtest.no_inline_rule.v1" + } + + fn payload_hash(&self, _hasher: &mut dyn Hasher) {} + + fn payload_eq(&self, other: &dyn ExtensionOp) -> bool { + other.as_any().downcast_ref::().is_some() + } + + fn clone_arc(&self) -> Arc { + Arc::new(self.clone()) + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn n_inputs(&self) -> usize { + 1 + } + + fn n_outputs(&self) -> usize { + 1 + } + + fn infer_output_meta( + &self, + input_dtypes: &[DType], + input_shapes: &[&[SymDim]], + ) -> Vec<(DType, Vec)> { + vec![(input_dtypes[0], input_shapes[0].to_vec())] + } + + fn eager_execute(&self, inputs: &[&Tensor]) -> tenferro_tensor::Result> { + Ok(vec![inputs[0].clone()]) + } +} + +impl ExtensionAdRule for CoverageRule { + fn family_id(&self) -> &'static str { + self.family + } + + fn linearize( + &self, + _op: &dyn ExtensionOp, + _builder: &mut FragmentBuilder, + _primal_in: &[GlobalValKey], + _primal_out: &[GlobalValKey], + tangent_in: &[Option], + _ctx: &mut ShapeGuardContext, + ) -> ADRuleResult>> { + Ok(vec![tangent_in[0]]) + } + + fn transpose_rule( + &self, + _op: &dyn ExtensionOp, + _emitter: &mut dyn OpEmitter, + cotangent_out: &[Option], + _inputs: &[ValRef], + _mode: &OpMode, + _ctx: &mut ShapeGuardContext, + ) -> ADRuleResult>> { + Ok(vec![cotangent_out[0]]) + } +} + #[test] fn default_instantiate_returns_none() { let factory: Arc = Arc::new(CoverageFamily { @@ -31,6 +127,89 @@ fn default_instantiate_returns_none() { assert!(factory.instantiate_default().is_none()); } +#[test] +fn extension_family_id_macro_generates_stable_const() { + assert_eq!(MacroRuleFamily::FAMILY_ID, "covtest.macro_rule.v1"); +} + +#[test] +fn register_and_lookup_rule_roundtrips() { + let family = "covtest.register_rule.v1"; + let rule: Arc = Arc::new(CoverageRule { family }); + register_extension_rule(rule).expect("first rule registration should succeed"); + + assert!(is_extension_rule_registered(family)); + let looked_up = lookup_extension_rule(family).expect("rule should be registered"); + assert_eq!(looked_up.family_id(), family); +} + +#[test] +fn register_rule_rejects_duplicate_family_id() { + let family = "covtest.duplicate_rule.v1"; + let first: Arc = Arc::new(CoverageRule { family }); + register_extension_rule(first).expect("first rule registration should succeed"); + + let second: Arc = Arc::new(CoverageRule { family }); + match register_extension_rule(second) { + Err(ExtensionRegistryError::DuplicateRule { family_id }) => { + assert_eq!(family_id, family); + } + other => panic!("expected DuplicateRule for {family:?}, got {other:?}"), + } +} + +#[test] +fn register_rule_rejects_malformed_family_id() { + let bad = "covtest.bad_rule"; + let rule: Arc = Arc::new(CoverageRule { family: bad }); + + match register_extension_rule(rule) { + Err(ExtensionRegistryError::MalformedFamilyId { family_id }) => { + assert_eq!(family_id, bad); + } + other => panic!("expected MalformedFamilyId for {bad:?}, got {other:?}"), + } +} + +#[test] +fn default_inline_rules_panic_with_registration_guidance() { + let op = NoInlineRuleOp; + + let mut builder = FragmentBuilder::::new(); + let mut ctx = ShapeGuardContext::default(); + let linearize_panic = catch_unwind(AssertUnwindSafe(|| { + let _ = op.linearize(&mut builder, &[], &[], &[], &mut ctx); + })); + assert!(linearize_panic.is_err()); + + let mut emitter = FragmentBuilder::::new(); + let mut ctx = ShapeGuardContext::default(); + let transpose_panic = catch_unwind(AssertUnwindSafe(|| { + let _ = op.transpose_rule(&mut emitter, &[], &[], &OpMode::Primal, &mut ctx); + })); + assert!(transpose_panic.is_err()); +} + +#[test] +fn missing_registered_rule_helpers_return_ad_rule_errors() { + let op = NoInlineRuleOp; + + let mut builder = FragmentBuilder::::new(); + let mut ctx = ShapeGuardContext::default(); + let linearize_err = + linearize_extension_rule(&op, &mut builder, &[], &[], &[], &mut ctx).unwrap_err(); + assert_eq!(linearize_err.rule(), ADRuleKind::Linearize); + assert!(linearize_err.to_string().contains(op.family_id())); + + let mut emitter = FragmentBuilder::::new(); + let mut ctx = ShapeGuardContext::default(); + let transpose_err = + transpose_extension_rule(&op, &mut emitter, &[], &[], &OpMode::Primal, &mut ctx) + .unwrap_err(); + assert_eq!(transpose_err.rule(), ADRuleKind::Transpose); + assert!(transpose_err.to_string().contains(op.family_id())); +} + #[test] fn register_rejects_malformed_family_ids() { // Each case targets a different reject branch in `is_valid_family_id`. diff --git a/tenferro-ops/src/tests/std_tensor_op_tests.rs b/tenferro-ops/src/tests/std_tensor_op_tests.rs index 899f3daa5..161a002ba 100644 --- a/tenferro-ops/src/tests/std_tensor_op_tests.rs +++ b/tenferro-ops/src/tests/std_tensor_op_tests.rs @@ -1,12 +1,16 @@ use crate::ad::context::ShapeGuardContext; use crate::dim_expr::DimExpr; +use crate::ext_op::{register_extension_rule, ExtensionAdRule, ExtensionOp}; use crate::std_tensor_op::StdTensorOp; use crate::{SymDim, TensorMeta}; -use chainrules_core::PrimitiveOp; +use chainrules_core::{ADRuleKind, ADRuleResult, PrimitiveOp}; use computegraph::fragment::{Fragment, FragmentBuilder}; use computegraph::types::{GlobalValKey, LocalValId, OpMode, ValRef}; -use computegraph::GraphOp; +use computegraph::{GraphOp, OpEmitter}; use num_complex::{Complex32, Complex64}; +use std::any::Any; +use std::hash::Hasher; +use std::sync::Arc; use tenferro_tensor::{ CompareDir, DType, DotGeneralConfig, GatherConfig, PadConfig, ScatterConfig, }; @@ -1345,3 +1349,140 @@ fn test_std_tensor_op_transpose_rule_panics_for_unimplemented_variant() { &mut ad_ctx, ); } + +#[derive(Clone, Debug)] +struct RuleOnlyExt { + family: &'static str, +} + +impl ExtensionOp for RuleOnlyExt { + fn family_id(&self) -> &'static str { + self.family + } + + fn payload_hash(&self, _hasher: &mut dyn Hasher) {} + + fn payload_eq(&self, other: &dyn ExtensionOp) -> bool { + other + .as_any() + .downcast_ref::() + .is_some_and(|rhs| rhs.family == self.family) + } + + fn clone_arc(&self) -> Arc { + Arc::new(self.clone()) + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn n_inputs(&self) -> usize { + 1 + } + + fn n_outputs(&self) -> usize { + 1 + } + + fn infer_output_meta( + &self, + input_dtypes: &[DType], + input_shapes: &[&[SymDim]], + ) -> Vec<(DType, Vec)> { + vec![(input_dtypes[0], input_shapes[0].to_vec())] + } + + fn eager_execute( + &self, + inputs: &[&tenferro_tensor::Tensor], + ) -> tenferro_tensor::Result> { + Ok(vec![inputs[0].clone()]) + } +} + +#[derive(Debug)] +struct RuleOnlyIdentityAd { + family: &'static str, +} + +impl ExtensionAdRule for RuleOnlyIdentityAd { + fn family_id(&self) -> &'static str { + self.family + } + + fn linearize( + &self, + _op: &dyn ExtensionOp, + _builder: &mut FragmentBuilder, + _primal_in: &[GlobalValKey], + _primal_out: &[GlobalValKey], + tangent_in: &[Option], + _ctx: &mut ShapeGuardContext, + ) -> ADRuleResult>> { + Ok(vec![tangent_in[0]]) + } + + fn transpose_rule( + &self, + _op: &dyn ExtensionOp, + _emitter: &mut dyn OpEmitter, + cotangent_out: &[Option], + _inputs: &[ValRef], + _mode: &OpMode, + _ctx: &mut ShapeGuardContext, + ) -> ADRuleResult>> { + Ok(vec![cotangent_out[0]]) + } +} + +#[test] +fn extension_try_linearize_uses_registered_rule() { + let family = "stdtensor.rule_only_identity.v1"; + let _ = register_extension_rule(Arc::new(RuleOnlyIdentityAd { family })); + let op = StdTensorOp::Extension(Arc::new(RuleOnlyExt { family })); + let mut builder = FragmentBuilder::::new(); + let mut ad_ctx = ShapeGuardContext::default(); + let dx = builder.add_input(tensor_input_key(900)); + let result = op + .try_linearize(&mut builder, &[], &[], &[Some(dx)], &mut ad_ctx) + .expect("registered extension rule should linearize"); + + assert_eq!(result, vec![Some(dx)]); +} + +#[test] +fn extension_try_transpose_uses_registered_rule() { + let family = "stdtensor.rule_only_transpose.v1"; + let _ = register_extension_rule(Arc::new(RuleOnlyIdentityAd { family })); + let op = StdTensorOp::Extension(Arc::new(RuleOnlyExt { family })); + let mut builder = FragmentBuilder::::new(); + let mut ad_ctx = ShapeGuardContext::default(); + let ct = builder.add_input(tensor_input_key(901)); + let result = op + .try_transpose_rule( + &mut builder, + &[Some(ct)], + &external_inputs(910, 1), + &linear_mode(&[true]), + &mut ad_ctx, + ) + .expect("registered extension rule should transpose"); + + assert_eq!(result, vec![Some(ct)]); +} + +#[test] +fn extension_try_linearize_reports_missing_rule() { + let family = "stdtensor.missing_rule.v1"; + let op = StdTensorOp::Extension(Arc::new(RuleOnlyExt { family })); + let mut builder = FragmentBuilder::::new(); + let mut ad_ctx = ShapeGuardContext::default(); + let dx = builder.add_input(tensor_input_key(920)); + let err = op + .try_linearize(&mut builder, &[], &[], &[Some(dx)], &mut ad_ctx) + .expect_err("missing extension rule should be an AD error"); + + assert_eq!(err.rule(), ADRuleKind::Linearize); + assert!(err.to_string().contains(family)); +} diff --git a/tenferro/src/eager.rs b/tenferro/src/eager.rs index 9ba50686e..15c1706b8 100644 --- a/tenferro/src/eager.rs +++ b/tenferro/src/eager.rs @@ -9,7 +9,7 @@ use tenferro_ops::ShapeGuardContext; use tenferro_tensor::cpu::CpuBackend; use tenferro_tensor::{Tensor, TensorBackend}; use tidu::{ - backward_dag, topo_sort_grad_dag, BackwardCallbacks, EagerOutput, EagerValue, GradNode, + topo_sort_grad_dag, try_backward_dag, BackwardCallbacks, EagerOutput, EagerValue, GradNode, LinearFragment, }; @@ -532,7 +532,7 @@ impl EagerTensor { backend: &mut *backend, }; let mut ad_ctx = ShapeGuardContext::with_global_metadata(); - let cotangents = backward_dag(&sorted, &self.key, seed, &mut callbacks, &mut ad_ctx); + let cotangents = try_backward_dag(&sorted, &self.key, seed, &mut callbacks, &mut ad_ctx)?; self.ctx.store_grads(&cotangents, &mut *backend)?; Ok(cotangents) } @@ -641,6 +641,35 @@ impl BackwardCallbacks for TenferroBackwardCallba .collect() } + fn try_eager_transpose( + &mut self, + linear: &LinearFragment, + cotangent_out: &[Option>], + external_data: &HashMap, Arc>, + ctx: &mut ShapeGuardContext, + ) -> chainrules_core::ADRuleResult>>> { + let mut emitter = EagerEmitter::new(self.backend); + emitter.external_data = external_data.clone(); + let cotangent_seed_ids = cotangent_out + .iter() + .map(|maybe_seed| { + maybe_seed + .as_ref() + .map(|seed| emitter.push_tensor(Arc::clone(seed))) + }) + .collect::>(); + + ctx.refresh_global_metadata(); + tidu::try_eager_transpose_fragment(linear, &mut emitter, &cotangent_seed_ids, ctx).map( + |cotangent_ids| { + cotangent_ids + .into_iter() + .map(|maybe_id| maybe_id.map(|id| emitter.tensor(id))) + .collect() + }, + ) + } + fn add_operands(&mut self, a: &Arc, b: &Arc) -> Arc { Arc::new( a.as_ref() diff --git a/tenferro/src/error.rs b/tenferro/src/error.rs index 8162146d6..3c352ae4d 100644 --- a/tenferro/src/error.rs +++ b/tenferro/src/error.rs @@ -42,6 +42,10 @@ pub enum Error { #[error(transparent)] TensorRuntime(#[from] tenferro_tensor::Error), + /// Automatic differentiation rule emission failed. + #[error(transparent)] + ADRule(#[from] chainrules_core::ADRuleError), + /// A `TracedTensor` passed to `eval_with_inputs` bindings is not a /// placeholder (has attached data). #[error( diff --git a/tenferro/src/extension.rs b/tenferro/src/extension.rs index 515a42a29..ed3e38856 100644 --- a/tenferro/src/extension.rs +++ b/tenferro/src/extension.rs @@ -3,8 +3,9 @@ //! This module exposes the Stage 6 `ExtensionOp` mechanism through the //! `tenferro` facade. External crates implement //! [`tenferro_ops::ext_op::ExtensionOp`], register an -//! [`ExtensionFactory`] through [`register_extension`], and build traced -//! graphs containing the extension via [`apply`]. +//! [`ExtensionFactory`] through [`register_extension`], register AD rules +//! through [`register_extension_rule`], and build traced or eager graphs +//! containing the extension via [`apply`] / [`apply_eager`]. //! //! See `docs/spec/extension-op.md` for the normative contract. //! @@ -29,23 +30,32 @@ use std::sync::Arc; use computegraph::fragment::FragmentBuilder; use computegraph::types::{OpMode, ValRef}; +use computegraph::GraphOp; use tenferro_ops::ext_op::ExtensionOp; use tenferro_ops::std_tensor_op::StdTensorOp; use tenferro_ops::SymDim; +use tenferro_tensor::{Tensor, TensorBackend}; use crate::checkpoint::CheckpointNode; +use crate::eager::{record_eager_outputs, EagerTensor}; +use crate::eager_exec::exec_op_on_tensors; +use crate::error::{Error, Result}; use crate::metadata::register_fragment_metadata; use crate::traced::{next_traced_id, TracedTensor}; pub use tenferro_ops::ext_op::{ - is_extension_registered, lookup_extension_factory, register_extension, ExtensionFactory, + is_extension_registered, is_extension_rule_registered, lookup_extension_factory, + lookup_extension_rule, register_extension, register_extension_rule, + ExtensionAdRule as _ExtensionAdRuleReexport, ExtensionFactory, ExtensionOp as _ExtensionOpReexport, ExtensionRegistryError, }; // Re-export under a canonical name (the `_ExtensionOpReexport` alias above // exists only so the macro-generated doc-test type bounds can find the // trait; downstream callers should use this name). +pub use tenferro_ops::ext_op::ExtensionAdRule as ExtensionAdRuleTrait; pub use tenferro_ops::ext_op::ExtensionOp as ExtensionOpTrait; +pub use tenferro_ops::ExtensionFamilyId; /// Apply an extension op in the traced graph. /// @@ -173,3 +183,78 @@ pub fn apply(op: Arc, inputs: &[&TracedTensor]) -> Vec( + op: Arc, + inputs: &[&EagerTensor], +) -> Result>> { + let Some(first) = inputs.first() else { + return Err(Error::Internal( + "extension::apply_eager requires at least one input tensor".to_string(), + )); + }; + if inputs.len() != op.n_inputs() { + return Err(Error::Internal(format!( + "extension::apply_eager: op family {:?} expects {} inputs, got {}", + op.family_id(), + op.n_inputs(), + inputs.len() + ))); + } + + 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 op = StdTensorOp::Extension(op); + let concrete_inputs: Vec<&Tensor> = inputs.iter().map(|tensor| tensor.data.as_ref()).collect(); + let outputs = { + let mut backend = ctx.backend.lock().unwrap(); + exec_op_on_tensors(&op, &concrete_inputs, &mut *backend)? + }; + if outputs.len() != op.n_outputs() { + return Err(Error::Internal(format!( + "expected {} eager outputs for {:?}, got {}", + op.n_outputs(), + op, + outputs.len() + ))); + } + + let outputs: Vec> = outputs.into_iter().map(Arc::new).collect(); + let traces = record_eager_outputs(&op, &outputs, inputs); + if traces.len() != outputs.len() { + return Err(Error::Internal(format!( + "expected {} eager traces for {:?}, got {}", + outputs.len(), + op, + traces.len() + ))); + } + + Ok(traces + .into_iter() + .zip(outputs) + .map(|(trace, output)| { + EagerTensor::new_result( + Arc::clone(&ctx), + trace.key, + output.as_ref().clone(), + trace.requires_grad, + trace.node, + ) + }) + .collect()) +} diff --git a/tenferro/src/traced.rs b/tenferro/src/traced.rs index c894ab833..352c5faf6 100644 --- a/tenferro/src/traced.rs +++ b/tenferro/src/traced.rs @@ -14,7 +14,7 @@ use tenferro_ops::input_key::TensorInputKey; use tenferro_ops::std_tensor_op::StdTensorOp; use tenferro_ops::ShapeGuardContext; use tenferro_tensor::{DType, DotGeneralConfig, Tensor, TensorBackend, TensorScalar, TypedTensor}; -use tidu::{differentiate, transpose}; +use tidu::{try_differentiate, try_transpose}; use super::compiler::compile_std_to_exec; use super::engine::Engine; @@ -677,7 +677,12 @@ impl TracedTensor { let ones = ones_tensor(self.dtype, vec![]); let seed = TracedTensor::from_tensor_concrete_shape(ones); - Ok(self.vjp(wrt, &seed)) + self.try_vjp_result(wrt, &seed)?.ok_or_else(|| { + Error::Internal(format!( + "grad output is inactive for {:?}", + leaf_input_key(wrt) + )) + }) } /// Like [`grad`](Self::grad) but returns `None` when the scalar output does @@ -697,7 +702,7 @@ impl TracedTensor { let ones = ones_tensor(self.dtype, vec![]); let seed = TracedTensor::from_tensor_concrete_shape(ones); - Ok(self.try_vjp(wrt, &seed)) + self.try_vjp_result(wrt, &seed) } /// Evaluate this tensor and replace its graph with a concrete leaf. @@ -775,6 +780,19 @@ impl TracedTensor { /// Like [`jvp`](Self::jvp) but returns `None` when the output does not /// depend on `wrt` (i.e. the tangent is structurally zero). pub fn try_jvp(&self, wrt: &TracedTensor, tangent: &TracedTensor) -> Option { + self.try_jvp_result(wrt, tangent) + .unwrap_or_else(|err| panic!("{err}")) + } + + /// Fallible variant of [`try_jvp`](Self::try_jvp). + /// + /// This returns an error when a primitive or extension cannot emit its + /// linearization rule. + pub fn try_jvp_result( + &self, + wrt: &TracedTensor, + tangent: &TracedTensor, + ) -> Result> { let wrt_input_key = leaf_input_key(wrt); let output_key = self.fragment.vals()[self.val].key.clone(); let aliases = self @@ -791,15 +809,17 @@ impl TracedTensor { roots.extend(checkpoint_fragments.iter().cloned()); let view = resolve(roots); let mut ad_ctx = ShapeGuardContext::with_global_metadata(); - let linear = differentiate( + let linear = try_differentiate( &view, std::slice::from_ref(&output_key), std::slice::from_ref(&wrt_input_key), next_pass_id(), &mut ad_ctx, &aliases, - ); - let tangent_output = linear.tangent_outputs[0]?; + )?; + let Some(tangent_output) = linear.tangent_outputs[0] else { + return Ok(None); + }; let tangent_input_key = linear_input_key(&linear.fragment, linear.tangent_inputs[0].1); register_fragment_metadata( &linear.fragment, @@ -831,7 +851,7 @@ impl TracedTensor { extra_roots.extend(checkpoint_fragments); extra_roots.extend(self.extra_roots.iter().cloned()); - Some(TracedTensor { + Ok(Some(TracedTensor { id: next_traced_id(), rank: self.rank, dtype: self.dtype, @@ -842,15 +862,27 @@ impl TracedTensor { inputs_map: Arc::new(inputs_map), extra_roots, checkpoint_chain: self.checkpoint_chain.clone(), - }) + })) } pub fn vjp(&self, wrt: &TracedTensor, cotangent: &TracedTensor) -> TracedTensor { - self.try_vjp(wrt, cotangent) - .unwrap_or_else(|| panic!("vjp output is inactive for {:?}", leaf_input_key(wrt))) + match self.try_vjp_result(wrt, cotangent) { + Ok(Some(vjp)) => vjp, + Ok(None) => panic!("vjp output is inactive for {:?}", leaf_input_key(wrt)), + Err(err) => panic!("{err}"), + } } - fn try_vjp(&self, wrt: &TracedTensor, cotangent: &TracedTensor) -> Option { + /// Fallible reverse-mode product helper. + /// + /// This returns `Ok(None)` when the cotangent for `wrt` is structurally + /// inactive, and returns an error when a primitive or extension is missing + /// the required AD rule. + pub fn try_vjp_result( + &self, + wrt: &TracedTensor, + cotangent: &TracedTensor, + ) -> Result> { let wrt_input_key = leaf_input_key(wrt); let output_key = self.fragment.vals()[self.val].key.clone(); let aliases = self @@ -867,15 +899,17 @@ impl TracedTensor { roots.extend(checkpoint_fragments.iter().cloned()); let view = resolve(roots); let mut ad_ctx = ShapeGuardContext::with_global_metadata(); - let linear = differentiate( + let linear = try_differentiate( &view, std::slice::from_ref(&output_key), std::slice::from_ref(&wrt_input_key), next_pass_id(), &mut ad_ctx, &aliases, - ); - linear.tangent_outputs[0]?; + )?; + if linear.tangent_outputs[0].is_none() { + return Ok(None); + } let linear_seed_key = linear_input_key(&linear.fragment, linear.tangent_inputs[0].1); register_fragment_metadata( &linear.fragment, @@ -890,7 +924,7 @@ impl TracedTensor { .iter() .map(|(_, local_id)| *local_id) .collect(); - let transposed = transpose(&linear, &mut ad_ctx); + let transposed = try_transpose(&linear, &mut ad_ctx)?; let cotangent_input_key = linear_input_key(&transposed.fragment, transposed.tangent_inputs[0].1); register_fragment_metadata( @@ -907,7 +941,9 @@ impl TracedTensor { )], ); let linear_fragment = Arc::new(linear.fragment); - let cotangent_output = transposed.tangent_outputs[0]?; + let Some(cotangent_output) = transposed.tangent_outputs[0] else { + return Ok(None); + }; let mut inputs_map = (*self.inputs_map).clone(); if let Some(chain) = &self.checkpoint_chain { @@ -950,7 +986,7 @@ impl TracedTensor { extra_roots.extend(checkpoint_fragments); extra_roots.extend(self.extra_roots.iter().cloned()); - Some(TracedTensor { + Ok(Some(TracedTensor { id: next_traced_id(), rank: wrt.rank, dtype: wrt.dtype, @@ -961,7 +997,7 @@ impl TracedTensor { inputs_map: Arc::new(inputs_map), extra_roots, checkpoint_chain: self.checkpoint_chain.clone(), - }) + })) } /// Elementwise addition with NumPy-style broadcasting. diff --git a/tenferro/tests/extension_op.rs b/tenferro/tests/extension_op.rs index 8031e88e7..732ae2bee 100644 --- a/tenferro/tests/extension_op.rs +++ b/tenferro/tests/extension_op.rs @@ -15,13 +15,18 @@ use std::any::Any; use std::hash::Hasher; +use std::panic::{catch_unwind, AssertUnwindSafe}; use std::sync::{Arc, Mutex, OnceLock}; +use chainrules_core::ADRuleResult; use computegraph::fragment::FragmentBuilder; use computegraph::types::{GlobalValKey, LocalValId, OpMode, ValRef}; use computegraph::OpEmitter; -use tenferro::extension::{apply, register_extension, ExtensionFactory}; -use tenferro::{CpuBackend, Engine, Tensor, TracedTensor}; +use tenferro::extension::{ + apply, apply_eager, register_extension, register_extension_rule, ExtensionAdRuleTrait, + ExtensionFactory, +}; +use tenferro::{CpuBackend, EagerContext, EagerTensor, Engine, Tensor, TracedTensor}; use tenferro_ops::ext_op::ExtensionOp; use tenferro_ops::std_tensor_op::StdTensorOp; use tenferro_ops::{ShapeGuardContext, SymDim}; @@ -50,6 +55,23 @@ fn register_once(factory: Arc) { ids.push(family_id); } +fn register_rule_once(rule: Arc) { + static REGISTERED: OnceLock>> = OnceLock::new(); + let guard = REGISTERED.get_or_init(|| Mutex::new(Vec::new())); + let mut ids = guard.lock().expect("test rule registry mutex"); + let family_id = rule.family_id(); + if ids.contains(&family_id) { + return; + } + if let Err(err) = register_extension_rule(rule) { + match err { + tenferro::extension::ExtensionRegistryError::DuplicateRule { .. } => {} + other => panic!("register_extension_rule failed: {other}"), + } + } + ids.push(family_id); +} + // ---------------------------------------------------------------------- // TestScaleBy2: single-input, single-output. y = x + x (= 2x). // ---------------------------------------------------------------------- @@ -181,6 +203,64 @@ impl ExtensionFactory for TestScaleBy2Factory { fn ensure_scale_by_2_registered() { register_once(Arc::new(TestScaleBy2Factory)); + register_rule_once(Arc::new(TestScaleBy2Rule)); +} + +#[derive(Debug)] +struct TestScaleBy2Rule; + +impl ExtensionAdRuleTrait for TestScaleBy2Rule { + fn family_id(&self) -> &'static str { + "tenferro-tests.scale_by_2.v1" + } + + fn linearize( + &self, + _op: &dyn ExtensionOp, + builder: &mut FragmentBuilder, + _primal_in: &[GlobalValKey], + _primal_out: &[GlobalValKey], + tangent_in: &[Option], + _ctx: &mut ShapeGuardContext, + ) -> ADRuleResult>> { + match tangent_in[0] { + Some(dx) => { + let sum = builder.add_op( + StdTensorOp::Add, + vec![ValRef::Local(dx), ValRef::Local(dx)], + OpMode::Linear { + active_mask: vec![true, true], + }, + ); + Ok(vec![Some(sum[0])]) + } + None => Ok(vec![None]), + } + } + + fn transpose_rule( + &self, + _op: &dyn ExtensionOp, + emitter: &mut dyn OpEmitter, + cotangent_out: &[Option], + _inputs: &[ValRef], + _mode: &OpMode, + _ctx: &mut ShapeGuardContext, + ) -> ADRuleResult>> { + match cotangent_out[0] { + Some(ct) => { + let sum = emitter.add_op( + StdTensorOp::Add, + vec![ValRef::Local(ct), ValRef::Local(ct)], + OpMode::Linear { + active_mask: vec![true, true], + }, + ); + Ok(vec![Some(sum[0])]) + } + None => Ok(vec![None]), + } + } } // ---------------------------------------------------------------------- @@ -277,6 +357,137 @@ impl ExtensionFactory for TestSwapFactory { fn ensure_swap_registered() { register_once(Arc::new(TestSwapFactory)); + register_rule_once(Arc::new(TestSwapRule)); +} + +#[derive(Debug)] +struct TestSwapRule; + +impl ExtensionAdRuleTrait for TestSwapRule { + fn family_id(&self) -> &'static str { + "tenferro-tests.swap.v1" + } + + fn linearize( + &self, + _op: &dyn ExtensionOp, + _builder: &mut FragmentBuilder, + _primal_in: &[GlobalValKey], + _primal_out: &[GlobalValKey], + tangent_in: &[Option], + _ctx: &mut ShapeGuardContext, + ) -> ADRuleResult>> { + Ok(vec![tangent_in[1], tangent_in[0]]) + } + + fn transpose_rule( + &self, + _op: &dyn ExtensionOp, + _emitter: &mut dyn OpEmitter, + cotangent_out: &[Option], + _inputs: &[ValRef], + _mode: &OpMode, + _ctx: &mut ShapeGuardContext, + ) -> ADRuleResult>> { + Ok(vec![cotangent_out[1], cotangent_out[0]]) + } +} + +// ---------------------------------------------------------------------- +// TestNoAd: forward-only extension. Missing AD must be reported as Error. +// ---------------------------------------------------------------------- + +#[derive(Clone, Debug, PartialEq)] +struct TestNoAd; + +impl ExtensionOp for TestNoAd { + fn family_id(&self) -> &'static str { + "tenferro-tests.no_ad.v1" + } + + fn payload_hash(&self, _hasher: &mut dyn Hasher) {} + + fn payload_eq(&self, other: &dyn ExtensionOp) -> bool { + other.as_any().downcast_ref::().is_some() + } + + fn clone_arc(&self) -> Arc { + Arc::new(self.clone()) + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn n_inputs(&self) -> usize { + 1 + } + + fn n_outputs(&self) -> usize { + 1 + } + + fn infer_output_meta( + &self, + input_dtypes: &[DType], + input_shapes: &[&[SymDim]], + ) -> Vec<(DType, Vec)> { + vec![(input_dtypes[0], input_shapes[0].to_vec())] + } + + fn eager_execute(&self, inputs: &[&Tensor]) -> tenferro_tensor::Result> { + Ok(vec![inputs[0].clone()]) + } +} + +// ---------------------------------------------------------------------- +// TestBadOutputCount: malformed extension for facade validation paths. +// ---------------------------------------------------------------------- + +#[derive(Clone, Debug, PartialEq)] +struct TestBadOutputCount; + +impl ExtensionOp for TestBadOutputCount { + fn family_id(&self) -> &'static str { + "tenferro-tests.bad_output_count.v1" + } + + fn payload_hash(&self, _hasher: &mut dyn Hasher) {} + + fn payload_eq(&self, other: &dyn ExtensionOp) -> bool { + other + .as_any() + .downcast_ref::() + .is_some() + } + + fn clone_arc(&self) -> Arc { + Arc::new(self.clone()) + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn n_inputs(&self) -> usize { + 1 + } + + fn n_outputs(&self) -> usize { + 2 + } + + fn infer_output_meta( + &self, + input_dtypes: &[DType], + input_shapes: &[&[SymDim]], + ) -> Vec<(DType, Vec)> { + vec![(input_dtypes[0], input_shapes[0].to_vec())] + } + + fn eager_execute(&self, inputs: &[&Tensor]) -> tenferro_tensor::Result> { + Ok(vec![inputs[0].clone()]) + } } // ---------------------------------------------------------------------- @@ -330,6 +541,136 @@ fn scale_by_2_grad_against_reduce_sum() { assert_eq!(f64_slice(grad_out), &[2.0, 2.0, 2.0, 2.0]); } +#[test] +fn scale_by_2_eager_backward_uses_registered_rule() { + ensure_scale_by_2_registered(); + + let ctx = EagerContext::with_backend(CpuBackend::new()); + let x = + EagerTensor::requires_grad_in(Tensor::from_vec(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0]), ctx); + let scaled = apply_eager(Arc::new(TestScaleBy2), &[&x]) + .expect("eager extension apply") + .into_iter() + .next() + .expect("single extension output"); + let loss = scaled.reduce_sum(&[0]).expect("loss"); + + let _ = loss.backward().expect("eager backward"); + + assert_eq!( + x.grad().unwrap().as_slice::().unwrap(), + &[2.0, 2.0, 2.0, 2.0] + ); +} + +#[test] +fn missing_extension_rule_errors_in_traced_grad() { + let x = TracedTensor::from_vec(vec![2], vec![1.0_f64, 2.0]); + let y = apply(Arc::new(TestNoAd), &[&x]) + .into_iter() + .next() + .expect("single output"); + let loss = y.reduce_sum(&[0]); + + let err = match loss.grad(&x) { + Ok(_) => panic!("missing extension AD rule unexpectedly succeeded"), + Err(err) => err, + }; + + assert!(err.to_string().contains("tenferro-tests.no_ad.v1")); +} + +#[test] +fn missing_extension_rule_errors_in_eager_backward() { + let ctx = EagerContext::with_backend(CpuBackend::new()); + let x = EagerTensor::requires_grad_in(Tensor::from_vec(vec![2], vec![1.0_f64, 2.0]), ctx); + let y = apply_eager(Arc::new(TestNoAd), &[&x]) + .expect("forward-only eager extension apply") + .into_iter() + .next() + .expect("single output"); + let loss = y.reduce_sum(&[0]).expect("loss"); + + let err = loss + .backward() + .expect_err("missing extension AD rule should error"); + + assert!(err.to_string().contains("tenferro-tests.no_ad.v1")); +} + +#[test] +fn apply_rejects_wrong_input_count() { + let panic = catch_unwind(AssertUnwindSafe(|| { + let _ = apply(Arc::new(TestScaleBy2), &[]); + })); + + assert!(panic.is_err()); +} + +#[test] +fn apply_rejects_mismatched_output_metadata_count() { + let x = TracedTensor::from_vec(vec![2], vec![1.0_f64, 2.0]); + let panic = catch_unwind(AssertUnwindSafe(|| { + let _ = apply(Arc::new(TestBadOutputCount), &[&x]); + })); + + assert!(panic.is_err()); +} + +#[test] +fn apply_eager_rejects_empty_input_list() { + let err = match apply_eager::(Arc::new(TestScaleBy2), &[]) { + Ok(_) => panic!("empty eager extension input list unexpectedly succeeded"), + Err(err) => err, + }; + + assert!(err.to_string().contains("requires at least one input")); +} + +#[test] +fn apply_eager_rejects_wrong_input_count() { + let ctx = EagerContext::with_backend(CpuBackend::new()); + let x = EagerTensor::requires_grad_in(Tensor::from_vec(vec![1], vec![1.0_f64]), ctx); + + let err = match apply_eager(Arc::new(TestSwap), &[&x]) { + Ok(_) => panic!("wrong eager extension input count unexpectedly succeeded"), + Err(err) => err, + }; + + assert!(err.to_string().contains("expects 2 inputs, got 1")); +} + +#[test] +fn apply_eager_rejects_cross_context_inputs() { + let lhs_ctx = EagerContext::with_backend(CpuBackend::new()); + let rhs_ctx = EagerContext::with_backend(CpuBackend::new()); + let lhs = EagerTensor::requires_grad_in(Tensor::from_vec(vec![1], vec![1.0_f64]), lhs_ctx); + let rhs = EagerTensor::requires_grad_in(Tensor::from_vec(vec![1], vec![2.0_f64]), rhs_ctx); + + let err = match apply_eager(Arc::new(TestSwap), &[&lhs, &rhs]) { + Ok(_) => panic!("cross-context eager extension inputs unexpectedly succeeded"), + Err(err) => err, + }; + + assert!(matches!( + err, + tenferro::error::Error::ContextMismatch { .. } + )); +} + +#[test] +fn apply_eager_rejects_mismatched_output_count() { + let ctx = EagerContext::with_backend(CpuBackend::new()); + let x = EagerTensor::requires_grad_in(Tensor::from_vec(vec![1], vec![1.0_f64]), ctx); + + let err = match apply_eager(Arc::new(TestBadOutputCount), &[&x]) { + Ok(_) => panic!("bad eager extension output count unexpectedly succeeded"), + Err(err) => err, + }; + + assert!(err.to_string().contains("expected 2 eager outputs")); +} + #[test] fn scale_by_2_grad_through_symbolic_placeholder() { ensure_scale_by_2_registered();