diff --git a/llm-refactor-plan.md b/llm-refactor-plan.md new file mode 100644 index 0000000..ddb1b8f --- /dev/null +++ b/llm-refactor-plan.md @@ -0,0 +1,264 @@ +# LLM Refactor Plan — petite-ad + +> **Status:** In progress on branch `llm-refactor`. +> Last updated: 2026-05-11 +> Commits: `ffe4984` (annotations/re-exports), `7e4c7e1` (split compiled.rs), `84dd468` (extract expr/tape) + +Generated by the LLM AI Coding Agent audit checklist. Each finding is evaluated +against the principles in the skill: predictable structure, flat explicit +architecture, linear functions with explicit state, regenerability, and comment +hygiene. + +--- + +## High Priority + +### 1. ✅ Split `src/multi/compiled.rs` into backend modules + +**File:** `src/multi/compiled.rs` (4,615 lines) + +This single file holds: WGPU backend, SIMD x2/x4 backends, scalar backend, +device planning types, batch buffer types, CompiledGraph, BackendKind dispatch, +OpCode/FlatInstruction, and all execution logic. Five distinct concerns are +interleaved in one file. + +**Proposed split:** + +``` +src/multi/ + backend/ + mod.rs → pub mod + re-exports + types.rs → OpCode, FlatInstruction, BackendCapabilities, Instruction + scalar.rs → ScalarBackend + simd.rs → SimdBackend (f64x2, f64x4) + wgpu.rs → WgpuBackend, WgpuBuffer, WgpuBufferSet (feature-gated) + device.rs → DeviceBatchPlan, DeviceBuffer, DeviceBufferSet, + DeviceMemoryLocation, DeviceTransferKind, etc. + dispatch.rs → BackendKind, BackendSupportReport, auto-dispatch logic + mock.rs → MockDeviceBackend, DeviceExecutionTrace + compiled.rs → keep: CompiledGraph, CompiledWorkspace, BatchInputs, + BatchValues, BatchGradients, metadata +``` + +**Benefit:** Each backend can be regenerated independently without +understanding the others. A future LLM working on WGPU won't need to parse +SIMD dispatch logic. + +### 2. ✅ Decompose `gradient_batch_auto_into` + +**File:** `src/multi/compiled.rs`, line ~1,852 + +The single largest function in the codebase. It dispatches to Scalar, SIMD x2, +SIMD x4, WGPU paths with inline per-backend logic. + +**Proposed fix:** +```rust +fn gradient_batch_auto_into(compiled, batch, buf) -> Result<()> { + let backend = compiled.recommended_batch_gradient_backend(); + match backend { + BackendKind::Scalar => gradient_batch_scalar_into(compiled, batch, buf), + BackendKind::SimdF64x2 => gradient_batch_simd_f64x2_into(compiled, batch, buf), + // ... + } +} +// Each arm becomes a standalone 20-50 line function. +``` + +**Benefit:** A future LLM replacing the SIMD backend only needs to read one +small function, not parse 1,000 lines of unrelated dispatch code. + +### 3. ✅ Split `src/multi/graph.rs` into concerns + +**File:** `src/multi/graph.rs` (6,315 lines) + +Contains: `Graph` (core API, 1,550-line impl), `Tape`/`TapeWorkspace`, +`ExprGraph`/`ExprNode` (expression graph with operator overloading), +`GradientCheckEntry`/`GradientCheckReport`, `GraphNode`, `NodeId`, +`DomainPolicy`, `GraphStats`, WGPU helpers (cfg-gated), serde support +(cfg-gated). + +**Proposed split:** + +``` +src/multi/ + graph/ + mod.rs → pub mod + re-exports + core.rs → Graph, GraphNode, NodeId, DomainPolicy, GraphStats + tape.rs → Tape, TapeWorkspace + expr.rs → ExprGraph, ExprNode, operator overloads + check.rs → GradientCheckEntry, GradientCheckReport + graph.rs → keep if preferred, or delete +``` + +**Benefit:** The `Graph` impl shrinks from 1,550 lines to ~400 lines of core +logic. ExprGraph changes won't touch Graph internals. + +### 4. ✅ Remove unused import lint suppressions in `multi.rs` + +**File:** `src/multi.rs`, lines 51-55 + +```rust +#[allow(unused_imports)] +pub use multi_ad_fr::MultiAD2FR; +#[allow(unused_imports)] +pub use multi_ad_rf::MultiAD2RF; +#[allow(unused_imports)] +pub use multi_ad_rr::MultiAD2RR; +``` + +These types are re-exported but never used internally in `multi.rs`. The +`#[allow(unused_imports)]` is noise. Replace with a module-level +`#![allow(unused_imports)]` or better: re-export them through lib.rs only and +remove from multi.rs since they're already in lib.rs's `pub use` block. + +### 5. ✅ Clean up `#[allow(dead_code)]` on public API methods + +**Affected files:** `src/mono/mono_fn.rs`, `src/multi/multi_fn.rs`, +`src/mono/types.rs`, `src/forward.rs`, `src/multi/builder.rs` + +Public library methods tagged `#[allow(dead_code)]` signal that no internal +code calls them. This is expected for a library, but the annotations clutter +the code. Options: +- Move to a single `#![allow(dead_code)]` at crate root +- Or remove the annotations (dead_code doesn't apply to `pub` items when + building as a library, only as a binary) + +--- + +## Medium Priority + +### 6. Unify duplicate `Dual` and `OpKind` types + +**Files:** `src/mono/mono_hessian_common.rs` and +`src/multi/multi_hessian_common.rs` + +Both define: +```rust +struct Dual { val: f64, tan: f64 } +enum OpKind { Sin, Cos, Tan, Exp, Neg, Ln, Sqrt, Abs, ... } +``` + +The two `Dual` types are nearly identical (mono adds `Dual::constant`, +multi adds `Add/Sub/Mul/Div` impls). The `OpKind` enums have different +variants (mono has 8, multi has 16+). + +**Proposal:** Move `Dual` to `src/types.rs` with both `constant` and arithmetic +impls. Keep `OpKind` separate since the variants differ, but name them +`MonoOpKind` and `MultiOpKind` to clarify which is which. + +### 7. Consolidate re-export blocks + +**Files:** `src/lib.rs` (30-line `pub use` block), `src/multi.rs` (30-line +`pub use` block) + +Half the types are re-exported in both files. A future LLM needs to read both +to understand what's publicly available. + +**Proposal:** Use `pub use multi::*;` in lib.rs for most types, with explicit +re-exports only for items that need renaming or feature-gating. Keep multi.rs +as the single source of truth for what `multi` exports. + +### 8. Inline `#[cfg(test)]` example function files + +**Files:** `src/mono/mf1.rs` through `mf4.rs` (4 files), `src/multi/f1.rs` +through `f3.rs` (3 files) + +Each file is ~30 lines: one struct + one trait impl, used in exactly one test. +Seven extra files for 210 lines of test code. + +**Proposal:** Move the struct definitions into the test module that uses them +(`src/mono/tests.rs`, `src/multi/tests.rs`) as module-level items, or +consolidate into a single `src/mono/tests/examples.rs`. + +### 9. Split `src/forward.rs` by variable arity + +**File:** `src/forward.rs` (840 lines) + +Contains both `ForwardAD::differentiate` (mono) and +`ForwardAD::directional_derivative` (multi), plus graph-based forward mode. +Three related but distinct APIs in one file. + +**Proposal:** Keep `src/forward.rs` for the `ForwardAD` struct and shared +types. Move mono differentiation to `src/mono/forward.rs` and multi +directional derivative to `src/multi/forward.rs`. + +### 10. Remove architecture notes from test code + +**File:** `src/tests_comprehensive.rs`, lines ~630-640 + +```rust +// Note: F1, F2, F3 are private modules, tested through multi::tests +// since MF1-4 are private implementation details +``` + +These comments document where other types are tested — information that belongs +in those test modules, not in the comprehensive test file. Remove or move. + +### 11. Extract per-backend plan constructors from `batch_plan` + +**File:** `src/multi/compiled.rs`, `fn batch_plan` (135 lines) + +This function builds a `DeviceBatchPlan` for each backend variant with inline +code for every variant. Extract: + +```rust +fn plan_scalar(graph, batch_size) -> DeviceBatchPlan { ... } +fn plan_simd_f64x2(graph, batch_size) -> DeviceBatchPlan { ... } +fn plan_wgpu(graph, batch_size) -> DeviceBatchPlan { ... } +``` + +--- + +## Lower Priority + +### 12. Trailing comma in macro invocation + +**File:** `src/mono/mf4.rs`, line 19 + +```rust +&mono_ops![sin, neg,] +``` + +The trailing comma after `neg` is technically valid in Rust macros but +inconsistent with other `mono_ops!` calls. Remove for consistency. + +### 13. Module-level docs missing in `src/macros.rs` + +**File:** `src/macros.rs` (107 lines) + +The file defines `mono_ops!` and `multi_ops!` macros but has no module-level +doc comment. Add a brief explanation of the macro API and syntax. + +### 14. Ambiguous module name `types.rs` in both mono and multi + +**Files:** `src/mono/types.rs`, `src/multi/types.rs` + +Both define different types but share the same name. When reading code via +`use super::types::*`, it's ambiguous which module is being referenced. + +**Proposal:** Rename to `mono_types.rs` and `multi_types.rs`, or use +`mono::types` / `multi::types` consistently in imports (already done, but +the filename convention would help). + +### 15. Redundant `#[cfg(test)]` on items inside `#[cfg(test)]` modules + +**Files:** `src/mono/mf1.rs` through `mf4.rs`, `src/multi/f1.rs` through `f3.rs` + +These files are declared as `#[cfg(test)] mod mf1;` in `mono.rs`, but each +struct inside also has `#[cfg(test)]`. The outer annotation already gates +the entire module. Remove the inner annotations. + +--- + +## Summary + +| Priority | Items | Effort Estimate | +|----------|-------|-----------------| +| High | 5 | 2-3 days (split compiled.rs + graph.rs) | +| Medium | 6 | 1-2 days | +| Lower | 4 | 0.5 day | + +The highest-value change is splitting `compiled.rs` into backend modules. +This alone addresses the core violation of the "predictable structure" and +"flat, explicit architecture" principles. A future LLM can work on the +WGPU backend without touching SIMD logic, and vice versa. diff --git a/src/lib.rs b/src/lib.rs index fd98968..a1fbdf7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -144,21 +144,17 @@ pub use mono::MonoAD; // Higher-order autodiff methods (exact Hessian computation) pub use mono::{MonoAD2FR, MonoAD2RF, MonoAD2RR}; -pub use multi::builder::GraphBuilder; -pub use multi::multi_ad_fr::MultiAD2FR; -pub use multi::multi_ad_rf::MultiAD2RF; -pub use multi::multi_ad_rr::MultiAD2RR; -pub use multi::MultiAD; pub use multi::{ - AcceleratorDeviceContext, AcceleratorDeviceKind, BackendCapabilities, BackendKind, - BackendRejectionReason, BackendSupportReport, BatchGradients, BatchGradientsBuffer, - BatchInputs, BatchLayout, BatchValues, BatchValuesBuffer, CompiledGraph, CompiledGraphMetadata, - CompiledWorkspace, DeviceBackend, DeviceBatchPlan, DeviceBuffer, DeviceBufferHandle, - DeviceBufferKind, DeviceBufferLayout, DeviceBufferSet, DeviceExecutionMode, + builder::GraphBuilder, multi_ad_fr::MultiAD2FR, multi_ad_rf::MultiAD2RF, + multi_ad_rr::MultiAD2RR, AcceleratorDeviceContext, AcceleratorDeviceKind, BackendCapabilities, + BackendKind, BackendRejectionReason, BackendSupportReport, BatchGradients, + BatchGradientsBuffer, BatchInputs, BatchLayout, BatchValues, BatchValuesBuffer, CompiledGraph, + CompiledGraphMetadata, CompiledWorkspace, DeviceBackend, DeviceBatchPlan, DeviceBuffer, + DeviceBufferHandle, DeviceBufferKind, DeviceBufferLayout, DeviceBufferSet, DeviceExecutionMode, DeviceExecutionTrace, DeviceMemoryLocation, DeviceTransferKind, DeviceTransferPlan, DeviceTransferPolicy, DomainPolicy, ExecutionBackend, ExprGraph, ExprNode, FlatInstruction, GpuBackendBoundary, GradientCheckEntry, GradientCheckReport, Graph, GraphNode, GraphStats, - Instruction, MockDeviceBackend, NodeId, OpCode, ScalarBackend, SimdBackend, Tape, + Instruction, MockDeviceBackend, MultiAD, NodeId, OpCode, ScalarBackend, SimdBackend, Tape, TapeWorkspace, UNUSED_NODE_ID, }; #[cfg(feature = "backend-wgpu")] diff --git a/src/macros.rs b/src/macros.rs index 012c0ae..457c2ce 100644 --- a/src/macros.rs +++ b/src/macros.rs @@ -1,3 +1,10 @@ +//! Convenience macros for building single-variable and multi-variable +//! computational graphs concisely. +//! +//! - [`mono_ops!`] — Build a `[MonoAD; N]` array of unary operations (sin, cos, tan, exp, etc.). +//! - [`multi_ops!`] — Build a `[(MultiAD, Vec); N]` array of multi-variable operations +//! with input variable markers. + /// Macro to convert function names to MonoAD enum at compile time. /// This avoids the function pointer comparison issue across library boundaries. /// diff --git a/src/mono/mf4.rs b/src/mono/mf4.rs index 96ec7bb..3e4f6ec 100644 --- a/src/mono/mf4.rs +++ b/src/mono/mf4.rs @@ -13,7 +13,7 @@ impl MonoFn for MF4 { } fn graph(&self) -> &'static GraphType { - &mono_ops![sin, neg,] + &mono_ops![sin, neg] } fn expected_value(&self) -> f64 { diff --git a/src/mono/mono_fn.rs b/src/mono/mono_fn.rs index 97cf0eb..b48de14 100644 --- a/src/mono/mono_fn.rs +++ b/src/mono/mono_fn.rs @@ -2,7 +2,6 @@ pub use super::mono_ad::MonoAD; pub use super::types::BackwardResultBox; /// Type alias for a graph of mono operations (slice of MonoAD) -#[allow(dead_code)] // Public API for library extension pub type GraphType = [MonoAD]; /// Trait for single-variable functions with analytical gradients. @@ -12,7 +11,6 @@ pub type GraphType = [MonoAD]; /// /// This trait is primarily intended for testing and demonstration purposes. /// Most users will work directly with the `MonoAD` enum. -#[allow(dead_code)] // Public API for library extension pub trait MonoFn { /// Returns the input value for this function. fn input(&self) -> f64; diff --git a/src/multi.rs b/src/multi.rs index 12bed78..29026b1 100644 --- a/src/multi.rs +++ b/src/multi.rs @@ -3,19 +3,21 @@ //! This module provides functionality for computing gradients of //! multi-variable functions using computational graphs. -// Example implementations - not part of public API mod f1; mod f2; mod f3; +pub mod backend; pub mod builder; pub mod compiled; +pub mod expr; pub mod graph; mod multi_ad; pub mod multi_ad_fr; pub mod multi_ad_rf; pub mod multi_ad_rr; mod parser; +pub mod tape; // Shared internal modules for multivariate derivative rules and Hessian computation. mod multi_hessian_common; @@ -26,31 +28,34 @@ mod multi_fn; mod tests; pub mod types; -pub use compiled::{ +pub use backend::{ AcceleratorDeviceContext, AcceleratorDeviceKind, BackendCapabilities, BackendKind, - BackendRejectionReason, BackendSupportReport, BatchGradients, BatchGradientsBuffer, - BatchInputs, BatchLayout, BatchValues, BatchValuesBuffer, CompiledGraph, CompiledGraphMetadata, - CompiledWorkspace, DeviceBackend, DeviceBatchPlan, DeviceBuffer, DeviceBufferHandle, - DeviceBufferKind, DeviceBufferLayout, DeviceBufferSet, DeviceExecutionMode, - DeviceExecutionTrace, DeviceMemoryLocation, DeviceTransferKind, DeviceTransferPlan, - DeviceTransferPolicy, ExecutionBackend, FlatInstruction, GpuBackendBoundary, Instruction, - MockDeviceBackend, OpCode, ScalarBackend, SimdBackend, UNUSED_NODE_ID, + BackendRejectionReason, BackendSupportReport, BatchLayout, DeviceBackend, DeviceBatchPlan, + DeviceBuffer, DeviceBufferHandle, DeviceBufferKind, DeviceBufferLayout, DeviceBufferSet, + DeviceExecutionMode, DeviceExecutionTrace, DeviceMemoryLocation, DeviceTransferKind, + DeviceTransferPlan, DeviceTransferPolicy, ExecutionBackend, FlatInstruction, + GpuBackendBoundary, Instruction, MockDeviceBackend, OpCode, ScalarBackend, SimdBackend, + UNUSED_NODE_ID, }; #[cfg(feature = "backend-wgpu")] -pub use compiled::{ +pub use backend::{ WgpuBackend, WgpuBuffer, WgpuBufferSet, WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES, }; +pub use compiled::{ + BatchGradients, BatchGradientsBuffer, BatchInputs, BatchValues, BatchValuesBuffer, + CompiledGraph, CompiledGraphMetadata, CompiledWorkspace, +}; +pub use expr::{ExprGraph, ExprNode}; pub use graph::{ - DomainPolicy, ExprGraph, ExprNode, GradientCheckEntry, GradientCheckReport, Graph, GraphNode, - GraphStats, NodeId, Tape, TapeWorkspace, + DomainPolicy, GradientCheckEntry, GradientCheckReport, Graph, GraphNode, GraphStats, NodeId, }; pub use multi_ad::MultiAD; pub use multi_fn::MultiFn; +pub use tape::{Tape, TapeWorkspace}; -// Re-exported at crate root via lib.rs — suppress unused-import warnings in this module -#[allow(unused_imports)] +#[allow(unused_imports)] // Used via lib.rs re-export path pub use multi_ad_fr::MultiAD2FR; -#[allow(unused_imports)] +#[allow(unused_imports)] // Used via lib.rs re-export path pub use multi_ad_rf::MultiAD2RF; -#[allow(unused_imports)] +#[allow(unused_imports)] // Used via lib.rs re-export path pub use multi_ad_rr::MultiAD2RR; diff --git a/src/multi/backend/device.rs b/src/multi/backend/device.rs new file mode 100644 index 0000000..08b400c --- /dev/null +++ b/src/multi/backend/device.rs @@ -0,0 +1,390 @@ +//! Device-oriented buffer types, batch planning, and the DeviceBackend trait. + +use crate::multi::backend::dispatch::BackendKind; +use crate::multi::compiled::CompiledGraph; +use crate::{AutodiffError, Result}; + +#[cfg(feature = "backend-wgpu")] +use crate::multi::backend::wgpu::WgpuBackend; + +/// Batch memory layout exposed for backend/device planning. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BatchLayout { + RowMajor, +} + +/// Logical device buffer role for batch execution planning. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DeviceBufferKind { + Inputs, + Values, + Outputs, + PrimaryValues, + Gradients, +} + +/// One logical device buffer description. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DeviceBufferLayout { + pub kind: DeviceBufferKind, + pub len: usize, + pub element_size: usize, +} + +/// Logical memory location for a planned backend buffer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DeviceMemoryLocation { + Host, + Device, +} + +/// Handle-like description for one planned backend buffer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DeviceBufferHandle { + pub kind: DeviceBufferKind, + pub location: DeviceMemoryLocation, + pub offset: usize, + pub len: usize, +} + +/// Logical transfer direction for backend execution planning. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DeviceTransferKind { + HostToDevice, + DeviceToHost, +} + +/// One logical host/device transfer needed by a batch plan. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DeviceTransferPlan { + pub kind: DeviceTransferKind, + pub buffer: DeviceBufferKind, + pub len: usize, +} + +/// Device buffer owned by a planned backend execution. +#[derive(Debug, Clone, PartialEq)] +pub struct DeviceBuffer { + handle: DeviceBufferHandle, + data: Vec, +} + +impl DeviceBuffer { + fn new(handle: DeviceBufferHandle) -> Self { + Self { + handle, + data: vec![0.0; handle.len], + } + } + + /// Return the immutable planned buffer handle. + #[must_use] + pub fn handle(&self) -> DeviceBufferHandle { + self.handle + } + + /// Return the owned buffer data. + #[must_use] + pub fn data(&self) -> &[f64] { + &self.data + } +} + +/// Owned buffer set allocated from a [`DeviceBatchPlan`]. +#[derive(Debug, Clone, PartialEq)] +pub struct DeviceBufferSet { + pub(crate) plan: DeviceBatchPlan, + buffers: Vec, +} + +impl DeviceBufferSet { + /// Allocate zero-initialized buffers for a batch plan. + #[must_use] + pub fn new(plan: DeviceBatchPlan) -> Self { + let buffers = plan + .buffer_handles + .iter() + .copied() + .map(DeviceBuffer::new) + .collect(); + Self { plan, buffers } + } + + /// Return the plan used to allocate this buffer set. + #[must_use] + pub fn plan(&self) -> &DeviceBatchPlan { + &self.plan + } + + /// Return all allocated buffers in plan order. + #[must_use] + pub fn buffers(&self) -> &[DeviceBuffer] { + &self.buffers + } + + /// Return an immutable buffer by logical kind. + pub fn buffer(&self, kind: DeviceBufferKind) -> Result<&DeviceBuffer> { + self.buffers + .iter() + .find(|buffer| buffer.handle.kind == kind) + .ok_or(AutodiffError::InvalidGraph { + reason: "device buffer kind is not in the plan", + }) + } + + /// Return a mutable buffer by logical kind. + pub(crate) fn buffer_mut(&mut self, kind: DeviceBufferKind) -> Result<&mut DeviceBuffer> { + self.buffers + .iter_mut() + .find(|buffer| buffer.handle.kind == kind) + .ok_or(AutodiffError::InvalidGraph { + reason: "device buffer kind is not in the plan", + }) + } + + /// Upload host data into a planned buffer. + pub fn upload(&mut self, kind: DeviceBufferKind, data: &[f64]) -> Result<()> { + let buffer = self.buffer_mut(kind)?; + if buffer.data.len() != data.len() { + return Err(AutodiffError::InvalidGraph { + reason: "upload length must match planned buffer length", + }); + } + buffer.data.copy_from_slice(data); + Ok(()) + } + + /// Download a planned buffer into an owned vector. + pub fn download(&self, kind: DeviceBufferKind) -> Result> { + Ok(self.buffer(kind)?.data.clone()) + } +} + +/// Device-oriented batch execution plan for a backend. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DeviceBatchPlan { + pub backend: BackendKind, + pub layout: BatchLayout, + pub batch_size: usize, + pub input_dim: usize, + pub output_dim: usize, + pub gradient_dim: usize, + pub value_count: usize, + pub buffers: Vec, + pub buffer_handles: Vec, + pub compute_transfer_plan: Vec, + pub gradient_transfer_plan: Vec, +} + +/// Batch execution mode for explicit device transfer planning. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DeviceExecutionMode { + ComputeBatch, + GradientBatch, +} + +/// Trace returned by explicit device-style execution. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DeviceExecutionTrace { + pub backend: BackendKind, + pub mode: DeviceExecutionMode, + pub transfers: Vec, + /// Whether this execution used a native accelerator kernel instead of host fallback logic. + pub used_native_kernel: bool, +} + +/// Feature-neutral accelerator device family. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AcceleratorDeviceKind { + MockCpu, + Cuda, + Wgpu, +} + +/// Feature-neutral accelerator context descriptor for future GPU backends. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AcceleratorDeviceContext { + pub kind: AcceleratorDeviceKind, + pub device_id: usize, + pub name: String, +} + +impl AcceleratorDeviceContext { + /// Return a context descriptor for the CPU mock-device backend. + #[must_use] + pub fn mock_cpu() -> Self { + Self { + kind: AcceleratorDeviceKind::MockCpu, + device_id: 0, + name: "mock-device-cpu".to_string(), + } + } + + /// Return a context descriptor for a future CUDA backend. + #[must_use] + pub fn cuda(device_id: usize) -> Self { + Self { + kind: AcceleratorDeviceKind::Cuda, + device_id, + name: format!("cuda:{device_id}"), + } + } + + /// Return a context descriptor for a future WGPU backend. + #[must_use] + pub fn wgpu(device_id: usize) -> Self { + Self { + kind: AcceleratorDeviceKind::Wgpu, + device_id, + name: format!("wgpu:{device_id}"), + } + } +} + +/// Host/device transfer policy for future accelerator backends. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DeviceTransferPolicy { + Explicit, + Automatic, +} + +/// Boundary object for future GPU backends. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GpuBackendBoundary { + pub context: AcceleratorDeviceContext, + pub transfer_policy: DeviceTransferPolicy, +} + +impl GpuBackendBoundary { + /// Create a boundary descriptor for a future accelerator backend. + #[must_use] + pub fn new(context: AcceleratorDeviceContext, transfer_policy: DeviceTransferPolicy) -> Self { + Self { + context, + transfer_policy, + } + } + + #[cfg(feature = "backend-wgpu")] + /// Initialize the real WGPU backend skeleton from this boundary. + pub fn initialize_wgpu(&self) -> Result { + WgpuBackend::from_boundary(self.clone()) + } + + /// Return an explicit error because generic GPU execution still requires a concrete backend. + pub fn unsupported_execution_error(&self) -> Result { + Err(AutodiffError::InvalidGraph { + reason: + "real GPU execution requires a concrete backend instance; initialize WgpuBackend", + }) + } +} + +/// Device-oriented batch backend planning interface. +pub trait DeviceBackend { + fn backend_kind(&self, graph: &CompiledGraph) -> BackendKind; + + fn batch_layout(&self) -> BatchLayout { + BatchLayout::RowMajor + } + + fn batch_plan(&self, graph: &CompiledGraph, batch_size: usize) -> DeviceBatchPlan { + let metadata = graph.metadata(); + let value_count = metadata.value_count.saturating_mul(batch_size); + let input_count = metadata.num_inputs.saturating_mul(batch_size); + let output_count = metadata.num_outputs.saturating_mul(batch_size); + let gradient_count = metadata.num_inputs.saturating_mul(batch_size); + let backend = self.backend_kind(graph); + let buffer_location = backend.memory_location(); + let buffers = vec![ + DeviceBufferLayout { + kind: DeviceBufferKind::Inputs, + len: input_count, + element_size: std::mem::size_of::(), + }, + DeviceBufferLayout { + kind: DeviceBufferKind::Values, + len: value_count, + element_size: std::mem::size_of::(), + }, + DeviceBufferLayout { + kind: DeviceBufferKind::Outputs, + len: output_count, + element_size: std::mem::size_of::(), + }, + DeviceBufferLayout { + kind: DeviceBufferKind::PrimaryValues, + len: batch_size, + element_size: std::mem::size_of::(), + }, + DeviceBufferLayout { + kind: DeviceBufferKind::Gradients, + len: gradient_count, + element_size: std::mem::size_of::(), + }, + ]; + let mut offset = 0; + let buffer_handles = buffers + .iter() + .map(|buffer| { + let handle = DeviceBufferHandle { + kind: buffer.kind, + location: buffer_location, + offset, + len: buffer.len, + }; + offset = offset.saturating_add(buffer.len); + handle + }) + .collect(); + let (compute_transfer_plan, gradient_transfer_plan) = + if buffer_location == DeviceMemoryLocation::Device { + ( + vec![ + DeviceTransferPlan { + kind: DeviceTransferKind::HostToDevice, + buffer: DeviceBufferKind::Inputs, + len: input_count, + }, + DeviceTransferPlan { + kind: DeviceTransferKind::DeviceToHost, + buffer: DeviceBufferKind::Outputs, + len: output_count, + }, + ], + vec![ + DeviceTransferPlan { + kind: DeviceTransferKind::HostToDevice, + buffer: DeviceBufferKind::Inputs, + len: input_count, + }, + DeviceTransferPlan { + kind: DeviceTransferKind::DeviceToHost, + buffer: DeviceBufferKind::PrimaryValues, + len: batch_size, + }, + DeviceTransferPlan { + kind: DeviceTransferKind::DeviceToHost, + buffer: DeviceBufferKind::Gradients, + len: gradient_count, + }, + ], + ) + } else { + (Vec::new(), Vec::new()) + }; + DeviceBatchPlan { + backend, + layout: self.batch_layout(), + batch_size, + input_dim: metadata.num_inputs, + output_dim: metadata.num_outputs, + gradient_dim: metadata.num_inputs, + value_count: metadata.value_count, + buffers, + buffer_handles, + compute_transfer_plan, + gradient_transfer_plan, + } + } +} diff --git a/src/multi/backend/dispatch.rs b/src/multi/backend/dispatch.rs new file mode 100644 index 0000000..c1d413e --- /dev/null +++ b/src/multi/backend/dispatch.rs @@ -0,0 +1,265 @@ +//! Backend dispatch types, ExecutionBackend trait, and auto-dispatch logic. + +use crate::multi::backend::device::{DeviceBackend, DeviceMemoryLocation}; +use crate::multi::backend::scalar::ScalarBackend; +use crate::multi::backend::simd::{ + compute_batch_simd_f64x2, compute_batch_simd_f64x4, gradient_batch_simd_f64x2, + gradient_batch_simd_f64x4, +}; +use crate::multi::backend::types::{ + supports_simd_f64x2_runtime, supports_simd_f64x4_runtime, supports_wgpu_runtime, + BackendCapabilities, OpCode, +}; +use crate::multi::compiled::{BatchGradientsBuffer, BatchInputs, BatchValuesBuffer, CompiledGraph}; +use crate::{AutodiffError, Result}; + +#[cfg(feature = "backend-wgpu")] +use crate::multi::backend::wgpu::WgpuBackend; + +/// Backend selected by automatic compiled-graph dispatch. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BackendKind { + Scalar, + MockDeviceCpu, + Wgpu, + SimdF64x4, + SimdF64x2, +} + +impl BackendKind { + /// Return the number of f64 lanes processed by this backend. + #[must_use] + pub fn lane_width(&self) -> usize { + match self { + BackendKind::Scalar | BackendKind::MockDeviceCpu | BackendKind::Wgpu => 1, + BackendKind::SimdF64x4 => 4, + BackendKind::SimdF64x2 => 2, + } + } + + /// Return runtime CPU features required by this backend. + #[must_use] + pub fn required_runtime_features(&self) -> &'static [&'static str] { + match self { + BackendKind::Scalar | BackendKind::MockDeviceCpu => &[], + BackendKind::Wgpu => &["backend-wgpu"], + BackendKind::SimdF64x4 => &["x86_64-avx"], + BackendKind::SimdF64x2 => &["x86_64-sse2"], + } + } + + /// Return whether the backend is available on this runtime target. + #[must_use] + pub fn runtime_available(&self) -> bool { + match self { + BackendKind::Scalar | BackendKind::MockDeviceCpu => true, + BackendKind::Wgpu => supports_wgpu_runtime(), + BackendKind::SimdF64x4 => supports_simd_f64x4_runtime(), + BackendKind::SimdF64x2 => supports_simd_f64x2_runtime(), + } + } + + /// Return runtime CPU features that are required but unavailable. + #[must_use] + pub fn unavailable_runtime_features(&self) -> Vec<&'static str> { + if self.runtime_available() { + Vec::new() + } else { + self.required_runtime_features().to_vec() + } + } + + /// Return a stable backend name. + #[must_use] + pub fn name(&self) -> &'static str { + match self { + BackendKind::Scalar => "scalar", + BackendKind::MockDeviceCpu => "mock-device-cpu", + BackendKind::Wgpu => "wgpu", + BackendKind::SimdF64x4 => "simd-f64x4", + BackendKind::SimdF64x2 => "simd-f64x2", + } + } + + /// Return the logical memory location used by this backend's batch plan. + #[must_use] + pub fn memory_location(&self) -> DeviceMemoryLocation { + match self { + BackendKind::MockDeviceCpu | BackendKind::Wgpu => DeviceMemoryLocation::Device, + BackendKind::Scalar | BackendKind::SimdF64x4 | BackendKind::SimdF64x2 => { + DeviceMemoryLocation::Host + } + } + } + + /// Return declared capabilities for this backend. + #[must_use] + pub fn capabilities(&self) -> BackendCapabilities { + match self { + BackendKind::Scalar | BackendKind::MockDeviceCpu => BackendCapabilities::scalar_f64(), + BackendKind::Wgpu => BackendCapabilities::wgpu_f64(), + BackendKind::SimdF64x4 => BackendCapabilities::simd_f64x4(), + BackendKind::SimdF64x2 => BackendCapabilities::simd_f64x2(), + } + } + + /// Execute batch value computation with this backend. + /// + /// # Performance + /// + /// When `self` is [`BackendKind::Wgpu`] this method calls + /// [`WgpuBackend::new_default`] on every invocation, which initialises a + /// new GPU device and queue each time. For repeated batch calls, create a + /// [`WgpuBackend`] once and call its [`ExecutionBackend::compute_batch`] + /// method directly instead of routing through `BackendKind::Wgpu`. + pub fn compute_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result<()> { + if matches!(self, BackendKind::Scalar | BackendKind::MockDeviceCpu) { + return ScalarBackend.compute_batch(graph, batch, buffer); + } + if matches!(self, BackendKind::Wgpu) { + #[cfg(feature = "backend-wgpu")] + { + return WgpuBackend::new_default()?.compute_batch(graph, batch, buffer); + } + #[cfg(not(feature = "backend-wgpu"))] + { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu backend requires the backend-wgpu cargo feature", + }); + } + } + let capabilities = self.capabilities(); + if !capabilities.supports_batch_compute { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support batch compute on this target", + }); + } + graph.validate_backend_capabilities(&capabilities)?; + match self { + BackendKind::Scalar | BackendKind::MockDeviceCpu | BackendKind::Wgpu => unreachable!(), + BackendKind::SimdF64x4 => compute_batch_simd_f64x4(graph, batch, buffer), + BackendKind::SimdF64x2 => compute_batch_simd_f64x2(graph, batch, buffer), + } + } + + /// Execute batch gradient computation with this backend. + /// + /// # Performance + /// + /// When `self` is [`BackendKind::Wgpu`] this method calls + /// [`WgpuBackend::new_default`] on every invocation, which initialises a + /// new GPU device and queue each time. For repeated batch calls, create a + /// [`WgpuBackend`] once and call its [`ExecutionBackend::gradient_batch`] + /// method directly instead of routing through `BackendKind::Wgpu`. + pub fn gradient_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result<()> { + if matches!(self, BackendKind::Scalar | BackendKind::MockDeviceCpu) { + return ScalarBackend.gradient_batch(graph, batch, buffer); + } + if matches!(self, BackendKind::Wgpu) { + #[cfg(feature = "backend-wgpu")] + { + return WgpuBackend::new_default()?.gradient_batch(graph, batch, buffer); + } + #[cfg(not(feature = "backend-wgpu"))] + { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu backend requires the backend-wgpu cargo feature", + }); + } + } + let capabilities = self.capabilities(); + if !capabilities.supports_batch_gradient { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support batch gradients on this target", + }); + } + graph.validate_backend_capabilities(&capabilities)?; + match self { + BackendKind::Scalar | BackendKind::MockDeviceCpu | BackendKind::Wgpu => unreachable!(), + BackendKind::SimdF64x4 => gradient_batch_simd_f64x4(graph, batch, buffer), + BackendKind::SimdF64x2 => gradient_batch_simd_f64x2(graph, batch, buffer), + } + } +} + +/// Reason a backend cannot execute a graph for a requested batch mode. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum BackendRejectionReason { + MissingF64, + UnsupportedOutputs, + UnsupportedOpcodes, + UnsupportedArities, + UnavailableRuntime, + NoBatchCompute, + NoBatchGradient, +} + +/// Static backend compatibility report for a compiled graph. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BackendSupportReport { + pub backend: BackendKind, + pub supports_f64: bool, + pub supports_required_outputs: bool, + pub supports_required_opcodes: bool, + pub supports_required_arities: bool, + pub supports_batch_compute: bool, + pub supports_batch_gradient: bool, + pub lane_width: usize, + pub runtime_available: bool, + pub required_runtime_features: Vec<&'static str>, + pub unavailable_runtime_features: Vec<&'static str>, + pub missing_opcodes: Vec, + pub batch_compute_rejection_reasons: Vec, + pub batch_gradient_rejection_reasons: Vec, +} + +impl BackendSupportReport { + /// Return whether this backend can compute batch values for the graph. + #[must_use] + pub fn can_compute_batch(&self) -> bool { + self.batch_compute_rejection_reasons.is_empty() + } + + /// Return whether this backend can compute batch gradients for the graph. + #[must_use] + pub fn can_gradient_batch(&self) -> bool { + self.batch_gradient_rejection_reasons.is_empty() + } +} + +/// Common execution-backend interface for future scalar, SIMD, and GPU backends. +pub trait ExecutionBackend { + fn name(&self) -> &'static str; + fn capabilities(&self) -> BackendCapabilities; + + fn compute(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result; + fn compute_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result<()>; + fn gradient(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result<(f64, Vec)>; + fn gradient_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result<()>; +} + +impl DeviceBackend for BackendKind { + fn backend_kind(&self, _graph: &CompiledGraph) -> BackendKind { + *self + } +} diff --git a/src/multi/backend/mock.rs b/src/multi/backend/mock.rs new file mode 100644 index 0000000..46d497f --- /dev/null +++ b/src/multi/backend/mock.rs @@ -0,0 +1,195 @@ +//! Mock device-style backend that executes on CPU while using device-oriented plans. + +use crate::multi::backend::device::{ + DeviceBackend, DeviceBufferKind, DeviceBufferSet, DeviceExecutionMode, DeviceExecutionTrace, + DeviceTransferKind, +}; +use crate::multi::backend::dispatch::{BackendKind, ExecutionBackend}; +use crate::multi::backend::scalar::ScalarBackend; +use crate::multi::backend::types::BackendCapabilities; +use crate::multi::compiled::{BatchGradientsBuffer, BatchInputs, BatchValuesBuffer, CompiledGraph}; +use crate::{AutodiffError, Result}; + +/// Mock device-style backend that executes on CPU while using device-oriented plans. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct MockDeviceBackend; + +impl DeviceBackend for MockDeviceBackend { + fn backend_kind(&self, _graph: &CompiledGraph) -> BackendKind { + BackendKind::MockDeviceCpu + } +} + +impl MockDeviceBackend { + /// Allocate a mock-device buffer set for a compiled graph and batch size. + #[must_use] + pub fn allocate_batch_buffers( + &self, + graph: &CompiledGraph, + batch_size: usize, + ) -> DeviceBufferSet { + DeviceBufferSet::new(self.batch_plan(graph, batch_size)) + } + + /// Execute batch value computation through explicit mock-device transfers. + pub fn compute_batch_with_buffers( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffers: &mut DeviceBufferSet, + output: &mut BatchValuesBuffer, + ) -> Result { + graph.check_batch(batch)?; + if buffers.plan.backend != BackendKind::MockDeviceCpu { + return Err(AutodiffError::InvalidGraph { + reason: "mock execution requires a mock-device buffer plan", + }); + } + if buffers.plan.batch_size != batch.batch_size || buffers.plan.input_dim != batch.input_dim + { + return Err(AutodiffError::InvalidGraph { + reason: "batch shape must match mock-device buffer plan", + }); + } + + let transfers = buffers.plan.compute_transfer_plan.clone(); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::HostToDevice { + if transfer.buffer != DeviceBufferKind::Inputs { + return Err(AutodiffError::InvalidGraph { + reason: "mock compute supports host-to-device input transfers only", + }); + } + buffers.upload(DeviceBufferKind::Inputs, batch.data)?; + } + } + + let device_inputs = buffers.download(DeviceBufferKind::Inputs)?; + let device_batch = BatchInputs::new(&device_inputs, batch.batch_size, batch.input_dim)?; + let mut scratch = BatchValuesBuffer::new(); + ScalarBackend.compute_batch(graph, device_batch, &mut scratch)?; + buffers.upload(DeviceBufferKind::Outputs, &scratch.data)?; + output.reset(batch.batch_size, graph.output_nodes.len()); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::DeviceToHost { + if transfer.buffer != DeviceBufferKind::Outputs { + return Err(AutodiffError::InvalidGraph { + reason: "mock compute supports device-to-host output transfers only", + }); + } + output + .data + .extend_from_slice(&buffers.download(DeviceBufferKind::Outputs)?); + } + } + + Ok(DeviceExecutionTrace { + backend: BackendKind::MockDeviceCpu, + mode: DeviceExecutionMode::ComputeBatch, + transfers, + used_native_kernel: false, + }) + } + + /// Execute batch gradient computation through explicit mock-device transfers. + pub fn gradient_batch_with_buffers( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffers: &mut DeviceBufferSet, + output: &mut BatchGradientsBuffer, + ) -> Result { + graph.check_batch(batch)?; + if buffers.plan.backend != BackendKind::MockDeviceCpu { + return Err(AutodiffError::InvalidGraph { + reason: "mock execution requires a mock-device buffer plan", + }); + } + if buffers.plan.batch_size != batch.batch_size || buffers.plan.input_dim != batch.input_dim + { + return Err(AutodiffError::InvalidGraph { + reason: "batch shape must match mock-device buffer plan", + }); + } + + let transfers = buffers.plan.gradient_transfer_plan.clone(); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::HostToDevice { + if transfer.buffer != DeviceBufferKind::Inputs { + return Err(AutodiffError::InvalidGraph { + reason: "mock gradient supports host-to-device input transfers only", + }); + } + buffers.upload(DeviceBufferKind::Inputs, batch.data)?; + } + } + + let device_inputs = buffers.download(DeviceBufferKind::Inputs)?; + let device_batch = BatchInputs::new(&device_inputs, batch.batch_size, batch.input_dim)?; + let mut scratch = BatchGradientsBuffer::new(); + ScalarBackend.gradient_batch(graph, device_batch, &mut scratch)?; + buffers.upload(DeviceBufferKind::PrimaryValues, &scratch.values)?; + buffers.upload(DeviceBufferKind::Gradients, &scratch.gradients)?; + output.reset(batch.batch_size, graph.num_inputs); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::DeviceToHost { + match transfer.buffer { + DeviceBufferKind::PrimaryValues => output + .values + .extend_from_slice(&buffers.download(DeviceBufferKind::PrimaryValues)?), + DeviceBufferKind::Gradients => output + .gradients + .extend_from_slice(&buffers.download(DeviceBufferKind::Gradients)?), + _ => { + return Err(AutodiffError::InvalidGraph { + reason: + "mock gradient supports primary-value and gradient downloads only", + }); + } + } + } + } + + Ok(DeviceExecutionTrace { + backend: BackendKind::MockDeviceCpu, + mode: DeviceExecutionMode::GradientBatch, + transfers, + used_native_kernel: false, + }) + } +} +impl ExecutionBackend for MockDeviceBackend { + fn name(&self) -> &'static str { + "mock-device-cpu" + } + + fn capabilities(&self) -> BackendCapabilities { + BackendCapabilities::scalar_f64() + } + + fn compute(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result { + ScalarBackend.compute(graph, inputs) + } + + fn compute_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result<()> { + ScalarBackend.compute_batch(graph, batch, buffer) + } + + fn gradient(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result<(f64, Vec)> { + ScalarBackend.gradient(graph, inputs) + } + + fn gradient_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result<()> { + ScalarBackend.gradient_batch(graph, batch, buffer) + } +} diff --git a/src/multi/backend/mod.rs b/src/multi/backend/mod.rs new file mode 100644 index 0000000..5f6c2cf --- /dev/null +++ b/src/multi/backend/mod.rs @@ -0,0 +1,22 @@ +//! Backend abstraction module for compiled graph execution. +//! +//! This module contains the types, traits, and backend implementations +//! for executing compiled graphs on different backends (scalar, SIMD, WGPU). + +pub(crate) mod device; +pub(crate) mod dispatch; +pub(crate) mod mock; +pub(crate) mod scalar; +pub(crate) mod simd; +pub(crate) mod types; +#[cfg(feature = "backend-wgpu")] +pub(crate) mod wgpu; + +pub use device::*; +pub use dispatch::*; +pub use mock::*; +pub use scalar::*; +pub use simd::*; +pub use types::*; +#[cfg(feature = "backend-wgpu")] +pub use wgpu::*; diff --git a/src/multi/backend/scalar.rs b/src/multi/backend/scalar.rs new file mode 100644 index 0000000..032e245 --- /dev/null +++ b/src/multi/backend/scalar.rs @@ -0,0 +1,75 @@ +//! Reference scalar f64 backend. + +use crate::multi::backend::device::DeviceBackend; +use crate::multi::backend::dispatch::{BackendKind, ExecutionBackend}; +use crate::multi::backend::types::BackendCapabilities; +use crate::multi::compiled::{BatchGradientsBuffer, BatchInputs, BatchValuesBuffer, CompiledGraph}; +use crate::{AutodiffError, Result}; + +/// Reference scalar backend for the execution-backend abstraction. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ScalarBackend; + +impl DeviceBackend for ScalarBackend { + fn backend_kind(&self, _graph: &CompiledGraph) -> BackendKind { + BackendKind::Scalar + } +} + +impl ExecutionBackend for ScalarBackend { + fn name(&self) -> &'static str { + "scalar" + } + + fn capabilities(&self) -> BackendCapabilities { + BackendCapabilities::scalar_f64() + } + + fn compute(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result { + graph.validate_backend_capabilities(&self.capabilities())?; + graph.compute(inputs) + } + + fn compute_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result<()> { + let capabilities = self.capabilities(); + if !capabilities.supports_batch_compute { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support batch compute", + }); + } + graph.validate_backend_capabilities(&capabilities)?; + graph.compute_batch_into(batch, buffer) + } + + fn gradient(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result<(f64, Vec)> { + let capabilities = self.capabilities(); + if !capabilities.supports_reverse_gradient { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support reverse gradients", + }); + } + graph.validate_backend_capabilities(&capabilities)?; + graph.gradient(inputs) + } + + fn gradient_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result<()> { + let capabilities = self.capabilities(); + if !capabilities.supports_batch_gradient { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support batch gradients", + }); + } + graph.validate_backend_capabilities(&capabilities)?; + graph.gradient_batch_into(batch, buffer) + } +} diff --git a/src/multi/backend/simd.rs b/src/multi/backend/simd.rs new file mode 100644 index 0000000..f082427 --- /dev/null +++ b/src/multi/backend/simd.rs @@ -0,0 +1,1163 @@ +//! Prototype f64x2/f64x4 SIMD backend for batch compute and batch gradients. + +use crate::multi::backend::device::DeviceBackend; +use crate::multi::backend::dispatch::{BackendKind, ExecutionBackend}; +use crate::multi::backend::types::{ + supports_simd_f64x4_runtime, BackendCapabilities, FlatInstruction, OpCode, +}; +use crate::multi::compiled::{BatchGradientsBuffer, BatchInputs, BatchValuesBuffer, CompiledGraph}; +use crate::multi::op_rules; +use crate::{AutodiffError, NodeId, Result}; + +#[cfg(target_arch = "x86_64")] +use std::arch::x86_64::{__m128d, __m256d}; + +/// Prototype f64x2 SIMD backend for batch compute and batch gradients. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct SimdBackend; + +impl DeviceBackend for SimdBackend { + fn backend_kind(&self, graph: &CompiledGraph) -> BackendKind { + graph + .simd_support_report() + .map(|report| report.backend) + .unwrap_or(BackendKind::SimdF64x2) + } +} + +impl ExecutionBackend for SimdBackend { + fn name(&self) -> &'static str { + "simd" + } + + fn capabilities(&self) -> BackendCapabilities { + BackendCapabilities::simd_f64x2() + } + + fn compute(&self, _graph: &CompiledGraph, _inputs: &[f64]) -> Result { + Err(AutodiffError::InvalidGraph { + reason: "simd backend currently supports batch compute only", + }) + } + + fn compute_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result<()> { + if let Ok(report) = graph.backend_support_report(BackendKind::SimdF64x4) { + if report.can_compute_batch() { + return compute_batch_simd_f64x4(graph, batch, buffer); + } + } + let capabilities = self.capabilities(); + if !capabilities.supports_batch_compute { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support batch compute on this target", + }); + } + graph.validate_backend_capabilities(&capabilities)?; + compute_batch_simd_f64x2(graph, batch, buffer) + } + + fn gradient(&self, _graph: &CompiledGraph, _inputs: &[f64]) -> Result<(f64, Vec)> { + Err(AutodiffError::InvalidGraph { + reason: "simd backend does not support reverse gradients yet", + }) + } + + fn gradient_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result<()> { + if let Ok(report) = graph.backend_support_report(BackendKind::SimdF64x4) { + if report.can_gradient_batch() { + return gradient_batch_simd_f64x4(graph, batch, buffer); + } + } + let capabilities = self.capabilities(); + if !capabilities.supports_batch_gradient { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support batch gradients on this target", + }); + } + graph.validate_backend_capabilities(&capabilities)?; + gradient_batch_simd_f64x2(graph, batch, buffer) + } +} + +fn simd_unsupported_opcode_error() -> AutodiffError { + AutodiffError::InvalidGraph { + reason: "simd backend does not support a required opcode", + } +} + +fn simd_scalar_value(opcode: OpCode, args: &[f64]) -> Result { + let op = opcode + .to_multi_ad() + .ok_or_else(simd_unsupported_opcode_error)?; + op_rules::forward_value(op, args) +} + +// Returns `[d0, d1]` where d1 is 0.0 for unary ops (arity ≤ 2, heap-free). +fn simd_scalar_first_derivatives(opcode: OpCode, args: &[f64], value: f64) -> Result<[f64; 2]> { + let op = opcode + .to_multi_ad() + .ok_or_else(simd_unsupported_opcode_error)?; + let derivs = op_rules::first_derivatives(op, args, value)?; + Ok([ + derivs.first().copied().unwrap_or(0.0), + derivs.get(1).copied().unwrap_or(0.0), + ]) +} + +fn append_scalar_compute_tail( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + start_row: usize, + buffer: &mut BatchValuesBuffer, +) -> Result<()> { + let mut workspace = graph.workspace(); + for row_index in start_row..batch.batch_size { + let outputs = + graph.compute_many_with_workspace(batch.try_row(row_index)?, &mut workspace)?; + buffer.data.extend_from_slice(outputs); + } + Ok(()) +} + +fn append_scalar_gradient_tail( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + start_row: usize, + buffer: &mut BatchGradientsBuffer, +) -> Result<()> { + let mut workspace = graph.workspace(); + for row_index in start_row..batch.batch_size { + let (value, gradient) = + graph.gradient_with_workspace(batch.try_row(row_index)?, &mut workspace)?; + buffer.values.push(value); + buffer.gradients.extend_from_slice(gradient); + } + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +fn checked_m128_lane(values: &[__m128d], index: NodeId) -> Result<__m128d> { + values + .get(index) + .copied() + .ok_or(AutodiffError::IndexOutOfBounds { + index, + max_index: values.len().saturating_sub(1), + }) +} + +#[cfg(target_arch = "x86_64")] +unsafe fn add_m128_cotangent( + cotangents: &mut [__m128d], + index: NodeId, + contribution: __m128d, +) -> Result<()> { + use std::arch::x86_64::_mm_add_pd; + + let max_index = cotangents.len().saturating_sub(1); + let target = cotangents + .get_mut(index) + .ok_or(AutodiffError::IndexOutOfBounds { index, max_index })?; + *target = _mm_add_pd(*target, contribution); + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +unsafe fn active_m128_contribution(current: __m128d, contribution: __m128d) -> __m128d { + use std::arch::x86_64::{_mm_and_pd, _mm_cmpneq_pd, _mm_setzero_pd}; + + let active = _mm_cmpneq_pd(current, _mm_setzero_pd()); + _mm_and_pd(contribution, active) +} + +#[cfg(target_arch = "x86_64")] +unsafe fn simd_f64x2_scalar_unary(input: __m128d, opcode: OpCode) -> Result<__m128d> { + use std::arch::x86_64::{_mm_set_pd, _mm_storeu_pd}; + + let mut stored = [0.0_f64; 2]; + _mm_storeu_pd(stored.as_mut_ptr(), input); + let first = simd_scalar_value(opcode, &[stored[0]])?; + let second = simd_scalar_value(opcode, &[stored[1]])?; + Ok(_mm_set_pd(second, first)) +} + +#[cfg(target_arch = "x86_64")] +unsafe fn simd_f64x2_scalar_unary_derivative( + input: __m128d, + output: __m128d, + opcode: OpCode, +) -> Result<__m128d> { + use std::arch::x86_64::{_mm_set_pd, _mm_storeu_pd}; + + let mut input_values = [0.0_f64; 2]; + let mut output_values = [0.0_f64; 2]; + _mm_storeu_pd(input_values.as_mut_ptr(), input); + _mm_storeu_pd(output_values.as_mut_ptr(), output); + let first = simd_scalar_first_derivatives(opcode, &[input_values[0]], output_values[0])?[0]; + let second = simd_scalar_first_derivatives(opcode, &[input_values[1]], output_values[1])?[0]; + Ok(_mm_set_pd(second, first)) +} + +#[cfg(target_arch = "x86_64")] +unsafe fn simd_f64x2_scalar_binary( + left: __m128d, + right: __m128d, + opcode: OpCode, +) -> Result<__m128d> { + use std::arch::x86_64::{_mm_set_pd, _mm_storeu_pd}; + + let mut left_values = [0.0_f64; 2]; + let mut right_values = [0.0_f64; 2]; + _mm_storeu_pd(left_values.as_mut_ptr(), left); + _mm_storeu_pd(right_values.as_mut_ptr(), right); + let first = simd_scalar_value(opcode, &[left_values[0], right_values[0]])?; + let second = simd_scalar_value(opcode, &[left_values[1], right_values[1]])?; + Ok(_mm_set_pd(second, first)) +} + +#[cfg(target_arch = "x86_64")] +unsafe fn simd_f64x2_scalar_binary_derivatives( + left: __m128d, + right: __m128d, + output: __m128d, + opcode: OpCode, +) -> Result<(__m128d, __m128d)> { + use std::arch::x86_64::{_mm_set_pd, _mm_storeu_pd}; + + let mut left_values = [0.0_f64; 2]; + let mut right_values = [0.0_f64; 2]; + let mut output_values = [0.0_f64; 2]; + _mm_storeu_pd(left_values.as_mut_ptr(), left); + _mm_storeu_pd(right_values.as_mut_ptr(), right); + _mm_storeu_pd(output_values.as_mut_ptr(), output); + let first = simd_scalar_first_derivatives( + opcode, + &[left_values[0], right_values[0]], + output_values[0], + )?; + let second = simd_scalar_first_derivatives( + opcode, + &[left_values[1], right_values[1]], + output_values[1], + )?; + Ok(( + _mm_set_pd(second[0], first[0]), + _mm_set_pd(second[1], first[1]), + )) +} + +#[cfg(target_arch = "x86_64")] +unsafe fn simd_f64x2_forward_values( + flat: &[FlatInstruction], + values: &mut Vec<__m128d>, +) -> Result<()> { + use std::arch::x86_64::{ + _mm_add_pd, _mm_and_pd, _mm_andnot_pd, _mm_cmpgt_pd, _mm_div_pd, _mm_mul_pd, _mm_set1_pd, + _mm_setzero_pd, _mm_sqrt_pd, _mm_sub_pd, + }; + + for instruction in flat { + let value = match instruction.opcode { + OpCode::Constant => _mm_set1_pd(instruction.value), + OpCode::Add => _mm_add_pd( + checked_m128_lane(values, instruction.left)?, + checked_m128_lane(values, instruction.right)?, + ), + OpCode::Sub => _mm_sub_pd( + checked_m128_lane(values, instruction.left)?, + checked_m128_lane(values, instruction.right)?, + ), + OpCode::Mul => _mm_mul_pd( + checked_m128_lane(values, instruction.left)?, + checked_m128_lane(values, instruction.right)?, + ), + OpCode::Div => _mm_div_pd( + checked_m128_lane(values, instruction.left)?, + checked_m128_lane(values, instruction.right)?, + ), + OpCode::Pow | OpCode::LogAddExp => { + let left = checked_m128_lane(values, instruction.left)?; + let right = checked_m128_lane(values, instruction.right)?; + simd_f64x2_scalar_binary(left, right, instruction.opcode)? + } + OpCode::Neg => _mm_sub_pd( + _mm_setzero_pd(), + checked_m128_lane(values, instruction.left)?, + ), + OpCode::Sqrt => _mm_sqrt_pd(checked_m128_lane(values, instruction.left)?), + OpCode::Relu => { + let input = checked_m128_lane(values, instruction.left)?; + let mask = _mm_cmpgt_pd(input, _mm_setzero_pd()); + _mm_and_pd(input, mask) + } + OpCode::Abs => { + let input = checked_m128_lane(values, instruction.left)?; + _mm_andnot_pd(_mm_set1_pd(-0.0), input) + } + OpCode::Sin + | OpCode::Cos + | OpCode::Tan + | OpCode::Tanh + | OpCode::Log1pExp + | OpCode::Exp + | OpCode::Ln => { + let input = checked_m128_lane(values, instruction.left)?; + simd_f64x2_scalar_unary(input, instruction.opcode)? + } + }; + values.push(value); + } + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +fn checked_m256_lane(values: &[__m256d], index: NodeId) -> Result<__m256d> { + values + .get(index) + .copied() + .ok_or(AutodiffError::IndexOutOfBounds { + index, + max_index: values.len().saturating_sub(1), + }) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn add_m256_cotangent( + cotangents: &mut [__m256d], + index: NodeId, + contribution: __m256d, +) -> Result<()> { + use std::arch::x86_64::_mm256_add_pd; + + let max_index = cotangents.len().saturating_sub(1); + let target = cotangents + .get_mut(index) + .ok_or(AutodiffError::IndexOutOfBounds { index, max_index })?; + *target = _mm256_add_pd(*target, contribution); + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn active_m256_contribution(current: __m256d, contribution: __m256d) -> __m256d { + use std::arch::x86_64::{_mm256_and_pd, _mm256_cmp_pd, _mm256_setzero_pd, _CMP_NEQ_UQ}; + + let active = _mm256_cmp_pd(current, _mm256_setzero_pd(), _CMP_NEQ_UQ); + _mm256_and_pd(contribution, active) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn simd_f64x4_scalar_unary(input: __m256d, opcode: OpCode) -> Result<__m256d> { + use std::arch::x86_64::{_mm256_set_pd, _mm256_storeu_pd}; + + let mut stored = [0.0_f64; 4]; + _mm256_storeu_pd(stored.as_mut_ptr(), input); + let first = simd_scalar_value(opcode, &[stored[0]])?; + let second = simd_scalar_value(opcode, &[stored[1]])?; + let third = simd_scalar_value(opcode, &[stored[2]])?; + let fourth = simd_scalar_value(opcode, &[stored[3]])?; + Ok(_mm256_set_pd(fourth, third, second, first)) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn simd_f64x4_scalar_unary_derivative( + input: __m256d, + output: __m256d, + opcode: OpCode, +) -> Result<__m256d> { + use std::arch::x86_64::{_mm256_set_pd, _mm256_storeu_pd}; + + let mut input_values = [0.0_f64; 4]; + let mut output_values = [0.0_f64; 4]; + _mm256_storeu_pd(input_values.as_mut_ptr(), input); + _mm256_storeu_pd(output_values.as_mut_ptr(), output); + let first = simd_scalar_first_derivatives(opcode, &[input_values[0]], output_values[0])?[0]; + let second = simd_scalar_first_derivatives(opcode, &[input_values[1]], output_values[1])?[0]; + let third = simd_scalar_first_derivatives(opcode, &[input_values[2]], output_values[2])?[0]; + let fourth = simd_scalar_first_derivatives(opcode, &[input_values[3]], output_values[3])?[0]; + Ok(_mm256_set_pd(fourth, third, second, first)) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn simd_f64x4_scalar_binary( + left: __m256d, + right: __m256d, + opcode: OpCode, +) -> Result<__m256d> { + use std::arch::x86_64::{_mm256_set_pd, _mm256_storeu_pd}; + + let mut left_values = [0.0_f64; 4]; + let mut right_values = [0.0_f64; 4]; + _mm256_storeu_pd(left_values.as_mut_ptr(), left); + _mm256_storeu_pd(right_values.as_mut_ptr(), right); + let first = simd_scalar_value(opcode, &[left_values[0], right_values[0]])?; + let second = simd_scalar_value(opcode, &[left_values[1], right_values[1]])?; + let third = simd_scalar_value(opcode, &[left_values[2], right_values[2]])?; + let fourth = simd_scalar_value(opcode, &[left_values[3], right_values[3]])?; + Ok(_mm256_set_pd(fourth, third, second, first)) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn simd_f64x4_scalar_binary_derivatives( + left: __m256d, + right: __m256d, + output: __m256d, + opcode: OpCode, +) -> Result<(__m256d, __m256d)> { + use std::arch::x86_64::{_mm256_set_pd, _mm256_storeu_pd}; + + let mut left_values = [0.0_f64; 4]; + let mut right_values = [0.0_f64; 4]; + let mut output_values = [0.0_f64; 4]; + _mm256_storeu_pd(left_values.as_mut_ptr(), left); + _mm256_storeu_pd(right_values.as_mut_ptr(), right); + _mm256_storeu_pd(output_values.as_mut_ptr(), output); + let first = simd_scalar_first_derivatives( + opcode, + &[left_values[0], right_values[0]], + output_values[0], + )?; + let second = simd_scalar_first_derivatives( + opcode, + &[left_values[1], right_values[1]], + output_values[1], + )?; + let third = simd_scalar_first_derivatives( + opcode, + &[left_values[2], right_values[2]], + output_values[2], + )?; + let fourth = simd_scalar_first_derivatives( + opcode, + &[left_values[3], right_values[3]], + output_values[3], + )?; + Ok(( + _mm256_set_pd(fourth[0], third[0], second[0], first[0]), + _mm256_set_pd(fourth[1], third[1], second[1], first[1]), + )) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn simd_f64x4_forward_values( + flat: &[FlatInstruction], + values: &mut Vec<__m256d>, +) -> Result<()> { + use std::arch::x86_64::{ + _mm256_add_pd, _mm256_and_pd, _mm256_andnot_pd, _mm256_cmp_pd, _mm256_div_pd, + _mm256_mul_pd, _mm256_set1_pd, _mm256_setzero_pd, _mm256_sqrt_pd, _mm256_sub_pd, + _CMP_GT_OQ, + }; + + for instruction in flat { + let value = match instruction.opcode { + OpCode::Constant => _mm256_set1_pd(instruction.value), + OpCode::Add => _mm256_add_pd( + checked_m256_lane(values, instruction.left)?, + checked_m256_lane(values, instruction.right)?, + ), + OpCode::Sub => _mm256_sub_pd( + checked_m256_lane(values, instruction.left)?, + checked_m256_lane(values, instruction.right)?, + ), + OpCode::Mul => _mm256_mul_pd( + checked_m256_lane(values, instruction.left)?, + checked_m256_lane(values, instruction.right)?, + ), + OpCode::Div => _mm256_div_pd( + checked_m256_lane(values, instruction.left)?, + checked_m256_lane(values, instruction.right)?, + ), + OpCode::Pow | OpCode::LogAddExp => { + let left = checked_m256_lane(values, instruction.left)?; + let right = checked_m256_lane(values, instruction.right)?; + simd_f64x4_scalar_binary(left, right, instruction.opcode)? + } + OpCode::Neg => _mm256_sub_pd( + _mm256_setzero_pd(), + checked_m256_lane(values, instruction.left)?, + ), + OpCode::Sqrt => _mm256_sqrt_pd(checked_m256_lane(values, instruction.left)?), + OpCode::Relu => { + let input = checked_m256_lane(values, instruction.left)?; + let mask = _mm256_cmp_pd(input, _mm256_setzero_pd(), _CMP_GT_OQ); + _mm256_and_pd(input, mask) + } + OpCode::Abs => { + let input = checked_m256_lane(values, instruction.left)?; + _mm256_andnot_pd(_mm256_set1_pd(-0.0), input) + } + OpCode::Sin + | OpCode::Cos + | OpCode::Tan + | OpCode::Tanh + | OpCode::Log1pExp + | OpCode::Exp + | OpCode::Ln => { + let input = checked_m256_lane(values, instruction.left)?; + simd_f64x4_scalar_unary(input, instruction.opcode)? + } + }; + values.push(value); + } + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +pub(crate) fn compute_batch_simd_f64x2( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, +) -> Result<()> { + use std::arch::x86_64::{_mm_set_pd, _mm_storeu_pd}; + + graph.check_batch(batch)?; + let flat = graph.flat_instructions_slice(); + let output_dim = graph.output_nodes.len(); + let value_count = graph.num_inputs + graph.instructions.len(); + buffer.reset(batch.batch_size, output_dim); + + let mut values: Vec<__m128d> = Vec::with_capacity(value_count); + let mut pair_outputs: Vec<[f64; 2]> = Vec::with_capacity(output_dim); + let mut row_index = 0; + while row_index + 1 < batch.batch_size { + let first = batch.try_row(row_index)?; + let second = batch.try_row(row_index + 1)?; + values.clear(); + for input_index in 0..graph.num_inputs { + unsafe { + values.push(_mm_set_pd(second[input_index], first[input_index])); + } + } + + unsafe { + simd_f64x2_forward_values(flat, &mut values)?; + } + + pair_outputs.clear(); + for &output in &graph.output_nodes { + let lane = checked_m128_lane(&values, output)?; + let mut stored = [0.0_f64; 2]; + unsafe { + _mm_storeu_pd(stored.as_mut_ptr(), lane); + } + pair_outputs.push(stored); + } + for output in &pair_outputs { + buffer.data.push(output[0]); + } + for output in &pair_outputs { + buffer.data.push(output[1]); + } + + row_index += 2; + } + + append_scalar_compute_tail(graph, batch, row_index, buffer)?; + + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn compute_batch_simd_f64x4_impl( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, +) -> Result<()> { + use std::arch::x86_64::{_mm256_set_pd, _mm256_storeu_pd}; + + graph.check_batch(batch)?; + let flat = graph.flat_instructions_slice(); + let output_dim = graph.output_nodes.len(); + let value_count = graph.num_inputs + graph.instructions.len(); + buffer.reset(batch.batch_size, output_dim); + + let mut values: Vec<__m256d> = Vec::with_capacity(value_count); + let mut quad_outputs: Vec<[f64; 4]> = Vec::with_capacity(output_dim); + let mut row_index = 0; + while row_index + 3 < batch.batch_size { + let first = batch.try_row(row_index)?; + let second = batch.try_row(row_index + 1)?; + let third = batch.try_row(row_index + 2)?; + let fourth = batch.try_row(row_index + 3)?; + values.clear(); + for input_index in 0..graph.num_inputs { + values.push(_mm256_set_pd( + fourth[input_index], + third[input_index], + second[input_index], + first[input_index], + )); + } + + simd_f64x4_forward_values(flat, &mut values)?; + + quad_outputs.clear(); + for &output in &graph.output_nodes { + let lane = checked_m256_lane(&values, output)?; + let mut stored = [0.0_f64; 4]; + _mm256_storeu_pd(stored.as_mut_ptr(), lane); + quad_outputs.push(stored); + } + for lane_index in 0..4 { + for output in &quad_outputs { + buffer.data.push(output[lane_index]); + } + } + + row_index += 4; + } + + append_scalar_compute_tail(graph, batch, row_index, buffer)?; + + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +pub(crate) fn compute_batch_simd_f64x4( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, +) -> Result<()> { + if !supports_simd_f64x4_runtime() { + return Err(AutodiffError::InvalidGraph { + reason: "simd f64x4 backend requires x86_64 AVX support", + }); + } + unsafe { compute_batch_simd_f64x4_impl(graph, batch, buffer) } +} + +#[cfg(not(target_arch = "x86_64"))] +pub(crate) fn compute_batch_simd_f64x4( + _graph: &CompiledGraph, + _batch: BatchInputs<'_>, + _buffer: &mut BatchValuesBuffer, +) -> Result<()> { + Err(AutodiffError::InvalidGraph { + reason: "simd f64x4 backend requires x86_64 AVX support", + }) +} + +#[cfg(not(target_arch = "x86_64"))] +pub(crate) fn compute_batch_simd_f64x2( + _graph: &CompiledGraph, + _batch: BatchInputs<'_>, + _buffer: &mut BatchValuesBuffer, +) -> Result<()> { + Err(AutodiffError::InvalidGraph { + reason: "simd backend requires x86_64 SSE2 support", + }) +} + +#[cfg(target_arch = "x86_64")] +pub(crate) fn gradient_batch_simd_f64x2( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, +) -> Result<()> { + use std::arch::x86_64::{ + _mm_add_pd, _mm_and_pd, _mm_cmpgt_pd, _mm_cmplt_pd, _mm_div_pd, _mm_mul_pd, _mm_set1_pd, + _mm_set_pd, _mm_setzero_pd, _mm_storeu_pd, _mm_sub_pd, + }; + + graph.check_batch(batch)?; + let flat = graph.flat_instructions_slice(); + buffer.reset(batch.batch_size, graph.num_inputs); + + let value_count = graph.num_inputs + graph.instructions.len(); + let mut values: Vec<__m128d> = Vec::with_capacity(value_count); + let mut cotangents: Vec<__m128d> = vec![unsafe { _mm_setzero_pd() }; value_count]; + let mut gradient_pairs = Vec::with_capacity(graph.num_inputs); + let mut row_index = 0; + while row_index + 1 < batch.batch_size { + let first = batch.try_row(row_index)?; + let second = batch.try_row(row_index + 1)?; + values.clear(); + for input_index in 0..graph.num_inputs { + unsafe { + values.push(_mm_set_pd(second[input_index], first[input_index])); + } + } + + unsafe { + simd_f64x2_forward_values(flat, &mut values)?; + } + + let Some(&output) = graph.output_nodes.first() else { + buffer.values.extend_from_slice(&[0.0, 0.0]); + buffer + .gradients + .resize(buffer.gradients.len() + graph.num_inputs * 2, 0.0); + row_index += 2; + continue; + }; + + // SAFETY: _mm_setzero_pd is safe; fill zeros before each backward pass. + cotangents.fill(unsafe { _mm_setzero_pd() }); + let max_index = cotangents.len().saturating_sub(1); + *cotangents + .get_mut(output) + .ok_or(AutodiffError::IndexOutOfBounds { + index: output, + max_index, + })? = unsafe { _mm_set1_pd(1.0) }; + + for instruction in flat.iter().rev() { + let current = checked_m128_lane(&cotangents, instruction.output)?; + unsafe { + match instruction.opcode { + OpCode::Constant => {} + OpCode::Add => { + let contribution = active_m128_contribution(current, current); + add_m128_cotangent(&mut cotangents, instruction.left, contribution)?; + add_m128_cotangent(&mut cotangents, instruction.right, contribution)?; + } + OpCode::Sub => { + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, current), + )?; + add_m128_cotangent( + &mut cotangents, + instruction.right, + active_m128_contribution( + current, + _mm_sub_pd(_mm_setzero_pd(), current), + ), + )?; + } + OpCode::Mul => { + let left = checked_m128_lane(&values, instruction.left)?; + let right = checked_m128_lane(&values, instruction.right)?; + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, _mm_mul_pd(current, right)), + )?; + add_m128_cotangent( + &mut cotangents, + instruction.right, + active_m128_contribution(current, _mm_mul_pd(current, left)), + )?; + } + OpCode::Div => { + let left = checked_m128_lane(&values, instruction.left)?; + let right = checked_m128_lane(&values, instruction.right)?; + let right_squared = _mm_mul_pd(right, right); + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, _mm_div_pd(current, right)), + )?; + let right_contribution = _mm_sub_pd( + _mm_setzero_pd(), + _mm_div_pd(_mm_mul_pd(current, left), right_squared), + ); + add_m128_cotangent( + &mut cotangents, + instruction.right, + active_m128_contribution(current, right_contribution), + )?; + } + OpCode::Pow | OpCode::LogAddExp => { + let left = checked_m128_lane(&values, instruction.left)?; + let right = checked_m128_lane(&values, instruction.right)?; + let output_value = checked_m128_lane(&values, instruction.output)?; + let (left_derivative, right_derivative) = + simd_f64x2_scalar_binary_derivatives( + left, + right, + output_value, + instruction.opcode, + )?; + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, _mm_mul_pd(current, left_derivative)), + )?; + add_m128_cotangent( + &mut cotangents, + instruction.right, + active_m128_contribution( + current, + _mm_mul_pd(current, right_derivative), + ), + )?; + } + OpCode::Neg => { + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution( + current, + _mm_sub_pd(_mm_setzero_pd(), current), + ), + )?; + } + OpCode::Sqrt => { + let output_value = checked_m128_lane(&values, instruction.output)?; + let contribution = + _mm_div_pd(_mm_mul_pd(current, _mm_set1_pd(0.5)), output_value); + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, contribution), + )?; + } + OpCode::Relu => { + let input_value = checked_m128_lane(&values, instruction.left)?; + let mask = _mm_cmpgt_pd(input_value, _mm_setzero_pd()); + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, _mm_and_pd(current, mask)), + )?; + } + OpCode::Sin + | OpCode::Cos + | OpCode::Tan + | OpCode::Tanh + | OpCode::Log1pExp + | OpCode::Exp + | OpCode::Ln => { + let input_value = checked_m128_lane(&values, instruction.left)?; + let output_value = checked_m128_lane(&values, instruction.output)?; + let derivative = simd_f64x2_scalar_unary_derivative( + input_value, + output_value, + instruction.opcode, + )?; + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, _mm_mul_pd(current, derivative)), + )?; + } + OpCode::Abs => { + let input_value = checked_m128_lane(&values, instruction.left)?; + let positive = _mm_cmpgt_pd(input_value, _mm_setzero_pd()); + let negative = _mm_cmplt_pd(input_value, _mm_setzero_pd()); + let sign = _mm_add_pd( + _mm_and_pd(_mm_set1_pd(1.0), positive), + _mm_and_pd(_mm_set1_pd(-1.0), negative), + ); + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, _mm_mul_pd(current, sign)), + )?; + } + } + } + } + + let output_value = checked_m128_lane(&values, output)?; + let mut value_pair = [0.0_f64; 2]; + unsafe { + _mm_storeu_pd(value_pair.as_mut_ptr(), output_value); + } + buffer.values.extend_from_slice(&value_pair); + + gradient_pairs.clear(); + for input_index in 0..graph.num_inputs { + let gradient_lane = checked_m128_lane(&cotangents, input_index)?; + let mut stored = [0.0_f64; 2]; + unsafe { + _mm_storeu_pd(stored.as_mut_ptr(), gradient_lane); + } + gradient_pairs.push(stored); + } + for pair in &gradient_pairs { + buffer.gradients.push(pair[0]); + } + for pair in &gradient_pairs { + buffer.gradients.push(pair[1]); + } + + row_index += 2; + } + + append_scalar_gradient_tail(graph, batch, row_index, buffer)?; + + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn gradient_batch_simd_f64x4_impl( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, +) -> Result<()> { + use std::arch::x86_64::{ + _mm256_add_pd, _mm256_and_pd, _mm256_cmp_pd, _mm256_div_pd, _mm256_mul_pd, _mm256_set1_pd, + _mm256_set_pd, _mm256_setzero_pd, _mm256_storeu_pd, _mm256_sub_pd, _CMP_GT_OQ, _CMP_LT_OQ, + }; + + graph.check_batch(batch)?; + let flat = graph.flat_instructions_slice(); + buffer.reset(batch.batch_size, graph.num_inputs); + + let value_count = graph.num_inputs + graph.instructions.len(); + let mut values: Vec<__m256d> = Vec::with_capacity(value_count); + let mut cotangents: Vec<__m256d> = vec![_mm256_setzero_pd(); value_count]; + let mut gradient_quads = Vec::with_capacity(graph.num_inputs); + let mut row_index = 0; + while row_index + 3 < batch.batch_size { + let first = batch.try_row(row_index)?; + let second = batch.try_row(row_index + 1)?; + let third = batch.try_row(row_index + 2)?; + let fourth = batch.try_row(row_index + 3)?; + values.clear(); + for input_index in 0..graph.num_inputs { + values.push(_mm256_set_pd( + fourth[input_index], + third[input_index], + second[input_index], + first[input_index], + )); + } + + simd_f64x4_forward_values(flat, &mut values)?; + + let Some(&output) = graph.output_nodes.first() else { + buffer.values.extend_from_slice(&[0.0; 4]); + buffer + .gradients + .resize(buffer.gradients.len() + graph.num_inputs * 4, 0.0); + row_index += 4; + continue; + }; + + cotangents.fill(_mm256_setzero_pd()); + let max_index = cotangents.len().saturating_sub(1); + *cotangents + .get_mut(output) + .ok_or(AutodiffError::IndexOutOfBounds { + index: output, + max_index, + })? = _mm256_set1_pd(1.0); + + for instruction in flat.iter().rev() { + let current = checked_m256_lane(&cotangents, instruction.output)?; + match instruction.opcode { + OpCode::Constant => {} + OpCode::Add => { + let contribution = active_m256_contribution(current, current); + add_m256_cotangent(&mut cotangents, instruction.left, contribution)?; + add_m256_cotangent(&mut cotangents, instruction.right, contribution)?; + } + OpCode::Sub => { + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, current), + )?; + add_m256_cotangent( + &mut cotangents, + instruction.right, + active_m256_contribution( + current, + _mm256_sub_pd(_mm256_setzero_pd(), current), + ), + )?; + } + OpCode::Mul => { + let left = checked_m256_lane(&values, instruction.left)?; + let right = checked_m256_lane(&values, instruction.right)?; + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, _mm256_mul_pd(current, right)), + )?; + add_m256_cotangent( + &mut cotangents, + instruction.right, + active_m256_contribution(current, _mm256_mul_pd(current, left)), + )?; + } + OpCode::Div => { + let left = checked_m256_lane(&values, instruction.left)?; + let right = checked_m256_lane(&values, instruction.right)?; + let right_squared = _mm256_mul_pd(right, right); + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, _mm256_div_pd(current, right)), + )?; + let right_contribution = _mm256_sub_pd( + _mm256_setzero_pd(), + _mm256_div_pd(_mm256_mul_pd(current, left), right_squared), + ); + add_m256_cotangent( + &mut cotangents, + instruction.right, + active_m256_contribution(current, right_contribution), + )?; + } + OpCode::Pow | OpCode::LogAddExp => { + let left = checked_m256_lane(&values, instruction.left)?; + let right = checked_m256_lane(&values, instruction.right)?; + let output_value = checked_m256_lane(&values, instruction.output)?; + let (left_derivative, right_derivative) = simd_f64x4_scalar_binary_derivatives( + left, + right, + output_value, + instruction.opcode, + )?; + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, _mm256_mul_pd(current, left_derivative)), + )?; + add_m256_cotangent( + &mut cotangents, + instruction.right, + active_m256_contribution(current, _mm256_mul_pd(current, right_derivative)), + )?; + } + OpCode::Neg => { + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution( + current, + _mm256_sub_pd(_mm256_setzero_pd(), current), + ), + )?; + } + OpCode::Sqrt => { + let output_value = checked_m256_lane(&values, instruction.output)?; + let contribution = + _mm256_div_pd(_mm256_mul_pd(current, _mm256_set1_pd(0.5)), output_value); + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, contribution), + )?; + } + OpCode::Relu => { + let input_value = checked_m256_lane(&values, instruction.left)?; + let mask = _mm256_cmp_pd(input_value, _mm256_setzero_pd(), _CMP_GT_OQ); + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, _mm256_and_pd(current, mask)), + )?; + } + OpCode::Sin + | OpCode::Cos + | OpCode::Tan + | OpCode::Tanh + | OpCode::Log1pExp + | OpCode::Exp + | OpCode::Ln => { + let input_value = checked_m256_lane(&values, instruction.left)?; + let output_value = checked_m256_lane(&values, instruction.output)?; + let derivative = simd_f64x4_scalar_unary_derivative( + input_value, + output_value, + instruction.opcode, + )?; + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, _mm256_mul_pd(current, derivative)), + )?; + } + OpCode::Abs => { + let input_value = checked_m256_lane(&values, instruction.left)?; + let positive = _mm256_cmp_pd(input_value, _mm256_setzero_pd(), _CMP_GT_OQ); + let negative = _mm256_cmp_pd(input_value, _mm256_setzero_pd(), _CMP_LT_OQ); + let sign = _mm256_add_pd( + _mm256_and_pd(_mm256_set1_pd(1.0), positive), + _mm256_and_pd(_mm256_set1_pd(-1.0), negative), + ); + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, _mm256_mul_pd(current, sign)), + )?; + } + } + } + + let output_value = checked_m256_lane(&values, output)?; + let mut value_quad = [0.0_f64; 4]; + _mm256_storeu_pd(value_quad.as_mut_ptr(), output_value); + buffer.values.extend_from_slice(&value_quad); + + gradient_quads.clear(); + for input_index in 0..graph.num_inputs { + let gradient_lane = checked_m256_lane(&cotangents, input_index)?; + let mut stored = [0.0_f64; 4]; + _mm256_storeu_pd(stored.as_mut_ptr(), gradient_lane); + gradient_quads.push(stored); + } + for lane_index in 0..4 { + for quad in &gradient_quads { + buffer.gradients.push(quad[lane_index]); + } + } + + row_index += 4; + } + + append_scalar_gradient_tail(graph, batch, row_index, buffer)?; + + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +pub(crate) fn gradient_batch_simd_f64x4( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, +) -> Result<()> { + if !supports_simd_f64x4_runtime() { + return Err(AutodiffError::InvalidGraph { + reason: "simd f64x4 backend requires x86_64 AVX support", + }); + } + unsafe { gradient_batch_simd_f64x4_impl(graph, batch, buffer) } +} + +#[cfg(not(target_arch = "x86_64"))] +pub(crate) fn gradient_batch_simd_f64x4( + _graph: &CompiledGraph, + _batch: BatchInputs<'_>, + _buffer: &mut BatchGradientsBuffer, +) -> Result<()> { + Err(AutodiffError::InvalidGraph { + reason: "simd f64x4 backend requires x86_64 AVX support", + }) +} + +#[cfg(not(target_arch = "x86_64"))] +pub(crate) fn gradient_batch_simd_f64x2( + _graph: &CompiledGraph, + _batch: BatchInputs<'_>, + _buffer: &mut BatchGradientsBuffer, +) -> Result<()> { + Err(AutodiffError::InvalidGraph { + reason: "simd backend requires x86_64 SSE2 support", + }) +} diff --git a/src/multi/backend/types.rs b/src/multi/backend/types.rs new file mode 100644 index 0000000..3032085 --- /dev/null +++ b/src/multi/backend/types.rs @@ -0,0 +1,312 @@ +//! Core type definitions for the backend abstraction. + +use crate::multi::multi_ad::MultiAD; +use crate::{AutodiffError, NodeId, Result}; + +/// One closure-free instruction in a compiled scalar graph. +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum Instruction { + /// Literal constant node. + Constant(f64), + /// Unary operation with one input node id. + Unary { op: MultiAD, arg: NodeId }, + /// Binary operation with two input node ids. + Binary { + op: MultiAD, + left: NodeId, + right: NodeId, + }, +} + +/// Sentinel node id used by flat backend instructions when an argument is absent. +pub const UNUSED_NODE_ID: NodeId = usize::MAX; + +/// Compact operation code used by flat backend instructions. +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum OpCode { + Constant, + Add, + Sub, + Mul, + Div, + Pow, + Sin, + Cos, + Tan, + Tanh, + Relu, + Log1pExp, + LogAddExp, + Neg, + Exp, + Ln, + Sqrt, + Abs, +} + +impl OpCode { + pub(crate) fn from_multi_ad(op: MultiAD) -> Result { + Ok(match op { + MultiAD::Add => OpCode::Add, + MultiAD::Sub => OpCode::Sub, + MultiAD::Mul => OpCode::Mul, + MultiAD::Div => OpCode::Div, + MultiAD::Pow => OpCode::Pow, + MultiAD::Sin => OpCode::Sin, + MultiAD::Cos => OpCode::Cos, + MultiAD::Tan => OpCode::Tan, + MultiAD::Tanh => OpCode::Tanh, + MultiAD::Relu => OpCode::Relu, + MultiAD::Log1pExp => OpCode::Log1pExp, + MultiAD::LogAddExp => OpCode::LogAddExp, + MultiAD::Neg => OpCode::Neg, + MultiAD::Exp => OpCode::Exp, + MultiAD::Ln => OpCode::Ln, + MultiAD::Sqrt => OpCode::Sqrt, + MultiAD::Abs => OpCode::Abs, + MultiAD::Inp => { + return Err(AutodiffError::InvalidGraph { + reason: "input markers are not backend instructions", + }); + } + }) + } + + pub(crate) fn to_multi_ad(self) -> Option { + Some(match self { + OpCode::Constant => return None, + OpCode::Add => MultiAD::Add, + OpCode::Sub => MultiAD::Sub, + OpCode::Mul => MultiAD::Mul, + OpCode::Div => MultiAD::Div, + OpCode::Pow => MultiAD::Pow, + OpCode::Sin => MultiAD::Sin, + OpCode::Cos => MultiAD::Cos, + OpCode::Tan => MultiAD::Tan, + OpCode::Tanh => MultiAD::Tanh, + OpCode::Relu => MultiAD::Relu, + OpCode::Log1pExp => MultiAD::Log1pExp, + OpCode::LogAddExp => MultiAD::LogAddExp, + OpCode::Neg => MultiAD::Neg, + OpCode::Exp => MultiAD::Exp, + OpCode::Ln => MultiAD::Ln, + OpCode::Sqrt => MultiAD::Sqrt, + OpCode::Abs => MultiAD::Abs, + }) + } + + /// Return the number of input slots consumed by this opcode. + #[must_use] + pub fn arity(&self) -> usize { + match self { + OpCode::Constant => 0, + OpCode::Sin + | OpCode::Cos + | OpCode::Tan + | OpCode::Tanh + | OpCode::Relu + | OpCode::Log1pExp + | OpCode::Neg + | OpCode::Exp + | OpCode::Ln + | OpCode::Sqrt + | OpCode::Abs => 1, + OpCode::Add + | OpCode::Sub + | OpCode::Mul + | OpCode::Div + | OpCode::Pow + | OpCode::LogAddExp => 2, + } + } +} + +/// Fully explicit instruction shape for SIMD/GPU backend lowering. +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct FlatInstruction { + pub opcode: OpCode, + pub output: NodeId, + pub left: NodeId, + pub right: NodeId, + pub value: f64, +} + +#[cfg(target_arch = "x86_64")] +#[inline] +pub(crate) fn supports_simd_f64x2_runtime() -> bool { + true +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +pub(crate) fn supports_simd_f64x2_runtime() -> bool { + false +} + +#[cfg(target_arch = "x86_64")] +#[inline] +pub(crate) fn supports_simd_f64x4_runtime() -> bool { + std::is_x86_feature_detected!("avx") +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +pub(crate) fn supports_simd_f64x4_runtime() -> bool { + false +} + +#[inline] +pub(crate) fn supports_wgpu_runtime() -> bool { + cfg!(feature = "backend-wgpu") +} + +/// Backend feature declaration used before dispatching compiled graphs. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BackendCapabilities { + pub supports_f64: bool, + pub supports_f32: bool, + pub supports_constants: bool, + pub supports_unary: bool, + pub supports_binary: bool, + pub supports_multi_output: bool, + pub supports_reverse_gradient: bool, + pub supports_batch_compute: bool, + pub supports_batch_gradient: bool, + pub supported_opcodes: Vec, +} + +impl BackendCapabilities { + /// Capabilities for the reference scalar f64 backend. + #[must_use] + pub fn scalar_f64() -> Self { + Self { + supports_f64: true, + supports_f32: false, + supports_constants: true, + supports_unary: true, + supports_binary: true, + supports_multi_output: true, + supports_reverse_gradient: true, + supports_batch_compute: true, + supports_batch_gradient: true, + supported_opcodes: vec![ + OpCode::Constant, + OpCode::Add, + OpCode::Sub, + OpCode::Mul, + OpCode::Div, + OpCode::Pow, + OpCode::Sin, + OpCode::Cos, + OpCode::Tan, + OpCode::Tanh, + OpCode::Relu, + OpCode::Log1pExp, + OpCode::LogAddExp, + OpCode::Neg, + OpCode::Exp, + OpCode::Ln, + OpCode::Sqrt, + OpCode::Abs, + ], + } + } + + /// Capabilities for the prototype f64x2 SIMD batch backend. + #[must_use] + pub fn simd_f64x2() -> Self { + Self { + supports_f64: true, + supports_f32: false, + supports_constants: true, + supports_unary: true, + supports_binary: true, + supports_multi_output: true, + supports_reverse_gradient: false, + supports_batch_compute: supports_simd_f64x2_runtime(), + supports_batch_gradient: supports_simd_f64x2_runtime(), + supported_opcodes: vec![ + OpCode::Constant, + OpCode::Add, + OpCode::Sub, + OpCode::Mul, + OpCode::Div, + OpCode::Pow, + OpCode::Sin, + OpCode::Cos, + OpCode::Tan, + OpCode::Tanh, + OpCode::Log1pExp, + OpCode::LogAddExp, + OpCode::Neg, + OpCode::Exp, + OpCode::Ln, + OpCode::Sqrt, + OpCode::Relu, + OpCode::Abs, + ], + } + } + + /// Capabilities for the prototype f64x4 SIMD batch backend. + #[must_use] + pub fn simd_f64x4() -> Self { + Self { + supports_f64: true, + supports_f32: false, + supports_constants: true, + supports_unary: true, + supports_binary: true, + supports_multi_output: true, + supports_reverse_gradient: false, + supports_batch_compute: supports_simd_f64x4_runtime(), + supports_batch_gradient: supports_simd_f64x4_runtime(), + supported_opcodes: vec![ + OpCode::Constant, + OpCode::Add, + OpCode::Sub, + OpCode::Mul, + OpCode::Div, + OpCode::Pow, + OpCode::Sin, + OpCode::Cos, + OpCode::Tan, + OpCode::Tanh, + OpCode::Log1pExp, + OpCode::LogAddExp, + OpCode::Neg, + OpCode::Exp, + OpCode::Ln, + OpCode::Sqrt, + OpCode::Relu, + OpCode::Abs, + ], + } + } + + /// Capabilities for the feature-gated WGPU batch backend skeleton. + #[must_use] + pub fn wgpu_f64() -> Self { + Self { + supports_f64: true, + supports_f32: false, + supports_constants: true, + supports_unary: true, + supports_binary: true, + supports_multi_output: true, + supports_reverse_gradient: true, + supports_batch_compute: supports_wgpu_runtime(), + supports_batch_gradient: supports_wgpu_runtime(), + supported_opcodes: BackendCapabilities::scalar_f64().supported_opcodes, + } + } + + /// Return whether the backend declares support for an opcode. + #[must_use] + pub fn supports_opcode(&self, opcode: OpCode) -> bool { + self.supported_opcodes.contains(&opcode) + } +} diff --git a/src/multi/backend/wgpu.rs b/src/multi/backend/wgpu.rs new file mode 100644 index 0000000..2b57109 --- /dev/null +++ b/src/multi/backend/wgpu.rs @@ -0,0 +1,974 @@ +//! WGPU backend skeleton with real device allocation and transfer plumbing. + +#[cfg(feature = "backend-wgpu")] +use std::sync::mpsc; + +#[cfg(feature = "backend-wgpu")] +use pollster::block_on; +#[cfg(feature = "backend-wgpu")] +use wgpu::{self, BufferUsages}; + +#[cfg(feature = "backend-wgpu")] +use crate::multi::backend::device::{ + AcceleratorDeviceContext, AcceleratorDeviceKind, DeviceBackend, DeviceBatchPlan, + DeviceBufferHandle, DeviceBufferKind, DeviceExecutionMode, DeviceExecutionTrace, + DeviceTransferKind, DeviceTransferPolicy, GpuBackendBoundary, +}; +#[cfg(feature = "backend-wgpu")] +use crate::multi::backend::dispatch::{BackendKind, ExecutionBackend}; +#[cfg(feature = "backend-wgpu")] +use crate::multi::backend::scalar::ScalarBackend; +#[cfg(feature = "backend-wgpu")] +use crate::multi::backend::types::{BackendCapabilities, OpCode, UNUSED_NODE_ID}; +#[cfg(feature = "backend-wgpu")] +use crate::multi::compiled::{BatchGradientsBuffer, BatchInputs, BatchValuesBuffer, CompiledGraph}; +#[cfg(feature = "backend-wgpu")] +use crate::{AutodiffError, Result}; + +#[cfg(feature = "backend-wgpu")] +/// Initialized WGPU backend skeleton with real device allocation and transfer plumbing. +#[derive(Debug, Clone)] +pub struct WgpuBackend { + pub(crate) boundary: GpuBackendBoundary, + pub(crate) device: wgpu::Device, + pub(crate) queue: wgpu::Queue, + pub(crate) adapter_info: wgpu::AdapterInfo, +} + +#[cfg(feature = "backend-wgpu")] +/// One real WGPU buffer allocated for a logical batch-plan role. +#[derive(Debug, Clone)] +pub struct WgpuBuffer { + pub(crate) handle: DeviceBufferHandle, + pub(crate) buffer: wgpu::Buffer, +} + +#[cfg(feature = "backend-wgpu")] +impl WgpuBuffer { + /// Return the immutable planned buffer handle. + #[must_use] + pub fn handle(&self) -> DeviceBufferHandle { + self.handle + } +} + +#[cfg(feature = "backend-wgpu")] +/// Owned WGPU buffers allocated from a [`DeviceBatchPlan`]. +#[derive(Debug, Clone)] +pub struct WgpuBufferSet { + pub(crate) plan: DeviceBatchPlan, + pub(crate) buffers: Vec, +} + +#[cfg(feature = "backend-wgpu")] +impl WgpuBufferSet { + /// Return the plan used to allocate this buffer set. + #[must_use] + pub fn plan(&self) -> &DeviceBatchPlan { + &self.plan + } + + /// Return all allocated GPU buffers in plan order. + #[must_use] + pub fn buffers(&self) -> &[WgpuBuffer] { + &self.buffers + } + + /// Return an immutable GPU buffer by logical kind. + pub fn buffer(&self, kind: DeviceBufferKind) -> Result<&WgpuBuffer> { + self.buffers + .iter() + .find(|buffer| buffer.handle.kind == kind) + .ok_or(AutodiffError::InvalidGraph { + reason: "wgpu buffer kind is not in the plan", + }) + } + + /// Upload host data into a planned GPU buffer. + pub fn upload( + &self, + backend: &WgpuBackend, + kind: DeviceBufferKind, + data: &[f64], + ) -> Result<()> { + backend.upload_buffer(self, kind, data) + } + + /// Download a planned GPU buffer into an owned vector. + pub fn download(&self, backend: &WgpuBackend, kind: DeviceBufferKind) -> Result> { + backend.download_buffer(self, kind) + } +} + +#[cfg(feature = "backend-wgpu")] +#[inline] +fn wgpu_buffer_size_bytes(len: usize) -> u64 { + let logical = len.saturating_mul(std::mem::size_of::()) as u64; + logical.max(8) +} + +#[cfg(feature = "backend-wgpu")] +fn encode_f64_slice(data: &[f64]) -> Vec { + let mut bytes = Vec::with_capacity(data.len().saturating_mul(std::mem::size_of::())); + for value in data { + bytes.extend_from_slice(&value.to_le_bytes()); + } + bytes +} + +#[cfg(feature = "backend-wgpu")] +fn decode_f64_bytes(bytes: &[u8], len: usize) -> Result> { + let expected_len = len.saturating_mul(std::mem::size_of::()); + if bytes.len() < expected_len { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu readback length does not match planned buffer length", + }); + } + let mut values = Vec::with_capacity(len); + for chunk in bytes[..expected_len].chunks_exact(std::mem::size_of::()) { + let mut array = [0_u8; std::mem::size_of::()]; + array.copy_from_slice(chunk); + values.push(f64::from_le_bytes(array)); + } + Ok(values) +} + +#[cfg(feature = "backend-wgpu")] +#[inline] +fn wgpu_buffer_size_bytes_f32(len: usize) -> u64 { + let logical = len.saturating_mul(std::mem::size_of::()) as u64; + logical.max(4) +} + +#[cfg(feature = "backend-wgpu")] +fn encode_f32_slice(data: &[f32]) -> Vec { + let mut bytes = Vec::with_capacity(data.len().saturating_mul(std::mem::size_of::())); + for value in data { + bytes.extend_from_slice(&value.to_le_bytes()); + } + bytes +} + +#[cfg(feature = "backend-wgpu")] +fn decode_f32_bytes(bytes: &[u8], len: usize) -> Result> { + let expected_len = len.saturating_mul(std::mem::size_of::()); + if bytes.len() < expected_len { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu readback length does not match requested f32 buffer length", + }); + } + let mut values = Vec::with_capacity(len); + for chunk in bytes[..expected_len].chunks_exact(std::mem::size_of::()) { + let mut array = [0_u8; std::mem::size_of::()]; + array.copy_from_slice(chunk); + values.push(f32::from_le_bytes(array)); + } + Ok(values) +} + +#[cfg(feature = "backend-wgpu")] +fn encode_u32_slice(data: &[u32]) -> Vec { + let mut bytes = Vec::with_capacity(data.len().saturating_mul(std::mem::size_of::())); + for value in data { + bytes.extend_from_slice(&value.to_le_bytes()); + } + bytes +} + +#[cfg(feature = "backend-wgpu")] +/// Conservative opcode subset currently supported by the exact-safe native WGPU batch-compute path. +pub const WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES: [OpCode; 4] = + [OpCode::Constant, OpCode::Neg, OpCode::Relu, OpCode::Abs]; + +#[cfg(feature = "backend-wgpu")] +fn checked_u32_from_usize(value: usize, reason: &'static str) -> Result { + u32::try_from(value).map_err(|_| AutodiffError::InvalidGraph { reason }) +} + +#[cfg(feature = "backend-wgpu")] +#[inline] +fn is_exact_f32_roundtrip(value: f64) -> bool { + if value.is_nan() { + return false; + } + let narrowed = value as f32; + let widened = f64::from(narrowed); + widened == value && (value != 0.0 || widened.is_sign_negative() == value.is_sign_negative()) +} + +#[cfg(feature = "backend-wgpu")] +fn wgpu_native_exact_safe_supports_opcode(opcode: OpCode) -> bool { + WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES.contains(&opcode) +} + +#[cfg(feature = "backend-wgpu")] +fn wgpu_native_exact_safe_graph(graph: &CompiledGraph) -> bool { + !graph.output_nodes().is_empty() + && graph.flat_instructions_slice().iter().all(|instruction| { + wgpu_native_exact_safe_supports_opcode(instruction.opcode) + && (instruction.opcode != OpCode::Constant + || is_exact_f32_roundtrip(instruction.value)) + }) +} + +#[cfg(feature = "backend-wgpu")] +fn wgpu_native_exact_safe_batch(batch: BatchInputs<'_>) -> bool { + batch.data.iter().copied().all(is_exact_f32_roundtrip) +} + +#[cfg(feature = "backend-wgpu")] +const WGPU_NATIVE_WORDS_PER_INSTRUCTION: usize = 8; + +#[cfg(feature = "backend-wgpu")] +const WGPU_NATIVE_SHADER: &str = r#" +@group(0) @binding(0) var input_data: array; +@group(0) @binding(1) var instruction_words: array; +@group(0) @binding(2) var output_nodes: array; +@group(0) @binding(3) var value_data: array; +@group(0) @binding(4) var output_data: array; +@group(0) @binding(5) var kernel_meta: array; + +fn relu_scalar(x: f32) -> f32 { + if x > 0.0 { + return x; + } + return 0.0; +} + +fn log1p_exp_scalar(x: f32) -> f32 { + if x > 0.0 { + return x + log(1.0 + exp(-x)); + } + return log(1.0 + exp(x)); +} + +fn log_add_exp_scalar(a: f32, b: f32) -> f32 { + var max_value = a; + var min_value = b; + if a < b { + max_value = b; + min_value = a; + } + return max_value + log(1.0 + exp(min_value - max_value)); +} + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) gid: vec3) { + let row = gid.x; + let num_inputs = kernel_meta[0u]; + let num_instructions = kernel_meta[1u]; + let num_outputs = kernel_meta[2u]; + let value_count = kernel_meta[3u]; + let batch_size = kernel_meta[4u]; + if row >= batch_size { + return; + } + + let input_base = row * num_inputs; + let value_base = row * value_count; + let output_base = row * num_outputs; + + var input_index = 0u; + loop { + if input_index >= num_inputs { + break; + } + value_data[value_base + input_index] = input_data[input_base + input_index]; + input_index = input_index + 1u; + } + + var instruction_index = 0u; + loop { + if instruction_index >= num_instructions { + break; + } + let word_base = instruction_index * 8u; + let opcode = instruction_words[word_base + 0u]; + let output_index = instruction_words[word_base + 1u]; + let left_index = instruction_words[word_base + 2u]; + let right_index = instruction_words[word_base + 3u]; + let value_bits = instruction_words[word_base + 4u]; + var left_value: f32 = 0.0; + if left_index != 4294967295u { + left_value = value_data[value_base + left_index]; + } + var right_value: f32 = 0.0; + if right_index != 4294967295u { + right_value = value_data[value_base + right_index]; + } + var result_value: f32 = 0.0; + switch opcode { + case 0u: { + result_value = bitcast(value_bits); + } + case 1u: { + result_value = left_value + right_value; + } + case 2u: { + result_value = left_value - right_value; + } + case 3u: { + result_value = left_value * right_value; + } + case 4u: { + result_value = left_value / right_value; + } + case 5u: { + result_value = pow(left_value, right_value); + } + case 6u: { + result_value = sin(left_value); + } + case 7u: { + result_value = cos(left_value); + } + case 8u: { + result_value = tan(left_value); + } + case 9u: { + result_value = tanh(left_value); + } + case 10u: { + result_value = relu_scalar(left_value); + } + case 11u: { + result_value = log1p_exp_scalar(left_value); + } + case 12u: { + result_value = log_add_exp_scalar(left_value, right_value); + } + case 13u: { + result_value = -left_value; + } + case 14u: { + result_value = exp(left_value); + } + case 15u: { + result_value = log(left_value); + } + case 16u: { + result_value = sqrt(left_value); + } + case 17u: { + result_value = abs(left_value); + } + default: { + result_value = 0.0; + } + } + value_data[value_base + output_index] = result_value; + instruction_index = instruction_index + 1u; + } + + var output_index = 0u; + loop { + if output_index >= num_outputs { + break; + } + let node = output_nodes[output_index]; + output_data[output_base + output_index] = value_data[value_base + node]; + output_index = output_index + 1u; + } +} +"#; + +#[cfg(feature = "backend-wgpu")] +impl WgpuBackend { + /// Create the default WGPU backend using device id 0 and automatic transfers. + pub fn new_default() -> Result { + Self::new( + AcceleratorDeviceContext::wgpu(0), + DeviceTransferPolicy::Automatic, + ) + } + + /// Create the WGPU backend skeleton from a context and transfer policy. + pub fn new( + context: AcceleratorDeviceContext, + transfer_policy: DeviceTransferPolicy, + ) -> Result { + Self::from_boundary(GpuBackendBoundary::new(context, transfer_policy)) + } + + /// Create the WGPU backend skeleton from a boundary descriptor. + pub fn from_boundary(boundary: GpuBackendBoundary) -> Result { + if boundary.context.kind != AcceleratorDeviceKind::Wgpu { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu backend requires a WGPU accelerator context", + }); + } + let instance = wgpu::Instance::default(); + let adapter = block_on(instance.request_adapter(&wgpu::RequestAdapterOptions { + power_preference: wgpu::PowerPreference::default(), + force_fallback_adapter: false, + compatible_surface: None, + })) + .map_err(|_| AutodiffError::InvalidGraph { + reason: "wgpu adapter request failed", + })?; + let adapter_info = adapter.get_info(); + let (device, queue) = block_on(adapter.request_device(&wgpu::DeviceDescriptor { + label: Some("petite-ad-wgpu"), + required_features: wgpu::Features::empty(), + required_limits: wgpu::Limits::default(), + memory_hints: wgpu::MemoryHints::Performance, + trace: wgpu::Trace::Off, + })) + .map_err(|_| AutodiffError::InvalidGraph { + reason: "wgpu device request failed", + })?; + Ok(Self { + boundary, + device, + queue, + adapter_info, + }) + } + + /// Return the boundary descriptor used for this backend. + #[must_use] + pub fn boundary(&self) -> &GpuBackendBoundary { + &self.boundary + } + + /// Return the accelerator context used for this backend. + #[must_use] + pub fn context(&self) -> &AcceleratorDeviceContext { + &self.boundary.context + } + + /// Return the configured transfer policy. + #[must_use] + pub fn transfer_policy(&self) -> DeviceTransferPolicy { + self.boundary.transfer_policy + } + + /// Return the resolved adapter name. + #[must_use] + pub fn adapter_name(&self) -> &str { + &self.adapter_info.name + } + + /// Return the conservative exact-safe opcode subset used by the native WGPU batch-compute path. + #[must_use] + pub fn native_batch_compute_supported_opcodes() -> &'static [OpCode] { + &WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES + } + + /// Return whether this graph is eligible for the restricted exact-safe native WGPU path. + #[must_use] + pub fn supports_native_batch_compute(&self, graph: &CompiledGraph) -> bool { + wgpu_native_exact_safe_graph(graph) + } + + /// Return whether a concrete batch is eligible for the exact-safe native WGPU path. + #[must_use] + pub fn supports_native_batch_compute_for_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + ) -> bool { + self.supports_native_batch_compute(graph) && wgpu_native_exact_safe_batch(batch) + } + + fn create_initialized_storage_buffer( + &self, + label: &'static str, + bytes: &[u8], + min_size: u64, + ) -> wgpu::Buffer { + let buffer = self.device.create_buffer(&wgpu::BufferDescriptor { + label: Some(label), + size: min_size.max(bytes.len() as u64).max(4), + usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + if !bytes.is_empty() { + self.queue.write_buffer(&buffer, 0, bytes); + } + buffer + } + + fn download_raw_buffer(&self, buffer: &wgpu::Buffer, size: u64) -> Result> { + let staging = self.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("petite-ad-wgpu-readback"), + size, + usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ, + mapped_at_creation: false, + }); + let mut encoder = self + .device + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("petite-ad-wgpu-readback-encoder"), + }); + encoder.copy_buffer_to_buffer(buffer, 0, &staging, 0, size); + let submission = self.queue.submit([encoder.finish()]); + let slice = staging.slice(..); + let (tx, rx) = mpsc::channel(); + slice.map_async(wgpu::MapMode::Read, move |result| { + let _ = tx.send(result); + }); + let _ = self.device.poll(wgpu::PollType::wait_for(submission)); + match rx.recv() { + Ok(Ok(())) => {} + _ => { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu readback mapping failed", + }); + } + } + let view = slice.get_mapped_range(); + let bytes = view.to_vec(); + drop(view); + staging.unmap(); + Ok(bytes) + } + + fn compute_batch_native( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + ) -> Result> { + let metadata = graph.metadata(); + let input_f32: Vec = batch.data.iter().map(|value| *value as f32).collect(); + let output_count = batch.batch_size.saturating_mul(metadata.num_outputs); + let input_buffer = self.create_initialized_storage_buffer( + "petite-ad-wgpu-native-inputs", + &encode_f32_slice(&input_f32), + wgpu_buffer_size_bytes_f32(input_f32.len()), + ); + let instruction_words = encode_u32_slice(&self.native_instruction_words(graph)?); + let instruction_buffer = self.create_initialized_storage_buffer( + "petite-ad-wgpu-native-instructions", + &instruction_words, + instruction_words.len() as u64, + ); + let output_nodes: Result> = graph + .output_nodes() + .iter() + .map(|node| { + checked_u32_from_usize(*node, "wgpu native output node index exceeds u32 range") + }) + .collect(); + let output_nodes = output_nodes?; + let output_node_buffer = self.create_initialized_storage_buffer( + "petite-ad-wgpu-native-output-nodes", + &encode_u32_slice(&output_nodes), + (output_nodes.len().max(1) * std::mem::size_of::()) as u64, + ); + let value_buffer = self.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("petite-ad-wgpu-native-values"), + size: wgpu_buffer_size_bytes_f32(batch.batch_size.saturating_mul(metadata.value_count)), + usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + let output_buffer = self.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("petite-ad-wgpu-native-outputs"), + size: wgpu_buffer_size_bytes_f32(output_count), + usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + let meta = [ + checked_u32_from_usize( + metadata.num_inputs, + "wgpu native input dimension exceeds u32 range", + )?, + checked_u32_from_usize( + metadata.num_instructions, + "wgpu native instruction count exceeds u32 range", + )?, + checked_u32_from_usize( + metadata.num_outputs, + "wgpu native output dimension exceeds u32 range", + )?, + checked_u32_from_usize( + metadata.value_count, + "wgpu native value count exceeds u32 range", + )?, + checked_u32_from_usize(batch.batch_size, "wgpu native batch size exceeds u32 range")?, + ]; + let meta_buffer = self.create_initialized_storage_buffer( + "petite-ad-wgpu-native-meta", + &encode_u32_slice(&meta), + (meta.len() * std::mem::size_of::()) as u64, + ); + self.queue.submit([]); + + let shader = self + .device + .create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("petite-ad-wgpu-native-compute-shader"), + source: wgpu::ShaderSource::Wgsl(WGPU_NATIVE_SHADER.into()), + }); + let pipeline = self + .device + .create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("petite-ad-wgpu-native-compute-pipeline"), + layout: None, + module: &shader, + entry_point: Some("main"), + compilation_options: Default::default(), + cache: None, + }); + let layout = pipeline.get_bind_group_layout(0); + let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("petite-ad-wgpu-native-bind-group"), + layout: &layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: input_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: instruction_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 2, + resource: output_node_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 3, + resource: value_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 4, + resource: output_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 5, + resource: meta_buffer.as_entire_binding(), + }, + ], + }); + let mut encoder = self + .device + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("petite-ad-wgpu-native-compute-encoder"), + }); + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("petite-ad-wgpu-native-compute-pass"), + timestamp_writes: None, + }); + pass.set_pipeline(&pipeline); + pass.set_bind_group(0, &bind_group, &[]); + let workgroups = checked_u32_from_usize( + batch.batch_size.div_ceil(64), + "wgpu native workgroup count exceeds u32 range", + )?; + pass.dispatch_workgroups(workgroups.max(1), 1, 1); + } + self.queue.submit([encoder.finish()]); + let raw_output = + self.download_raw_buffer(&output_buffer, wgpu_buffer_size_bytes_f32(output_count))?; + let output_f32 = decode_f32_bytes(&raw_output, output_count)?; + Ok(output_f32.into_iter().map(f64::from).collect()) + } + + fn native_instruction_words(&self, graph: &CompiledGraph) -> Result> { + let mut words = Vec::with_capacity( + graph.flat_instructions_slice().len() * WGPU_NATIVE_WORDS_PER_INSTRUCTION, + ); + for instruction in graph.flat_instructions_slice() { + words.push(match instruction.opcode { + OpCode::Constant => 0, + OpCode::Add => 1, + OpCode::Sub => 2, + OpCode::Mul => 3, + OpCode::Div => 4, + OpCode::Pow => 5, + OpCode::Sin => 6, + OpCode::Cos => 7, + OpCode::Tan => 8, + OpCode::Tanh => 9, + OpCode::Relu => 10, + OpCode::Log1pExp => 11, + OpCode::LogAddExp => 12, + OpCode::Neg => 13, + OpCode::Exp => 14, + OpCode::Ln => 15, + OpCode::Sqrt => 16, + OpCode::Abs => 17, + }); + words.push(checked_u32_from_usize( + instruction.output, + "wgpu native instruction output index exceeds u32 range", + )?); + words.push(if instruction.left == UNUSED_NODE_ID { + u32::MAX + } else { + checked_u32_from_usize( + instruction.left, + "wgpu native instruction left index exceeds u32 range", + )? + }); + words.push(if instruction.right == UNUSED_NODE_ID { + u32::MAX + } else { + checked_u32_from_usize( + instruction.right, + "wgpu native instruction right index exceeds u32 range", + )? + }); + words.push((instruction.value as f32).to_bits()); + words.extend_from_slice(&[0, 0, 0]); + } + Ok(words) + } + + /// Allocate a real WGPU buffer set for a compiled graph and batch size. + pub fn allocate_batch_buffers( + &self, + graph: &CompiledGraph, + batch_size: usize, + ) -> Result { + let plan = self.batch_plan(graph, batch_size); + let mut buffers = Vec::with_capacity(plan.buffer_handles.len()); + for handle in &plan.buffer_handles { + let label = format!("petite-ad-wgpu-{:?}", handle.kind); + let buffer = self.device.create_buffer(&wgpu::BufferDescriptor { + label: Some(label.as_str()), + size: wgpu_buffer_size_bytes(handle.len), + usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + buffers.push(WgpuBuffer { + handle: *handle, + buffer, + }); + } + Ok(WgpuBufferSet { plan, buffers }) + } + + fn upload_buffer( + &self, + buffers: &WgpuBufferSet, + kind: DeviceBufferKind, + data: &[f64], + ) -> Result<()> { + let buffer = buffers.buffer(kind)?; + if buffer.handle.len != data.len() { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu upload length must match planned buffer length", + }); + } + if data.is_empty() { + return Ok(()); + } + let bytes = encode_f64_slice(data); + self.queue.write_buffer(&buffer.buffer, 0, &bytes); + self.queue.submit([]); + Ok(()) + } + + fn download_buffer(&self, buffers: &WgpuBufferSet, kind: DeviceBufferKind) -> Result> { + let buffer = buffers.buffer(kind)?; + if buffer.handle.len == 0 { + return Ok(Vec::new()); + } + let raw = + self.download_raw_buffer(&buffer.buffer, wgpu_buffer_size_bytes(buffer.handle.len))?; + decode_f64_bytes(&raw, buffer.handle.len) + } + + /// Execute batch value computation through explicit WGPU transfers. + pub fn compute_batch_with_buffers( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffers: &mut WgpuBufferSet, + output: &mut BatchValuesBuffer, + ) -> Result { + graph.check_batch(batch)?; + if buffers.plan.backend != BackendKind::Wgpu { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu execution requires a WGPU buffer plan", + }); + } + if buffers.plan.batch_size != batch.batch_size || buffers.plan.input_dim != batch.input_dim + { + return Err(AutodiffError::InvalidGraph { + reason: "batch shape must match WGPU buffer plan", + }); + } + + let transfers = buffers.plan.compute_transfer_plan.clone(); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::HostToDevice { + if transfer.buffer != DeviceBufferKind::Inputs { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu compute supports host-to-device input transfers only", + }); + } + buffers.upload(self, DeviceBufferKind::Inputs, batch.data)?; + } + } + + let used_native_kernel = self.supports_native_batch_compute_for_batch(graph, batch); + let scratch_data = if used_native_kernel { + self.compute_batch_native(graph, batch)? + } else { + let device_inputs = buffers.download(self, DeviceBufferKind::Inputs)?; + let device_batch = BatchInputs::new(&device_inputs, batch.batch_size, batch.input_dim)?; + let mut scratch = BatchValuesBuffer::new(); + ScalarBackend.compute_batch(graph, device_batch, &mut scratch)?; + scratch.data + }; + buffers.upload(self, DeviceBufferKind::Outputs, &scratch_data)?; + output.reset(batch.batch_size, graph.output_nodes.len()); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::DeviceToHost { + if transfer.buffer != DeviceBufferKind::Outputs { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu compute supports device-to-host output transfers only", + }); + } + output + .data + .extend_from_slice(&buffers.download(self, DeviceBufferKind::Outputs)?); + } + } + + Ok(DeviceExecutionTrace { + backend: BackendKind::Wgpu, + mode: DeviceExecutionMode::ComputeBatch, + transfers, + used_native_kernel, + }) + } + + /// Execute batch gradient computation through explicit WGPU transfers. + pub fn gradient_batch_with_buffers( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffers: &mut WgpuBufferSet, + output: &mut BatchGradientsBuffer, + ) -> Result { + graph.check_batch(batch)?; + if buffers.plan.backend != BackendKind::Wgpu { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu execution requires a WGPU buffer plan", + }); + } + if buffers.plan.batch_size != batch.batch_size || buffers.plan.input_dim != batch.input_dim + { + return Err(AutodiffError::InvalidGraph { + reason: "batch shape must match WGPU buffer plan", + }); + } + + let transfers = buffers.plan.gradient_transfer_plan.clone(); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::HostToDevice { + if transfer.buffer != DeviceBufferKind::Inputs { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu gradient supports host-to-device input transfers only", + }); + } + buffers.upload(self, DeviceBufferKind::Inputs, batch.data)?; + } + } + + let device_inputs = buffers.download(self, DeviceBufferKind::Inputs)?; + let device_batch = BatchInputs::new(&device_inputs, batch.batch_size, batch.input_dim)?; + let mut scratch = BatchGradientsBuffer::new(); + ScalarBackend.gradient_batch(graph, device_batch, &mut scratch)?; + buffers.upload(self, DeviceBufferKind::PrimaryValues, &scratch.values)?; + buffers.upload(self, DeviceBufferKind::Gradients, &scratch.gradients)?; + output.reset(batch.batch_size, graph.num_inputs); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::DeviceToHost { + match transfer.buffer { + DeviceBufferKind::PrimaryValues => output.values.extend_from_slice( + &buffers.download(self, DeviceBufferKind::PrimaryValues)?, + ), + DeviceBufferKind::Gradients => output + .gradients + .extend_from_slice(&buffers.download(self, DeviceBufferKind::Gradients)?), + _ => { + return Err(AutodiffError::InvalidGraph { + reason: + "wgpu gradient supports primary-value and gradient downloads only", + }); + } + } + } + } + + Ok(DeviceExecutionTrace { + backend: BackendKind::Wgpu, + mode: DeviceExecutionMode::GradientBatch, + transfers, + used_native_kernel: false, + }) + } +} + +impl DeviceBackend for WgpuBackend { + fn backend_kind(&self, _graph: &CompiledGraph) -> BackendKind { + BackendKind::Wgpu + } +} + +impl ExecutionBackend for WgpuBackend { + fn name(&self) -> &'static str { + "wgpu" + } + + fn capabilities(&self) -> BackendCapabilities { + BackendCapabilities::wgpu_f64() + } + + fn compute(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result { + let batch = BatchInputs::new(inputs, 1, inputs.len())?; + let mut values = BatchValuesBuffer::new(); + self.compute_batch(graph, batch, &mut values)?; + values + .data + .first() + .copied() + .ok_or(AutodiffError::InvalidGraph { + reason: "wgpu compute did not produce an output value", + }) + } + + fn compute_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result<()> { + graph.validate_backend_capabilities(&self.capabilities())?; + let mut gpu_buffers = self.allocate_batch_buffers(graph, batch.batch_size)?; + self.compute_batch_with_buffers(graph, batch, &mut gpu_buffers, buffer)?; + Ok(()) + } + + fn gradient(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result<(f64, Vec)> { + let batch = BatchInputs::new(inputs, 1, inputs.len())?; + let mut gradients = BatchGradientsBuffer::new(); + self.gradient_batch(graph, batch, &mut gradients)?; + let value = gradients + .values + .first() + .copied() + .ok_or(AutodiffError::InvalidGraph { + reason: "wgpu gradient did not produce a primary output value", + })?; + Ok((value, gradients.gradients)) + } + + fn gradient_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result<()> { + graph.validate_backend_capabilities(&self.capabilities())?; + let mut gpu_buffers = self.allocate_batch_buffers(graph, batch.batch_size)?; + self.gradient_batch_with_buffers(graph, batch, &mut gpu_buffers, buffer)?; + Ok(()) + } +} diff --git a/src/multi/builder.rs b/src/multi/builder.rs index 66d7454..86e67ab 100644 --- a/src/multi/builder.rs +++ b/src/multi/builder.rs @@ -42,7 +42,6 @@ use super::multi_ad::MultiAD; #[derive(Debug, Clone)] pub struct GraphBuilder { /// Number of input variables - #[allow(dead_code)] num_inputs: usize, /// Operations in the legacy tuple computation graph operations: Vec<(MultiAD, Vec)>, diff --git a/src/multi/compiled.rs b/src/multi/compiled.rs index 6b8640c..1e683ab 100644 --- a/src/multi/compiled.rs +++ b/src/multi/compiled.rs @@ -1,4130 +1,863 @@ //! Closure-free compiled instruction IR for acceleration-ready graph execution. -use super::multi_ad::MultiAD; use super::op_rules; use crate::{AutodiffError, NodeId, Result}; +use super::backend::{ + BackendCapabilities, BackendKind, BackendRejectionReason, BackendSupportReport, DeviceBackend, + DeviceBatchPlan, DeviceBufferSet, DeviceExecutionTrace, FlatInstruction, Instruction, + MockDeviceBackend, OpCode, UNUSED_NODE_ID, +}; +#[cfg(test)] +use super::backend::{DeviceBufferKind, DeviceMemoryLocation}; #[cfg(feature = "backend-wgpu")] -use std::sync::mpsc; - -#[cfg(target_arch = "x86_64")] -use std::arch::x86_64::{__m128d, __m256d}; - -#[cfg(feature = "backend-wgpu")] -use pollster::block_on; -#[cfg(feature = "backend-wgpu")] -use wgpu::{self, BufferUsages}; +use super::backend::{WgpuBackend, WgpuBufferSet}; +#[cfg(test)] +use super::multi_ad::MultiAD; -/// One closure-free instruction in a compiled scalar graph. -#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +/// Flat row-major batch input view. #[derive(Debug, Clone, Copy, PartialEq)] -pub enum Instruction { - /// Literal constant node. - Constant(f64), - /// Unary operation with one input node id. - Unary { op: MultiAD, arg: NodeId }, - /// Binary operation with two input node ids. - Binary { - op: MultiAD, - left: NodeId, - right: NodeId, - }, -} - -/// Sentinel node id used by flat backend instructions when an argument is absent. -pub const UNUSED_NODE_ID: NodeId = usize::MAX; - -/// Compact operation code used by flat backend instructions. -#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub enum OpCode { - Constant, - Add, - Sub, - Mul, - Div, - Pow, - Sin, - Cos, - Tan, - Tanh, - Relu, - Log1pExp, - LogAddExp, - Neg, - Exp, - Ln, - Sqrt, - Abs, +pub struct BatchInputs<'a> { + /// Flat row-major input data. + pub data: &'a [f64], + /// Number of rows in the batch. + pub batch_size: usize, + /// Number of inputs per row. + pub input_dim: usize, } -impl OpCode { - fn from_multi_ad(op: MultiAD) -> Result { - Ok(match op { - MultiAD::Add => OpCode::Add, - MultiAD::Sub => OpCode::Sub, - MultiAD::Mul => OpCode::Mul, - MultiAD::Div => OpCode::Div, - MultiAD::Pow => OpCode::Pow, - MultiAD::Sin => OpCode::Sin, - MultiAD::Cos => OpCode::Cos, - MultiAD::Tan => OpCode::Tan, - MultiAD::Tanh => OpCode::Tanh, - MultiAD::Relu => OpCode::Relu, - MultiAD::Log1pExp => OpCode::Log1pExp, - MultiAD::LogAddExp => OpCode::LogAddExp, - MultiAD::Neg => OpCode::Neg, - MultiAD::Exp => OpCode::Exp, - MultiAD::Ln => OpCode::Ln, - MultiAD::Sqrt => OpCode::Sqrt, - MultiAD::Abs => OpCode::Abs, - MultiAD::Inp => { - return Err(AutodiffError::InvalidGraph { - reason: "input markers are not backend instructions", - }); - } +impl<'a> BatchInputs<'a> { + /// Create a validated row-major input view. + pub fn new(data: &'a [f64], batch_size: usize, input_dim: usize) -> Result { + if data.len() != batch_size.saturating_mul(input_dim) { + return Err(AutodiffError::InvalidGraph { + reason: "batch data length must equal batch_size * input_dim", + }); + } + Ok(Self { + data, + batch_size, + input_dim, }) } - fn to_multi_ad(self) -> Option { - Some(match self { - OpCode::Constant => return None, - OpCode::Add => MultiAD::Add, - OpCode::Sub => MultiAD::Sub, - OpCode::Mul => MultiAD::Mul, - OpCode::Div => MultiAD::Div, - OpCode::Pow => MultiAD::Pow, - OpCode::Sin => MultiAD::Sin, - OpCode::Cos => MultiAD::Cos, - OpCode::Tan => MultiAD::Tan, - OpCode::Tanh => MultiAD::Tanh, - OpCode::Relu => MultiAD::Relu, - OpCode::Log1pExp => MultiAD::Log1pExp, - OpCode::LogAddExp => MultiAD::LogAddExp, - OpCode::Neg => MultiAD::Neg, - OpCode::Exp => MultiAD::Exp, - OpCode::Ln => MultiAD::Ln, - OpCode::Sqrt => MultiAD::Sqrt, - OpCode::Abs => MultiAD::Abs, - }) + /// Return one input row. + /// + /// Prefer [`BatchInputs::try_row`] when invalid row indices should return an error. + #[must_use] + pub fn row(&self, index: usize) -> &[f64] { + self.try_row(index).expect("batch row index out of bounds") } - /// Return the number of input slots consumed by this opcode. - #[must_use] - pub fn arity(&self) -> usize { - match self { - OpCode::Constant => 0, - OpCode::Sin - | OpCode::Cos - | OpCode::Tan - | OpCode::Tanh - | OpCode::Relu - | OpCode::Log1pExp - | OpCode::Neg - | OpCode::Exp - | OpCode::Ln - | OpCode::Sqrt - | OpCode::Abs => 1, - OpCode::Add - | OpCode::Sub - | OpCode::Mul - | OpCode::Div - | OpCode::Pow - | OpCode::LogAddExp => 2, + /// Return one input row, or an error when `index` is out of range. + pub fn try_row(&self, index: usize) -> Result<&[f64]> { + if index >= self.batch_size { + return Err(AutodiffError::IndexOutOfBounds { + index, + max_index: self.batch_size.saturating_sub(1), + }); } + let start = index * self.input_dim; + Ok(&self.data[start..start + self.input_dim]) } } -/// Fully explicit instruction shape for SIMD/GPU backend lowering. -#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] -#[derive(Debug, Clone, Copy, PartialEq)] -pub struct FlatInstruction { - pub opcode: OpCode, - pub output: NodeId, - pub left: NodeId, - pub right: NodeId, - pub value: f64, +/// Flat row-major batch values. +#[derive(Debug, Clone, PartialEq)] +pub struct BatchValues { + pub data: Vec, + pub batch_size: usize, + pub output_dim: usize, } -#[cfg(target_arch = "x86_64")] -#[inline] -fn supports_simd_f64x2_runtime() -> bool { - true +/// Reusable flat row-major batch value buffer. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct BatchValuesBuffer { + pub data: Vec, + pub batch_size: usize, + pub output_dim: usize, } -#[cfg(not(target_arch = "x86_64"))] -#[inline] -fn supports_simd_f64x2_runtime() -> bool { - false -} +impl BatchValuesBuffer { + /// Create an empty reusable output buffer. + #[must_use] + pub fn new() -> Self { + Self::default() + } -#[cfg(target_arch = "x86_64")] -#[inline] -fn supports_simd_f64x4_runtime() -> bool { - std::is_x86_feature_detected!("avx") -} + pub(crate) fn reset(&mut self, batch_size: usize, output_dim: usize) { + self.data.clear(); + self.data.reserve(batch_size.saturating_mul(output_dim)); + self.batch_size = batch_size; + self.output_dim = output_dim; + } -#[cfg(not(target_arch = "x86_64"))] -#[inline] -fn supports_simd_f64x4_runtime() -> bool { - false + /// Clone the current buffer contents into an owned result value. + #[must_use] + pub fn to_values(&self) -> BatchValues { + BatchValues { + data: self.data.clone(), + batch_size: self.batch_size, + output_dim: self.output_dim, + } + } } -#[inline] -fn supports_wgpu_runtime() -> bool { - cfg!(feature = "backend-wgpu") +/// Flat row-major batch scalar-output gradients. +#[derive(Debug, Clone, PartialEq)] +pub struct BatchGradients { + pub values: Vec, + pub gradients: Vec, + pub batch_size: usize, + pub input_dim: usize, } -/// Backend feature declaration used before dispatching compiled graphs. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct BackendCapabilities { - pub supports_f64: bool, - pub supports_f32: bool, - pub supports_constants: bool, - pub supports_unary: bool, - pub supports_binary: bool, - pub supports_multi_output: bool, - pub supports_reverse_gradient: bool, - pub supports_batch_compute: bool, - pub supports_batch_gradient: bool, - pub supported_opcodes: Vec, +/// Reusable flat row-major batch scalar-output gradient buffer. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct BatchGradientsBuffer { + pub values: Vec, + pub gradients: Vec, + pub batch_size: usize, + pub input_dim: usize, } -impl BackendCapabilities { - /// Capabilities for the reference scalar f64 backend. - #[must_use] - pub fn scalar_f64() -> Self { - Self { - supports_f64: true, - supports_f32: false, - supports_constants: true, - supports_unary: true, - supports_binary: true, - supports_multi_output: true, - supports_reverse_gradient: true, - supports_batch_compute: true, - supports_batch_gradient: true, - supported_opcodes: vec![ - OpCode::Constant, - OpCode::Add, - OpCode::Sub, - OpCode::Mul, - OpCode::Div, - OpCode::Pow, - OpCode::Sin, - OpCode::Cos, - OpCode::Tan, - OpCode::Tanh, - OpCode::Relu, - OpCode::Log1pExp, - OpCode::LogAddExp, - OpCode::Neg, - OpCode::Exp, - OpCode::Ln, - OpCode::Sqrt, - OpCode::Abs, - ], - } - } - - /// Capabilities for the prototype f64x2 SIMD batch backend. +impl BatchGradientsBuffer { + /// Create an empty reusable gradient buffer. #[must_use] - pub fn simd_f64x2() -> Self { - Self { - supports_f64: true, - supports_f32: false, - supports_constants: true, - supports_unary: true, - supports_binary: true, - supports_multi_output: true, - supports_reverse_gradient: false, - supports_batch_compute: supports_simd_f64x2_runtime(), - supports_batch_gradient: supports_simd_f64x2_runtime(), - supported_opcodes: vec![ - OpCode::Constant, - OpCode::Add, - OpCode::Sub, - OpCode::Mul, - OpCode::Div, - OpCode::Pow, - OpCode::Sin, - OpCode::Cos, - OpCode::Tan, - OpCode::Tanh, - OpCode::Log1pExp, - OpCode::LogAddExp, - OpCode::Neg, - OpCode::Exp, - OpCode::Ln, - OpCode::Sqrt, - OpCode::Relu, - OpCode::Abs, - ], - } + pub fn new() -> Self { + Self::default() } - /// Capabilities for the prototype f64x4 SIMD batch backend. - #[must_use] - pub fn simd_f64x4() -> Self { - Self { - supports_f64: true, - supports_f32: false, - supports_constants: true, - supports_unary: true, - supports_binary: true, - supports_multi_output: true, - supports_reverse_gradient: false, - supports_batch_compute: supports_simd_f64x4_runtime(), - supports_batch_gradient: supports_simd_f64x4_runtime(), - supported_opcodes: vec![ - OpCode::Constant, - OpCode::Add, - OpCode::Sub, - OpCode::Mul, - OpCode::Div, - OpCode::Pow, - OpCode::Sin, - OpCode::Cos, - OpCode::Tan, - OpCode::Tanh, - OpCode::Log1pExp, - OpCode::LogAddExp, - OpCode::Neg, - OpCode::Exp, - OpCode::Ln, - OpCode::Sqrt, - OpCode::Relu, - OpCode::Abs, - ], - } + pub(crate) fn reset(&mut self, batch_size: usize, input_dim: usize) { + self.values.clear(); + self.gradients.clear(); + self.values.reserve(batch_size); + self.gradients.reserve(batch_size.saturating_mul(input_dim)); + self.batch_size = batch_size; + self.input_dim = input_dim; } - /// Capabilities for the feature-gated WGPU batch backend skeleton. + /// Clone the current buffer contents into an owned result value. #[must_use] - pub fn wgpu_f64() -> Self { - Self { - supports_f64: true, - supports_f32: false, - supports_constants: true, - supports_unary: true, - supports_binary: true, - supports_multi_output: true, - supports_reverse_gradient: true, - supports_batch_compute: supports_wgpu_runtime(), - supports_batch_gradient: supports_wgpu_runtime(), - supported_opcodes: BackendCapabilities::scalar_f64().supported_opcodes, + pub fn to_gradients(&self) -> BatchGradients { + BatchGradients { + values: self.values.clone(), + gradients: self.gradients.clone(), + batch_size: self.batch_size, + input_dim: self.input_dim, } } - - /// Return whether the backend declares support for an opcode. - #[must_use] - pub fn supports_opcode(&self, opcode: OpCode) -> bool { - self.supported_opcodes.contains(&opcode) - } } -/// Reference scalar backend for the execution-backend abstraction. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct ScalarBackend; - -/// Mock device-style backend that executes on CPU while using device-oriented plans. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct MockDeviceBackend; +/// Static metadata for backend selection and workspace planning. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CompiledGraphMetadata { + pub num_inputs: usize, + pub num_outputs: usize, + pub num_instructions: usize, + pub num_constants: usize, + pub num_unary: usize, + pub num_binary: usize, + pub value_count: usize, + pub is_scalar_output: bool, +} -/// Prototype f64x2 SIMD backend for batch compute and batch gradients. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct SimdBackend; +/// Reusable buffers for [`CompiledGraph`] execution. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct CompiledWorkspace { + values: Vec, + cotangents: Vec, + gradients: Vec, + outputs: Vec, +} -/// Backend selected by automatic compiled-graph dispatch. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum BackendKind { - Scalar, - MockDeviceCpu, - Wgpu, - SimdF64x4, - SimdF64x2, +/// Closure-free compiled graph representation. +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[derive(Debug, Clone, PartialEq)] +pub struct CompiledGraph { + pub(crate) num_inputs: usize, + pub(crate) instructions: Vec, + pub(crate) output_nodes: Vec, + flat_instructions: Vec, } -impl BackendKind { - /// Return the number of f64 lanes processed by this backend. - #[must_use] - pub fn lane_width(&self) -> usize { - match self { - BackendKind::Scalar | BackendKind::MockDeviceCpu | BackendKind::Wgpu => 1, - BackendKind::SimdF64x4 => 4, - BackendKind::SimdF64x2 => 2, - } +impl CompiledGraph { + pub(crate) fn new( + num_inputs: usize, + instructions: Vec, + output_nodes: Vec, + ) -> Result { + let flat_instructions = Self::lower_flat_instructions(num_inputs, &instructions)?; + Ok(Self { + num_inputs, + instructions, + output_nodes, + flat_instructions, + }) } - /// Return runtime CPU features required by this backend. + /// Return number of graph inputs. #[must_use] - pub fn required_runtime_features(&self) -> &'static [&'static str] { - match self { - BackendKind::Scalar | BackendKind::MockDeviceCpu => &[], - BackendKind::Wgpu => &["backend-wgpu"], - BackendKind::SimdF64x4 => &["x86_64-avx"], - BackendKind::SimdF64x2 => &["x86_64-sse2"], - } + pub fn num_inputs(&self) -> usize { + self.num_inputs } - /// Return whether the backend is available on this runtime target. + /// Return compiled instructions. #[must_use] - pub fn runtime_available(&self) -> bool { - match self { - BackendKind::Scalar | BackendKind::MockDeviceCpu => true, - BackendKind::Wgpu => supports_wgpu_runtime(), - BackendKind::SimdF64x4 => supports_simd_f64x4_runtime(), - BackendKind::SimdF64x2 => supports_simd_f64x2_runtime(), - } + pub fn instructions(&self) -> &[Instruction] { + &self.instructions } - /// Return runtime CPU features that are required but unavailable. + /// Return selected output nodes. #[must_use] - pub fn unavailable_runtime_features(&self) -> Vec<&'static str> { - if self.runtime_available() { - Vec::new() - } else { - self.required_runtime_features().to_vec() - } + pub fn output_nodes(&self) -> &[NodeId] { + &self.output_nodes } - /// Return a stable backend name. - #[must_use] - pub fn name(&self) -> &'static str { - match self { - BackendKind::Scalar => "scalar", - BackendKind::MockDeviceCpu => "mock-device-cpu", - BackendKind::Wgpu => "wgpu", - BackendKind::SimdF64x4 => "simd-f64x4", - BackendKind::SimdF64x2 => "simd-f64x2", + fn lower_flat_instructions( + num_inputs: usize, + instructions: &[Instruction], + ) -> Result> { + let mut flat = Vec::with_capacity(instructions.len()); + for (offset, instruction) in instructions.iter().enumerate() { + let output = num_inputs + offset; + let item = match *instruction { + Instruction::Constant(value) => FlatInstruction { + opcode: OpCode::Constant, + output, + left: UNUSED_NODE_ID, + right: UNUSED_NODE_ID, + value, + }, + Instruction::Unary { op, arg } => FlatInstruction { + opcode: OpCode::from_multi_ad(op)?, + output, + left: arg, + right: UNUSED_NODE_ID, + value: 0.0, + }, + Instruction::Binary { op, left, right } => FlatInstruction { + opcode: OpCode::from_multi_ad(op)?, + output, + left, + right, + value: 0.0, + }, + }; + flat.push(item); } + Ok(flat) } - /// Return the logical memory location used by this backend's batch plan. + /// Return a zero-copy view of flat instructions with explicit output and argument slots. #[must_use] - pub fn memory_location(&self) -> DeviceMemoryLocation { - match self { - BackendKind::MockDeviceCpu | BackendKind::Wgpu => DeviceMemoryLocation::Device, - BackendKind::Scalar | BackendKind::SimdF64x4 | BackendKind::SimdF64x2 => { - DeviceMemoryLocation::Host - } - } + pub fn flat_instructions_slice(&self) -> &[FlatInstruction] { + &self.flat_instructions } - /// Return declared capabilities for this backend. - #[must_use] - pub fn capabilities(&self) -> BackendCapabilities { - match self { - BackendKind::Scalar | BackendKind::MockDeviceCpu => BackendCapabilities::scalar_f64(), - BackendKind::Wgpu => BackendCapabilities::wgpu_f64(), - BackendKind::SimdF64x4 => BackendCapabilities::simd_f64x4(), - BackendKind::SimdF64x2 => BackendCapabilities::simd_f64x2(), - } + /// Return owned flat instructions with explicit output and argument slots. + pub fn flat_instructions(&self) -> Result> { + Ok(self.flat_instructions_slice().to_vec()) } - /// Execute batch value computation with this backend. - pub fn compute_batch( - &self, - graph: &CompiledGraph, - batch: BatchInputs<'_>, - buffer: &mut BatchValuesBuffer, - ) -> Result<()> { - if matches!(self, BackendKind::Scalar | BackendKind::MockDeviceCpu) { - return ScalarBackend.compute_batch(graph, batch, buffer); - } - if matches!(self, BackendKind::Wgpu) { - #[cfg(feature = "backend-wgpu")] - { - return WgpuBackend::new_default()?.compute_batch(graph, batch, buffer); - } - #[cfg(not(feature = "backend-wgpu"))] - { - return Err(AutodiffError::InvalidGraph { - reason: "wgpu backend requires the backend-wgpu cargo feature", - }); - } - } - let capabilities = self.capabilities(); - if !capabilities.supports_batch_compute { + /// Validate whether static graph requirements fit backend capabilities. + pub fn validate_backend_capabilities(&self, capabilities: &BackendCapabilities) -> Result<()> { + if !capabilities.supports_f64 { return Err(AutodiffError::InvalidGraph { - reason: "backend does not support batch compute on this target", + reason: "backend must support f64 execution", }); } - graph.validate_backend_capabilities(&capabilities)?; - match self { - BackendKind::Scalar | BackendKind::MockDeviceCpu | BackendKind::Wgpu => unreachable!(), - BackendKind::SimdF64x4 => compute_batch_simd_f64x4(graph, batch, buffer), - BackendKind::SimdF64x2 => compute_batch_simd_f64x2(graph, batch, buffer), + if self.output_nodes.len() > 1 && !capabilities.supports_multi_output { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support multi-output graphs", + }); } - } - - /// Execute batch gradient computation with this backend. - pub fn gradient_batch( - &self, - graph: &CompiledGraph, - batch: BatchInputs<'_>, - buffer: &mut BatchGradientsBuffer, - ) -> Result<()> { - if matches!(self, BackendKind::Scalar | BackendKind::MockDeviceCpu) { - return ScalarBackend.gradient_batch(graph, batch, buffer); - } - if matches!(self, BackendKind::Wgpu) { - #[cfg(feature = "backend-wgpu")] - { - return WgpuBackend::new_default()?.gradient_batch(graph, batch, buffer); - } - #[cfg(not(feature = "backend-wgpu"))] - { - return Err(AutodiffError::InvalidGraph { - reason: "wgpu backend requires the backend-wgpu cargo feature", - }); - } - } - let capabilities = self.capabilities(); - if !capabilities.supports_batch_gradient { - return Err(AutodiffError::InvalidGraph { - reason: "backend does not support batch gradients on this target", - }); - } - graph.validate_backend_capabilities(&capabilities)?; - match self { - BackendKind::Scalar | BackendKind::MockDeviceCpu | BackendKind::Wgpu => unreachable!(), - BackendKind::SimdF64x4 => gradient_batch_simd_f64x4(graph, batch, buffer), - BackendKind::SimdF64x2 => gradient_batch_simd_f64x2(graph, batch, buffer), - } - } -} - -/// Batch memory layout exposed for backend/device planning. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum BatchLayout { - RowMajor, -} - -/// Logical device buffer role for batch execution planning. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum DeviceBufferKind { - Inputs, - Values, - Outputs, - PrimaryValues, - Gradients, -} - -/// One logical device buffer description. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct DeviceBufferLayout { - pub kind: DeviceBufferKind, - pub len: usize, - pub element_size: usize, -} - -/// Logical memory location for a planned backend buffer. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum DeviceMemoryLocation { - Host, - Device, -} - -/// Handle-like description for one planned backend buffer. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct DeviceBufferHandle { - pub kind: DeviceBufferKind, - pub location: DeviceMemoryLocation, - pub offset: usize, - pub len: usize, -} - -/// Logical transfer direction for backend execution planning. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum DeviceTransferKind { - HostToDevice, - DeviceToHost, -} - -/// One logical host/device transfer needed by a batch plan. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct DeviceTransferPlan { - pub kind: DeviceTransferKind, - pub buffer: DeviceBufferKind, - pub len: usize, -} - -/// Device buffer owned by a planned backend execution. -#[derive(Debug, Clone, PartialEq)] -pub struct DeviceBuffer { - handle: DeviceBufferHandle, - data: Vec, -} - -impl DeviceBuffer { - fn new(handle: DeviceBufferHandle) -> Self { - Self { - handle, - data: vec![0.0; handle.len], - } - } - - /// Return the immutable planned buffer handle. - #[must_use] - pub fn handle(&self) -> DeviceBufferHandle { - self.handle - } - - /// Return the owned buffer data. - #[must_use] - pub fn data(&self) -> &[f64] { - &self.data - } -} - -/// Owned buffer set allocated from a [`DeviceBatchPlan`]. -#[derive(Debug, Clone, PartialEq)] -pub struct DeviceBufferSet { - plan: DeviceBatchPlan, - buffers: Vec, -} - -impl DeviceBufferSet { - /// Allocate zero-initialized buffers for a batch plan. - #[must_use] - pub fn new(plan: DeviceBatchPlan) -> Self { - let buffers = plan - .buffer_handles - .iter() - .copied() - .map(DeviceBuffer::new) - .collect(); - Self { plan, buffers } - } - - /// Return the plan used to allocate this buffer set. - #[must_use] - pub fn plan(&self) -> &DeviceBatchPlan { - &self.plan - } - - /// Return all allocated buffers in plan order. - #[must_use] - pub fn buffers(&self) -> &[DeviceBuffer] { - &self.buffers - } - - /// Return an immutable buffer by logical kind. - pub fn buffer(&self, kind: DeviceBufferKind) -> Result<&DeviceBuffer> { - self.buffers - .iter() - .find(|buffer| buffer.handle.kind == kind) - .ok_or(AutodiffError::InvalidGraph { - reason: "device buffer kind is not in the plan", - }) - } - - /// Return a mutable buffer by logical kind. - fn buffer_mut(&mut self, kind: DeviceBufferKind) -> Result<&mut DeviceBuffer> { - self.buffers - .iter_mut() - .find(|buffer| buffer.handle.kind == kind) - .ok_or(AutodiffError::InvalidGraph { - reason: "device buffer kind is not in the plan", - }) - } - - /// Upload host data into a planned buffer. - pub fn upload(&mut self, kind: DeviceBufferKind, data: &[f64]) -> Result<()> { - let buffer = self.buffer_mut(kind)?; - if buffer.data.len() != data.len() { - return Err(AutodiffError::InvalidGraph { - reason: "upload length must match planned buffer length", - }); - } - buffer.data.copy_from_slice(data); - Ok(()) - } - - /// Download a planned buffer into an owned vector. - pub fn download(&self, kind: DeviceBufferKind) -> Result> { - Ok(self.buffer(kind)?.data.clone()) - } -} - -/// Device-oriented batch execution plan for a backend. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct DeviceBatchPlan { - pub backend: BackendKind, - pub layout: BatchLayout, - pub batch_size: usize, - pub input_dim: usize, - pub output_dim: usize, - pub gradient_dim: usize, - pub value_count: usize, - pub buffers: Vec, - pub buffer_handles: Vec, - pub compute_transfer_plan: Vec, - pub gradient_transfer_plan: Vec, -} - -/// Batch execution mode for explicit device transfer planning. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum DeviceExecutionMode { - ComputeBatch, - GradientBatch, -} - -/// Trace returned by explicit device-style execution. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct DeviceExecutionTrace { - pub backend: BackendKind, - pub mode: DeviceExecutionMode, - pub transfers: Vec, - /// Whether this execution used a native accelerator kernel instead of host fallback logic. - pub used_native_kernel: bool, -} - -/// Feature-neutral accelerator device family. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum AcceleratorDeviceKind { - MockCpu, - Cuda, - Wgpu, -} - -/// Feature-neutral accelerator context descriptor for future GPU backends. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AcceleratorDeviceContext { - pub kind: AcceleratorDeviceKind, - pub device_id: usize, - pub name: String, -} - -impl AcceleratorDeviceContext { - /// Return a context descriptor for the CPU mock-device backend. - #[must_use] - pub fn mock_cpu() -> Self { - Self { - kind: AcceleratorDeviceKind::MockCpu, - device_id: 0, - name: "mock-device-cpu".to_string(), - } - } - - /// Return a context descriptor for a future CUDA backend. - #[must_use] - pub fn cuda(device_id: usize) -> Self { - Self { - kind: AcceleratorDeviceKind::Cuda, - device_id, - name: format!("cuda:{device_id}"), - } - } - - /// Return a context descriptor for a future WGPU backend. - #[must_use] - pub fn wgpu(device_id: usize) -> Self { - Self { - kind: AcceleratorDeviceKind::Wgpu, - device_id, - name: format!("wgpu:{device_id}"), - } - } -} - -/// Host/device transfer policy for future accelerator backends. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum DeviceTransferPolicy { - Explicit, - Automatic, -} - -/// Boundary object for future GPU backends. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct GpuBackendBoundary { - pub context: AcceleratorDeviceContext, - pub transfer_policy: DeviceTransferPolicy, -} - -impl GpuBackendBoundary { - /// Create a boundary descriptor for a future accelerator backend. - #[must_use] - pub fn new(context: AcceleratorDeviceContext, transfer_policy: DeviceTransferPolicy) -> Self { - Self { - context, - transfer_policy, - } - } - - #[cfg(feature = "backend-wgpu")] - /// Initialize the real WGPU backend skeleton from this boundary. - pub fn initialize_wgpu(&self) -> Result { - WgpuBackend::from_boundary(self.clone()) - } - - /// Return an explicit error because generic GPU execution still requires a concrete backend. - pub fn unsupported_execution_error(&self) -> Result { - Err(AutodiffError::InvalidGraph { - reason: - "real GPU execution requires a concrete backend instance; initialize WgpuBackend", - }) - } -} - -#[cfg(feature = "backend-wgpu")] -/// Initialized WGPU backend skeleton with real device allocation and transfer plumbing. -#[derive(Debug, Clone)] -pub struct WgpuBackend { - boundary: GpuBackendBoundary, - device: wgpu::Device, - queue: wgpu::Queue, - adapter_info: wgpu::AdapterInfo, -} - -#[cfg(feature = "backend-wgpu")] -/// One real WGPU buffer allocated for a logical batch-plan role. -#[derive(Debug, Clone)] -pub struct WgpuBuffer { - handle: DeviceBufferHandle, - buffer: wgpu::Buffer, -} - -#[cfg(feature = "backend-wgpu")] -impl WgpuBuffer { - /// Return the immutable planned buffer handle. - #[must_use] - pub fn handle(&self) -> DeviceBufferHandle { - self.handle - } -} - -#[cfg(feature = "backend-wgpu")] -/// Owned WGPU buffers allocated from a [`DeviceBatchPlan`]. -#[derive(Debug, Clone)] -pub struct WgpuBufferSet { - plan: DeviceBatchPlan, - buffers: Vec, -} - -#[cfg(feature = "backend-wgpu")] -impl WgpuBufferSet { - /// Return the plan used to allocate this buffer set. - #[must_use] - pub fn plan(&self) -> &DeviceBatchPlan { - &self.plan - } - - /// Return all allocated GPU buffers in plan order. - #[must_use] - pub fn buffers(&self) -> &[WgpuBuffer] { - &self.buffers - } - - /// Return an immutable GPU buffer by logical kind. - pub fn buffer(&self, kind: DeviceBufferKind) -> Result<&WgpuBuffer> { - self.buffers - .iter() - .find(|buffer| buffer.handle.kind == kind) - .ok_or(AutodiffError::InvalidGraph { - reason: "wgpu buffer kind is not in the plan", - }) - } - - /// Upload host data into a planned GPU buffer. - pub fn upload( - &self, - backend: &WgpuBackend, - kind: DeviceBufferKind, - data: &[f64], - ) -> Result<()> { - backend.upload_buffer(self, kind, data) - } - - /// Download a planned GPU buffer into an owned vector. - pub fn download(&self, backend: &WgpuBackend, kind: DeviceBufferKind) -> Result> { - backend.download_buffer(self, kind) - } -} - -/// Device-oriented batch backend planning interface. -pub trait DeviceBackend { - fn backend_kind(&self, graph: &CompiledGraph) -> BackendKind; - - fn batch_layout(&self) -> BatchLayout { - BatchLayout::RowMajor - } - - fn batch_plan(&self, graph: &CompiledGraph, batch_size: usize) -> DeviceBatchPlan { - let metadata = graph.metadata(); - let value_count = metadata.value_count.saturating_mul(batch_size); - let input_count = metadata.num_inputs.saturating_mul(batch_size); - let output_count = metadata.num_outputs.saturating_mul(batch_size); - let gradient_count = metadata.num_inputs.saturating_mul(batch_size); - let backend = self.backend_kind(graph); - let buffer_location = backend.memory_location(); - let buffers = vec![ - DeviceBufferLayout { - kind: DeviceBufferKind::Inputs, - len: input_count, - element_size: std::mem::size_of::(), - }, - DeviceBufferLayout { - kind: DeviceBufferKind::Values, - len: value_count, - element_size: std::mem::size_of::(), - }, - DeviceBufferLayout { - kind: DeviceBufferKind::Outputs, - len: output_count, - element_size: std::mem::size_of::(), - }, - DeviceBufferLayout { - kind: DeviceBufferKind::PrimaryValues, - len: batch_size, - element_size: std::mem::size_of::(), - }, - DeviceBufferLayout { - kind: DeviceBufferKind::Gradients, - len: gradient_count, - element_size: std::mem::size_of::(), - }, - ]; - let mut offset = 0; - let buffer_handles = buffers - .iter() - .map(|buffer| { - let handle = DeviceBufferHandle { - kind: buffer.kind, - location: buffer_location, - offset, - len: buffer.len, - }; - offset = offset.saturating_add(buffer.len); - handle - }) - .collect(); - let (compute_transfer_plan, gradient_transfer_plan) = - if buffer_location == DeviceMemoryLocation::Device { - ( - vec![ - DeviceTransferPlan { - kind: DeviceTransferKind::HostToDevice, - buffer: DeviceBufferKind::Inputs, - len: input_count, - }, - DeviceTransferPlan { - kind: DeviceTransferKind::DeviceToHost, - buffer: DeviceBufferKind::Outputs, - len: output_count, - }, - ], - vec![ - DeviceTransferPlan { - kind: DeviceTransferKind::HostToDevice, - buffer: DeviceBufferKind::Inputs, - len: input_count, - }, - DeviceTransferPlan { - kind: DeviceTransferKind::DeviceToHost, - buffer: DeviceBufferKind::PrimaryValues, - len: batch_size, - }, - DeviceTransferPlan { - kind: DeviceTransferKind::DeviceToHost, - buffer: DeviceBufferKind::Gradients, - len: gradient_count, - }, - ], - ) - } else { - (Vec::new(), Vec::new()) - }; - DeviceBatchPlan { - backend, - layout: self.batch_layout(), - batch_size, - input_dim: metadata.num_inputs, - output_dim: metadata.num_outputs, - gradient_dim: metadata.num_inputs, - value_count: metadata.value_count, - buffers, - buffer_handles, - compute_transfer_plan, - gradient_transfer_plan, - } - } -} - -/// Reason a backend cannot execute a graph for a requested batch mode. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub enum BackendRejectionReason { - MissingF64, - UnsupportedOutputs, - UnsupportedOpcodes, - UnsupportedArities, - UnavailableRuntime, - NoBatchCompute, - NoBatchGradient, -} - -/// Static backend compatibility report for a compiled graph. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct BackendSupportReport { - pub backend: BackendKind, - pub supports_f64: bool, - pub supports_required_outputs: bool, - pub supports_required_opcodes: bool, - pub supports_required_arities: bool, - pub supports_batch_compute: bool, - pub supports_batch_gradient: bool, - pub lane_width: usize, - pub runtime_available: bool, - pub required_runtime_features: Vec<&'static str>, - pub unavailable_runtime_features: Vec<&'static str>, - pub missing_opcodes: Vec, - pub batch_compute_rejection_reasons: Vec, - pub batch_gradient_rejection_reasons: Vec, -} - -impl BackendSupportReport { - /// Return whether this backend can compute batch values for the graph. - #[must_use] - pub fn can_compute_batch(&self) -> bool { - self.batch_compute_rejection_reasons.is_empty() - } - - /// Return whether this backend can compute batch gradients for the graph. - #[must_use] - pub fn can_gradient_batch(&self) -> bool { - self.batch_gradient_rejection_reasons.is_empty() - } -} - -/// Common execution-backend interface for future scalar, SIMD, and GPU backends. -pub trait ExecutionBackend { - fn name(&self) -> &'static str; - fn capabilities(&self) -> BackendCapabilities; - - fn compute(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result; - fn compute_batch( - &self, - graph: &CompiledGraph, - batch: BatchInputs<'_>, - buffer: &mut BatchValuesBuffer, - ) -> Result<()>; - fn gradient(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result<(f64, Vec)>; - fn gradient_batch( - &self, - graph: &CompiledGraph, - batch: BatchInputs<'_>, - buffer: &mut BatchGradientsBuffer, - ) -> Result<()>; -} - -/// Flat row-major batch input view. -#[derive(Debug, Clone, Copy, PartialEq)] -pub struct BatchInputs<'a> { - /// Flat row-major input data. - pub data: &'a [f64], - /// Number of rows in the batch. - pub batch_size: usize, - /// Number of inputs per row. - pub input_dim: usize, -} - -impl<'a> BatchInputs<'a> { - /// Create a validated row-major input view. - pub fn new(data: &'a [f64], batch_size: usize, input_dim: usize) -> Result { - if data.len() != batch_size.saturating_mul(input_dim) { - return Err(AutodiffError::InvalidGraph { - reason: "batch data length must equal batch_size * input_dim", - }); - } - Ok(Self { - data, - batch_size, - input_dim, - }) - } - - /// Return one input row. - /// - /// Prefer [`BatchInputs::try_row`] when invalid row indices should return an error. - #[must_use] - pub fn row(&self, index: usize) -> &[f64] { - self.try_row(index).expect("batch row index out of bounds") - } - - /// Return one input row, or an error when `index` is out of range. - pub fn try_row(&self, index: usize) -> Result<&[f64]> { - if index >= self.batch_size { - return Err(AutodiffError::IndexOutOfBounds { - index, - max_index: self.batch_size.saturating_sub(1), - }); - } - let start = index * self.input_dim; - Ok(&self.data[start..start + self.input_dim]) - } -} - -/// Flat row-major batch values. -#[derive(Debug, Clone, PartialEq)] -pub struct BatchValues { - pub data: Vec, - pub batch_size: usize, - pub output_dim: usize, -} - -/// Reusable flat row-major batch value buffer. -#[derive(Debug, Clone, Default, PartialEq)] -pub struct BatchValuesBuffer { - pub data: Vec, - pub batch_size: usize, - pub output_dim: usize, -} - -impl BatchValuesBuffer { - /// Create an empty reusable output buffer. - #[must_use] - pub fn new() -> Self { - Self::default() - } - - fn reset(&mut self, batch_size: usize, output_dim: usize) { - self.data.clear(); - self.data.reserve(batch_size.saturating_mul(output_dim)); - self.batch_size = batch_size; - self.output_dim = output_dim; - } - - /// Clone the current buffer contents into an owned result value. - #[must_use] - pub fn to_values(&self) -> BatchValues { - BatchValues { - data: self.data.clone(), - batch_size: self.batch_size, - output_dim: self.output_dim, - } - } -} - -/// Flat row-major batch scalar-output gradients. -#[derive(Debug, Clone, PartialEq)] -pub struct BatchGradients { - pub values: Vec, - pub gradients: Vec, - pub batch_size: usize, - pub input_dim: usize, -} - -/// Reusable flat row-major batch scalar-output gradient buffer. -#[derive(Debug, Clone, Default, PartialEq)] -pub struct BatchGradientsBuffer { - pub values: Vec, - pub gradients: Vec, - pub batch_size: usize, - pub input_dim: usize, -} - -impl BatchGradientsBuffer { - /// Create an empty reusable gradient buffer. - #[must_use] - pub fn new() -> Self { - Self::default() - } - - fn reset(&mut self, batch_size: usize, input_dim: usize) { - self.values.clear(); - self.gradients.clear(); - self.values.reserve(batch_size); - self.gradients.reserve(batch_size.saturating_mul(input_dim)); - self.batch_size = batch_size; - self.input_dim = input_dim; - } - - /// Clone the current buffer contents into an owned result value. - #[must_use] - pub fn to_gradients(&self) -> BatchGradients { - BatchGradients { - values: self.values.clone(), - gradients: self.gradients.clone(), - batch_size: self.batch_size, - input_dim: self.input_dim, - } - } -} - -/// Static metadata for backend selection and workspace planning. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct CompiledGraphMetadata { - pub num_inputs: usize, - pub num_outputs: usize, - pub num_instructions: usize, - pub num_constants: usize, - pub num_unary: usize, - pub num_binary: usize, - pub value_count: usize, - pub is_scalar_output: bool, -} - -/// Reusable buffers for [`CompiledGraph`] execution. -#[derive(Debug, Clone, Default, PartialEq)] -pub struct CompiledWorkspace { - values: Vec, - cotangents: Vec, - gradients: Vec, - outputs: Vec, -} - -/// Closure-free compiled graph representation. -#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] -#[derive(Debug, Clone, PartialEq)] -pub struct CompiledGraph { - pub(crate) num_inputs: usize, - pub(crate) instructions: Vec, - pub(crate) output_nodes: Vec, - flat_instructions: Vec, -} - -impl CompiledGraph { - pub(crate) fn new( - num_inputs: usize, - instructions: Vec, - output_nodes: Vec, - ) -> Result { - let flat_instructions = Self::lower_flat_instructions(num_inputs, &instructions)?; - Ok(Self { - num_inputs, - instructions, - output_nodes, - flat_instructions, - }) - } - - /// Return number of graph inputs. - #[must_use] - pub fn num_inputs(&self) -> usize { - self.num_inputs - } - - /// Return compiled instructions. - #[must_use] - pub fn instructions(&self) -> &[Instruction] { - &self.instructions - } - - /// Return selected output nodes. - #[must_use] - pub fn output_nodes(&self) -> &[NodeId] { - &self.output_nodes - } - - fn lower_flat_instructions( - num_inputs: usize, - instructions: &[Instruction], - ) -> Result> { - let mut flat = Vec::with_capacity(instructions.len()); - for (offset, instruction) in instructions.iter().enumerate() { - let output = num_inputs + offset; - let item = match *instruction { - Instruction::Constant(value) => FlatInstruction { - opcode: OpCode::Constant, - output, - left: UNUSED_NODE_ID, - right: UNUSED_NODE_ID, - value, - }, - Instruction::Unary { op, arg } => FlatInstruction { - opcode: OpCode::from_multi_ad(op)?, - output, - left: arg, - right: UNUSED_NODE_ID, - value: 0.0, - }, - Instruction::Binary { op, left, right } => FlatInstruction { - opcode: OpCode::from_multi_ad(op)?, - output, - left, - right, - value: 0.0, - }, - }; - flat.push(item); - } - Ok(flat) - } - - /// Return a zero-copy view of flat instructions with explicit output and argument slots. - #[must_use] - pub fn flat_instructions_slice(&self) -> &[FlatInstruction] { - &self.flat_instructions - } - - /// Return owned flat instructions with explicit output and argument slots. - pub fn flat_instructions(&self) -> Result> { - Ok(self.flat_instructions_slice().to_vec()) - } - - /// Validate whether static graph requirements fit backend capabilities. - pub fn validate_backend_capabilities(&self, capabilities: &BackendCapabilities) -> Result<()> { - if !capabilities.supports_f64 { - return Err(AutodiffError::InvalidGraph { - reason: "backend must support f64 execution", - }); - } - if self.output_nodes.len() > 1 && !capabilities.supports_multi_output { - return Err(AutodiffError::InvalidGraph { - reason: "backend does not support multi-output graphs", - }); - } - for instruction in self.flat_instructions_slice() { - if !capabilities.supports_opcode(instruction.opcode) { - return Err(AutodiffError::InvalidGraph { - reason: "backend does not support a required opcode", - }); - } - match instruction.opcode.arity() { - 0 if !capabilities.supports_constants => { - return Err(AutodiffError::InvalidGraph { - reason: "backend does not support constants", - }); - } - 1 if !capabilities.supports_unary => { - return Err(AutodiffError::InvalidGraph { - reason: "backend does not support unary operations", - }); - } - 2 if !capabilities.supports_binary => { - return Err(AutodiffError::InvalidGraph { - reason: "backend does not support binary operations", - }); - } - _ => {} - } - } - Ok(()) - } - - /// Return static execution metadata for backend planning. - #[must_use] - pub fn metadata(&self) -> CompiledGraphMetadata { - let mut num_constants = 0; - let mut num_unary = 0; - let mut num_binary = 0; - for instruction in &self.instructions { - match instruction { - Instruction::Constant(_) => num_constants += 1, - Instruction::Unary { .. } => num_unary += 1, - Instruction::Binary { .. } => num_binary += 1, - } - } - CompiledGraphMetadata { - num_inputs: self.num_inputs, - num_outputs: self.output_nodes.len(), - num_instructions: self.instructions.len(), - num_constants, - num_unary, - num_binary, - value_count: self.num_inputs + self.instructions.len(), - is_scalar_output: self.output_nodes.len() == 1, - } - } - - /// Create a reusable workspace sized for this compiled graph. - #[must_use] - pub fn workspace(&self) -> CompiledWorkspace { - let value_count = self.num_inputs + self.instructions.len(); - CompiledWorkspace { - values: Vec::with_capacity(value_count), - cotangents: Vec::with_capacity(value_count), - gradients: Vec::with_capacity(self.num_inputs), - outputs: Vec::with_capacity(self.output_nodes.len()), - } - } - - fn check_input_len(&self, inputs: &[f64]) -> Result<()> { - if inputs.len() == self.num_inputs { - Ok(()) - } else { - Err(AutodiffError::InvalidGraph { - reason: "input length must match compiled graph input count", - }) - } - } - - fn check_batch(&self, batch: BatchInputs<'_>) -> Result<()> { - if batch.input_dim == self.num_inputs { - Ok(()) - } else { - Err(AutodiffError::InvalidGraph { - reason: "batch input_dim must match compiled graph input count", - }) - } - } - - fn gather_unary(values: &[f64], arg: NodeId) -> Result<[f64; 1]> { - let value = values - .get(arg) - .copied() - .ok_or(AutodiffError::IndexOutOfBounds { - index: arg, - max_index: values.len().saturating_sub(1), - })?; - Ok([value]) - } - - fn gather_binary(values: &[f64], left: NodeId, right: NodeId) -> Result<[f64; 2]> { - let left_value = values - .get(left) - .copied() - .ok_or(AutodiffError::IndexOutOfBounds { - index: left, - max_index: values.len().saturating_sub(1), - })?; - let right_value = values - .get(right) - .copied() - .ok_or(AutodiffError::IndexOutOfBounds { - index: right, - max_index: values.len().saturating_sub(1), - })?; - Ok([left_value, right_value]) - } - - fn fill_values(&self, inputs: &[f64], workspace: &mut CompiledWorkspace) -> Result<()> { - self.check_input_len(inputs)?; - workspace.values.clear(); - workspace - .values - .reserve(self.num_inputs + self.instructions.len()); - workspace.values.extend_from_slice(inputs); - - for instruction in &self.instructions { - let value = match *instruction { - Instruction::Constant(value) => value, - Instruction::Unary { op, arg } => { - let args = Self::gather_unary(&workspace.values, arg)?; - op_rules::forward_value(op, &args)? - } - Instruction::Binary { op, left, right } => { - let args = Self::gather_binary(&workspace.values, left, right)?; - op_rules::forward_value(op, &args)? - } - }; - workspace.values.push(value); - } - Ok(()) - } - - /// Compute primary output. - pub fn compute(&self, inputs: &[f64]) -> Result { - let mut workspace = self.workspace(); - self.compute_with_workspace(inputs, &mut workspace) - } - - /// Compute all selected outputs. - pub fn compute_many(&self, inputs: &[f64]) -> Result> { - let mut workspace = self.workspace(); - Ok(self - .compute_many_with_workspace(inputs, &mut workspace)? - .to_vec()) - } - - /// Compute primary output with a reusable workspace. - pub fn compute_with_workspace( - &self, - inputs: &[f64], - workspace: &mut CompiledWorkspace, - ) -> Result { - let outputs = self.compute_many_with_workspace(inputs, workspace)?; - Ok(outputs.first().copied().unwrap_or(0.0)) - } - - /// Compute all selected outputs with a reusable workspace. - pub fn compute_many_with_workspace<'a>( - &self, - inputs: &[f64], - workspace: &'a mut CompiledWorkspace, - ) -> Result<&'a [f64]> { - self.fill_values(inputs, workspace)?; - workspace.outputs.clear(); - for &output in &self.output_nodes { - let value = - workspace - .values - .get(output) - .copied() - .ok_or(AutodiffError::IndexOutOfBounds { - index: output, - max_index: workspace.values.len().saturating_sub(1), - })?; - workspace.outputs.push(value); - } - Ok(&workspace.outputs) - } - - fn reverse_for_output(&self, output: NodeId, workspace: &mut CompiledWorkspace) -> Result<()> { - workspace.cotangents.clear(); - workspace.cotangents.resize(workspace.values.len(), 0.0); - if output >= workspace.cotangents.len() { - return Err(AutodiffError::IndexOutOfBounds { - index: output, - max_index: workspace.cotangents.len().saturating_sub(1), - }); - } - workspace.cotangents[output] = 1.0; - - for (offset, instruction) in self.instructions.iter().enumerate().rev() { - let node_id = self.num_inputs + offset; - let current = workspace.cotangents[node_id]; - if current == 0.0 { - continue; - } - match *instruction { - Instruction::Constant(_) => {} - Instruction::Unary { op, arg } => { - let args = Self::gather_unary(&workspace.values, arg)?; - let value = workspace.values[node_id]; - if let op_rules::LocalRule::Unary { dy, .. } = - op_rules::local_rule(op, &args, value)? - { - workspace.cotangents[arg] += current * dy; - } - } - Instruction::Binary { op, left, right } => { - let args = Self::gather_binary(&workspace.values, left, right)?; - let value = workspace.values[node_id]; - if let op_rules::LocalRule::Binary { - dy_left, dy_right, .. - } = op_rules::local_rule(op, &args, value)? - { - workspace.cotangents[left] += current * dy_left; - workspace.cotangents[right] += current * dy_right; - } - } - } - } - Ok(()) - } - - /// Compute primary output and gradient. - pub fn gradient(&self, inputs: &[f64]) -> Result<(f64, Vec)> { - let mut workspace = self.workspace(); - let (value, gradient) = self.gradient_with_workspace(inputs, &mut workspace)?; - Ok((value, gradient.to_vec())) - } - - /// Compute primary output and gradient with a reusable workspace. - pub fn gradient_with_workspace<'a>( - &self, - inputs: &[f64], - workspace: &'a mut CompiledWorkspace, - ) -> Result<(f64, &'a [f64])> { - self.fill_values(inputs, workspace)?; - workspace.gradients.clear(); - workspace.gradients.resize(self.num_inputs, 0.0); - let Some(&output) = self.output_nodes.first() else { - return Ok((0.0, &workspace.gradients)); - }; - self.reverse_for_output(output, workspace)?; - workspace - .gradients - .copy_from_slice(&workspace.cotangents[..self.num_inputs]); - Ok((workspace.values[output], &workspace.gradients)) - } - - /// Compute Jacobian for all selected outputs. - pub fn jacobian(&self, inputs: &[f64]) -> Result>> { - let mut workspace = self.workspace(); - self.jacobian_with_workspace(inputs, &mut workspace) - } - - /// Compute Jacobian for all selected outputs with a reusable workspace. - pub fn jacobian_with_workspace( - &self, - inputs: &[f64], - workspace: &mut CompiledWorkspace, - ) -> Result>> { - self.fill_values(inputs, workspace)?; - let mut jacobian = Vec::with_capacity(self.output_nodes.len()); - for &output in &self.output_nodes { - self.reverse_for_output(output, workspace)?; - jacobian.push(workspace.cotangents[..self.num_inputs].to_vec()); - } - Ok(jacobian) - } - - /// Return static compatibility details for a backend. - pub fn backend_support_report(&self, backend: BackendKind) -> Result { - let capabilities = backend.capabilities(); - let supports_required_outputs = - self.output_nodes.len() <= 1 || capabilities.supports_multi_output; - let mut missing_opcodes = Vec::new(); - let mut supports_required_arities = true; - for instruction in self.flat_instructions_slice() { - if !capabilities.supports_opcode(instruction.opcode) - && !missing_opcodes.contains(&instruction.opcode) - { - missing_opcodes.push(instruction.opcode); - } - match instruction.opcode.arity() { - 0 if !capabilities.supports_constants => supports_required_arities = false, - 1 if !capabilities.supports_unary => supports_required_arities = false, - 2 if !capabilities.supports_binary => supports_required_arities = false, - _ => {} - } - } - let runtime_available = backend.runtime_available(); - let supports_required_opcodes = missing_opcodes.is_empty(); - let mut common_reasons = Vec::new(); - if !capabilities.supports_f64 { - common_reasons.push(BackendRejectionReason::MissingF64); - } - if !supports_required_outputs { - common_reasons.push(BackendRejectionReason::UnsupportedOutputs); - } - if !supports_required_opcodes { - common_reasons.push(BackendRejectionReason::UnsupportedOpcodes); - } - if !supports_required_arities { - common_reasons.push(BackendRejectionReason::UnsupportedArities); - } - if !runtime_available { - common_reasons.push(BackendRejectionReason::UnavailableRuntime); - } - - let mut batch_compute_rejection_reasons = common_reasons.clone(); - if !capabilities.supports_batch_compute { - batch_compute_rejection_reasons.push(BackendRejectionReason::NoBatchCompute); - } - let mut batch_gradient_rejection_reasons = common_reasons; - if !capabilities.supports_batch_gradient { - batch_gradient_rejection_reasons.push(BackendRejectionReason::NoBatchGradient); - } - - Ok(BackendSupportReport { - backend, - supports_f64: capabilities.supports_f64, - supports_required_outputs, - supports_required_opcodes, - supports_required_arities, - supports_batch_compute: capabilities.supports_batch_compute, - supports_batch_gradient: capabilities.supports_batch_gradient, - lane_width: backend.lane_width(), - runtime_available, - required_runtime_features: backend.required_runtime_features().to_vec(), - unavailable_runtime_features: backend.unavailable_runtime_features(), - missing_opcodes, - batch_compute_rejection_reasons, - batch_gradient_rejection_reasons, - }) - } - - /// Return static compatibility details for all built-in backends. - pub fn backend_support_reports(&self) -> Result> { - Ok(vec![ - self.backend_support_report(BackendKind::Scalar)?, - self.backend_support_report(BackendKind::MockDeviceCpu)?, - self.backend_support_report(BackendKind::Wgpu)?, - self.backend_support_report(BackendKind::SimdF64x4)?, - self.backend_support_report(BackendKind::SimdF64x2)?, - ]) - } - - /// Return a device-oriented batch buffer plan for a backend. - #[must_use] - pub fn device_batch_plan(&self, backend: BackendKind, batch_size: usize) -> DeviceBatchPlan { - backend.batch_plan(self, batch_size) - } - - /// Allocate mock-device buffers for this compiled graph and batch size. - #[must_use] - pub fn allocate_mock_device_buffers(&self, batch_size: usize) -> DeviceBufferSet { - MockDeviceBackend.allocate_batch_buffers(self, batch_size) - } - - /// Execute batch value computation through mock-device buffers. - pub fn compute_batch_mock_device_into( - &self, - batch: BatchInputs<'_>, - buffers: &mut DeviceBufferSet, - output: &mut BatchValuesBuffer, - ) -> Result { - MockDeviceBackend.compute_batch_with_buffers(self, batch, buffers, output) - } - - /// Execute batch gradient computation through mock-device buffers. - pub fn gradient_batch_mock_device_into( - &self, - batch: BatchInputs<'_>, - buffers: &mut DeviceBufferSet, - output: &mut BatchGradientsBuffer, - ) -> Result { - MockDeviceBackend.gradient_batch_with_buffers(self, batch, buffers, output) - } - - #[cfg(feature = "backend-wgpu")] - /// Allocate real WGPU buffers for this compiled graph and batch size. - pub fn allocate_wgpu_buffers( - &self, - backend: &WgpuBackend, - batch_size: usize, - ) -> Result { - backend.allocate_batch_buffers(self, batch_size) - } - - #[cfg(feature = "backend-wgpu")] - /// Execute batch value computation through real WGPU buffers. - pub fn compute_batch_wgpu_into( - &self, - backend: &WgpuBackend, - batch: BatchInputs<'_>, - buffers: &mut WgpuBufferSet, - output: &mut BatchValuesBuffer, - ) -> Result { - backend.compute_batch_with_buffers(self, batch, buffers, output) - } - - #[cfg(feature = "backend-wgpu")] - /// Execute batch gradient computation through real WGPU buffers. - pub fn gradient_batch_wgpu_into( - &self, - backend: &WgpuBackend, - batch: BatchInputs<'_>, - buffers: &mut WgpuBufferSet, - output: &mut BatchGradientsBuffer, - ) -> Result { - backend.gradient_batch_with_buffers(self, batch, buffers, output) - } - - #[cfg(feature = "backend-wgpu")] - /// Return whether this graph is statically eligible for the exact-safe native WGPU compute path. - #[must_use] - pub fn supports_native_wgpu_batch_compute(&self, backend: &WgpuBackend) -> bool { - backend.supports_native_batch_compute(self) - } - - #[cfg(feature = "backend-wgpu")] - /// Return whether this graph and concrete batch can use the exact-safe native WGPU compute path. - #[must_use] - pub fn supports_native_wgpu_batch_compute_for_batch( - &self, - backend: &WgpuBackend, - batch: BatchInputs<'_>, - ) -> bool { - backend.supports_native_batch_compute_for_batch(self, batch) - } - - /// Return static compatibility details for the preferred SIMD backend. - pub fn simd_support_report(&self) -> Result { - let f64x4 = self.backend_support_report(BackendKind::SimdF64x4)?; - if f64x4.supports_batch_compute || f64x4.supports_batch_gradient { - return Ok(f64x4); - } - self.backend_support_report(BackendKind::SimdF64x2) - } - - /// Return the preferred backend for batch value computation. - #[must_use] - pub fn recommended_batch_compute_backend(&self) -> BackendKind { - match self.backend_support_report(BackendKind::SimdF64x4) { - Ok(report) if report.can_compute_batch() => BackendKind::SimdF64x4, - _ => match self.backend_support_report(BackendKind::SimdF64x2) { - Ok(report) if report.can_compute_batch() => BackendKind::SimdF64x2, - _ => BackendKind::Scalar, - }, - } - } - - /// Return the preferred backend for batch gradient computation. - #[must_use] - pub fn recommended_batch_gradient_backend(&self) -> BackendKind { - match self.backend_support_report(BackendKind::SimdF64x4) { - Ok(report) if report.can_gradient_batch() => BackendKind::SimdF64x4, - _ => match self.backend_support_report(BackendKind::SimdF64x2) { - Ok(report) if report.can_gradient_batch() => BackendKind::SimdF64x2, - _ => BackendKind::Scalar, - }, - } - } - - /// Compute all selected outputs for a batch of rows. - pub fn compute_batch(&self, batch: BatchInputs<'_>) -> Result { - let mut buffer = BatchValuesBuffer::new(); - self.compute_batch_into(batch, &mut buffer)?; - Ok(BatchValues { - data: buffer.data, - batch_size: buffer.batch_size, - output_dim: buffer.output_dim, - }) - } - - /// Compute all selected outputs for a batch with automatic backend dispatch. - pub fn compute_batch_auto(&self, batch: BatchInputs<'_>) -> Result<(BackendKind, BatchValues)> { - let mut buffer = BatchValuesBuffer::new(); - let backend = self.compute_batch_auto_into(batch, &mut buffer)?; - Ok(( - backend, - BatchValues { - data: buffer.data, - batch_size: buffer.batch_size, - output_dim: buffer.output_dim, - }, - )) - } - - /// Compute all selected outputs into a reusable output buffer. - pub fn compute_batch_into( - &self, - batch: BatchInputs<'_>, - buffer: &mut BatchValuesBuffer, - ) -> Result<()> { - self.check_batch(batch)?; - let mut workspace = self.workspace(); - let output_dim = self.output_nodes.len(); - buffer.reset(batch.batch_size, output_dim); - for row_index in 0..batch.batch_size { - let outputs = - self.compute_many_with_workspace(batch.try_row(row_index)?, &mut workspace)?; - buffer.data.extend_from_slice(outputs); - } - Ok(()) - } - - /// Compute all selected outputs into a reusable buffer with automatic backend dispatch. - pub fn compute_batch_auto_into( - &self, - batch: BatchInputs<'_>, - buffer: &mut BatchValuesBuffer, - ) -> Result { - let backend = self.recommended_batch_compute_backend(); - backend.compute_batch(self, batch, buffer)?; - Ok(backend) - } - - /// Compute primary-output values and gradients for a batch of rows. - pub fn gradient_batch(&self, batch: BatchInputs<'_>) -> Result { - let mut buffer = BatchGradientsBuffer::new(); - self.gradient_batch_into(batch, &mut buffer)?; - Ok(BatchGradients { - values: buffer.values, - gradients: buffer.gradients, - batch_size: buffer.batch_size, - input_dim: buffer.input_dim, - }) - } - - /// Compute primary-output values and gradients for a batch with automatic backend dispatch. - pub fn gradient_batch_auto( - &self, - batch: BatchInputs<'_>, - ) -> Result<(BackendKind, BatchGradients)> { - let mut buffer = BatchGradientsBuffer::new(); - let backend = self.gradient_batch_auto_into(batch, &mut buffer)?; - Ok(( - backend, - BatchGradients { - values: buffer.values, - gradients: buffer.gradients, - batch_size: buffer.batch_size, - input_dim: buffer.input_dim, - }, - )) - } - - /// Compute primary-output values and gradients into a reusable gradient buffer. - pub fn gradient_batch_into( - &self, - batch: BatchInputs<'_>, - buffer: &mut BatchGradientsBuffer, - ) -> Result<()> { - self.check_batch(batch)?; - let mut workspace = self.workspace(); - buffer.reset(batch.batch_size, self.num_inputs); - for row_index in 0..batch.batch_size { - let (value, gradient) = - self.gradient_with_workspace(batch.try_row(row_index)?, &mut workspace)?; - buffer.values.push(value); - buffer.gradients.extend_from_slice(gradient); - } - Ok(()) - } - - /// Compute primary-output values and gradients into a reusable buffer with automatic backend dispatch. - pub fn gradient_batch_auto_into( - &self, - batch: BatchInputs<'_>, - buffer: &mut BatchGradientsBuffer, - ) -> Result { - let backend = self.recommended_batch_gradient_backend(); - backend.gradient_batch(self, batch, buffer)?; - Ok(backend) - } -} - -fn simd_unsupported_opcode_error() -> AutodiffError { - AutodiffError::InvalidGraph { - reason: "simd backend does not support a required opcode", - } -} - -fn simd_scalar_value(opcode: OpCode, args: &[f64]) -> Result { - let op = opcode - .to_multi_ad() - .ok_or_else(simd_unsupported_opcode_error)?; - op_rules::forward_value(op, args) -} - -fn simd_scalar_first_derivatives(opcode: OpCode, args: &[f64], value: f64) -> Result> { - let op = opcode - .to_multi_ad() - .ok_or_else(simd_unsupported_opcode_error)?; - op_rules::first_derivatives(op, args, value) -} - -fn append_scalar_compute_tail( - graph: &CompiledGraph, - batch: BatchInputs<'_>, - start_row: usize, - buffer: &mut BatchValuesBuffer, -) -> Result<()> { - let mut workspace = graph.workspace(); - for row_index in start_row..batch.batch_size { - let outputs = - graph.compute_many_with_workspace(batch.try_row(row_index)?, &mut workspace)?; - buffer.data.extend_from_slice(outputs); - } - Ok(()) -} - -fn append_scalar_gradient_tail( - graph: &CompiledGraph, - batch: BatchInputs<'_>, - start_row: usize, - buffer: &mut BatchGradientsBuffer, -) -> Result<()> { - let mut workspace = graph.workspace(); - for row_index in start_row..batch.batch_size { - let (value, gradient) = - graph.gradient_with_workspace(batch.try_row(row_index)?, &mut workspace)?; - buffer.values.push(value); - buffer.gradients.extend_from_slice(gradient); - } - Ok(()) -} - -#[cfg(target_arch = "x86_64")] -fn checked_m128_lane(values: &[__m128d], index: NodeId) -> Result<__m128d> { - values - .get(index) - .copied() - .ok_or(AutodiffError::IndexOutOfBounds { - index, - max_index: values.len().saturating_sub(1), - }) -} - -#[cfg(target_arch = "x86_64")] -unsafe fn add_m128_cotangent( - cotangents: &mut [__m128d], - index: NodeId, - contribution: __m128d, -) -> Result<()> { - use std::arch::x86_64::_mm_add_pd; - - let max_index = cotangents.len().saturating_sub(1); - let target = cotangents - .get_mut(index) - .ok_or(AutodiffError::IndexOutOfBounds { index, max_index })?; - *target = _mm_add_pd(*target, contribution); - Ok(()) -} - -#[cfg(target_arch = "x86_64")] -unsafe fn active_m128_contribution(current: __m128d, contribution: __m128d) -> __m128d { - use std::arch::x86_64::{_mm_and_pd, _mm_cmpneq_pd, _mm_setzero_pd}; - - let active = _mm_cmpneq_pd(current, _mm_setzero_pd()); - _mm_and_pd(contribution, active) -} - -#[cfg(target_arch = "x86_64")] -unsafe fn simd_f64x2_scalar_unary(input: __m128d, opcode: OpCode) -> Result<__m128d> { - use std::arch::x86_64::{_mm_set_pd, _mm_storeu_pd}; - - let mut stored = [0.0_f64; 2]; - _mm_storeu_pd(stored.as_mut_ptr(), input); - let first = simd_scalar_value(opcode, &[stored[0]])?; - let second = simd_scalar_value(opcode, &[stored[1]])?; - Ok(_mm_set_pd(second, first)) -} - -#[cfg(target_arch = "x86_64")] -unsafe fn simd_f64x2_scalar_unary_derivative( - input: __m128d, - output: __m128d, - opcode: OpCode, -) -> Result<__m128d> { - use std::arch::x86_64::{_mm_set_pd, _mm_storeu_pd}; - - let mut input_values = [0.0_f64; 2]; - let mut output_values = [0.0_f64; 2]; - _mm_storeu_pd(input_values.as_mut_ptr(), input); - _mm_storeu_pd(output_values.as_mut_ptr(), output); - let first = simd_scalar_first_derivatives(opcode, &[input_values[0]], output_values[0])?[0]; - let second = simd_scalar_first_derivatives(opcode, &[input_values[1]], output_values[1])?[0]; - Ok(_mm_set_pd(second, first)) -} - -#[cfg(target_arch = "x86_64")] -unsafe fn simd_f64x2_scalar_binary( - left: __m128d, - right: __m128d, - opcode: OpCode, -) -> Result<__m128d> { - use std::arch::x86_64::{_mm_set_pd, _mm_storeu_pd}; - - let mut left_values = [0.0_f64; 2]; - let mut right_values = [0.0_f64; 2]; - _mm_storeu_pd(left_values.as_mut_ptr(), left); - _mm_storeu_pd(right_values.as_mut_ptr(), right); - let first = simd_scalar_value(opcode, &[left_values[0], right_values[0]])?; - let second = simd_scalar_value(opcode, &[left_values[1], right_values[1]])?; - Ok(_mm_set_pd(second, first)) -} - -#[cfg(target_arch = "x86_64")] -unsafe fn simd_f64x2_scalar_binary_derivatives( - left: __m128d, - right: __m128d, - output: __m128d, - opcode: OpCode, -) -> Result<(__m128d, __m128d)> { - use std::arch::x86_64::{_mm_set_pd, _mm_storeu_pd}; - - let mut left_values = [0.0_f64; 2]; - let mut right_values = [0.0_f64; 2]; - let mut output_values = [0.0_f64; 2]; - _mm_storeu_pd(left_values.as_mut_ptr(), left); - _mm_storeu_pd(right_values.as_mut_ptr(), right); - _mm_storeu_pd(output_values.as_mut_ptr(), output); - let first = simd_scalar_first_derivatives( - opcode, - &[left_values[0], right_values[0]], - output_values[0], - )?; - let second = simd_scalar_first_derivatives( - opcode, - &[left_values[1], right_values[1]], - output_values[1], - )?; - Ok(( - _mm_set_pd(second[0], first[0]), - _mm_set_pd(second[1], first[1]), - )) -} - -#[cfg(target_arch = "x86_64")] -unsafe fn simd_f64x2_forward_values( - flat: &[FlatInstruction], - values: &mut Vec<__m128d>, -) -> Result<()> { - use std::arch::x86_64::{ - _mm_add_pd, _mm_and_pd, _mm_andnot_pd, _mm_cmpgt_pd, _mm_div_pd, _mm_mul_pd, _mm_set1_pd, - _mm_setzero_pd, _mm_sqrt_pd, _mm_sub_pd, - }; - - for instruction in flat { - let value = match instruction.opcode { - OpCode::Constant => _mm_set1_pd(instruction.value), - OpCode::Add => _mm_add_pd( - checked_m128_lane(values, instruction.left)?, - checked_m128_lane(values, instruction.right)?, - ), - OpCode::Sub => _mm_sub_pd( - checked_m128_lane(values, instruction.left)?, - checked_m128_lane(values, instruction.right)?, - ), - OpCode::Mul => _mm_mul_pd( - checked_m128_lane(values, instruction.left)?, - checked_m128_lane(values, instruction.right)?, - ), - OpCode::Div => _mm_div_pd( - checked_m128_lane(values, instruction.left)?, - checked_m128_lane(values, instruction.right)?, - ), - OpCode::Pow | OpCode::LogAddExp => { - let left = checked_m128_lane(values, instruction.left)?; - let right = checked_m128_lane(values, instruction.right)?; - simd_f64x2_scalar_binary(left, right, instruction.opcode)? - } - OpCode::Neg => _mm_sub_pd( - _mm_setzero_pd(), - checked_m128_lane(values, instruction.left)?, - ), - OpCode::Sqrt => _mm_sqrt_pd(checked_m128_lane(values, instruction.left)?), - OpCode::Relu => { - let input = checked_m128_lane(values, instruction.left)?; - let mask = _mm_cmpgt_pd(input, _mm_setzero_pd()); - _mm_and_pd(input, mask) - } - OpCode::Abs => { - let input = checked_m128_lane(values, instruction.left)?; - _mm_andnot_pd(_mm_set1_pd(-0.0), input) - } - OpCode::Sin - | OpCode::Cos - | OpCode::Tan - | OpCode::Tanh - | OpCode::Log1pExp - | OpCode::Exp - | OpCode::Ln => { - let input = checked_m128_lane(values, instruction.left)?; - simd_f64x2_scalar_unary(input, instruction.opcode)? - } - }; - values.push(value); - } - Ok(()) -} - -#[cfg(target_arch = "x86_64")] -fn checked_m256_lane(values: &[__m256d], index: NodeId) -> Result<__m256d> { - values - .get(index) - .copied() - .ok_or(AutodiffError::IndexOutOfBounds { - index, - max_index: values.len().saturating_sub(1), - }) -} - -#[cfg(target_arch = "x86_64")] -#[target_feature(enable = "avx")] -unsafe fn add_m256_cotangent( - cotangents: &mut [__m256d], - index: NodeId, - contribution: __m256d, -) -> Result<()> { - use std::arch::x86_64::_mm256_add_pd; - - let max_index = cotangents.len().saturating_sub(1); - let target = cotangents - .get_mut(index) - .ok_or(AutodiffError::IndexOutOfBounds { index, max_index })?; - *target = _mm256_add_pd(*target, contribution); - Ok(()) -} - -#[cfg(target_arch = "x86_64")] -#[target_feature(enable = "avx")] -unsafe fn active_m256_contribution(current: __m256d, contribution: __m256d) -> __m256d { - use std::arch::x86_64::{_mm256_and_pd, _mm256_cmp_pd, _mm256_setzero_pd, _CMP_NEQ_UQ}; - - let active = _mm256_cmp_pd(current, _mm256_setzero_pd(), _CMP_NEQ_UQ); - _mm256_and_pd(contribution, active) -} - -#[cfg(target_arch = "x86_64")] -#[target_feature(enable = "avx")] -unsafe fn simd_f64x4_scalar_unary(input: __m256d, opcode: OpCode) -> Result<__m256d> { - use std::arch::x86_64::{_mm256_set_pd, _mm256_storeu_pd}; - - let mut stored = [0.0_f64; 4]; - _mm256_storeu_pd(stored.as_mut_ptr(), input); - let first = simd_scalar_value(opcode, &[stored[0]])?; - let second = simd_scalar_value(opcode, &[stored[1]])?; - let third = simd_scalar_value(opcode, &[stored[2]])?; - let fourth = simd_scalar_value(opcode, &[stored[3]])?; - Ok(_mm256_set_pd(fourth, third, second, first)) -} - -#[cfg(target_arch = "x86_64")] -#[target_feature(enable = "avx")] -unsafe fn simd_f64x4_scalar_unary_derivative( - input: __m256d, - output: __m256d, - opcode: OpCode, -) -> Result<__m256d> { - use std::arch::x86_64::{_mm256_set_pd, _mm256_storeu_pd}; - - let mut input_values = [0.0_f64; 4]; - let mut output_values = [0.0_f64; 4]; - _mm256_storeu_pd(input_values.as_mut_ptr(), input); - _mm256_storeu_pd(output_values.as_mut_ptr(), output); - let first = simd_scalar_first_derivatives(opcode, &[input_values[0]], output_values[0])?[0]; - let second = simd_scalar_first_derivatives(opcode, &[input_values[1]], output_values[1])?[0]; - let third = simd_scalar_first_derivatives(opcode, &[input_values[2]], output_values[2])?[0]; - let fourth = simd_scalar_first_derivatives(opcode, &[input_values[3]], output_values[3])?[0]; - Ok(_mm256_set_pd(fourth, third, second, first)) -} - -#[cfg(target_arch = "x86_64")] -#[target_feature(enable = "avx")] -unsafe fn simd_f64x4_scalar_binary( - left: __m256d, - right: __m256d, - opcode: OpCode, -) -> Result<__m256d> { - use std::arch::x86_64::{_mm256_set_pd, _mm256_storeu_pd}; - - let mut left_values = [0.0_f64; 4]; - let mut right_values = [0.0_f64; 4]; - _mm256_storeu_pd(left_values.as_mut_ptr(), left); - _mm256_storeu_pd(right_values.as_mut_ptr(), right); - let first = simd_scalar_value(opcode, &[left_values[0], right_values[0]])?; - let second = simd_scalar_value(opcode, &[left_values[1], right_values[1]])?; - let third = simd_scalar_value(opcode, &[left_values[2], right_values[2]])?; - let fourth = simd_scalar_value(opcode, &[left_values[3], right_values[3]])?; - Ok(_mm256_set_pd(fourth, third, second, first)) -} - -#[cfg(target_arch = "x86_64")] -#[target_feature(enable = "avx")] -unsafe fn simd_f64x4_scalar_binary_derivatives( - left: __m256d, - right: __m256d, - output: __m256d, - opcode: OpCode, -) -> Result<(__m256d, __m256d)> { - use std::arch::x86_64::{_mm256_set_pd, _mm256_storeu_pd}; - - let mut left_values = [0.0_f64; 4]; - let mut right_values = [0.0_f64; 4]; - let mut output_values = [0.0_f64; 4]; - _mm256_storeu_pd(left_values.as_mut_ptr(), left); - _mm256_storeu_pd(right_values.as_mut_ptr(), right); - _mm256_storeu_pd(output_values.as_mut_ptr(), output); - let first = simd_scalar_first_derivatives( - opcode, - &[left_values[0], right_values[0]], - output_values[0], - )?; - let second = simd_scalar_first_derivatives( - opcode, - &[left_values[1], right_values[1]], - output_values[1], - )?; - let third = simd_scalar_first_derivatives( - opcode, - &[left_values[2], right_values[2]], - output_values[2], - )?; - let fourth = simd_scalar_first_derivatives( - opcode, - &[left_values[3], right_values[3]], - output_values[3], - )?; - Ok(( - _mm256_set_pd(fourth[0], third[0], second[0], first[0]), - _mm256_set_pd(fourth[1], third[1], second[1], first[1]), - )) -} - -#[cfg(target_arch = "x86_64")] -#[target_feature(enable = "avx")] -unsafe fn simd_f64x4_forward_values( - flat: &[FlatInstruction], - values: &mut Vec<__m256d>, -) -> Result<()> { - use std::arch::x86_64::{ - _mm256_add_pd, _mm256_and_pd, _mm256_andnot_pd, _mm256_cmp_pd, _mm256_div_pd, - _mm256_mul_pd, _mm256_set1_pd, _mm256_setzero_pd, _mm256_sqrt_pd, _mm256_sub_pd, - _CMP_GT_OQ, - }; - - for instruction in flat { - let value = match instruction.opcode { - OpCode::Constant => _mm256_set1_pd(instruction.value), - OpCode::Add => _mm256_add_pd( - checked_m256_lane(values, instruction.left)?, - checked_m256_lane(values, instruction.right)?, - ), - OpCode::Sub => _mm256_sub_pd( - checked_m256_lane(values, instruction.left)?, - checked_m256_lane(values, instruction.right)?, - ), - OpCode::Mul => _mm256_mul_pd( - checked_m256_lane(values, instruction.left)?, - checked_m256_lane(values, instruction.right)?, - ), - OpCode::Div => _mm256_div_pd( - checked_m256_lane(values, instruction.left)?, - checked_m256_lane(values, instruction.right)?, - ), - OpCode::Pow | OpCode::LogAddExp => { - let left = checked_m256_lane(values, instruction.left)?; - let right = checked_m256_lane(values, instruction.right)?; - simd_f64x4_scalar_binary(left, right, instruction.opcode)? - } - OpCode::Neg => _mm256_sub_pd( - _mm256_setzero_pd(), - checked_m256_lane(values, instruction.left)?, - ), - OpCode::Sqrt => _mm256_sqrt_pd(checked_m256_lane(values, instruction.left)?), - OpCode::Relu => { - let input = checked_m256_lane(values, instruction.left)?; - let mask = _mm256_cmp_pd(input, _mm256_setzero_pd(), _CMP_GT_OQ); - _mm256_and_pd(input, mask) - } - OpCode::Abs => { - let input = checked_m256_lane(values, instruction.left)?; - _mm256_andnot_pd(_mm256_set1_pd(-0.0), input) - } - OpCode::Sin - | OpCode::Cos - | OpCode::Tan - | OpCode::Tanh - | OpCode::Log1pExp - | OpCode::Exp - | OpCode::Ln => { - let input = checked_m256_lane(values, instruction.left)?; - simd_f64x4_scalar_unary(input, instruction.opcode)? - } - }; - values.push(value); - } - Ok(()) -} - -#[cfg(target_arch = "x86_64")] -fn compute_batch_simd_f64x2( - graph: &CompiledGraph, - batch: BatchInputs<'_>, - buffer: &mut BatchValuesBuffer, -) -> Result<()> { - use std::arch::x86_64::{_mm_set_pd, _mm_storeu_pd}; - - graph.check_batch(batch)?; - let flat = graph.flat_instructions_slice(); - let output_dim = graph.output_nodes.len(); - let value_count = graph.num_inputs + graph.instructions.len(); - buffer.reset(batch.batch_size, output_dim); - - let mut values: Vec<__m128d> = Vec::with_capacity(value_count); - let mut pair_outputs: Vec<[f64; 2]> = Vec::with_capacity(output_dim); - let mut row_index = 0; - while row_index + 1 < batch.batch_size { - let first = batch.try_row(row_index)?; - let second = batch.try_row(row_index + 1)?; - values.clear(); - for input_index in 0..graph.num_inputs { - unsafe { - values.push(_mm_set_pd(second[input_index], first[input_index])); - } - } - - unsafe { - simd_f64x2_forward_values(flat, &mut values)?; - } - - pair_outputs.clear(); - for &output in &graph.output_nodes { - let lane = checked_m128_lane(&values, output)?; - let mut stored = [0.0_f64; 2]; - unsafe { - _mm_storeu_pd(stored.as_mut_ptr(), lane); - } - pair_outputs.push(stored); - } - for output in &pair_outputs { - buffer.data.push(output[0]); - } - for output in &pair_outputs { - buffer.data.push(output[1]); - } - - row_index += 2; - } - - append_scalar_compute_tail(graph, batch, row_index, buffer)?; - - Ok(()) -} - -#[cfg(target_arch = "x86_64")] -#[target_feature(enable = "avx")] -unsafe fn compute_batch_simd_f64x4_impl( - graph: &CompiledGraph, - batch: BatchInputs<'_>, - buffer: &mut BatchValuesBuffer, -) -> Result<()> { - use std::arch::x86_64::{_mm256_set_pd, _mm256_storeu_pd}; - - graph.check_batch(batch)?; - let flat = graph.flat_instructions_slice(); - let output_dim = graph.output_nodes.len(); - let value_count = graph.num_inputs + graph.instructions.len(); - buffer.reset(batch.batch_size, output_dim); - - let mut values: Vec<__m256d> = Vec::with_capacity(value_count); - let mut quad_outputs: Vec<[f64; 4]> = Vec::with_capacity(output_dim); - let mut row_index = 0; - while row_index + 3 < batch.batch_size { - let first = batch.try_row(row_index)?; - let second = batch.try_row(row_index + 1)?; - let third = batch.try_row(row_index + 2)?; - let fourth = batch.try_row(row_index + 3)?; - values.clear(); - for input_index in 0..graph.num_inputs { - values.push(_mm256_set_pd( - fourth[input_index], - third[input_index], - second[input_index], - first[input_index], - )); - } - - simd_f64x4_forward_values(flat, &mut values)?; - - quad_outputs.clear(); - for &output in &graph.output_nodes { - let lane = checked_m256_lane(&values, output)?; - let mut stored = [0.0_f64; 4]; - _mm256_storeu_pd(stored.as_mut_ptr(), lane); - quad_outputs.push(stored); - } - for lane_index in 0..4 { - for output in &quad_outputs { - buffer.data.push(output[lane_index]); - } - } - - row_index += 4; - } - - append_scalar_compute_tail(graph, batch, row_index, buffer)?; - - Ok(()) -} - -#[cfg(target_arch = "x86_64")] -fn compute_batch_simd_f64x4( - graph: &CompiledGraph, - batch: BatchInputs<'_>, - buffer: &mut BatchValuesBuffer, -) -> Result<()> { - if !supports_simd_f64x4_runtime() { - return Err(AutodiffError::InvalidGraph { - reason: "simd f64x4 backend requires x86_64 AVX support", - }); - } - unsafe { compute_batch_simd_f64x4_impl(graph, batch, buffer) } -} - -#[cfg(not(target_arch = "x86_64"))] -fn compute_batch_simd_f64x4( - _graph: &CompiledGraph, - _batch: BatchInputs<'_>, - _buffer: &mut BatchValuesBuffer, -) -> Result<()> { - Err(AutodiffError::InvalidGraph { - reason: "simd f64x4 backend requires x86_64 AVX support", - }) -} - -#[cfg(not(target_arch = "x86_64"))] -fn compute_batch_simd_f64x2( - _graph: &CompiledGraph, - _batch: BatchInputs<'_>, - _buffer: &mut BatchValuesBuffer, -) -> Result<()> { - Err(AutodiffError::InvalidGraph { - reason: "simd backend requires x86_64 SSE2 support", - }) -} - -#[cfg(target_arch = "x86_64")] -fn gradient_batch_simd_f64x2( - graph: &CompiledGraph, - batch: BatchInputs<'_>, - buffer: &mut BatchGradientsBuffer, -) -> Result<()> { - use std::arch::x86_64::{ - _mm_add_pd, _mm_and_pd, _mm_cmpgt_pd, _mm_cmplt_pd, _mm_div_pd, _mm_mul_pd, _mm_set1_pd, - _mm_set_pd, _mm_setzero_pd, _mm_storeu_pd, _mm_sub_pd, - }; - - graph.check_batch(batch)?; - let flat = graph.flat_instructions_slice(); - buffer.reset(batch.batch_size, graph.num_inputs); - - let value_count = graph.num_inputs + graph.instructions.len(); - let mut values: Vec<__m128d> = Vec::with_capacity(value_count); - let mut cotangents: Vec<__m128d> = Vec::with_capacity(value_count); - let mut gradient_pairs = Vec::with_capacity(graph.num_inputs); - let mut row_index = 0; - while row_index + 1 < batch.batch_size { - let first = batch.try_row(row_index)?; - let second = batch.try_row(row_index + 1)?; - values.clear(); - for input_index in 0..graph.num_inputs { - unsafe { - values.push(_mm_set_pd(second[input_index], first[input_index])); - } - } - - unsafe { - simd_f64x2_forward_values(flat, &mut values)?; - } - - let Some(&output) = graph.output_nodes.first() else { - buffer.values.extend_from_slice(&[0.0, 0.0]); - buffer - .gradients - .resize(buffer.gradients.len() + graph.num_inputs * 2, 0.0); - row_index += 2; - continue; - }; - - cotangents.clear(); - cotangents.resize_with(value_count, || unsafe { _mm_setzero_pd() }); - let max_index = cotangents.len().saturating_sub(1); - *cotangents - .get_mut(output) - .ok_or(AutodiffError::IndexOutOfBounds { - index: output, - max_index, - })? = unsafe { _mm_set1_pd(1.0) }; - - for instruction in flat.iter().rev() { - let current = checked_m128_lane(&cotangents, instruction.output)?; - unsafe { - match instruction.opcode { - OpCode::Constant => {} - OpCode::Add => { - let contribution = active_m128_contribution(current, current); - add_m128_cotangent(&mut cotangents, instruction.left, contribution)?; - add_m128_cotangent(&mut cotangents, instruction.right, contribution)?; - } - OpCode::Sub => { - add_m128_cotangent( - &mut cotangents, - instruction.left, - active_m128_contribution(current, current), - )?; - add_m128_cotangent( - &mut cotangents, - instruction.right, - active_m128_contribution( - current, - _mm_sub_pd(_mm_setzero_pd(), current), - ), - )?; - } - OpCode::Mul => { - let left = checked_m128_lane(&values, instruction.left)?; - let right = checked_m128_lane(&values, instruction.right)?; - add_m128_cotangent( - &mut cotangents, - instruction.left, - active_m128_contribution(current, _mm_mul_pd(current, right)), - )?; - add_m128_cotangent( - &mut cotangents, - instruction.right, - active_m128_contribution(current, _mm_mul_pd(current, left)), - )?; - } - OpCode::Div => { - let left = checked_m128_lane(&values, instruction.left)?; - let right = checked_m128_lane(&values, instruction.right)?; - let right_squared = _mm_mul_pd(right, right); - add_m128_cotangent( - &mut cotangents, - instruction.left, - active_m128_contribution(current, _mm_div_pd(current, right)), - )?; - let right_contribution = _mm_sub_pd( - _mm_setzero_pd(), - _mm_div_pd(_mm_mul_pd(current, left), right_squared), - ); - add_m128_cotangent( - &mut cotangents, - instruction.right, - active_m128_contribution(current, right_contribution), - )?; - } - OpCode::Pow | OpCode::LogAddExp => { - let left = checked_m128_lane(&values, instruction.left)?; - let right = checked_m128_lane(&values, instruction.right)?; - let output_value = checked_m128_lane(&values, instruction.output)?; - let (left_derivative, right_derivative) = - simd_f64x2_scalar_binary_derivatives( - left, - right, - output_value, - instruction.opcode, - )?; - add_m128_cotangent( - &mut cotangents, - instruction.left, - active_m128_contribution(current, _mm_mul_pd(current, left_derivative)), - )?; - add_m128_cotangent( - &mut cotangents, - instruction.right, - active_m128_contribution( - current, - _mm_mul_pd(current, right_derivative), - ), - )?; - } - OpCode::Neg => { - add_m128_cotangent( - &mut cotangents, - instruction.left, - active_m128_contribution( - current, - _mm_sub_pd(_mm_setzero_pd(), current), - ), - )?; - } - OpCode::Sqrt => { - let output_value = checked_m128_lane(&values, instruction.output)?; - let contribution = - _mm_div_pd(_mm_mul_pd(current, _mm_set1_pd(0.5)), output_value); - add_m128_cotangent( - &mut cotangents, - instruction.left, - active_m128_contribution(current, contribution), - )?; - } - OpCode::Relu => { - let input_value = checked_m128_lane(&values, instruction.left)?; - let mask = _mm_cmpgt_pd(input_value, _mm_setzero_pd()); - add_m128_cotangent( - &mut cotangents, - instruction.left, - active_m128_contribution(current, _mm_and_pd(current, mask)), - )?; - } - OpCode::Sin - | OpCode::Cos - | OpCode::Tan - | OpCode::Tanh - | OpCode::Log1pExp - | OpCode::Exp - | OpCode::Ln => { - let input_value = checked_m128_lane(&values, instruction.left)?; - let output_value = checked_m128_lane(&values, instruction.output)?; - let derivative = simd_f64x2_scalar_unary_derivative( - input_value, - output_value, - instruction.opcode, - )?; - add_m128_cotangent( - &mut cotangents, - instruction.left, - active_m128_contribution(current, _mm_mul_pd(current, derivative)), - )?; - } - OpCode::Abs => { - let input_value = checked_m128_lane(&values, instruction.left)?; - let positive = _mm_cmpgt_pd(input_value, _mm_setzero_pd()); - let negative = _mm_cmplt_pd(input_value, _mm_setzero_pd()); - let sign = _mm_add_pd( - _mm_and_pd(_mm_set1_pd(1.0), positive), - _mm_and_pd(_mm_set1_pd(-1.0), negative), - ); - add_m128_cotangent( - &mut cotangents, - instruction.left, - active_m128_contribution(current, _mm_mul_pd(current, sign)), - )?; - } - } - } - } - - let output_value = checked_m128_lane(&values, output)?; - let mut value_pair = [0.0_f64; 2]; - unsafe { - _mm_storeu_pd(value_pair.as_mut_ptr(), output_value); - } - buffer.values.extend_from_slice(&value_pair); - - gradient_pairs.clear(); - for input_index in 0..graph.num_inputs { - let gradient_lane = checked_m128_lane(&cotangents, input_index)?; - let mut stored = [0.0_f64; 2]; - unsafe { - _mm_storeu_pd(stored.as_mut_ptr(), gradient_lane); - } - gradient_pairs.push(stored); - } - for pair in &gradient_pairs { - buffer.gradients.push(pair[0]); - } - for pair in &gradient_pairs { - buffer.gradients.push(pair[1]); - } - - row_index += 2; - } - - append_scalar_gradient_tail(graph, batch, row_index, buffer)?; - - Ok(()) -} - -#[cfg(target_arch = "x86_64")] -#[target_feature(enable = "avx")] -unsafe fn gradient_batch_simd_f64x4_impl( - graph: &CompiledGraph, - batch: BatchInputs<'_>, - buffer: &mut BatchGradientsBuffer, -) -> Result<()> { - use std::arch::x86_64::{ - _mm256_add_pd, _mm256_and_pd, _mm256_cmp_pd, _mm256_div_pd, _mm256_mul_pd, _mm256_set1_pd, - _mm256_set_pd, _mm256_setzero_pd, _mm256_storeu_pd, _mm256_sub_pd, _CMP_GT_OQ, _CMP_LT_OQ, - }; - - graph.check_batch(batch)?; - let flat = graph.flat_instructions_slice(); - buffer.reset(batch.batch_size, graph.num_inputs); - - let value_count = graph.num_inputs + graph.instructions.len(); - let mut values: Vec<__m256d> = Vec::with_capacity(value_count); - let mut cotangents: Vec<__m256d> = Vec::with_capacity(value_count); - let mut gradient_quads = Vec::with_capacity(graph.num_inputs); - let mut row_index = 0; - while row_index + 3 < batch.batch_size { - let first = batch.try_row(row_index)?; - let second = batch.try_row(row_index + 1)?; - let third = batch.try_row(row_index + 2)?; - let fourth = batch.try_row(row_index + 3)?; - values.clear(); - for input_index in 0..graph.num_inputs { - values.push(_mm256_set_pd( - fourth[input_index], - third[input_index], - second[input_index], - first[input_index], - )); - } - - simd_f64x4_forward_values(flat, &mut values)?; - - let Some(&output) = graph.output_nodes.first() else { - buffer.values.extend_from_slice(&[0.0; 4]); - buffer - .gradients - .resize(buffer.gradients.len() + graph.num_inputs * 4, 0.0); - row_index += 4; - continue; - }; - - cotangents.clear(); - cotangents.resize_with(value_count, || _mm256_setzero_pd()); - let max_index = cotangents.len().saturating_sub(1); - *cotangents - .get_mut(output) - .ok_or(AutodiffError::IndexOutOfBounds { - index: output, - max_index, - })? = _mm256_set1_pd(1.0); - - for instruction in flat.iter().rev() { - let current = checked_m256_lane(&cotangents, instruction.output)?; - match instruction.opcode { - OpCode::Constant => {} - OpCode::Add => { - let contribution = active_m256_contribution(current, current); - add_m256_cotangent(&mut cotangents, instruction.left, contribution)?; - add_m256_cotangent(&mut cotangents, instruction.right, contribution)?; - } - OpCode::Sub => { - add_m256_cotangent( - &mut cotangents, - instruction.left, - active_m256_contribution(current, current), - )?; - add_m256_cotangent( - &mut cotangents, - instruction.right, - active_m256_contribution( - current, - _mm256_sub_pd(_mm256_setzero_pd(), current), - ), - )?; - } - OpCode::Mul => { - let left = checked_m256_lane(&values, instruction.left)?; - let right = checked_m256_lane(&values, instruction.right)?; - add_m256_cotangent( - &mut cotangents, - instruction.left, - active_m256_contribution(current, _mm256_mul_pd(current, right)), - )?; - add_m256_cotangent( - &mut cotangents, - instruction.right, - active_m256_contribution(current, _mm256_mul_pd(current, left)), - )?; - } - OpCode::Div => { - let left = checked_m256_lane(&values, instruction.left)?; - let right = checked_m256_lane(&values, instruction.right)?; - let right_squared = _mm256_mul_pd(right, right); - add_m256_cotangent( - &mut cotangents, - instruction.left, - active_m256_contribution(current, _mm256_div_pd(current, right)), - )?; - let right_contribution = _mm256_sub_pd( - _mm256_setzero_pd(), - _mm256_div_pd(_mm256_mul_pd(current, left), right_squared), - ); - add_m256_cotangent( - &mut cotangents, - instruction.right, - active_m256_contribution(current, right_contribution), - )?; - } - OpCode::Pow | OpCode::LogAddExp => { - let left = checked_m256_lane(&values, instruction.left)?; - let right = checked_m256_lane(&values, instruction.right)?; - let output_value = checked_m256_lane(&values, instruction.output)?; - let (left_derivative, right_derivative) = simd_f64x4_scalar_binary_derivatives( - left, - right, - output_value, - instruction.opcode, - )?; - add_m256_cotangent( - &mut cotangents, - instruction.left, - active_m256_contribution(current, _mm256_mul_pd(current, left_derivative)), - )?; - add_m256_cotangent( - &mut cotangents, - instruction.right, - active_m256_contribution(current, _mm256_mul_pd(current, right_derivative)), - )?; - } - OpCode::Neg => { - add_m256_cotangent( - &mut cotangents, - instruction.left, - active_m256_contribution( - current, - _mm256_sub_pd(_mm256_setzero_pd(), current), - ), - )?; - } - OpCode::Sqrt => { - let output_value = checked_m256_lane(&values, instruction.output)?; - let contribution = - _mm256_div_pd(_mm256_mul_pd(current, _mm256_set1_pd(0.5)), output_value); - add_m256_cotangent( - &mut cotangents, - instruction.left, - active_m256_contribution(current, contribution), - )?; - } - OpCode::Relu => { - let input_value = checked_m256_lane(&values, instruction.left)?; - let mask = _mm256_cmp_pd(input_value, _mm256_setzero_pd(), _CMP_GT_OQ); - add_m256_cotangent( - &mut cotangents, - instruction.left, - active_m256_contribution(current, _mm256_and_pd(current, mask)), - )?; - } - OpCode::Sin - | OpCode::Cos - | OpCode::Tan - | OpCode::Tanh - | OpCode::Log1pExp - | OpCode::Exp - | OpCode::Ln => { - let input_value = checked_m256_lane(&values, instruction.left)?; - let output_value = checked_m256_lane(&values, instruction.output)?; - let derivative = simd_f64x4_scalar_unary_derivative( - input_value, - output_value, - instruction.opcode, - )?; - add_m256_cotangent( - &mut cotangents, - instruction.left, - active_m256_contribution(current, _mm256_mul_pd(current, derivative)), - )?; - } - OpCode::Abs => { - let input_value = checked_m256_lane(&values, instruction.left)?; - let positive = _mm256_cmp_pd(input_value, _mm256_setzero_pd(), _CMP_GT_OQ); - let negative = _mm256_cmp_pd(input_value, _mm256_setzero_pd(), _CMP_LT_OQ); - let sign = _mm256_add_pd( - _mm256_and_pd(_mm256_set1_pd(1.0), positive), - _mm256_and_pd(_mm256_set1_pd(-1.0), negative), - ); - add_m256_cotangent( - &mut cotangents, - instruction.left, - active_m256_contribution(current, _mm256_mul_pd(current, sign)), - )?; - } - } - } - - let output_value = checked_m256_lane(&values, output)?; - let mut value_quad = [0.0_f64; 4]; - _mm256_storeu_pd(value_quad.as_mut_ptr(), output_value); - buffer.values.extend_from_slice(&value_quad); - - gradient_quads.clear(); - for input_index in 0..graph.num_inputs { - let gradient_lane = checked_m256_lane(&cotangents, input_index)?; - let mut stored = [0.0_f64; 4]; - _mm256_storeu_pd(stored.as_mut_ptr(), gradient_lane); - gradient_quads.push(stored); - } - for lane_index in 0..4 { - for quad in &gradient_quads { - buffer.gradients.push(quad[lane_index]); + for instruction in self.flat_instructions_slice() { + if !capabilities.supports_opcode(instruction.opcode) { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support a required opcode", + }); } - } - - row_index += 4; - } - - append_scalar_gradient_tail(graph, batch, row_index, buffer)?; - - Ok(()) -} - -#[cfg(target_arch = "x86_64")] -fn gradient_batch_simd_f64x4( - graph: &CompiledGraph, - batch: BatchInputs<'_>, - buffer: &mut BatchGradientsBuffer, -) -> Result<()> { - if !supports_simd_f64x4_runtime() { - return Err(AutodiffError::InvalidGraph { - reason: "simd f64x4 backend requires x86_64 AVX support", - }); - } - unsafe { gradient_batch_simd_f64x4_impl(graph, batch, buffer) } -} - -#[cfg(not(target_arch = "x86_64"))] -fn gradient_batch_simd_f64x4( - _graph: &CompiledGraph, - _batch: BatchInputs<'_>, - _buffer: &mut BatchGradientsBuffer, -) -> Result<()> { - Err(AutodiffError::InvalidGraph { - reason: "simd f64x4 backend requires x86_64 AVX support", - }) -} - -#[cfg(not(target_arch = "x86_64"))] -fn gradient_batch_simd_f64x2( - _graph: &CompiledGraph, - _batch: BatchInputs<'_>, - _buffer: &mut BatchGradientsBuffer, -) -> Result<()> { - Err(AutodiffError::InvalidGraph { - reason: "simd backend requires x86_64 SSE2 support", - }) -} - -impl MockDeviceBackend { - /// Allocate a mock-device buffer set for a compiled graph and batch size. - #[must_use] - pub fn allocate_batch_buffers( - &self, - graph: &CompiledGraph, - batch_size: usize, - ) -> DeviceBufferSet { - DeviceBufferSet::new(self.batch_plan(graph, batch_size)) - } - - /// Execute batch value computation through explicit mock-device transfers. - pub fn compute_batch_with_buffers( - &self, - graph: &CompiledGraph, - batch: BatchInputs<'_>, - buffers: &mut DeviceBufferSet, - output: &mut BatchValuesBuffer, - ) -> Result { - graph.check_batch(batch)?; - if buffers.plan.backend != BackendKind::MockDeviceCpu { - return Err(AutodiffError::InvalidGraph { - reason: "mock execution requires a mock-device buffer plan", - }); - } - if buffers.plan.batch_size != batch.batch_size || buffers.plan.input_dim != batch.input_dim - { - return Err(AutodiffError::InvalidGraph { - reason: "batch shape must match mock-device buffer plan", - }); - } - - let transfers = buffers.plan.compute_transfer_plan.clone(); - for transfer in &transfers { - if transfer.kind == DeviceTransferKind::HostToDevice { - if transfer.buffer != DeviceBufferKind::Inputs { + match instruction.opcode.arity() { + 0 if !capabilities.supports_constants => { return Err(AutodiffError::InvalidGraph { - reason: "mock compute supports host-to-device input transfers only", + reason: "backend does not support constants", }); } - buffers.upload(DeviceBufferKind::Inputs, batch.data)?; - } - } - - let device_inputs = buffers.download(DeviceBufferKind::Inputs)?; - let device_batch = BatchInputs::new(&device_inputs, batch.batch_size, batch.input_dim)?; - let mut scratch = BatchValuesBuffer::new(); - ScalarBackend.compute_batch(graph, device_batch, &mut scratch)?; - buffers.upload(DeviceBufferKind::Outputs, &scratch.data)?; - output.reset(batch.batch_size, graph.output_nodes.len()); - for transfer in &transfers { - if transfer.kind == DeviceTransferKind::DeviceToHost { - if transfer.buffer != DeviceBufferKind::Outputs { + 1 if !capabilities.supports_unary => { return Err(AutodiffError::InvalidGraph { - reason: "mock compute supports device-to-host output transfers only", + reason: "backend does not support unary operations", }); } - output - .data - .extend_from_slice(&buffers.download(DeviceBufferKind::Outputs)?); - } - } - - Ok(DeviceExecutionTrace { - backend: BackendKind::MockDeviceCpu, - mode: DeviceExecutionMode::ComputeBatch, - transfers, - used_native_kernel: false, - }) - } - - /// Execute batch gradient computation through explicit mock-device transfers. - pub fn gradient_batch_with_buffers( - &self, - graph: &CompiledGraph, - batch: BatchInputs<'_>, - buffers: &mut DeviceBufferSet, - output: &mut BatchGradientsBuffer, - ) -> Result { - graph.check_batch(batch)?; - if buffers.plan.backend != BackendKind::MockDeviceCpu { - return Err(AutodiffError::InvalidGraph { - reason: "mock execution requires a mock-device buffer plan", - }); - } - if buffers.plan.batch_size != batch.batch_size || buffers.plan.input_dim != batch.input_dim - { - return Err(AutodiffError::InvalidGraph { - reason: "batch shape must match mock-device buffer plan", - }); - } - - let transfers = buffers.plan.gradient_transfer_plan.clone(); - for transfer in &transfers { - if transfer.kind == DeviceTransferKind::HostToDevice { - if transfer.buffer != DeviceBufferKind::Inputs { + 2 if !capabilities.supports_binary => { return Err(AutodiffError::InvalidGraph { - reason: "mock gradient supports host-to-device input transfers only", + reason: "backend does not support binary operations", }); } - buffers.upload(DeviceBufferKind::Inputs, batch.data)?; + _ => {} } } + Ok(()) + } - let device_inputs = buffers.download(DeviceBufferKind::Inputs)?; - let device_batch = BatchInputs::new(&device_inputs, batch.batch_size, batch.input_dim)?; - let mut scratch = BatchGradientsBuffer::new(); - ScalarBackend.gradient_batch(graph, device_batch, &mut scratch)?; - buffers.upload(DeviceBufferKind::PrimaryValues, &scratch.values)?; - buffers.upload(DeviceBufferKind::Gradients, &scratch.gradients)?; - output.reset(batch.batch_size, graph.num_inputs); - for transfer in &transfers { - if transfer.kind == DeviceTransferKind::DeviceToHost { - match transfer.buffer { - DeviceBufferKind::PrimaryValues => output - .values - .extend_from_slice(&buffers.download(DeviceBufferKind::PrimaryValues)?), - DeviceBufferKind::Gradients => output - .gradients - .extend_from_slice(&buffers.download(DeviceBufferKind::Gradients)?), - _ => { - return Err(AutodiffError::InvalidGraph { - reason: - "mock gradient supports primary-value and gradient downloads only", - }); - } - } + /// Return static execution metadata for backend planning. + #[must_use] + pub fn metadata(&self) -> CompiledGraphMetadata { + let mut num_constants = 0; + let mut num_unary = 0; + let mut num_binary = 0; + for instruction in &self.instructions { + match instruction { + Instruction::Constant(_) => num_constants += 1, + Instruction::Unary { .. } => num_unary += 1, + Instruction::Binary { .. } => num_binary += 1, } } - - Ok(DeviceExecutionTrace { - backend: BackendKind::MockDeviceCpu, - mode: DeviceExecutionMode::GradientBatch, - transfers, - used_native_kernel: false, - }) - } -} - -#[cfg(feature = "backend-wgpu")] -#[inline] -fn wgpu_buffer_size_bytes(len: usize) -> u64 { - let logical = len.saturating_mul(std::mem::size_of::()) as u64; - logical.max(8) -} - -#[cfg(feature = "backend-wgpu")] -fn encode_f64_slice(data: &[f64]) -> Vec { - let mut bytes = Vec::with_capacity(data.len().saturating_mul(std::mem::size_of::())); - for value in data { - bytes.extend_from_slice(&value.to_ne_bytes()); - } - bytes -} - -#[cfg(feature = "backend-wgpu")] -fn decode_f64_bytes(bytes: &[u8], len: usize) -> Result> { - let expected_len = len.saturating_mul(std::mem::size_of::()); - if bytes.len() < expected_len { - return Err(AutodiffError::InvalidGraph { - reason: "wgpu readback length does not match planned buffer length", - }); - } - let mut values = Vec::with_capacity(len); - for chunk in bytes[..expected_len].chunks_exact(std::mem::size_of::()) { - let mut array = [0_u8; std::mem::size_of::()]; - array.copy_from_slice(chunk); - values.push(f64::from_ne_bytes(array)); - } - Ok(values) -} - -#[cfg(feature = "backend-wgpu")] -#[inline] -fn wgpu_buffer_size_bytes_f32(len: usize) -> u64 { - let logical = len.saturating_mul(std::mem::size_of::()) as u64; - logical.max(4) -} - -#[cfg(feature = "backend-wgpu")] -fn encode_f32_slice(data: &[f32]) -> Vec { - let mut bytes = Vec::with_capacity(data.len().saturating_mul(std::mem::size_of::())); - for value in data { - bytes.extend_from_slice(&value.to_ne_bytes()); - } - bytes -} - -#[cfg(feature = "backend-wgpu")] -fn decode_f32_bytes(bytes: &[u8], len: usize) -> Result> { - let expected_len = len.saturating_mul(std::mem::size_of::()); - if bytes.len() < expected_len { - return Err(AutodiffError::InvalidGraph { - reason: "wgpu readback length does not match requested f32 buffer length", - }); - } - let mut values = Vec::with_capacity(len); - for chunk in bytes[..expected_len].chunks_exact(std::mem::size_of::()) { - let mut array = [0_u8; std::mem::size_of::()]; - array.copy_from_slice(chunk); - values.push(f32::from_ne_bytes(array)); - } - Ok(values) -} - -#[cfg(feature = "backend-wgpu")] -fn encode_u32_slice(data: &[u32]) -> Vec { - let mut bytes = Vec::with_capacity(data.len().saturating_mul(std::mem::size_of::())); - for value in data { - bytes.extend_from_slice(&value.to_ne_bytes()); - } - bytes -} - -#[cfg(feature = "backend-wgpu")] -/// Conservative opcode subset currently supported by the exact-safe native WGPU batch-compute path. -pub const WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES: [OpCode; 4] = - [OpCode::Constant, OpCode::Neg, OpCode::Relu, OpCode::Abs]; - -#[cfg(feature = "backend-wgpu")] -fn checked_u32_from_usize(value: usize, reason: &'static str) -> Result { - u32::try_from(value).map_err(|_| AutodiffError::InvalidGraph { reason }) -} - -#[cfg(feature = "backend-wgpu")] -#[inline] -fn is_exact_f32_roundtrip(value: f64) -> bool { - if value.is_nan() { - return false; - } - let narrowed = value as f32; - let widened = f64::from(narrowed); - widened == value && (value != 0.0 || widened.is_sign_negative() == value.is_sign_negative()) -} - -#[cfg(feature = "backend-wgpu")] -fn wgpu_native_exact_safe_supports_opcode(opcode: OpCode) -> bool { - WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES.contains(&opcode) -} - -#[cfg(feature = "backend-wgpu")] -fn wgpu_native_exact_safe_graph(graph: &CompiledGraph) -> bool { - !graph.output_nodes().is_empty() - && graph.flat_instructions_slice().iter().all(|instruction| { - wgpu_native_exact_safe_supports_opcode(instruction.opcode) - && (instruction.opcode != OpCode::Constant - || is_exact_f32_roundtrip(instruction.value)) - }) -} - -#[cfg(feature = "backend-wgpu")] -fn wgpu_native_exact_safe_batch(batch: BatchInputs<'_>) -> bool { - batch.data.iter().copied().all(is_exact_f32_roundtrip) -} - -#[cfg(feature = "backend-wgpu")] -const WGPU_NATIVE_WORDS_PER_INSTRUCTION: usize = 8; - -#[cfg(feature = "backend-wgpu")] -const WGPU_NATIVE_SHADER: &str = r#" -@group(0) @binding(0) var input_data: array; -@group(0) @binding(1) var instruction_words: array; -@group(0) @binding(2) var output_nodes: array; -@group(0) @binding(3) var value_data: array; -@group(0) @binding(4) var output_data: array; -@group(0) @binding(5) var kernel_meta: array; - -fn relu_scalar(x: f32) -> f32 { - if x > 0.0 { - return x; - } - return 0.0; -} - -fn log1p_exp_scalar(x: f32) -> f32 { - if x > 0.0 { - return x + log(1.0 + exp(-x)); - } - return log(1.0 + exp(x)); -} - -fn log_add_exp_scalar(a: f32, b: f32) -> f32 { - var max_value = a; - var min_value = b; - if a < b { - max_value = b; - min_value = a; - } - return max_value + log(1.0 + exp(min_value - max_value)); -} - -@compute @workgroup_size(64) -fn main(@builtin(global_invocation_id) gid: vec3) { - let row = gid.x; - let num_inputs = kernel_meta[0u]; - let num_instructions = kernel_meta[1u]; - let num_outputs = kernel_meta[2u]; - let value_count = kernel_meta[3u]; - let batch_size = kernel_meta[4u]; - if row >= batch_size { - return; - } - - let input_base = row * num_inputs; - let value_base = row * value_count; - let output_base = row * num_outputs; - - var input_index = 0u; - loop { - if input_index >= num_inputs { - break; + CompiledGraphMetadata { + num_inputs: self.num_inputs, + num_outputs: self.output_nodes.len(), + num_instructions: self.instructions.len(), + num_constants, + num_unary, + num_binary, + value_count: self.num_inputs + self.instructions.len(), + is_scalar_output: self.output_nodes.len() == 1, } - value_data[value_base + input_index] = input_data[input_base + input_index]; - input_index = input_index + 1u; } - var instruction_index = 0u; - loop { - if instruction_index >= num_instructions { - break; - } - let word_base = instruction_index * 8u; - let opcode = instruction_words[word_base + 0u]; - let output_index = instruction_words[word_base + 1u]; - let left_index = instruction_words[word_base + 2u]; - let right_index = instruction_words[word_base + 3u]; - let value_bits = instruction_words[word_base + 4u]; - var left_value: f32 = 0.0; - if left_index != 4294967295u { - left_value = value_data[value_base + left_index]; - } - var right_value: f32 = 0.0; - if right_index != 4294967295u { - right_value = value_data[value_base + right_index]; - } - var result_value: f32 = 0.0; - switch opcode { - case 0u: { - result_value = bitcast(value_bits); - } - case 1u: { - result_value = left_value + right_value; - } - case 2u: { - result_value = left_value - right_value; - } - case 3u: { - result_value = left_value * right_value; - } - case 4u: { - result_value = left_value / right_value; - } - case 5u: { - result_value = pow(left_value, right_value); - } - case 6u: { - result_value = sin(left_value); - } - case 7u: { - result_value = cos(left_value); - } - case 8u: { - result_value = tan(left_value); - } - case 9u: { - result_value = tanh(left_value); - } - case 10u: { - result_value = relu_scalar(left_value); - } - case 11u: { - result_value = log1p_exp_scalar(left_value); - } - case 12u: { - result_value = log_add_exp_scalar(left_value, right_value); - } - case 13u: { - result_value = -left_value; - } - case 14u: { - result_value = exp(left_value); - } - case 15u: { - result_value = log(left_value); - } - case 16u: { - result_value = sqrt(left_value); - } - case 17u: { - result_value = abs(left_value); - } - default: { - result_value = 0.0; - } + /// Create a reusable workspace sized for this compiled graph. + #[must_use] + pub fn workspace(&self) -> CompiledWorkspace { + let value_count = self.num_inputs + self.instructions.len(); + CompiledWorkspace { + values: Vec::with_capacity(value_count), + cotangents: Vec::with_capacity(value_count), + gradients: Vec::with_capacity(self.num_inputs), + outputs: Vec::with_capacity(self.output_nodes.len()), } - value_data[value_base + output_index] = result_value; - instruction_index = instruction_index + 1u; } - var output_index = 0u; - loop { - if output_index >= num_outputs { - break; + pub(crate) fn check_input_len(&self, inputs: &[f64]) -> Result<()> { + if inputs.len() == self.num_inputs { + Ok(()) + } else { + Err(AutodiffError::InvalidGraph { + reason: "input length must match compiled graph input count", + }) } - let node = output_nodes[output_index]; - output_data[output_base + output_index] = value_data[value_base + node]; - output_index = output_index + 1u; - } -} -"#; - -#[cfg(feature = "backend-wgpu")] -impl WgpuBackend { - /// Create the default WGPU backend using device id 0 and automatic transfers. - pub fn new_default() -> Result { - Self::new( - AcceleratorDeviceContext::wgpu(0), - DeviceTransferPolicy::Automatic, - ) } - /// Create the WGPU backend skeleton from a context and transfer policy. - pub fn new( - context: AcceleratorDeviceContext, - transfer_policy: DeviceTransferPolicy, - ) -> Result { - Self::from_boundary(GpuBackendBoundary::new(context, transfer_policy)) - } - - /// Create the WGPU backend skeleton from a boundary descriptor. - pub fn from_boundary(boundary: GpuBackendBoundary) -> Result { - if boundary.context.kind != AcceleratorDeviceKind::Wgpu { - return Err(AutodiffError::InvalidGraph { - reason: "wgpu backend requires a WGPU accelerator context", - }); + pub(crate) fn check_batch(&self, batch: BatchInputs<'_>) -> Result<()> { + if batch.input_dim == self.num_inputs { + Ok(()) + } else { + Err(AutodiffError::InvalidGraph { + reason: "batch input_dim must match compiled graph input count", + }) } - let instance = wgpu::Instance::default(); - let adapter = block_on(instance.request_adapter(&wgpu::RequestAdapterOptions { - power_preference: wgpu::PowerPreference::default(), - force_fallback_adapter: false, - compatible_surface: None, - })) - .map_err(|_| AutodiffError::InvalidGraph { - reason: "wgpu adapter request failed", - })?; - let adapter_info = adapter.get_info(); - let (device, queue) = block_on(adapter.request_device(&wgpu::DeviceDescriptor { - label: Some("petite-ad-wgpu"), - required_features: wgpu::Features::empty(), - required_limits: wgpu::Limits::default(), - memory_hints: wgpu::MemoryHints::Performance, - trace: wgpu::Trace::Off, - })) - .map_err(|_| AutodiffError::InvalidGraph { - reason: "wgpu device request failed", - })?; - Ok(Self { - boundary, - device, - queue, - adapter_info, - }) } - /// Return the boundary descriptor used for this backend. - #[must_use] - pub fn boundary(&self) -> &GpuBackendBoundary { - &self.boundary + fn gather_unary(values: &[f64], arg: NodeId) -> Result<[f64; 1]> { + let value = values + .get(arg) + .copied() + .ok_or(AutodiffError::IndexOutOfBounds { + index: arg, + max_index: values.len().saturating_sub(1), + })?; + Ok([value]) } - /// Return the accelerator context used for this backend. - #[must_use] - pub fn context(&self) -> &AcceleratorDeviceContext { - &self.boundary.context + fn gather_binary(values: &[f64], left: NodeId, right: NodeId) -> Result<[f64; 2]> { + let left_value = values + .get(left) + .copied() + .ok_or(AutodiffError::IndexOutOfBounds { + index: left, + max_index: values.len().saturating_sub(1), + })?; + let right_value = values + .get(right) + .copied() + .ok_or(AutodiffError::IndexOutOfBounds { + index: right, + max_index: values.len().saturating_sub(1), + })?; + Ok([left_value, right_value]) } - /// Return the configured transfer policy. - #[must_use] - pub fn transfer_policy(&self) -> DeviceTransferPolicy { - self.boundary.transfer_policy - } + pub(crate) fn fill_values( + &self, + inputs: &[f64], + workspace: &mut CompiledWorkspace, + ) -> Result<()> { + self.check_input_len(inputs)?; + workspace.values.clear(); + workspace + .values + .reserve(self.num_inputs + self.instructions.len()); + workspace.values.extend_from_slice(inputs); - /// Return the resolved adapter name. - #[must_use] - pub fn adapter_name(&self) -> &str { - &self.adapter_info.name + for instruction in &self.instructions { + let value = match *instruction { + Instruction::Constant(value) => value, + Instruction::Unary { op, arg } => { + let args = Self::gather_unary(&workspace.values, arg)?; + op_rules::forward_value(op, &args)? + } + Instruction::Binary { op, left, right } => { + let args = Self::gather_binary(&workspace.values, left, right)?; + op_rules::forward_value(op, &args)? + } + }; + workspace.values.push(value); + } + Ok(()) } - /// Return the conservative exact-safe opcode subset used by the native WGPU batch-compute path. - #[must_use] - pub fn native_batch_compute_supported_opcodes() -> &'static [OpCode] { - &WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES + /// Compute primary output. + pub fn compute(&self, inputs: &[f64]) -> Result { + let mut workspace = self.workspace(); + self.compute_with_workspace(inputs, &mut workspace) } - /// Return whether this graph is eligible for the restricted exact-safe native WGPU path. - #[must_use] - pub fn supports_native_batch_compute(&self, graph: &CompiledGraph) -> bool { - wgpu_native_exact_safe_graph(graph) + /// Compute all selected outputs. + pub fn compute_many(&self, inputs: &[f64]) -> Result> { + let mut workspace = self.workspace(); + Ok(self + .compute_many_with_workspace(inputs, &mut workspace)? + .to_vec()) } - /// Return whether a concrete batch is eligible for the exact-safe native WGPU path. - #[must_use] - pub fn supports_native_batch_compute_for_batch( + /// Compute primary output with a reusable workspace. + pub fn compute_with_workspace( &self, - graph: &CompiledGraph, - batch: BatchInputs<'_>, - ) -> bool { - self.supports_native_batch_compute(graph) && wgpu_native_exact_safe_batch(batch) + inputs: &[f64], + workspace: &mut CompiledWorkspace, + ) -> Result { + let outputs = self.compute_many_with_workspace(inputs, workspace)?; + Ok(outputs.first().copied().unwrap_or(0.0)) } - fn create_initialized_storage_buffer( + /// Compute all selected outputs with a reusable workspace. + pub fn compute_many_with_workspace<'a>( &self, - label: &'static str, - bytes: &[u8], - min_size: u64, - ) -> wgpu::Buffer { - let buffer = self.device.create_buffer(&wgpu::BufferDescriptor { - label: Some(label), - size: min_size.max(bytes.len() as u64).max(4), - usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, - mapped_at_creation: false, - }); - if !bytes.is_empty() { - self.queue.write_buffer(&buffer, 0, bytes); - } - buffer - } - - fn download_raw_buffer(&self, buffer: &wgpu::Buffer, size: u64) -> Result> { - let staging = self.device.create_buffer(&wgpu::BufferDescriptor { - label: Some("petite-ad-wgpu-readback"), - size, - usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ, - mapped_at_creation: false, - }); - let mut encoder = self - .device - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("petite-ad-wgpu-readback-encoder"), - }); - encoder.copy_buffer_to_buffer(buffer, 0, &staging, 0, size); - let submission = self.queue.submit([encoder.finish()]); - let slice = staging.slice(..); - let (tx, rx) = mpsc::channel(); - slice.map_async(wgpu::MapMode::Read, move |result| { - let _ = tx.send(result); - }); - let _ = self.device.poll(wgpu::PollType::wait_for(submission)); - match rx.recv() { - Ok(Ok(())) => {} - _ => { - return Err(AutodiffError::InvalidGraph { - reason: "wgpu readback mapping failed", - }); - } + inputs: &[f64], + workspace: &'a mut CompiledWorkspace, + ) -> Result<&'a [f64]> { + self.fill_values(inputs, workspace)?; + workspace.outputs.clear(); + for &output in &self.output_nodes { + let value = + workspace + .values + .get(output) + .copied() + .ok_or(AutodiffError::IndexOutOfBounds { + index: output, + max_index: workspace.values.len().saturating_sub(1), + })?; + workspace.outputs.push(value); } - let view = slice.get_mapped_range(); - let bytes = view.to_vec(); - drop(view); - staging.unmap(); - Ok(bytes) + Ok(&workspace.outputs) } - fn compute_batch_native( - &self, - graph: &CompiledGraph, - batch: BatchInputs<'_>, - ) -> Result> { - let metadata = graph.metadata(); - let input_f32: Vec = batch.data.iter().map(|value| *value as f32).collect(); - let output_count = batch.batch_size.saturating_mul(metadata.num_outputs); - let input_buffer = self.create_initialized_storage_buffer( - "petite-ad-wgpu-native-inputs", - &encode_f32_slice(&input_f32), - wgpu_buffer_size_bytes_f32(input_f32.len()), - ); - let instruction_words = encode_u32_slice(&self.native_instruction_words(graph)?); - let instruction_buffer = self.create_initialized_storage_buffer( - "petite-ad-wgpu-native-instructions", - &instruction_words, - instruction_words.len() as u64, - ); - let output_nodes: Result> = graph - .output_nodes() - .iter() - .map(|node| { - checked_u32_from_usize(*node, "wgpu native output node index exceeds u32 range") - }) - .collect(); - let output_nodes = output_nodes?; - let output_node_buffer = self.create_initialized_storage_buffer( - "petite-ad-wgpu-native-output-nodes", - &encode_u32_slice(&output_nodes), - (output_nodes.len().max(1) * std::mem::size_of::()) as u64, - ); - let value_buffer = self.device.create_buffer(&wgpu::BufferDescriptor { - label: Some("petite-ad-wgpu-native-values"), - size: wgpu_buffer_size_bytes_f32(batch.batch_size.saturating_mul(metadata.value_count)), - usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, - mapped_at_creation: false, - }); - let output_buffer = self.device.create_buffer(&wgpu::BufferDescriptor { - label: Some("petite-ad-wgpu-native-outputs"), - size: wgpu_buffer_size_bytes_f32(output_count), - usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, - mapped_at_creation: false, - }); - let meta = [ - checked_u32_from_usize( - metadata.num_inputs, - "wgpu native input dimension exceeds u32 range", - )?, - checked_u32_from_usize( - metadata.num_instructions, - "wgpu native instruction count exceeds u32 range", - )?, - checked_u32_from_usize( - metadata.num_outputs, - "wgpu native output dimension exceeds u32 range", - )?, - checked_u32_from_usize( - metadata.value_count, - "wgpu native value count exceeds u32 range", - )?, - checked_u32_from_usize(batch.batch_size, "wgpu native batch size exceeds u32 range")?, - ]; - let meta_buffer = self.create_initialized_storage_buffer( - "petite-ad-wgpu-native-meta", - &encode_u32_slice(&meta), - (meta.len() * std::mem::size_of::()) as u64, - ); - self.queue.submit([]); - - let shader = self - .device - .create_shader_module(wgpu::ShaderModuleDescriptor { - label: Some("petite-ad-wgpu-native-compute-shader"), - source: wgpu::ShaderSource::Wgsl(WGPU_NATIVE_SHADER.into()), - }); - let pipeline = self - .device - .create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some("petite-ad-wgpu-native-compute-pipeline"), - layout: None, - module: &shader, - entry_point: Some("main"), - compilation_options: Default::default(), - cache: None, - }); - let layout = pipeline.get_bind_group_layout(0); - let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor { - label: Some("petite-ad-wgpu-native-bind-group"), - layout: &layout, - entries: &[ - wgpu::BindGroupEntry { - binding: 0, - resource: input_buffer.as_entire_binding(), - }, - wgpu::BindGroupEntry { - binding: 1, - resource: instruction_buffer.as_entire_binding(), - }, - wgpu::BindGroupEntry { - binding: 2, - resource: output_node_buffer.as_entire_binding(), - }, - wgpu::BindGroupEntry { - binding: 3, - resource: value_buffer.as_entire_binding(), - }, - wgpu::BindGroupEntry { - binding: 4, - resource: output_buffer.as_entire_binding(), - }, - wgpu::BindGroupEntry { - binding: 5, - resource: meta_buffer.as_entire_binding(), - }, - ], - }); - let mut encoder = self - .device - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("petite-ad-wgpu-native-compute-encoder"), - }); - { - let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("petite-ad-wgpu-native-compute-pass"), - timestamp_writes: None, + fn reverse_for_output(&self, output: NodeId, workspace: &mut CompiledWorkspace) -> Result<()> { + workspace.cotangents.clear(); + workspace.cotangents.resize(workspace.values.len(), 0.0); + if output >= workspace.cotangents.len() { + return Err(AutodiffError::IndexOutOfBounds { + index: output, + max_index: workspace.cotangents.len().saturating_sub(1), }); - pass.set_pipeline(&pipeline); - pass.set_bind_group(0, &bind_group, &[]); - let workgroups = checked_u32_from_usize( - batch.batch_size.div_ceil(64), - "wgpu native workgroup count exceeds u32 range", - )?; - pass.dispatch_workgroups(workgroups.max(1), 1, 1); } - self.queue.submit([encoder.finish()]); - let raw_output = - self.download_raw_buffer(&output_buffer, wgpu_buffer_size_bytes_f32(output_count))?; - let output_f32 = decode_f32_bytes(&raw_output, output_count)?; - Ok(output_f32.into_iter().map(f64::from).collect()) - } + workspace.cotangents[output] = 1.0; - fn native_instruction_words(&self, graph: &CompiledGraph) -> Result> { - let mut words = Vec::with_capacity( - graph.flat_instructions_slice().len() * WGPU_NATIVE_WORDS_PER_INSTRUCTION, - ); - for instruction in graph.flat_instructions_slice() { - words.push(match instruction.opcode { - OpCode::Constant => 0, - OpCode::Add => 1, - OpCode::Sub => 2, - OpCode::Mul => 3, - OpCode::Div => 4, - OpCode::Pow => 5, - OpCode::Sin => 6, - OpCode::Cos => 7, - OpCode::Tan => 8, - OpCode::Tanh => 9, - OpCode::Relu => 10, - OpCode::Log1pExp => 11, - OpCode::LogAddExp => 12, - OpCode::Neg => 13, - OpCode::Exp => 14, - OpCode::Ln => 15, - OpCode::Sqrt => 16, - OpCode::Abs => 17, - }); - words.push(checked_u32_from_usize( - instruction.output, - "wgpu native instruction output index exceeds u32 range", - )?); - words.push(if instruction.left == UNUSED_NODE_ID { - u32::MAX - } else { - checked_u32_from_usize( - instruction.left, - "wgpu native instruction left index exceeds u32 range", - )? - }); - words.push(if instruction.right == UNUSED_NODE_ID { - u32::MAX - } else { - checked_u32_from_usize( - instruction.right, - "wgpu native instruction right index exceeds u32 range", - )? - }); - words.push((instruction.value as f32).to_bits()); - words.extend_from_slice(&[0, 0, 0]); + for (offset, instruction) in self.instructions.iter().enumerate().rev() { + let node_id = self.num_inputs + offset; + let current = workspace.cotangents[node_id]; + if current == 0.0 { + continue; + } + match *instruction { + Instruction::Constant(_) => {} + Instruction::Unary { op, arg } => { + let args = Self::gather_unary(&workspace.values, arg)?; + let value = workspace.values[node_id]; + if let op_rules::LocalRule::Unary { dy, .. } = + op_rules::local_rule(op, &args, value)? + { + workspace.cotangents[arg] += current * dy; + } + } + Instruction::Binary { op, left, right } => { + let args = Self::gather_binary(&workspace.values, left, right)?; + let value = workspace.values[node_id]; + if let op_rules::LocalRule::Binary { + dy_left, dy_right, .. + } = op_rules::local_rule(op, &args, value)? + { + workspace.cotangents[left] += current * dy_left; + workspace.cotangents[right] += current * dy_right; + } + } + } } - Ok(words) + Ok(()) } - /// Allocate a real WGPU buffer set for a compiled graph and batch size. - pub fn allocate_batch_buffers( - &self, - graph: &CompiledGraph, - batch_size: usize, - ) -> Result { - let plan = self.batch_plan(graph, batch_size); - let mut buffers = Vec::with_capacity(plan.buffer_handles.len()); - for handle in &plan.buffer_handles { - let label = format!("petite-ad-wgpu-{:?}", handle.kind); - let buffer = self.device.create_buffer(&wgpu::BufferDescriptor { - label: Some(label.as_str()), - size: wgpu_buffer_size_bytes(handle.len), - usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, - mapped_at_creation: false, - }); - buffers.push(WgpuBuffer { - handle: *handle, - buffer, - }); - } - Ok(WgpuBufferSet { plan, buffers }) + /// Compute primary output and gradient. + pub fn gradient(&self, inputs: &[f64]) -> Result<(f64, Vec)> { + let mut workspace = self.workspace(); + let (value, gradient) = self.gradient_with_workspace(inputs, &mut workspace)?; + Ok((value, gradient.to_vec())) } - fn upload_buffer( + /// Compute primary output and gradient with a reusable workspace. + pub fn gradient_with_workspace<'a>( &self, - buffers: &WgpuBufferSet, - kind: DeviceBufferKind, - data: &[f64], - ) -> Result<()> { - let buffer = buffers.buffer(kind)?; - if buffer.handle.len != data.len() { - return Err(AutodiffError::InvalidGraph { - reason: "wgpu upload length must match planned buffer length", - }); - } - if data.is_empty() { - return Ok(()); - } - let bytes = encode_f64_slice(data); - self.queue.write_buffer(&buffer.buffer, 0, &bytes); - self.queue.submit([]); - Ok(()) + inputs: &[f64], + workspace: &'a mut CompiledWorkspace, + ) -> Result<(f64, &'a [f64])> { + self.fill_values(inputs, workspace)?; + workspace.gradients.clear(); + workspace.gradients.resize(self.num_inputs, 0.0); + let Some(&output) = self.output_nodes.first() else { + return Ok((0.0, &workspace.gradients)); + }; + self.reverse_for_output(output, workspace)?; + workspace + .gradients + .copy_from_slice(&workspace.cotangents[..self.num_inputs]); + Ok((workspace.values[output], &workspace.gradients)) } - fn download_buffer(&self, buffers: &WgpuBufferSet, kind: DeviceBufferKind) -> Result> { - let buffer = buffers.buffer(kind)?; - if buffer.handle.len == 0 { - return Ok(Vec::new()); - } - let raw = - self.download_raw_buffer(&buffer.buffer, wgpu_buffer_size_bytes(buffer.handle.len))?; - decode_f64_bytes(&raw, buffer.handle.len) + /// Compute Jacobian for all selected outputs. + pub fn jacobian(&self, inputs: &[f64]) -> Result>> { + let mut workspace = self.workspace(); + self.jacobian_with_workspace(inputs, &mut workspace) } - /// Execute batch value computation through explicit WGPU transfers. - pub fn compute_batch_with_buffers( + /// Compute Jacobian for all selected outputs with a reusable workspace. + pub fn jacobian_with_workspace( &self, - graph: &CompiledGraph, - batch: BatchInputs<'_>, - buffers: &mut WgpuBufferSet, - output: &mut BatchValuesBuffer, - ) -> Result { - graph.check_batch(batch)?; - if buffers.plan.backend != BackendKind::Wgpu { - return Err(AutodiffError::InvalidGraph { - reason: "wgpu execution requires a WGPU buffer plan", - }); - } - if buffers.plan.batch_size != batch.batch_size || buffers.plan.input_dim != batch.input_dim - { - return Err(AutodiffError::InvalidGraph { - reason: "batch shape must match WGPU buffer plan", - }); + inputs: &[f64], + workspace: &mut CompiledWorkspace, + ) -> Result>> { + self.fill_values(inputs, workspace)?; + let mut jacobian = Vec::with_capacity(self.output_nodes.len()); + for &output in &self.output_nodes { + self.reverse_for_output(output, workspace)?; + jacobian.push(workspace.cotangents[..self.num_inputs].to_vec()); } + Ok(jacobian) + } - let transfers = buffers.plan.compute_transfer_plan.clone(); - for transfer in &transfers { - if transfer.kind == DeviceTransferKind::HostToDevice { - if transfer.buffer != DeviceBufferKind::Inputs { - return Err(AutodiffError::InvalidGraph { - reason: "wgpu compute supports host-to-device input transfers only", - }); - } - buffers.upload(self, DeviceBufferKind::Inputs, batch.data)?; + /// Return static compatibility details for a backend. + pub fn backend_support_report(&self, backend: BackendKind) -> Result { + let capabilities = backend.capabilities(); + let supports_required_outputs = + self.output_nodes.len() <= 1 || capabilities.supports_multi_output; + let mut missing_opcodes = Vec::new(); + let mut supports_required_arities = true; + for instruction in self.flat_instructions_slice() { + if !capabilities.supports_opcode(instruction.opcode) + && !missing_opcodes.contains(&instruction.opcode) + { + missing_opcodes.push(instruction.opcode); } - } - - let used_native_kernel = self.supports_native_batch_compute_for_batch(graph, batch); - let scratch_data = if used_native_kernel { - self.compute_batch_native(graph, batch)? - } else { - let device_inputs = buffers.download(self, DeviceBufferKind::Inputs)?; - let device_batch = BatchInputs::new(&device_inputs, batch.batch_size, batch.input_dim)?; - let mut scratch = BatchValuesBuffer::new(); - ScalarBackend.compute_batch(graph, device_batch, &mut scratch)?; - scratch.data - }; - buffers.upload(self, DeviceBufferKind::Outputs, &scratch_data)?; - output.reset(batch.batch_size, graph.output_nodes.len()); - for transfer in &transfers { - if transfer.kind == DeviceTransferKind::DeviceToHost { - if transfer.buffer != DeviceBufferKind::Outputs { - return Err(AutodiffError::InvalidGraph { - reason: "wgpu compute supports device-to-host output transfers only", - }); - } - output - .data - .extend_from_slice(&buffers.download(self, DeviceBufferKind::Outputs)?); + match instruction.opcode.arity() { + 0 if !capabilities.supports_constants => supports_required_arities = false, + 1 if !capabilities.supports_unary => supports_required_arities = false, + 2 if !capabilities.supports_binary => supports_required_arities = false, + _ => {} } } - - Ok(DeviceExecutionTrace { - backend: BackendKind::Wgpu, - mode: DeviceExecutionMode::ComputeBatch, - transfers, - used_native_kernel, - }) - } - - /// Execute batch gradient computation through explicit WGPU transfers. - pub fn gradient_batch_with_buffers( - &self, - graph: &CompiledGraph, - batch: BatchInputs<'_>, - buffers: &mut WgpuBufferSet, - output: &mut BatchGradientsBuffer, - ) -> Result { - graph.check_batch(batch)?; - if buffers.plan.backend != BackendKind::Wgpu { - return Err(AutodiffError::InvalidGraph { - reason: "wgpu execution requires a WGPU buffer plan", - }); + let runtime_available = backend.runtime_available(); + let supports_required_opcodes = missing_opcodes.is_empty(); + let mut common_reasons = Vec::new(); + if !capabilities.supports_f64 { + common_reasons.push(BackendRejectionReason::MissingF64); } - if buffers.plan.batch_size != batch.batch_size || buffers.plan.input_dim != batch.input_dim - { - return Err(AutodiffError::InvalidGraph { - reason: "batch shape must match WGPU buffer plan", - }); + if !supports_required_outputs { + common_reasons.push(BackendRejectionReason::UnsupportedOutputs); } - - let transfers = buffers.plan.gradient_transfer_plan.clone(); - for transfer in &transfers { - if transfer.kind == DeviceTransferKind::HostToDevice { - if transfer.buffer != DeviceBufferKind::Inputs { - return Err(AutodiffError::InvalidGraph { - reason: "wgpu gradient supports host-to-device input transfers only", - }); - } - buffers.upload(self, DeviceBufferKind::Inputs, batch.data)?; - } + if !supports_required_opcodes { + common_reasons.push(BackendRejectionReason::UnsupportedOpcodes); } - - let device_inputs = buffers.download(self, DeviceBufferKind::Inputs)?; - let device_batch = BatchInputs::new(&device_inputs, batch.batch_size, batch.input_dim)?; - let mut scratch = BatchGradientsBuffer::new(); - ScalarBackend.gradient_batch(graph, device_batch, &mut scratch)?; - buffers.upload(self, DeviceBufferKind::PrimaryValues, &scratch.values)?; - buffers.upload(self, DeviceBufferKind::Gradients, &scratch.gradients)?; - output.reset(batch.batch_size, graph.num_inputs); - for transfer in &transfers { - if transfer.kind == DeviceTransferKind::DeviceToHost { - match transfer.buffer { - DeviceBufferKind::PrimaryValues => output.values.extend_from_slice( - &buffers.download(self, DeviceBufferKind::PrimaryValues)?, - ), - DeviceBufferKind::Gradients => output - .gradients - .extend_from_slice(&buffers.download(self, DeviceBufferKind::Gradients)?), - _ => { - return Err(AutodiffError::InvalidGraph { - reason: - "wgpu gradient supports primary-value and gradient downloads only", - }); - } - } - } + if !supports_required_arities { + common_reasons.push(BackendRejectionReason::UnsupportedArities); + } + if !runtime_available { + common_reasons.push(BackendRejectionReason::UnavailableRuntime); } - Ok(DeviceExecutionTrace { - backend: BackendKind::Wgpu, - mode: DeviceExecutionMode::GradientBatch, - transfers, - used_native_kernel: false, - }) - } -} - -impl DeviceBackend for ScalarBackend { - fn backend_kind(&self, _graph: &CompiledGraph) -> BackendKind { - BackendKind::Scalar - } -} - -impl DeviceBackend for MockDeviceBackend { - fn backend_kind(&self, _graph: &CompiledGraph) -> BackendKind { - BackendKind::MockDeviceCpu - } -} - -#[cfg(feature = "backend-wgpu")] -impl DeviceBackend for WgpuBackend { - fn backend_kind(&self, _graph: &CompiledGraph) -> BackendKind { - BackendKind::Wgpu - } -} + let mut batch_compute_rejection_reasons = common_reasons.clone(); + if !capabilities.supports_batch_compute { + batch_compute_rejection_reasons.push(BackendRejectionReason::NoBatchCompute); + } + let mut batch_gradient_rejection_reasons = common_reasons; + if !capabilities.supports_batch_gradient { + batch_gradient_rejection_reasons.push(BackendRejectionReason::NoBatchGradient); + } -impl DeviceBackend for SimdBackend { - fn backend_kind(&self, graph: &CompiledGraph) -> BackendKind { - graph - .simd_support_report() - .map(|report| report.backend) - .unwrap_or(BackendKind::SimdF64x2) + Ok(BackendSupportReport { + backend, + supports_f64: capabilities.supports_f64, + supports_required_outputs, + supports_required_opcodes, + supports_required_arities, + supports_batch_compute: capabilities.supports_batch_compute, + supports_batch_gradient: capabilities.supports_batch_gradient, + lane_width: backend.lane_width(), + runtime_available, + required_runtime_features: backend.required_runtime_features().to_vec(), + unavailable_runtime_features: backend.unavailable_runtime_features(), + missing_opcodes, + batch_compute_rejection_reasons, + batch_gradient_rejection_reasons, + }) } -} -impl DeviceBackend for BackendKind { - fn backend_kind(&self, _graph: &CompiledGraph) -> BackendKind { - *self + /// Return static compatibility details for all built-in backends. + pub fn backend_support_reports(&self) -> Result> { + Ok(vec![ + self.backend_support_report(BackendKind::Scalar)?, + self.backend_support_report(BackendKind::MockDeviceCpu)?, + self.backend_support_report(BackendKind::Wgpu)?, + self.backend_support_report(BackendKind::SimdF64x4)?, + self.backend_support_report(BackendKind::SimdF64x2)?, + ]) } -} -impl ExecutionBackend for SimdBackend { - fn name(&self) -> &'static str { - "simd" + /// Return a device-oriented batch buffer plan for a backend. + #[must_use] + pub fn device_batch_plan(&self, backend: BackendKind, batch_size: usize) -> DeviceBatchPlan { + backend.batch_plan(self, batch_size) } - fn capabilities(&self) -> BackendCapabilities { - BackendCapabilities::simd_f64x2() + /// Allocate mock-device buffers for this compiled graph and batch size. + #[must_use] + pub fn allocate_mock_device_buffers(&self, batch_size: usize) -> DeviceBufferSet { + MockDeviceBackend.allocate_batch_buffers(self, batch_size) } - fn compute(&self, _graph: &CompiledGraph, _inputs: &[f64]) -> Result { - Err(AutodiffError::InvalidGraph { - reason: "simd backend currently supports batch compute only", - }) + /// Execute batch value computation through mock-device buffers. + pub fn compute_batch_mock_device_into( + &self, + batch: BatchInputs<'_>, + buffers: &mut DeviceBufferSet, + output: &mut BatchValuesBuffer, + ) -> Result { + MockDeviceBackend.compute_batch_with_buffers(self, batch, buffers, output) } - fn compute_batch( + /// Execute batch gradient computation through mock-device buffers. + pub fn gradient_batch_mock_device_into( &self, - graph: &CompiledGraph, batch: BatchInputs<'_>, - buffer: &mut BatchValuesBuffer, - ) -> Result<()> { - if let Ok(report) = graph.backend_support_report(BackendKind::SimdF64x4) { - if report.can_compute_batch() { - return compute_batch_simd_f64x4(graph, batch, buffer); - } - } - let capabilities = self.capabilities(); - if !capabilities.supports_batch_compute { - return Err(AutodiffError::InvalidGraph { - reason: "backend does not support batch compute on this target", - }); - } - graph.validate_backend_capabilities(&capabilities)?; - compute_batch_simd_f64x2(graph, batch, buffer) + buffers: &mut DeviceBufferSet, + output: &mut BatchGradientsBuffer, + ) -> Result { + MockDeviceBackend.gradient_batch_with_buffers(self, batch, buffers, output) } - fn gradient(&self, _graph: &CompiledGraph, _inputs: &[f64]) -> Result<(f64, Vec)> { - Err(AutodiffError::InvalidGraph { - reason: "simd backend does not support reverse gradients yet", - }) + #[cfg(feature = "backend-wgpu")] + /// Allocate real WGPU buffers for this compiled graph and batch size. + pub fn allocate_wgpu_buffers( + &self, + backend: &WgpuBackend, + batch_size: usize, + ) -> Result { + backend.allocate_batch_buffers(self, batch_size) } - fn gradient_batch( + #[cfg(feature = "backend-wgpu")] + /// Execute batch value computation through real WGPU buffers. + pub fn compute_batch_wgpu_into( &self, - graph: &CompiledGraph, + backend: &WgpuBackend, batch: BatchInputs<'_>, - buffer: &mut BatchGradientsBuffer, - ) -> Result<()> { - if let Ok(report) = graph.backend_support_report(BackendKind::SimdF64x4) { - if report.can_gradient_batch() { - return gradient_batch_simd_f64x4(graph, batch, buffer); - } - } - let capabilities = self.capabilities(); - if !capabilities.supports_batch_gradient { - return Err(AutodiffError::InvalidGraph { - reason: "backend does not support batch gradients on this target", - }); - } - graph.validate_backend_capabilities(&capabilities)?; - gradient_batch_simd_f64x2(graph, batch, buffer) - } -} - -impl ExecutionBackend for MockDeviceBackend { - fn name(&self) -> &'static str { - "mock-device-cpu" + buffers: &mut WgpuBufferSet, + output: &mut BatchValuesBuffer, + ) -> Result { + backend.compute_batch_with_buffers(self, batch, buffers, output) } - fn capabilities(&self) -> BackendCapabilities { - BackendCapabilities::scalar_f64() + #[cfg(feature = "backend-wgpu")] + /// Execute batch gradient computation through real WGPU buffers. + pub fn gradient_batch_wgpu_into( + &self, + backend: &WgpuBackend, + batch: BatchInputs<'_>, + buffers: &mut WgpuBufferSet, + output: &mut BatchGradientsBuffer, + ) -> Result { + backend.gradient_batch_with_buffers(self, batch, buffers, output) } - fn compute(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result { - ScalarBackend.compute(graph, inputs) + #[cfg(feature = "backend-wgpu")] + /// Return whether this graph is statically eligible for the exact-safe native WGPU compute path. + #[must_use] + pub fn supports_native_wgpu_batch_compute(&self, backend: &WgpuBackend) -> bool { + backend.supports_native_batch_compute(self) } - fn compute_batch( + #[cfg(feature = "backend-wgpu")] + /// Return whether this graph and concrete batch can use the exact-safe native WGPU compute path. + #[must_use] + pub fn supports_native_wgpu_batch_compute_for_batch( &self, - graph: &CompiledGraph, + backend: &WgpuBackend, batch: BatchInputs<'_>, - buffer: &mut BatchValuesBuffer, - ) -> Result<()> { - ScalarBackend.compute_batch(graph, batch, buffer) + ) -> bool { + backend.supports_native_batch_compute_for_batch(self, batch) } - fn gradient(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result<(f64, Vec)> { - ScalarBackend.gradient(graph, inputs) + /// Return static compatibility details for the preferred SIMD backend. + pub fn simd_support_report(&self) -> Result { + let f64x4 = self.backend_support_report(BackendKind::SimdF64x4)?; + if f64x4.supports_batch_compute || f64x4.supports_batch_gradient { + return Ok(f64x4); + } + self.backend_support_report(BackendKind::SimdF64x2) } - fn gradient_batch( - &self, - graph: &CompiledGraph, - batch: BatchInputs<'_>, - buffer: &mut BatchGradientsBuffer, - ) -> Result<()> { - ScalarBackend.gradient_batch(graph, batch, buffer) + /// Return the preferred backend for batch value computation. + #[must_use] + pub fn recommended_batch_compute_backend(&self) -> BackendKind { + match self.backend_support_report(BackendKind::SimdF64x4) { + Ok(report) if report.can_compute_batch() => BackendKind::SimdF64x4, + _ => match self.backend_support_report(BackendKind::SimdF64x2) { + Ok(report) if report.can_compute_batch() => BackendKind::SimdF64x2, + _ => BackendKind::Scalar, + }, + } } -} -#[cfg(feature = "backend-wgpu")] -impl ExecutionBackend for WgpuBackend { - fn name(&self) -> &'static str { - "wgpu" + /// Return the preferred backend for batch gradient computation. + #[must_use] + pub fn recommended_batch_gradient_backend(&self) -> BackendKind { + match self.backend_support_report(BackendKind::SimdF64x4) { + Ok(report) if report.can_gradient_batch() => BackendKind::SimdF64x4, + _ => match self.backend_support_report(BackendKind::SimdF64x2) { + Ok(report) if report.can_gradient_batch() => BackendKind::SimdF64x2, + _ => BackendKind::Scalar, + }, + } } - fn capabilities(&self) -> BackendCapabilities { - BackendCapabilities::wgpu_f64() + /// Compute all selected outputs for a batch of rows. + pub fn compute_batch(&self, batch: BatchInputs<'_>) -> Result { + let mut buffer = BatchValuesBuffer::new(); + self.compute_batch_into(batch, &mut buffer)?; + Ok(BatchValues { + data: buffer.data, + batch_size: buffer.batch_size, + output_dim: buffer.output_dim, + }) } - fn compute(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result { - let batch = BatchInputs::new(inputs, 1, inputs.len())?; - let mut values = BatchValuesBuffer::new(); - self.compute_batch(graph, batch, &mut values)?; - values - .data - .first() - .copied() - .ok_or(AutodiffError::InvalidGraph { - reason: "wgpu compute did not produce an output value", - }) + /// Compute all selected outputs for a batch with automatic backend dispatch. + pub fn compute_batch_auto(&self, batch: BatchInputs<'_>) -> Result<(BackendKind, BatchValues)> { + let mut buffer = BatchValuesBuffer::new(); + let backend = self.compute_batch_auto_into(batch, &mut buffer)?; + Ok(( + backend, + BatchValues { + data: buffer.data, + batch_size: buffer.batch_size, + output_dim: buffer.output_dim, + }, + )) } - fn compute_batch( + /// Compute all selected outputs into a reusable output buffer. + pub fn compute_batch_into( &self, - graph: &CompiledGraph, batch: BatchInputs<'_>, buffer: &mut BatchValuesBuffer, ) -> Result<()> { - graph.validate_backend_capabilities(&self.capabilities())?; - let mut gpu_buffers = self.allocate_batch_buffers(graph, batch.batch_size)?; - self.compute_batch_with_buffers(graph, batch, &mut gpu_buffers, buffer)?; + self.check_batch(batch)?; + let mut workspace = self.workspace(); + let output_dim = self.output_nodes.len(); + buffer.reset(batch.batch_size, output_dim); + for row_index in 0..batch.batch_size { + let outputs = + self.compute_many_with_workspace(batch.try_row(row_index)?, &mut workspace)?; + buffer.data.extend_from_slice(outputs); + } Ok(()) } - fn gradient(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result<(f64, Vec)> { - let batch = BatchInputs::new(inputs, 1, inputs.len())?; - let mut gradients = BatchGradientsBuffer::new(); - self.gradient_batch(graph, batch, &mut gradients)?; - let value = gradients - .values - .first() - .copied() - .ok_or(AutodiffError::InvalidGraph { - reason: "wgpu gradient did not produce a primary output value", - })?; - Ok((value, gradients.gradients)) - } - - fn gradient_batch( + /// Compute all selected outputs into a reusable buffer with automatic backend dispatch. + pub fn compute_batch_auto_into( &self, - graph: &CompiledGraph, batch: BatchInputs<'_>, - buffer: &mut BatchGradientsBuffer, - ) -> Result<()> { - graph.validate_backend_capabilities(&self.capabilities())?; - let mut gpu_buffers = self.allocate_batch_buffers(graph, batch.batch_size)?; - self.gradient_batch_with_buffers(graph, batch, &mut gpu_buffers, buffer)?; - Ok(()) - } -} - -impl ExecutionBackend for ScalarBackend { - fn name(&self) -> &'static str { - "scalar" + buffer: &mut BatchValuesBuffer, + ) -> Result { + let backend = self.recommended_batch_compute_backend(); + backend.compute_batch(self, batch, buffer)?; + Ok(backend) } - fn capabilities(&self) -> BackendCapabilities { - BackendCapabilities::scalar_f64() + /// Compute primary-output values and gradients for a batch of rows. + pub fn gradient_batch(&self, batch: BatchInputs<'_>) -> Result { + let mut buffer = BatchGradientsBuffer::new(); + self.gradient_batch_into(batch, &mut buffer)?; + Ok(BatchGradients { + values: buffer.values, + gradients: buffer.gradients, + batch_size: buffer.batch_size, + input_dim: buffer.input_dim, + }) } - fn compute(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result { - graph.validate_backend_capabilities(&self.capabilities())?; - graph.compute(inputs) + /// Compute primary-output values and gradients for a batch with automatic backend dispatch. + pub fn gradient_batch_auto( + &self, + batch: BatchInputs<'_>, + ) -> Result<(BackendKind, BatchGradients)> { + let mut buffer = BatchGradientsBuffer::new(); + let backend = self.gradient_batch_auto_into(batch, &mut buffer)?; + Ok(( + backend, + BatchGradients { + values: buffer.values, + gradients: buffer.gradients, + batch_size: buffer.batch_size, + input_dim: buffer.input_dim, + }, + )) } - fn compute_batch( + /// Compute primary-output values and gradients into a reusable gradient buffer. + pub fn gradient_batch_into( &self, - graph: &CompiledGraph, batch: BatchInputs<'_>, - buffer: &mut BatchValuesBuffer, + buffer: &mut BatchGradientsBuffer, ) -> Result<()> { - let capabilities = self.capabilities(); - if !capabilities.supports_batch_compute { - return Err(AutodiffError::InvalidGraph { - reason: "backend does not support batch compute", - }); - } - graph.validate_backend_capabilities(&capabilities)?; - graph.compute_batch_into(batch, buffer) - } - - fn gradient(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result<(f64, Vec)> { - let capabilities = self.capabilities(); - if !capabilities.supports_reverse_gradient { - return Err(AutodiffError::InvalidGraph { - reason: "backend does not support reverse gradients", - }); + self.check_batch(batch)?; + let mut workspace = self.workspace(); + buffer.reset(batch.batch_size, self.num_inputs); + for row_index in 0..batch.batch_size { + let (value, gradient) = + self.gradient_with_workspace(batch.try_row(row_index)?, &mut workspace)?; + buffer.values.push(value); + buffer.gradients.extend_from_slice(gradient); } - graph.validate_backend_capabilities(&capabilities)?; - graph.gradient(inputs) + Ok(()) } - fn gradient_batch( + /// Compute primary-output values and gradients into a reusable buffer with automatic backend dispatch. + pub fn gradient_batch_auto_into( &self, - graph: &CompiledGraph, batch: BatchInputs<'_>, buffer: &mut BatchGradientsBuffer, - ) -> Result<()> { - let capabilities = self.capabilities(); - if !capabilities.supports_batch_gradient { - return Err(AutodiffError::InvalidGraph { - reason: "backend does not support batch gradients", - }); - } - graph.validate_backend_capabilities(&capabilities)?; - graph.gradient_batch_into(batch, buffer) + ) -> Result { + let backend = self.recommended_batch_gradient_backend(); + backend.gradient_batch(self, batch, buffer)?; + Ok(backend) } } diff --git a/src/multi/expr.rs b/src/multi/expr.rs new file mode 100644 index 0000000..2737020 --- /dev/null +++ b/src/multi/expr.rs @@ -0,0 +1,215 @@ +//! Expression graph with operator-overloaded construction. + +use std::{ + cell::RefCell, + ops::{Add, Div, Mul, Neg, Sub}, + rc::Rc, +}; + +use super::graph::{Graph, NodeId}; +use super::multi_ad::MultiAD; +use crate::Result; + +/// Shared expression graph used for operator-overloaded graph construction. +#[derive(Debug, Clone)] +pub struct ExprGraph { + graph: Rc>, +} + +/// A node handle tied to an [`ExprGraph`]. +#[derive(Debug, Clone)] +pub struct ExprNode { + graph: Rc>, + node: NodeId, +} + +impl ExprGraph { + /// Create an expression graph with `num_inputs` input variables. + #[must_use] + pub fn new(num_inputs: usize) -> Self { + Self { + graph: Rc::new(RefCell::new(Graph::new(num_inputs))), + } + } + + /// Return an expression node for an input. + #[must_use] + pub fn input(&self, input_index: usize) -> ExprNode { + let node = self.graph.borrow().input(input_index); + ExprNode { + graph: Rc::clone(&self.graph), + node, + } + } + + /// Return an expression node for a literal constant. + pub fn constant(&self, value: f64) -> ExprNode { + let node = self.graph.borrow_mut().constant(value); + ExprNode { + graph: Rc::clone(&self.graph), + node, + } + } + + /// Select an expression node as the graph output. + pub fn set_output(&self, expr: &ExprNode) -> Result<()> { + self.graph.borrow_mut().set_output(expr.node)?; + Ok(()) + } + + /// Clone out the underlying reusable graph. + #[must_use] + pub fn graph(&self) -> Graph { + self.graph.borrow().clone() + } +} + +impl ExprNode { + /// Return the underlying node id. + #[must_use] + pub fn node_id(&self) -> NodeId { + self.node + } + + /// Check that `self` and `other` belong to the same [`ExprGraph`]. + /// + /// # Panics + /// + /// Panics if `self` and `other` were created from different [`ExprGraph`] + /// instances. All arithmetic operators (`+`, `-`, `*`, `/`) between two + /// [`ExprNode`] values delegate here, so mixing nodes from separate graphs + /// will panic at runtime. + fn same_graph(&self, other: &ExprNode) { + assert!( + Rc::ptr_eq(&self.graph, &other.graph), + "ExprNode graph mismatch" + ); + } + + fn unary(&self, op: MultiAD) -> ExprNode { + let node = self.graph.borrow_mut().push_operation(op, vec![self.node]); + ExprNode { + graph: Rc::clone(&self.graph), + node, + } + } + + fn binary(&self, op: MultiAD, other: &ExprNode) -> ExprNode { + self.same_graph(other); + let node = self + .graph + .borrow_mut() + .push_operation(op, vec![self.node, other.node]); + ExprNode { + graph: Rc::clone(&self.graph), + node, + } + } + + fn binary_const(&self, op: MultiAD, value: f64) -> ExprNode { + let mut graph = self.graph.borrow_mut(); + let constant = graph.constant(value); + let node = graph.push_operation(op, vec![self.node, constant]); + ExprNode { + graph: Rc::clone(&self.graph), + node, + } + } + + /// Append `sin(self)`. + pub fn sin(&self) -> ExprNode { + self.unary(MultiAD::Sin) + } + + /// Append `cos(self)`. + pub fn cos(&self) -> ExprNode { + self.unary(MultiAD::Cos) + } + + /// Append `exp(self)`. + pub fn exp(&self) -> ExprNode { + self.unary(MultiAD::Exp) + } + + /// Append `ln(self)`. + pub fn ln(&self) -> ExprNode { + self.unary(MultiAD::Ln) + } + + /// Append `sqrt(self)`. + pub fn sqrt(&self) -> ExprNode { + self.unary(MultiAD::Sqrt) + } +} + +impl Add for ExprNode { + type Output = ExprNode; + + fn add(self, rhs: ExprNode) -> Self::Output { + self.binary(MultiAD::Add, &rhs) + } +} + +impl Add for ExprNode { + type Output = ExprNode; + + fn add(self, rhs: f64) -> Self::Output { + self.binary_const(MultiAD::Add, rhs) + } +} + +impl Sub for ExprNode { + type Output = ExprNode; + + fn sub(self, rhs: ExprNode) -> Self::Output { + self.binary(MultiAD::Sub, &rhs) + } +} + +impl Sub for ExprNode { + type Output = ExprNode; + + fn sub(self, rhs: f64) -> Self::Output { + self.binary_const(MultiAD::Sub, rhs) + } +} + +impl Mul for ExprNode { + type Output = ExprNode; + + fn mul(self, rhs: ExprNode) -> Self::Output { + self.binary(MultiAD::Mul, &rhs) + } +} + +impl Mul for ExprNode { + type Output = ExprNode; + + fn mul(self, rhs: f64) -> Self::Output { + self.binary_const(MultiAD::Mul, rhs) + } +} + +impl Div for ExprNode { + type Output = ExprNode; + + fn div(self, rhs: ExprNode) -> Self::Output { + self.binary(MultiAD::Div, &rhs) + } +} + +impl Div for ExprNode { + type Output = ExprNode; + + fn div(self, rhs: f64) -> Self::Output { + self.binary_const(MultiAD::Div, rhs) + } +} + +impl Neg for ExprNode { + type Output = ExprNode; + + fn neg(self) -> Self::Output { + self.unary(MultiAD::Neg) + } +} diff --git a/src/multi/graph.rs b/src/multi/graph.rs index 5b82ee6..bd69aef 100644 --- a/src/multi/graph.rs +++ b/src/multi/graph.rs @@ -33,35 +33,29 @@ //! assert_eq!(jacobian.len(), 2); //! ``` -use std::{ - cell::RefCell, - fmt::Write, - ops::{Add, Div, Mul, Neg, Sub}, - rc::Rc, - sync::Arc, -}; +use std::{fmt::Write, sync::Arc}; +use super::backend::{ + BackendKind, BackendSupportReport, DeviceBatchPlan, DeviceBufferSet, DeviceExecutionTrace, + Instruction, +}; use super::compiled::{ - BackendKind, BackendSupportReport, BatchGradients, BatchGradientsBuffer, BatchInputs, - BatchValues, BatchValuesBuffer, CompiledGraph, CompiledGraphMetadata, CompiledWorkspace, - DeviceBatchPlan, DeviceBufferSet, DeviceExecutionTrace, Instruction, + BatchGradients, BatchGradientsBuffer, BatchInputs, BatchValues, BatchValuesBuffer, + CompiledGraph, CompiledGraphMetadata, CompiledWorkspace, }; +#[cfg(test)] +use super::expr::ExprGraph; use super::multi_ad::MultiAD; use super::multi_ad_fr::MultiAD2FR; use super::multi_ad_rf::MultiAD2RF; use super::op_rules; use super::parser; +#[cfg(test)] +use super::tape::TapeWorkspace; +use super::tape::{CompiledArgRange, CompiledNode, Tape}; use super::types::BackwardResultBox; use crate::{AutodiffError, Result}; -/// Reusable scratch buffers for repeated [`Tape`] evaluation. -#[derive(Debug, Clone, Default, PartialEq)] -pub struct TapeWorkspace { - values: Vec, - cotangent_values: Vec, - gradients: Vec, -} - /// A handle to an input or computed node in a graph. pub type NodeId = usize; @@ -82,8 +76,8 @@ pub enum GraphNode { #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] #[derive(Debug, Clone, PartialEq)] pub struct Graph { - num_inputs: usize, - nodes: Vec, + pub(super) num_inputs: usize, + pub(super) nodes: Vec, output_nodes: Vec, input_names: Vec>, output_names: Vec<(NodeId, String)>, @@ -97,27 +91,6 @@ pub struct Graph { /// /// Today this stores the graph structure directly and provides a stable place /// to grow future precomputation or buffer reuse without changing the public API. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct CompiledArgRange { - start: usize, - len: usize, -} - -#[derive(Debug, Clone, Copy, PartialEq)] -enum CompiledNode { - Constant(f64), - Operation { - op: MultiAD, - arg_range: CompiledArgRange, - }, -} - -#[derive(Debug, Clone, Copy, PartialEq)] -enum BackwardLocal { - Unary(f64), - Binary(f64, f64), -} - #[derive(Debug, Clone, PartialEq)] enum ExactLocal { None, @@ -137,14 +110,6 @@ enum ExactLocal { }, } -#[derive(Debug, Clone, PartialEq)] -pub struct Tape { - graph: Graph, - output_indices: Vec, - compiled_nodes: Arc<[CompiledNode]>, - arg_indices: Arc<[usize]>, -} - /// One component comparison from [`Graph::check_gradient`]. #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] #[derive(Debug, Clone, PartialEq)] @@ -197,19 +162,6 @@ pub struct GraphStats { pub op_counts: Vec<(MultiAD, usize)>, } -/// Shared expression graph used for operator-overloaded graph construction. -#[derive(Debug, Clone)] -pub struct ExprGraph { - graph: Rc>, -} - -/// A node handle tied to an [`ExprGraph`]. -#[derive(Debug, Clone)] -pub struct ExprNode { - graph: Rc>, - node: NodeId, -} - #[inline] fn op_name(op: MultiAD) -> &'static str { op_rules::op_name(op) @@ -220,106 +172,6 @@ fn expected_arity(op: MultiAD) -> usize { op_rules::expected_arity(op) } -#[inline] -fn compiled_arg_indices(arg_indices_storage: &[usize], arg_range: CompiledArgRange) -> &[usize] { - &arg_indices_storage[arg_range.start..arg_range.start + arg_range.len] -} - -#[inline] -fn with_compiled_arg_values( - arg_indices_storage: &[usize], - arg_range: CompiledArgRange, - values: &[f64], - f: F, -) -> Result -where - F: FnOnce(&[f64]) -> Result, -{ - let arg_indices = compiled_arg_indices(arg_indices_storage, arg_range); - match arg_indices.len() { - 0 => f(&[]), - 1 => { - MultiAD::check_value_index(arg_indices[0], values.len())?; - let args = [values[arg_indices[0]]]; - f(&args) - } - 2 => { - MultiAD::check_value_index(arg_indices[0], values.len())?; - MultiAD::check_value_index(arg_indices[1], values.len())?; - let args = [values[arg_indices[0]], values[arg_indices[1]]]; - f(&args) - } - _ => { - let arg_values = MultiAD::gather_arg_values(arg_indices, values)?; - f(&arg_values) - } - } -} - -#[inline] -fn backward_local(rule: op_rules::LocalRule) -> BackwardLocal { - match rule { - op_rules::LocalRule::Unary { dy, .. } => BackwardLocal::Unary(dy), - op_rules::LocalRule::Binary { - dy_left, dy_right, .. - } => BackwardLocal::Binary(dy_left, dy_right), - } -} - -fn reverse_accumulate_compiled( - num_inputs: usize, - compiled_nodes: &[CompiledNode], - arg_indices_storage: &[usize], - values: &[f64], - cotangent_values: &mut [f64], -) -> Result<()> { - for (offset, node) in compiled_nodes.iter().enumerate().rev() { - let node_id = num_inputs + offset; - let current_cotangent = cotangent_values[node_id]; - if current_cotangent == 0.0 { - continue; - } - - let CompiledNode::Operation { op, arg_range } = *node else { - continue; - }; - let input_indices = compiled_arg_indices(arg_indices_storage, arg_range); - let value = values[node_id]; - - with_compiled_arg_values(arg_indices_storage, arg_range, values, |args| { - match op_rules::local_rule(op, args, value)? { - op_rules::LocalRule::Unary { dy, .. } => { - cotangent_values[input_indices[0]] += current_cotangent * dy; - } - op_rules::LocalRule::Binary { - dy_left, dy_right, .. - } => { - cotangent_values[input_indices[0]] += current_cotangent * dy_left; - cotangent_values[input_indices[1]] += current_cotangent * dy_right; - } - } - Ok(()) - })?; - } - - Ok(()) -} - -impl TapeWorkspace { - /// Create an empty reusable workspace. - #[must_use] - pub fn new() -> Self { - Self::default() - } - - /// Clear all retained buffers while keeping their capacity. - pub fn clear(&mut self) { - self.values.clear(); - self.cotangent_values.clear(); - self.gradients.clear(); - } -} - impl Graph { #[inline] fn output_name_for(&self, node_id: NodeId) -> Option<&str> { @@ -2182,189 +2034,6 @@ impl Graph { } } -impl ExprGraph { - /// Create an expression graph with `num_inputs` input variables. - #[must_use] - pub fn new(num_inputs: usize) -> Self { - Self { - graph: Rc::new(RefCell::new(Graph::new(num_inputs))), - } - } - - /// Return an expression node for an input. - #[must_use] - pub fn input(&self, input_index: usize) -> ExprNode { - let node = self.graph.borrow().input(input_index); - ExprNode { - graph: Rc::clone(&self.graph), - node, - } - } - - /// Return an expression node for a literal constant. - pub fn constant(&self, value: f64) -> ExprNode { - let node = self.graph.borrow_mut().constant(value); - ExprNode { - graph: Rc::clone(&self.graph), - node, - } - } - - /// Select an expression node as the graph output. - pub fn set_output(&self, expr: &ExprNode) -> Result<()> { - self.graph.borrow_mut().set_output(expr.node)?; - Ok(()) - } - - /// Clone out the underlying reusable graph. - #[must_use] - pub fn graph(&self) -> Graph { - self.graph.borrow().clone() - } -} - -impl ExprNode { - /// Return the underlying node id. - #[must_use] - pub fn node_id(&self) -> NodeId { - self.node - } - - fn same_graph(&self, other: &ExprNode) { - assert!( - Rc::ptr_eq(&self.graph, &other.graph), - "ExprNode graph mismatch" - ); - } - - fn unary(&self, op: MultiAD) -> ExprNode { - let node = self.graph.borrow_mut().push_operation(op, vec![self.node]); - ExprNode { - graph: Rc::clone(&self.graph), - node, - } - } - - fn binary(&self, op: MultiAD, other: &ExprNode) -> ExprNode { - self.same_graph(other); - let node = self - .graph - .borrow_mut() - .push_operation(op, vec![self.node, other.node]); - ExprNode { - graph: Rc::clone(&self.graph), - node, - } - } - - fn binary_const(&self, op: MultiAD, value: f64) -> ExprNode { - let mut graph = self.graph.borrow_mut(); - let constant = graph.constant(value); - let node = graph.push_operation(op, vec![self.node, constant]); - ExprNode { - graph: Rc::clone(&self.graph), - node, - } - } - - /// Append `sin(self)`. - pub fn sin(&self) -> ExprNode { - self.unary(MultiAD::Sin) - } - - /// Append `cos(self)`. - pub fn cos(&self) -> ExprNode { - self.unary(MultiAD::Cos) - } - - /// Append `exp(self)`. - pub fn exp(&self) -> ExprNode { - self.unary(MultiAD::Exp) - } - - /// Append `ln(self)`. - pub fn ln(&self) -> ExprNode { - self.unary(MultiAD::Ln) - } - - /// Append `sqrt(self)`. - pub fn sqrt(&self) -> ExprNode { - self.unary(MultiAD::Sqrt) - } -} - -impl Add for ExprNode { - type Output = ExprNode; - - fn add(self, rhs: ExprNode) -> Self::Output { - self.binary(MultiAD::Add, &rhs) - } -} - -impl Add for ExprNode { - type Output = ExprNode; - - fn add(self, rhs: f64) -> Self::Output { - self.binary_const(MultiAD::Add, rhs) - } -} - -impl Sub for ExprNode { - type Output = ExprNode; - - fn sub(self, rhs: ExprNode) -> Self::Output { - self.binary(MultiAD::Sub, &rhs) - } -} - -impl Sub for ExprNode { - type Output = ExprNode; - - fn sub(self, rhs: f64) -> Self::Output { - self.binary_const(MultiAD::Sub, rhs) - } -} - -impl Mul for ExprNode { - type Output = ExprNode; - - fn mul(self, rhs: ExprNode) -> Self::Output { - self.binary(MultiAD::Mul, &rhs) - } -} - -impl Mul for ExprNode { - type Output = ExprNode; - - fn mul(self, rhs: f64) -> Self::Output { - self.binary_const(MultiAD::Mul, rhs) - } -} - -impl Div for ExprNode { - type Output = ExprNode; - - fn div(self, rhs: ExprNode) -> Self::Output { - self.binary(MultiAD::Div, &rhs) - } -} - -impl Div for ExprNode { - type Output = ExprNode; - - fn div(self, rhs: f64) -> Self::Output { - self.binary_const(MultiAD::Div, rhs) - } -} - -impl Neg for ExprNode { - type Output = ExprNode; - - fn neg(self) -> Self::Output { - self.unary(MultiAD::Neg) - } -} - impl TryFrom<&Graph> for Vec<(MultiAD, Vec)> { type Error = crate::AutodiffError; @@ -2381,425 +2050,6 @@ impl TryFrom for Vec<(MultiAD, Vec)> { } } -impl Tape { - #[inline] - fn check_input_len(&self, inputs: &[f64]) -> Result<()> { - if inputs.len() == self.graph.num_inputs { - Ok(()) - } else { - Err(AutodiffError::InvalidGraph { - reason: "input length must match graph.num_inputs()", - }) - } - } - - #[inline] - fn fill_values(&self, inputs: &[f64], workspace: &mut TapeWorkspace) -> Result<()> { - self.fill_values_inner(inputs, workspace, false) - } - - #[inline] - fn fill_values_checked(&self, inputs: &[f64], workspace: &mut TapeWorkspace) -> Result<()> { - self.fill_values_inner(inputs, workspace, true) - } - - fn fill_values_inner( - &self, - inputs: &[f64], - workspace: &mut TapeWorkspace, - checked: bool, - ) -> Result<()> { - self.check_input_len(inputs)?; - workspace.values.clear(); - workspace - .values - .reserve(self.graph.num_inputs + self.graph.nodes.len()); - workspace.values.extend_from_slice(inputs); - - for node in self.compiled_nodes.iter() { - match node { - CompiledNode::Constant(value) => workspace.values.push(*value), - CompiledNode::Operation { op, arg_range } => { - let value = with_compiled_arg_values( - &self.arg_indices, - *arg_range, - &workspace.values, - |args| { - if checked { - op.forward_checked(args) - } else { - op.forward(args) - } - }, - )?; - workspace.values.push(value); - } - } - } - - Ok(()) - } - - /// Return the underlying graph. - #[must_use] - pub fn graph(&self) -> &Graph { - &self.graph - } - - /// Create a reusable workspace sized for this tape. - #[must_use] - pub fn workspace(&self) -> TapeWorkspace { - TapeWorkspace { - values: Vec::with_capacity(self.graph.num_inputs + self.graph.nodes.len()), - cotangent_values: Vec::with_capacity(self.graph.num_inputs + self.graph.nodes.len()), - gradients: Vec::with_capacity(self.graph.num_inputs), - } - } - - /// Compute only the primary output value. - pub fn compute(&self, inputs: &[f64]) -> Result { - let mut workspace = self.workspace(); - self.compute_with_workspace(inputs, &mut workspace) - } - - /// Compute all selected output values. - pub fn compute_many(&self, inputs: &[f64]) -> Result> { - let mut workspace = self.workspace(); - self.compute_many_with_workspace(inputs, &mut workspace) - } - - /// Compute all selected output values with checked-domain validation. - pub fn compute_many_checked(&self, inputs: &[f64]) -> Result> { - let mut workspace = self.workspace(); - self.compute_many_with_workspace_checked(inputs, &mut workspace) - } - - /// Compute only the primary output value with checked-domain validation. - pub fn compute_checked(&self, inputs: &[f64]) -> Result { - let mut workspace = self.workspace(); - self.compute_with_workspace_checked(inputs, &mut workspace) - } - - /// Compute only the primary output value using a reusable workspace. - pub fn compute_with_workspace( - &self, - inputs: &[f64], - workspace: &mut TapeWorkspace, - ) -> Result { - self.fill_values(inputs, workspace)?; - Ok(self - .output_indices - .first() - .map(|&index| workspace.values[index]) - .unwrap_or(0.0)) - } - - /// Compute all selected output values using a reusable workspace. - pub fn compute_many_with_workspace( - &self, - inputs: &[f64], - workspace: &mut TapeWorkspace, - ) -> Result> { - self.fill_values(inputs, workspace)?; - Ok(self - .output_indices - .iter() - .map(|&index| workspace.values[index]) - .collect()) - } - - /// Compute all selected output values using a reusable workspace with checked-domain validation. - pub fn compute_many_with_workspace_checked( - &self, - inputs: &[f64], - workspace: &mut TapeWorkspace, - ) -> Result> { - self.fill_values_checked(inputs, workspace)?; - Ok(self - .output_indices - .iter() - .map(|&index| workspace.values[index]) - .collect()) - } - - /// Compute only the primary output value using a reusable workspace with checked-domain validation. - pub fn compute_with_workspace_checked( - &self, - inputs: &[f64], - workspace: &mut TapeWorkspace, - ) -> Result { - self.fill_values_checked(inputs, workspace)?; - Ok(self - .output_indices - .first() - .map(|&index| workspace.values[index]) - .unwrap_or(0.0)) - } - - /// Compute the output value and gradient closure. - pub fn compute_grad(&self, inputs: &[f64]) -> Result { - self.check_input_len(inputs)?; - let mut values: Vec = - Vec::with_capacity(self.graph.num_inputs + self.graph.nodes.len()); - values.extend_from_slice(inputs); - - let mut backward_locals: Vec> = - Vec::with_capacity(self.compiled_nodes.len()); - - for node in self.compiled_nodes.iter() { - match node { - CompiledNode::Constant(value) => { - values.push(*value); - backward_locals.push(None); - } - CompiledNode::Operation { op, arg_range } => { - let (value, backward_local) = - with_compiled_arg_values(&self.arg_indices, *arg_range, &values, |args| { - let value = op.forward(args)?; - let local_rule = op_rules::local_rule(*op, args, value)?; - Ok((value, backward_local(local_rule))) - })?; - values.push(value); - backward_locals.push(Some(backward_local)); - } - } - } - - let final_value = self - .output_indices - .first() - .map(|&index| values[index]) - .unwrap_or(0.0); - let num_inputs = self.graph.num_inputs; - let values_len = values.len(); - let output_index = self.output_indices.first().copied(); - let compiled_nodes = Arc::clone(&self.compiled_nodes); - let arg_indices = Arc::clone(&self.arg_indices); - - let backward_fn = Box::new(move |cotangent: f64| -> Vec { - let Some(final_output_index) = output_index else { - return Vec::new(); - }; - - let mut cotangent_values = vec![0.0; values_len]; - cotangent_values[final_output_index] = cotangent; - - for (offset, backward_local) in backward_locals.iter().enumerate().rev() { - let Some(local) = backward_local else { - continue; - }; - let node_id = num_inputs + offset; - let current_cotangent = cotangent_values[node_id]; - if current_cotangent == 0.0 { - continue; - } - - let CompiledNode::Operation { arg_range, .. } = compiled_nodes[offset] else { - continue; - }; - let input_indices = compiled_arg_indices(&arg_indices, arg_range); - match local { - BackwardLocal::Unary(dy) => { - cotangent_values[input_indices[0]] += current_cotangent * dy; - } - BackwardLocal::Binary(dy_left, dy_right) => { - cotangent_values[input_indices[0]] += current_cotangent * dy_left; - cotangent_values[input_indices[1]] += current_cotangent * dy_right; - } - } - } - - cotangent_values[..num_inputs.min(cotangent_values.len())].to_vec() - }); - - Ok((final_value, backward_fn)) - } - - /// Compute the output value and gradient eagerly using a reusable workspace. - pub fn gradient_with_workspace<'a>( - &self, - inputs: &[f64], - workspace: &'a mut TapeWorkspace, - ) -> Result<(f64, &'a [f64])> { - self.fill_values(inputs, workspace)?; - - workspace.cotangent_values.clear(); - workspace - .cotangent_values - .resize(workspace.values.len(), 0.0); - workspace.gradients.clear(); - workspace.gradients.resize(self.graph.num_inputs, 0.0); - - let Some(output_index) = self.output_indices.first().copied() else { - return Ok((0.0, &workspace.gradients)); - }; - - workspace.cotangent_values[output_index] = 1.0; - reverse_accumulate_compiled( - self.graph.num_inputs, - &self.compiled_nodes, - &self.arg_indices, - &workspace.values, - &mut workspace.cotangent_values, - )?; - - workspace - .gradients - .copy_from_slice(&workspace.cotangent_values[..self.graph.num_inputs]); - Ok((workspace.values[output_index], &workspace.gradients)) - } - - /// Compute the output value and gradient eagerly using a reusable workspace with checked-domain validation. - pub fn gradient_with_workspace_checked<'a>( - &self, - inputs: &[f64], - workspace: &'a mut TapeWorkspace, - ) -> Result<(f64, &'a [f64])> { - self.fill_values_checked(inputs, workspace)?; - - workspace.cotangent_values.clear(); - workspace - .cotangent_values - .resize(workspace.values.len(), 0.0); - workspace.gradients.clear(); - workspace.gradients.resize(self.graph.num_inputs, 0.0); - - let Some(output_index) = self.output_indices.first().copied() else { - return Ok((0.0, &workspace.gradients)); - }; - - workspace.cotangent_values[output_index] = 1.0; - reverse_accumulate_compiled( - self.graph.num_inputs, - &self.compiled_nodes, - &self.arg_indices, - &workspace.values, - &mut workspace.cotangent_values, - )?; - - workspace - .gradients - .copy_from_slice(&workspace.cotangent_values[..self.graph.num_inputs]); - Ok((workspace.values[output_index], &workspace.gradients)) - } - - /// Compute the output value and gradient eagerly. - pub fn gradient(&self, inputs: &[f64]) -> Result<(f64, Vec)> { - let mut workspace = self.workspace(); - let (value, gradient) = self.gradient_with_workspace(inputs, &mut workspace)?; - Ok((value, gradient.to_vec())) - } - - /// Compute the output value and gradient eagerly with checked-domain validation. - pub fn gradient_checked(&self, inputs: &[f64]) -> Result<(f64, Vec)> { - let mut workspace = self.workspace(); - let (value, gradient) = self.gradient_with_workspace_checked(inputs, &mut workspace)?; - Ok((value, gradient.to_vec())) - } - - /// Compute the Jacobian for all selected outputs. - pub fn jacobian(&self, inputs: &[f64]) -> Result>> { - let mut workspace = self.workspace(); - self.jacobian_with_workspace(inputs, &mut workspace) - } - - /// Compute the Jacobian for all selected outputs with checked-domain validation. - pub fn jacobian_checked(&self, inputs: &[f64]) -> Result>> { - let mut workspace = self.workspace(); - self.jacobian_with_workspace_checked(inputs, &mut workspace) - } - - /// Compute the Jacobian for all selected outputs using a reusable workspace. - pub fn jacobian_with_workspace( - &self, - inputs: &[f64], - workspace: &mut TapeWorkspace, - ) -> Result>> { - self.jacobian_with_workspace_inner(inputs, workspace, false) - } - - /// Compute the Jacobian for all selected outputs using a reusable workspace with checked-domain validation. - pub fn jacobian_with_workspace_checked( - &self, - inputs: &[f64], - workspace: &mut TapeWorkspace, - ) -> Result>> { - self.jacobian_with_workspace_inner(inputs, workspace, true) - } - - fn jacobian_with_workspace_inner( - &self, - inputs: &[f64], - workspace: &mut TapeWorkspace, - checked: bool, - ) -> Result>> { - if checked { - self.fill_values_checked(inputs, workspace)?; - } else { - self.fill_values(inputs, workspace)?; - } - let mut jacobian: Vec> = Vec::with_capacity(self.output_indices.len()); - - for &output_index in &self.output_indices { - workspace.cotangent_values.clear(); - workspace - .cotangent_values - .resize(workspace.values.len(), 0.0); - workspace.gradients.clear(); - workspace.gradients.resize(self.graph.num_inputs, 0.0); - workspace.cotangent_values[output_index] = 1.0; - reverse_accumulate_compiled( - self.graph.num_inputs, - &self.compiled_nodes, - &self.arg_indices, - &workspace.values, - &mut workspace.cotangent_values, - )?; - - workspace - .gradients - .copy_from_slice(&workspace.cotangent_values[..self.graph.num_inputs]); - jacobian.push(workspace.gradients.clone()); - } - - Ok(jacobian) - } - - /// Compute a finite-difference Hessian by differentiating eager gradients. - pub fn compute_hessian(&self, inputs: &[f64]) -> Result>> { - let num_inputs = self.graph.num_inputs; - let epsilon = 1e-5; - let mut hessian = vec![vec![0.0; num_inputs]; num_inputs]; - - if num_inputs == 0 { - let (_value, _grad) = self.gradient(inputs)?; - return Ok(hessian); - } - - let mut plus_workspace = self.workspace(); - let mut minus_workspace = self.workspace(); - - for j in 0..num_inputs { - let mut inputs_plus = inputs.to_vec(); - inputs_plus[j] += epsilon; - - let mut inputs_minus = inputs.to_vec(); - inputs_minus[j] -= epsilon; - - let (_value_plus, grad_plus) = - self.gradient_with_workspace(&inputs_plus, &mut plus_workspace)?; - let (_value_minus, grad_minus) = - self.gradient_with_workspace(&inputs_minus, &mut minus_workspace)?; - - for i in 0..num_inputs { - hessian[i][j] = (grad_plus[i] - grad_minus[i]) / (2.0 * epsilon); - } - } - - Ok(hessian) - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/multi/multi_fn.rs b/src/multi/multi_fn.rs index f00fef7..6fc91f8 100644 --- a/src/multi/multi_fn.rs +++ b/src/multi/multi_fn.rs @@ -3,7 +3,6 @@ pub use super::types::BackwardResultBox; use crate::error::Result; /// Type alias for a multi-variable computation graph -#[allow(dead_code)] // Public API for library extension pub type GraphType = [(MultiAD, Vec)]; /// Trait for multi-variable functions with analytical gradients. @@ -13,7 +12,6 @@ pub type GraphType = [(MultiAD, Vec)]; /// /// This trait is primarily intended for testing and demonstration purposes. /// Most users will work directly with the `MultiAD` enum. -#[allow(dead_code)] // Public API for library extension pub trait MultiFn { /// Returns the input values for this function. fn inputs(&self) -> Vec; diff --git a/src/multi/tape.rs b/src/multi/tape.rs new file mode 100644 index 0000000..134aef0 --- /dev/null +++ b/src/multi/tape.rs @@ -0,0 +1,565 @@ +//! Tape-based compiled graph evaluation. + +use std::sync::Arc; + +use super::graph::{Graph, NodeId}; +use super::multi_ad::MultiAD; +use super::op_rules; +use super::types::BackwardResultBox; +use crate::{AutodiffError, Result}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct CompiledArgRange { + pub(super) start: usize, + pub(super) len: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub(super) enum CompiledNode { + Constant(f64), + Operation { + op: MultiAD, + arg_range: CompiledArgRange, + }, +} + +#[derive(Debug, Clone, Copy, PartialEq)] +enum BackwardLocal { + Unary(f64), + Binary(f64, f64), +} + +/// Reusable scratch buffers for repeated [`Tape`] evaluation. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct TapeWorkspace { + pub(super) values: Vec, + pub(super) cotangent_values: Vec, + pub(super) gradients: Vec, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct Tape { + pub(super) graph: Graph, + pub(super) output_indices: Vec, + pub(super) compiled_nodes: Arc<[CompiledNode]>, + pub(super) arg_indices: Arc<[usize]>, +} + +#[inline] +fn compiled_arg_indices(arg_indices_storage: &[usize], arg_range: CompiledArgRange) -> &[usize] { + &arg_indices_storage[arg_range.start..arg_range.start + arg_range.len] +} + +#[inline] +fn with_compiled_arg_values( + arg_indices_storage: &[usize], + arg_range: CompiledArgRange, + values: &[f64], + f: F, +) -> Result +where + F: FnOnce(&[f64]) -> Result, +{ + let arg_indices = compiled_arg_indices(arg_indices_storage, arg_range); + match arg_indices.len() { + 0 => f(&[]), + 1 => { + MultiAD::check_value_index(arg_indices[0], values.len())?; + let args = [values[arg_indices[0]]]; + f(&args) + } + 2 => { + MultiAD::check_value_index(arg_indices[0], values.len())?; + MultiAD::check_value_index(arg_indices[1], values.len())?; + let args = [values[arg_indices[0]], values[arg_indices[1]]]; + f(&args) + } + _ => { + let arg_values = MultiAD::gather_arg_values(arg_indices, values)?; + f(&arg_values) + } + } +} + +#[inline] +fn backward_local(rule: op_rules::LocalRule) -> BackwardLocal { + match rule { + op_rules::LocalRule::Unary { dy, .. } => BackwardLocal::Unary(dy), + op_rules::LocalRule::Binary { + dy_left, dy_right, .. + } => BackwardLocal::Binary(dy_left, dy_right), + } +} + +fn reverse_accumulate_compiled( + num_inputs: usize, + compiled_nodes: &[CompiledNode], + arg_indices_storage: &[usize], + values: &[f64], + cotangent_values: &mut [f64], +) -> Result<()> { + for (offset, node) in compiled_nodes.iter().enumerate().rev() { + let node_id = num_inputs + offset; + let current_cotangent = cotangent_values[node_id]; + if current_cotangent == 0.0 { + continue; + } + + let CompiledNode::Operation { op, arg_range } = *node else { + continue; + }; + let input_indices = compiled_arg_indices(arg_indices_storage, arg_range); + let value = values[node_id]; + + with_compiled_arg_values(arg_indices_storage, arg_range, values, |args| { + match op_rules::local_rule(op, args, value)? { + op_rules::LocalRule::Unary { dy, .. } => { + cotangent_values[input_indices[0]] += current_cotangent * dy; + } + op_rules::LocalRule::Binary { + dy_left, dy_right, .. + } => { + cotangent_values[input_indices[0]] += current_cotangent * dy_left; + cotangent_values[input_indices[1]] += current_cotangent * dy_right; + } + } + Ok(()) + })?; + } + + Ok(()) +} + +impl TapeWorkspace { + /// Create an empty reusable workspace. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Clear all retained buffers while keeping their capacity. + pub fn clear(&mut self) { + self.values.clear(); + self.cotangent_values.clear(); + self.gradients.clear(); + } +} + +impl Tape { + #[inline] + fn check_input_len(&self, inputs: &[f64]) -> Result<()> { + if inputs.len() == self.graph.num_inputs { + Ok(()) + } else { + Err(AutodiffError::InvalidGraph { + reason: "input length must match graph.num_inputs()", + }) + } + } + + #[inline] + fn fill_values(&self, inputs: &[f64], workspace: &mut TapeWorkspace) -> Result<()> { + self.fill_values_inner(inputs, workspace, false) + } + + #[inline] + fn fill_values_checked(&self, inputs: &[f64], workspace: &mut TapeWorkspace) -> Result<()> { + self.fill_values_inner(inputs, workspace, true) + } + + fn fill_values_inner( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + checked: bool, + ) -> Result<()> { + self.check_input_len(inputs)?; + workspace.values.clear(); + workspace + .values + .reserve(self.graph.num_inputs + self.graph.nodes.len()); + workspace.values.extend_from_slice(inputs); + + for node in self.compiled_nodes.iter() { + match node { + CompiledNode::Constant(value) => workspace.values.push(*value), + CompiledNode::Operation { op, arg_range } => { + let value = with_compiled_arg_values( + &self.arg_indices, + *arg_range, + &workspace.values, + |args| { + if checked { + op.forward_checked(args) + } else { + op.forward(args) + } + }, + )?; + workspace.values.push(value); + } + } + } + + Ok(()) + } + + /// Return the underlying graph. + #[must_use] + pub fn graph(&self) -> &Graph { + &self.graph + } + + /// Create a reusable workspace sized for this tape. + #[must_use] + pub fn workspace(&self) -> TapeWorkspace { + TapeWorkspace { + values: Vec::with_capacity(self.graph.num_inputs + self.graph.nodes.len()), + cotangent_values: Vec::with_capacity(self.graph.num_inputs + self.graph.nodes.len()), + gradients: Vec::with_capacity(self.graph.num_inputs), + } + } + + /// Compute only the primary output value. + pub fn compute(&self, inputs: &[f64]) -> Result { + let mut workspace = self.workspace(); + self.compute_with_workspace(inputs, &mut workspace) + } + + /// Compute all selected output values. + pub fn compute_many(&self, inputs: &[f64]) -> Result> { + let mut workspace = self.workspace(); + self.compute_many_with_workspace(inputs, &mut workspace) + } + + /// Compute all selected output values with checked-domain validation. + pub fn compute_many_checked(&self, inputs: &[f64]) -> Result> { + let mut workspace = self.workspace(); + self.compute_many_with_workspace_checked(inputs, &mut workspace) + } + + /// Compute only the primary output value with checked-domain validation. + pub fn compute_checked(&self, inputs: &[f64]) -> Result { + let mut workspace = self.workspace(); + self.compute_with_workspace_checked(inputs, &mut workspace) + } + + /// Compute only the primary output value using a reusable workspace. + pub fn compute_with_workspace( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + ) -> Result { + self.fill_values(inputs, workspace)?; + Ok(self + .output_indices + .first() + .map(|&index| workspace.values[index]) + .unwrap_or(0.0)) + } + + /// Compute all selected output values using a reusable workspace. + pub fn compute_many_with_workspace( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + ) -> Result> { + self.fill_values(inputs, workspace)?; + Ok(self + .output_indices + .iter() + .map(|&index| workspace.values[index]) + .collect()) + } + + /// Compute all selected output values using a reusable workspace with checked-domain validation. + pub fn compute_many_with_workspace_checked( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + ) -> Result> { + self.fill_values_checked(inputs, workspace)?; + Ok(self + .output_indices + .iter() + .map(|&index| workspace.values[index]) + .collect()) + } + + /// Compute only the primary output value using a reusable workspace with checked-domain validation. + pub fn compute_with_workspace_checked( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + ) -> Result { + self.fill_values_checked(inputs, workspace)?; + Ok(self + .output_indices + .first() + .map(|&index| workspace.values[index]) + .unwrap_or(0.0)) + } + + /// Compute the output value and gradient closure. + pub fn compute_grad(&self, inputs: &[f64]) -> Result { + self.check_input_len(inputs)?; + let mut values: Vec = + Vec::with_capacity(self.graph.num_inputs + self.graph.nodes.len()); + values.extend_from_slice(inputs); + + let mut backward_locals: Vec> = + Vec::with_capacity(self.compiled_nodes.len()); + + for node in self.compiled_nodes.iter() { + match node { + CompiledNode::Constant(value) => { + values.push(*value); + backward_locals.push(None); + } + CompiledNode::Operation { op, arg_range } => { + let (value, backward_local) = + with_compiled_arg_values(&self.arg_indices, *arg_range, &values, |args| { + let value = op.forward(args)?; + let local_rule = op_rules::local_rule(*op, args, value)?; + Ok((value, backward_local(local_rule))) + })?; + values.push(value); + backward_locals.push(Some(backward_local)); + } + } + } + + let final_value = self + .output_indices + .first() + .map(|&index| values[index]) + .unwrap_or(0.0); + let num_inputs = self.graph.num_inputs; + let values_len = values.len(); + let output_index = self.output_indices.first().copied(); + let compiled_nodes = Arc::clone(&self.compiled_nodes); + let arg_indices = Arc::clone(&self.arg_indices); + + let backward_fn = Box::new(move |cotangent: f64| -> Vec { + let Some(final_output_index) = output_index else { + return Vec::new(); + }; + + let mut cotangent_values = vec![0.0; values_len]; + cotangent_values[final_output_index] = cotangent; + + for (offset, backward_local) in backward_locals.iter().enumerate().rev() { + let Some(local) = backward_local else { + continue; + }; + let node_id = num_inputs + offset; + let current_cotangent = cotangent_values[node_id]; + if current_cotangent == 0.0 { + continue; + } + + let CompiledNode::Operation { arg_range, .. } = compiled_nodes[offset] else { + continue; + }; + let input_indices = compiled_arg_indices(&arg_indices, arg_range); + match local { + BackwardLocal::Unary(dy) => { + cotangent_values[input_indices[0]] += current_cotangent * dy; + } + BackwardLocal::Binary(dy_left, dy_right) => { + cotangent_values[input_indices[0]] += current_cotangent * dy_left; + cotangent_values[input_indices[1]] += current_cotangent * dy_right; + } + } + } + + cotangent_values[..num_inputs.min(cotangent_values.len())].to_vec() + }); + + Ok((final_value, backward_fn)) + } + + /// Compute the output value and gradient eagerly using a reusable workspace. + pub fn gradient_with_workspace<'a>( + &self, + inputs: &[f64], + workspace: &'a mut TapeWorkspace, + ) -> Result<(f64, &'a [f64])> { + self.fill_values(inputs, workspace)?; + + workspace.cotangent_values.clear(); + workspace + .cotangent_values + .resize(workspace.values.len(), 0.0); + workspace.gradients.clear(); + workspace.gradients.resize(self.graph.num_inputs, 0.0); + + let Some(output_index) = self.output_indices.first().copied() else { + return Ok((0.0, &workspace.gradients)); + }; + + workspace.cotangent_values[output_index] = 1.0; + reverse_accumulate_compiled( + self.graph.num_inputs, + &self.compiled_nodes, + &self.arg_indices, + &workspace.values, + &mut workspace.cotangent_values, + )?; + + workspace + .gradients + .copy_from_slice(&workspace.cotangent_values[..self.graph.num_inputs]); + Ok((workspace.values[output_index], &workspace.gradients)) + } + + /// Compute the output value and gradient eagerly using a reusable workspace with checked-domain validation. + pub fn gradient_with_workspace_checked<'a>( + &self, + inputs: &[f64], + workspace: &'a mut TapeWorkspace, + ) -> Result<(f64, &'a [f64])> { + self.fill_values_checked(inputs, workspace)?; + + workspace.cotangent_values.clear(); + workspace + .cotangent_values + .resize(workspace.values.len(), 0.0); + workspace.gradients.clear(); + workspace.gradients.resize(self.graph.num_inputs, 0.0); + + let Some(output_index) = self.output_indices.first().copied() else { + return Ok((0.0, &workspace.gradients)); + }; + + workspace.cotangent_values[output_index] = 1.0; + reverse_accumulate_compiled( + self.graph.num_inputs, + &self.compiled_nodes, + &self.arg_indices, + &workspace.values, + &mut workspace.cotangent_values, + )?; + + workspace + .gradients + .copy_from_slice(&workspace.cotangent_values[..self.graph.num_inputs]); + Ok((workspace.values[output_index], &workspace.gradients)) + } + + /// Compute the output value and gradient eagerly. + pub fn gradient(&self, inputs: &[f64]) -> Result<(f64, Vec)> { + let mut workspace = self.workspace(); + let (value, gradient) = self.gradient_with_workspace(inputs, &mut workspace)?; + Ok((value, gradient.to_vec())) + } + + /// Compute the output value and gradient eagerly with checked-domain validation. + pub fn gradient_checked(&self, inputs: &[f64]) -> Result<(f64, Vec)> { + let mut workspace = self.workspace(); + let (value, gradient) = self.gradient_with_workspace_checked(inputs, &mut workspace)?; + Ok((value, gradient.to_vec())) + } + + /// Compute the Jacobian for all selected outputs. + pub fn jacobian(&self, inputs: &[f64]) -> Result>> { + let mut workspace = self.workspace(); + self.jacobian_with_workspace(inputs, &mut workspace) + } + + /// Compute the Jacobian for all selected outputs with checked-domain validation. + pub fn jacobian_checked(&self, inputs: &[f64]) -> Result>> { + let mut workspace = self.workspace(); + self.jacobian_with_workspace_checked(inputs, &mut workspace) + } + + /// Compute the Jacobian for all selected outputs using a reusable workspace. + pub fn jacobian_with_workspace( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + ) -> Result>> { + self.jacobian_with_workspace_inner(inputs, workspace, false) + } + + /// Compute the Jacobian for all selected outputs using a reusable workspace with checked-domain validation. + pub fn jacobian_with_workspace_checked( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + ) -> Result>> { + self.jacobian_with_workspace_inner(inputs, workspace, true) + } + + fn jacobian_with_workspace_inner( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + checked: bool, + ) -> Result>> { + if checked { + self.fill_values_checked(inputs, workspace)?; + } else { + self.fill_values(inputs, workspace)?; + } + let mut jacobian: Vec> = Vec::with_capacity(self.output_indices.len()); + + for &output_index in &self.output_indices { + workspace.cotangent_values.clear(); + workspace + .cotangent_values + .resize(workspace.values.len(), 0.0); + workspace.gradients.clear(); + workspace.gradients.resize(self.graph.num_inputs, 0.0); + workspace.cotangent_values[output_index] = 1.0; + reverse_accumulate_compiled( + self.graph.num_inputs, + &self.compiled_nodes, + &self.arg_indices, + &workspace.values, + &mut workspace.cotangent_values, + )?; + + workspace + .gradients + .copy_from_slice(&workspace.cotangent_values[..self.graph.num_inputs]); + jacobian.push(workspace.gradients.clone()); + } + + Ok(jacobian) + } + + /// Compute a finite-difference Hessian by differentiating eager gradients. + pub fn compute_hessian(&self, inputs: &[f64]) -> Result>> { + let num_inputs = self.graph.num_inputs; + let epsilon = 1e-5; + let mut hessian = vec![vec![0.0; num_inputs]; num_inputs]; + + if num_inputs == 0 { + let (_value, _grad) = self.gradient(inputs)?; + return Ok(hessian); + } + + let mut plus_workspace = self.workspace(); + let mut minus_workspace = self.workspace(); + + for j in 0..num_inputs { + let mut inputs_plus = inputs.to_vec(); + inputs_plus[j] += epsilon; + + let mut inputs_minus = inputs.to_vec(); + inputs_minus[j] -= epsilon; + + let (_value_plus, grad_plus) = + self.gradient_with_workspace(&inputs_plus, &mut plus_workspace)?; + let (_value_minus, grad_minus) = + self.gradient_with_workspace(&inputs_minus, &mut minus_workspace)?; + + for i in 0..num_inputs { + hessian[i][j] = (grad_plus[i] - grad_minus[i]) / (2.0 * epsilon); + } + } + + Ok(hessian) + } +} diff --git a/src/tests_comprehensive.rs b/src/tests_comprehensive.rs index 163c797..d4663ad 100644 --- a/src/tests_comprehensive.rs +++ b/src/tests_comprehensive.rs @@ -627,19 +627,6 @@ fn test_dual_constant() { assert!(approx_eq(d.tan, 0.0, 1e-10)); } -// ============================================================================ -// MonoFn Trait Tests - covered in mono::tests module -// since MF1-4 are private implementation details - -// ============================================================================ -// MultiFn Trait Tests -// ============================================================================ - -// Note: F1, F2, F3 are private modules, tested through multi::tests - -// MultiFn trait tests are covered in multi::tests module -// since F1, F2, F3 are private implementation details - // ============================================================================ // Arity Error Tests // ============================================================================