atomr-accel-cuda 0.10.0

//! `CudnnActor` — Phase 2 cuDNN slice. Wraps a [`cudarc::cudnn::Cudnn`]
//! handle and exposes the v9 frontend graph API plus legacy
//! `ConvForward` / `Activation` / `Softmax` shims for back-compat.
//!
//! # Module layout
//!
//! ```text
//! kernel/cudnn/
//! ├── mod.rs           — CudnnActor, CudnnMsg, CudnnInner, props
//! ├── graph.rs         — v9 frontend graph spec (TensorSpec, OpSpec,
//! │                     OperationGraphSpec) + plan cache
//! ├── conv.rs          — ConvFwdRequest<T>, ConvBwdDataRequest<T>,
//! │                     ConvBwdFilterRequest<T>
//! ├── norm.rs          — BatchNormRequest<T>, LayerNormRequest<T>,
//! │                     InstanceNormRequest<T>, GroupNormRequest<T>,
//! │                     NormBwdRequest<T>
//! ├── pool.rs          — PoolFwdRequest<T>, PoolBwdRequest<T>
//! ├── attention.rs     — MultiHeadAttnFwdRequest<T>,
//! │                     MultiHeadAttnBwdRequest<T>
//! ├── rnn.rs           — RnnFwdRequest<T>, RnnBwdRequest<T>
//! └── activation.rs    — ActivationFwdRequest<T>, SoftmaxFwdRequest<T>,
//!                       DropoutFwdRequest<T>, LrnFwdRequest<T>
//! ```
//!
//! # Op coverage
//!
//! | Family       | Forward | Backward | Notes                                         |
//! |--------------|:-------:|:--------:|-----------------------------------------------|
//! | Conv         | ✓       | ✓ data + filter | 1D/2D/3D, NCHW + NHWC, groups, dilation |
//! | Pool         | ✓       | ✓        | max, avg, avg-exclude-padding                 |
//! | BatchNorm    | ✓       | ✓        | training, inference, persistent               |
//! | LayerNorm    | ✓       | ✓        | training, inference                           |
//! | InstanceNorm | ✓       | ✓        |                                               |
//! | GroupNorm    | ✓       | ✓        |                                               |
//! | Activation   | ✓       | (fused with conv epilogue) | relu, sigmoid, tanh, gelu, gelu_approx, swish, elu, softplus, identity |
//! | Softmax      | ✓       | (planned) | instance + channel mode                      |
//! | Dropout      | ✓       | (planned) |                                              |
//! | LRN          | ✓       | (planned) |                                              |
//! | Attention    | ✓       | ✓        | causal, sliding-window, GQA/MQA, dropout     |
//! | RNN/LSTM/GRU | ✓       | ✓        | uni + bi, multi-layer, dropout              |
//!
//! # Dtype matrix
//!
//! Every request type is generic over `T: crate::dtype::CudnnSupported`.
//! Implementations cover `f32`, `f64`, `i8`, plus `half::f16` and
//! `half::bf16` under the `f16` feature.

#![allow(dead_code)]

pub mod activation;
pub mod attention;
pub mod conv;
pub mod graph;
pub mod norm;
pub mod pool;
pub mod rnn;

use std::sync::Arc;

use async_trait::async_trait;
use atomr_core::actor::{Actor, Context, Props};
use cudarc::cudnn::Cudnn;
use cudarc::driver::CudaSlice;
use parking_lot::Mutex;
use tokio::sync::oneshot;

use crate::completion::CompletionStrategy;
use crate::device::DeviceState;
use crate::error::GpuError;
use crate::gpu_ref::GpuRef;
use crate::kernel::dispatch::{CudnnDispatch, CudnnDispatchCtx};
use crate::stream::StreamAllocator;

pub use activation::{
    ActivationFwdRequest, ActivationKind, DropoutFwdRequest, LrnFwdRequest, LrnParams,
    SoftmaxFwdRequest, SoftmaxMode,
};
pub use attention::{
    AttentionMask, AttentionParams, MultiHeadAttnBwdRequest, MultiHeadAttnFwdRequest,
};
pub use conv::{
    ConvBwdDataRequest, ConvBwdFilterRequest, ConvDescParams, ConvFwdRequest, EpilogueKind,
};
pub use graph::{
    cache_key, CachedPlan, DtypeTag, NormMode, NormPhase, OpSpec, OperationGraphSpec, PlanCache,
    PlanCacheKey, PointwiseMode, PoolKind, ReduceOp, TensorLayout, TensorSpec,
    DEFAULT_PLAN_CACHE_SIZE,
};
pub use norm::{
    BatchNormRequest, GroupNormRequest, InstanceNormRequest, LayerNormRequest, NormBwdRequest,
};
pub use pool::{PoolBwdRequest, PoolFwdRequest, PoolMode, PoolParams};
pub use rnn::{RnnBwdRequest, RnnDirection, RnnFwdRequest, RnnMode, RnnParams};

const LIB: &str = "cudnn";

// ----- Legacy back-compat parameter / request types -------------------

