singe-cublas 0.1.0-alpha.5

use std::{
    ffi::{self, CStr},
    fmt::{self, Display, Formatter},
    result,
};

use num_enum::{IntoPrimitive, TryFromPrimitive};
use singe_core::impl_enum_conversion;
use singe_cuda::error::Error as CudaError;
use thiserror::Error;

use singe_cublas_sys as sys;

#[derive(Error, Debug)]
pub enum Error {
    #[error("cuda error: {0}")]
    Cuda(#[from] CudaError),

    #[error("cublas error ({code}): {message}")]
    Cublas { code: Status, message: String },

    #[error("string contains interior nul byte")]
    InteriorNul,

    #[error("unexpected null handle")]
    NullHandle,

    #[error("{name} is out of range")]
    OutOfRange { name: String },

    #[error("unexpected attribute size: expected {expected} bytes, got {actual}")]
    AttributeSizeMismatch { expected: usize, actual: usize },

    #[error("{name} has mismatched length")]
    MismatchedLength { name: String },

    #[error("invalid vector increment")]
    InvalidIncrement,

    #[error("invalid matrix leading dimension")]
    InvalidLeadingDimension,

    #[error("invalid matrix shape")]
    InvalidMatrixShape,

    #[error("invalid vector shape")]
    InvalidVectorShape,

    #[error("stream belongs to a different cuda context")]
    StreamContextMismatch,

    #[error("operation requires host pointer mode")]
    RequiresHostPointerMode,

    #[error("scalar pointer modes do not match")]
    ScalarPointerModeMismatch,
}

pub type Result<T> = result::Result<T, Error>;

/// cuBLAS status returned by library calls.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, TryFromPrimitive, IntoPrimitive)]
#[repr(u32)]
pub enum Status {
    /// The operation completed successfully.
    Success = sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS as _,
    /// The cuBLAS library was not initialized.
    /// Common causes include a missing prior [`Context::create`](crate::context::Context::create) call, an error in the CUDA runtime called by cuBLAS, or an error in the hardware setup.
    /// To correct: call [`Context::create`](crate::context::Context::create) before the operation and check that the hardware, driver, and cuBLAS library are correctly installed.
    NotInitialized = sys::cublasStatus_t::CUBLAS_STATUS_NOT_INITIALIZED as _,
    /// Resource allocation failed inside the cuBLAS library.
    /// Usually caused by a device memory allocation failure.
    /// To correct: before the operation, deallocate previously allocated memory as much as possible.
    AllocFailed = sys::cublasStatus_t::CUBLAS_STATUS_ALLOC_FAILED as _,
    /// An unsupported value or parameter was passed to the operation, such as a negative vector size.
    /// To correct: ensure that all the parameters being passed have valid values.
    InvalidValue = sys::cublasStatus_t::CUBLAS_STATUS_INVALID_VALUE as _,
    /// The operation requires a feature absent from the device architecture; usually caused by compute capability lower than 5.0.
    /// To correct: compile and run the application on a device with appropriate compute capability.
    ArchMismatch = sys::cublasStatus_t::CUBLAS_STATUS_ARCH_MISMATCH as _,
    /// An access to GPU memory space failed, which is usually caused by a failure to bind a texture.
    /// To correct: before the operation, unbind any previously bound textures.
    MappingError = sys::cublasStatus_t::CUBLAS_STATUS_MAPPING_ERROR as _,
    /// The GPU program failed to execute.
    /// A kernel launch failure on the GPU is a common cause.
    /// To correct: check that the hardware, an appropriate version of the driver, and the cuBLAS library are correctly installed.
    ExecutionFailed = sys::cublasStatus_t::CUBLAS_STATUS_EXECUTION_FAILED as _,
    /// An internal cuBLAS operation failed.
    /// Usually caused by an asynchronous memory copy failure.
    /// To correct: check that the hardware, an appropriate version of the driver, and the cuBLAS library are correctly installed.
    /// Also check that memory passed to the operation is not released before the operation completes.
    InternalError = sys::cublasStatus_t::CUBLAS_STATUS_INTERNAL_ERROR as _,
    /// The requested operation is not supported.
    NotSupported = sys::cublasStatus_t::CUBLAS_STATUS_NOT_SUPPORTED as _,
    /// The requested operation requires a license, and an error was detected
    /// when checking the current licensing.
    /// This error can happen if the license is not present or expired, or if
    /// `NVIDIA_LICENSE_FILE` is not set properly.
    LicenseError = sys::cublasStatus_t::CUBLAS_STATUS_LICENSE_ERROR as _,
}

impl_enum_conversion!(sys::cublasStatus_t, Status);

impl Status {
    pub const fn description(self) -> &'static str {
        match self {
            Self::Success => "success",
            Self::NotInitialized => "not initialized",
            Self::AllocFailed => "allocation failed",
            Self::InvalidValue => "invalid value",
            Self::ArchMismatch => "architecture mismatch",
            Self::MappingError => "mapping error",
            Self::ExecutionFailed => "execution failed",
            Self::InternalError => "internal error",
            Self::NotSupported => "not supported",
            Self::LicenseError => "license error",
        }
    }
}

impl Display for Status {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            Self::Success => write!(f, "CUBLAS_STATUS_SUCCESS"),
            Self::NotInitialized => write!(f, "CUBLAS_STATUS_NOT_INITIALIZED"),
            Self::AllocFailed => write!(f, "CUBLAS_STATUS_ALLOC_FAILED"),
            Self::InvalidValue => write!(f, "CUBLAS_STATUS_INVALID_VALUE"),
            Self::ArchMismatch => write!(f, "CUBLAS_STATUS_ARCH_MISMATCH"),
            Self::MappingError => write!(f, "CUBLAS_STATUS_MAPPING_ERROR"),
            Self::ExecutionFailed => write!(f, "CUBLAS_STATUS_EXECUTION_FAILED"),
            Self::InternalError => write!(f, "CUBLAS_STATUS_INTERNAL_ERROR"),
            Self::NotSupported => write!(f, "CUBLAS_STATUS_NOT_SUPPORTED"),
            Self::LicenseError => write!(f, "CUBLAS_STATUS_LICENSE_ERROR"),
        }
    }
}

impl From<sys::cublasStatus_t> for Error {
    fn from(status: sys::cublasStatus_t) -> Self {
        debug_assert_ne!(status, sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS);

        let message = unsafe {
            let c_ptr = sys::cublasGetStatusString(status);
            if c_ptr.is_null() {
                String::from("unknown cublas error")
            } else {
                CStr::from_ptr(c_ptr).to_string_lossy().into_owned()
            }
        };

        Self::Cublas {
            code: status.into(),
            message,
        }
    }
}

impl From<Status> for Error {
    fn from(status: Status) -> Self {
        sys::cublasStatus_t::from(status).into()
    }
}

impl From<ffi::NulError> for Error {
    fn from(_: ffi::NulError) -> Self {
        Self::InteriorNul
    }
}

#[macro_export]
macro_rules! try_ffi {
    ($expr:expr) => {{
        let status = { $expr };
        if status != singe_cublas_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
            Err($crate::error::Error::from(status))
        } else {
            Ok(())
        }
    }};
}