onnx-runtime-ep-cuda 0.1.0-dev.2

//! `MatMul` on the GPU via cuBLASLt (`docs/ORT2.md` §15.3).
//!
//! Supports **2-D** `A[M,K] · B[K,N]` and **3-D batched** `A[b,M,K] · B[b,K,N]`
//! (equal batch on both operands) for f32 / f16 / bf16, all in true fp32
//! accumulation. The row-major → column-major mapping lives in [`crate::blas`].
//!
//! ## Phase-2a limits (all reported as actionable errors, never panics)
//!
//! * rank < 2, rank > 3, or 1-D operand promotion → "not implemented in Phase 2a"
//! * broadcast / unequal batch dims → "not implemented in Phase 2a"
//! * non-contiguous (strided) device inputs → "not implemented in Phase 2a"
//! * dtypes other than f32 / f16 / bf16 → "not implemented in Phase 2a"
//! * mismatched inner dims / dtypes → a plain kernel error (a real mistake, not
//!   a missing feature)

use std::sync::Arc;

use onnx_runtime_ep_api::{EpError, Kernel, KernelFactory, Result, TensorMut, TensorView};
use onnx_runtime_ir::{DataType, Node};

use crate::blas::{self, GemmDtype, GemmParams, WORKSPACE_BYTES};
use crate::error::not_implemented;
use crate::runtime::{cuptr, CudaRuntime};

/// Factory for [`MatMulKernel`]; carries the shared CUDA runtime.
pub struct MatMulFactory {
    pub runtime: Arc<CudaRuntime>,
}

impl KernelFactory for MatMulFactory {
    fn create(&self, _node: &Node, _input_shapes: &[Vec<usize>]) -> Result<Box<dyn Kernel>> {
        Ok(Box::new(MatMulKernel {
            runtime: self.runtime.clone(),
        }))
    }
}

/// cuBLASLt-backed f32/f16/bf16 MatMul kernel.
pub struct MatMulKernel {
    runtime: Arc<CudaRuntime>,
}

/// Map an ONNX element type to a cuBLASLt GEMM dtype, or an actionable
/// Phase-2a "unsupported dtype" error.
fn gemm_dtype(dt: DataType) -> Result<GemmDtype> {
    match dt {
        DataType::Float32 => Ok(GemmDtype::F32),
        DataType::Float16 => Ok(GemmDtype::F16),
        DataType::BFloat16 => Ok(GemmDtype::Bf16),
        other => Err(not_implemented(format!("MatMul with dtype {other:?}"))),
    }
}

/// Resolve `(batch, M, K, N)` from the two operand shapes, enforcing the
/// Phase-2a shape support. Inner-dim disagreement is a genuine error; everything
/// outside the supported rank/batch set is a deferred-feature error.
fn gemm_dims(a: &[usize], b: &[usize]) -> Result<(usize, usize, usize, usize)> {
    match (a.len(), b.len()) {
        (2, 2) => {
            let (m, k, n) = (a[0], a[1], b[1]);
            if b[0] != k {
                return Err(inner_mismatch(a, b));
            }
            Ok((1, m, k, n))
        }
        (3, 3) => {
            if a[0] != b[0] {
                return Err(not_implemented(format!(
                    "batched MatMul with unequal/broadcast batch dims (A batch {}, B batch {})",
                    a[0], b[0]
                )));
            }
            let (batch, m, k, n) = (a[0], a[1], a[2], b[2]);
            if b[1] != k {
                return Err(inner_mismatch(a, b));
            }
            Ok((batch, m, k, n))
        }
        _ => Err(not_implemented(format!(
            "MatMul with operand ranks {}D x {}D (only 2-D and equal-batch 3-D are supported)",
            a.len(),
            b.len()
        ))),
    }
}

fn inner_mismatch(a: &[usize], b: &[usize]) -> EpError {
    EpError::KernelFailed(format!(
        "cuda_ep MatMul: inner dimensions disagree between A {a:?} and B {b:?}"
    ))
}