/// Convolution parameters (cuDNN 2D conv subset).
///
/// **Deprecated** — kept for back-compat with the F2 ConvForward API.
/// New code should construct [`ConvDescParams`] directly.
#[derive(Debug, Clone, Copy)]
pub struct ConvParams {
    pub pad: [i32; 2],
    pub stride: [i32; 2],
    pub dilation: [i32; 2],
}

/// Legacy F2 ConvForward request (NCHW, f32 only).
///
/// **Deprecated** — use [`ConvFwdRequest<T>`] under
/// [`CudnnMsg::Op`] for new code.
pub struct ConvForwardRequest {
    pub x: GpuRef<f32>,
    pub x_dims: [i32; 4],
    pub w: GpuRef<f32>,
    pub w_dims: [i32; 4],
    pub y: GpuRef<f32>,
    pub y_dims: [i32; 4],
    pub conv: ConvParams,
    pub alpha: f32,
    pub beta: f32,
    pub reply: oneshot::Sender<Result<(), GpuError>>,
}

/// Legacy F2 Activation request.
pub struct ActivationRequest {
    pub kind: ActivationKind,
    pub x: GpuRef<f32>,
    pub y: GpuRef<f32>,
    pub dims: [i32; 4],
    pub alpha: f32,
    pub beta: f32,
    pub reply: oneshot::Sender<Result<(), GpuError>>,
}

/// Legacy F2 Softmax request.
pub struct SoftmaxRequest {
    pub x: GpuRef<f32>,
    pub y: GpuRef<f32>,
    pub dims: [i32; 4],
    pub alpha: f32,
    pub beta: f32,
    pub reply: oneshot::Sender<Result<(), GpuError>>,
}

// ----- CudnnMsg + actor ----------------------------------------------

/// Mailbox message for [`CudnnActor`].
///
/// Modern callers send `CudnnMsg::Op(Box<dyn CudnnDispatch>)` with a
/// typed request struct (e.g. `ConvFwdRequest<f16>`). The legacy
/// `ConvForward` / `Activation` / `Softmax` variants are retained for
/// back-compat and are slated for removal once downstream users
/// migrate.
pub enum CudnnMsg {
    /// Generic typed cuDNN op (canonical form). The boxed trait object
    /// carries dtype + op kind for telemetry and dispatches via
    /// [`CudnnDispatch::dispatch`].
    Op(Box<dyn CudnnDispatch>),

    /// **Deprecated** — use [`CudnnMsg::Op`] with [`ConvFwdRequest<f32>`].
    #[deprecated(note = "use CudnnMsg::Op with ConvFwdRequest<f32>")]
    ConvForward(Box<ConvForwardRequest>),

    /// **Deprecated** — use [`CudnnMsg::Op`] with [`ActivationFwdRequest<f32>`].
    #[deprecated(note = "use CudnnMsg::Op with ActivationFwdRequest<f32>")]
    Activation(Box<ActivationRequest>),

    /// **Deprecated** — use [`CudnnMsg::Op`] with [`SoftmaxFwdRequest<f32>`].
    #[deprecated(note = "use CudnnMsg::Op with SoftmaxFwdRequest<f32>")]
    Softmax(Box<SoftmaxRequest>),
}

pub struct CudnnActor {
    inner: CudnnInner,
}

struct SendCudnn(Arc<Cudnn>);
unsafe impl Send for SendCudnn {}
unsafe impl Sync for SendCudnn {}

enum CudnnInner {
    Real {
        handle: SendCudnn,
        stream: Arc<cudarc::driver::CudaStream>,
        completion: Arc<dyn CompletionStrategy>,
        plan_cache: Mutex<PlanCache>,
        workspace: Mutex<Option<CudaSlice<u8>>>,
        #[allow(dead_code)]
        state: Arc<DeviceState>,
    },
    Mock,
}

impl CudnnActor {
    pub fn props(
        stream: Arc<cudarc::driver::CudaStream>,
        _allocator: Arc<dyn StreamAllocator>,
        completion: Arc<dyn CompletionStrategy>,
        state: Arc<DeviceState>,
    ) -> Props<Self> {
        Props::create(move || {
            let handle = match Cudnn::new(stream.clone()) {
                Ok(h) => h,
                Err(e) => panic!("ContextPoisoned: Cudnn::new failed: {e}"),
            };
            CudnnActor {
                inner: CudnnInner::Real {
                    handle: SendCudnn(handle),
                    stream: stream.clone(),
                    completion: completion.clone(),
                    plan_cache: Mutex::new(PlanCache::new(DEFAULT_PLAN_CACHE_SIZE)),
                    workspace: Mutex::new(None),
                    state: state.clone(),
                },
            }
        })
    }

    pub fn mock_props() -> Props<Self> {
        Props::create(|| CudnnActor {
            inner: CudnnInner::Mock,
        })
    }
}

#[async_trait]
impl Actor for CudnnActor {
    type Msg = CudnnMsg;

    async fn handle(&mut self, _ctx: &mut Context<Self>, msg: CudnnMsg) {
        match &self.inner {
            CudnnInner::Mock => reply_mock(msg),
            CudnnInner::Real {
                handle,
                stream,
                completion,
                plan_cache,
                workspace,
                ..
            } => match msg {
                CudnnMsg::Op(op) => {
                    let ctx = CudnnDispatchCtx {
                        handle: handle.0.clone(),
                        stream: stream.clone(),
                        completion: completion.clone(),
                        plan_cache,
                        workspace,
                    };
                    op.dispatch(&ctx);
                }
                #[allow(deprecated)]
                CudnnMsg::ConvForward(req) => {
                    handle_legacy_conv_fwd(*req);
                }
                #[allow(deprecated)]
                CudnnMsg::Activation(req) => {
                    handle_legacy_activation(*req);
                }
                #[allow(deprecated)]
                CudnnMsg::Softmax(req) => {
                    handle_legacy_softmax(*req);
                }
            },
        }
    }
}

fn reply_mock(msg: CudnnMsg) {
    let err = || GpuError::Unrecoverable("CudnnActor in mock mode".into());
    match msg {
        CudnnMsg::Op(_) => {
            // The op's reply channel is owned inside the box; we
            // dispatch to a no-op variant of the dispatcher to send a
            // clear error. The dispatch impls all check
            // `mock`-equivalent state and reply with LibraryError, so
            // we just drop the box here — that closes any oneshot
            // senders inside (which the receivers see as Closed).
        }
        #[allow(deprecated)]
        CudnnMsg::ConvForward(r) => {
            let _ = r.reply.send(Err(err()));
        }
        #[allow(deprecated)]
        CudnnMsg::Activation(r) => {
            let _ = r.reply.send(Err(err()));
        }
        #[allow(deprecated)]
        CudnnMsg::Softmax(r) => {
            let _ = r.reply.send(Err(err()));
        }
    }
}

#[allow(deprecated)]
fn handle_legacy_conv_fwd(req: ConvForwardRequest) {
    // The legacy launch path lived in cudnn_actor.rs; for the v2
    // skeleton we reply with a clear migration message. Real callers
    // should use CudnnMsg::Op(ConvFwdRequest<f32>) which routes
    // through the v9 frontend graph builder.
    let _ = req.reply.send(Err(GpuError::LibraryError {
        lib: LIB,
        msg: "ConvForward (legacy) is deprecated; send CudnnMsg::Op(ConvFwdRequest<f32>) \
              for v9 frontend dispatch"
            .to_string(),
    }));
}

#[allow(deprecated)]
fn handle_legacy_activation(req: ActivationRequest) {
    let _ = req.reply.send(Err(GpuError::LibraryError {
        lib: LIB,
        msg: "Activation (legacy) is deprecated; send CudnnMsg::Op(ActivationFwdRequest<f32>)"
            .to_string(),
    }));
}

#[allow(deprecated)]
fn handle_legacy_softmax(req: SoftmaxRequest) {
    let _ = req.reply.send(Err(GpuError::LibraryError {
        lib: LIB,
        msg: "Softmax (legacy) is deprecated; send CudnnMsg::Op(SoftmaxFwdRequest<f32>)"
            .to_string(),
    }));
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The deprecated `ConvForward` variant still constructs and the
    /// boxed legacy request still carries its inner fields.
    #[test]
    #[allow(deprecated)]
    fn deprecated_conv_forward_alias_still_constructs() {
        let (tx, _rx) = oneshot::channel();
        // We can't construct a real GpuRef here, but constructing the
        // request struct itself is sufficient for the alias check —
        // the field types compile.
        // Skip building the GpuRefs (would need a real CudaSlice);
        // instead exercise the related ConvParams + flag round-trip.
        let p = ConvParams {
            pad: [0, 0],
            stride: [1, 1],
            dilation: [1, 1],
        };
        assert_eq!(p.pad, [0, 0]);
        assert_eq!(p.stride, [1, 1]);
        // Verify variant tag construction by way of a dummy enum
        // pattern — we don't build the boxed variant (no GpuRef
        // available), but we confirm the variant exists by reference.
        fn _accepts_legacy(_: &CudnnMsg) {}
        // Build a fresh Op variant from a tiny dispatcher to confirm
        // CudnnMsg::Op carries Box<dyn CudnnDispatch>.
        struct Probe(oneshot::Sender<Result<(), GpuError>>);
        impl CudnnDispatch for Probe {
            fn dtype_name(&self) -> &'static str {
                "f32"
            }
            fn op_kind(&self) -> &'static str {
                "probe"
            }
            fn dispatch(self: Box<Self>, _ctx: &CudnnDispatchCtx<'_>) {
                let _ = self.0.send(Ok(()));
            }
        }
        let msg = CudnnMsg::Op(Box::new(Probe(tx)));
        _accepts_legacy(&msg);
    }

    #[test]
    fn cudnn_dispatch_is_object_safe() {
        // Verifies the trait is dyn-safe (compile-only check).
        fn _accept(_: Box<dyn CudnnDispatch>) {}
    }

    #[test]
    fn plan_cache_default_size_matches_constant() {
        let pc = PlanCache::default();
        assert_eq!(pc.cap(), DEFAULT_PLAN_CACHE_SIZE);
    }
}