impl MatMulKernel {
    fn run(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
        if inputs.len() != 2 || outputs.len() != 1 {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep MatMul: expected 2 inputs and 1 output, got {} and {}",
                inputs.len(),
                outputs.len()
            )));
        }
        let a = &inputs[0];
        let b = &inputs[1];

        // All operands must share one supported element type.
        let dtype = gemm_dtype(a.dtype)?;
        if b.dtype != a.dtype || outputs[0].dtype != a.dtype {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep MatMul: mixed dtypes A={:?} B={:?} C={:?} (all must match)",
                a.dtype, b.dtype, outputs[0].dtype
            )));
        }

        // Phase 2a requires dense, row-major device buffers. Strided views (e.g.
        // a transposed input) are deferred; the graph must materialise them.
        if !a.is_contiguous() || !b.is_contiguous() {
            return Err(not_implemented(
                "MatMul with a non-contiguous (strided) input; \
                 insert an explicit copy/transpose before the MatMul",
            ));
        }
        if !outputs[0].is_contiguous() {
            return Err(not_implemented("MatMul with a non-contiguous output"));
        }

        let (batch, m, k, n) = gemm_dims(a.shape, b.shape)?;

        // Sanity-check the output extent matches the computed product shape.
        let expected_out = batch * m * n;
        if outputs[0].numel() != expected_out {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep MatMul: output has {} elements, expected {} ({}x{}x{})",
                outputs[0].numel(),
                expected_out,
                batch,
                m,
                n
            )));
        }

        // Device pointers (byte_offset applied). These are opaque CUDA
        // addresses, never dereferenced on the host.
        let a_ptr = cuptr(a.data_ptr::<u8>() as *const std::ffi::c_void);
        let b_ptr = cuptr(b.data_ptr::<u8>() as *const std::ffi::c_void);
        let c_ptr = cuptr(outputs[0].data_ptr_mut::<u8>() as *const std::ffi::c_void);

        let params = GemmParams {
            dtype,
            a: a_ptr,
            b: b_ptr,
            c: c_ptr,
            m,
            k,
            n,
            batch,
        };

        // Per-call workspace (Phase 2b: pool this on the stream).
        let workspace = self.runtime.alloc_raw(WORKSPACE_BYTES)?;

        // SAFETY: pointers are live device allocations sized for the shapes
        // above (bounds validated by the session before dispatch); the context
        // is bound by `alloc_raw`; `workspace` is live for the whole call; C
        // aliases neither A nor B (distinct output buffer).
        let result = unsafe {
            blas::gemm(
                self.runtime.blas(),
                self.runtime.stream_ptr(),
                &params,
                workspace,
                WORKSPACE_BYTES,
            )
        }
        .and_then(|()| self.runtime.synchronize());

        // Always release the workspace, even on failure.
        // SAFETY: `workspace` came from the `alloc_raw` above and is freed once.
        let free = unsafe { self.runtime.free_raw(workspace) };
        result.and(free)
    }
}

impl Kernel for MatMulKernel {
    fn execute(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
        self.run(inputs, outputs)
    }

    fn supports_strided_input(&self, _input_idx: usize) -> bool {
        // Phase 2a: dense inputs only (see `run`).
        false
    }

    fn cuda_graph_compatible(&self) -> bool {
        // The per-call workspace alloc/free and heuristic query are not graph-
        // capturable as written; revisit when the workspace is pooled (Phase 2b).
        false
    }
}

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

    #[test]
    fn dims_2d_ok() {
        assert_eq!(gemm_dims(&[2, 3], &[3, 4]).unwrap(), (1, 2, 3, 4));
    }

    #[test]
    fn dims_3d_equal_batch_ok() {
        assert_eq!(gemm_dims(&[5, 2, 3], &[5, 3, 4]).unwrap(), (5, 2, 3, 4));
    }

    #[test]
    fn dims_inner_mismatch_is_plain_error() {
        let e = gemm_dims(&[2, 3], &[4, 5]).unwrap_err();
        let msg = format!("{e}");
        assert!(msg.contains("inner dimensions disagree"), "{msg}");
        // A genuine mistake, not a deferred feature.
        assert!(!msg.contains("not yet implemented"), "{msg}");
    }

    #[test]
    fn dims_broadcast_batch_is_deferred() {
        let e = gemm_dims(&[1, 2, 3], &[5, 3, 4]).unwrap_err();
        assert!(format!("{e}").contains("not yet implemented"));
    }

    #[test]
    fn dims_high_rank_is_deferred() {
        let e = gemm_dims(&[2, 2, 2, 2], &[2, 2, 2, 2]).unwrap_err();
        let msg = format!("{e}");
        assert!(msg.contains("4D x 4D"), "{msg}");
        assert!(msg.contains("not yet implemented"), "{msg}");
    }

    #[test]
    fn dtype_mapping_and_unsupported() {
        assert_eq!(gemm_dtype(DataType::Float32).unwrap(), GemmDtype::F32);
        assert_eq!(gemm_dtype(DataType::Float16).unwrap(), GemmDtype::F16);
        assert_eq!(gemm_dtype(DataType::BFloat16).unwrap(), GemmDtype::Bf16);
        let e = gemm_dtype(DataType::Int64).unwrap_err();
        let msg = format!("{e}");
        assert!(msg.contains("dtype Int64"), "{msg}");
        assert!(msg.contains("not yet implemented"), "{msg}");
    }
}