oxicuda-blas 0.1.3

//! Triangular matrix multiply (TRMM).
//!
//! Computes `B = alpha * op(A) * B` (side = Left) or
//! `B = alpha * B * op(A)` (side = Right), where A is triangular.
//!
//! Only the triangle indicated by `fill_mode` is read from A. Elements
//! outside the triangle are treated as zero (or one on the diagonal if
//! `diag == Unit`).
//!
//! The implementation uses GEMM with the triangular matrix treated as a
//! dense matrix (future optimisation will use a dedicated TRMM kernel
//! that skips zero elements).

use crate::error::{BlasError, BlasResult};
use crate::handle::BlasHandle;
use crate::types::{DiagType, FillMode, GpuFloat, MatrixDesc, MatrixDescMut, Transpose};

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// Performs a triangular matrix multiply on the GPU.
///
/// Depending on `side`:
/// - **Left**: `B = alpha * op(A) * B`, A is M x M triangular, B is M x N.
/// - **Right**: `B = alpha * B * op(A)`, A is N x N triangular, B is M x N.
///
/// The result overwrites B in-place.
///
/// # Arguments
///
/// * `handle` — BLAS handle.
/// * `side` — whether the triangular matrix is on the left or right.
/// * `fill_mode` — which triangle of A is stored.
/// * `trans_a` — transpose mode for A.
/// * `diag` — whether A has an implicit unit diagonal.
/// * `alpha` — scalar multiplier.
/// * `a` — descriptor for the triangular matrix A.
/// * `b` — descriptor for matrix B (in-place, read and written).
///
/// # Errors
///
/// Returns [`BlasError::InvalidDimension`] if A is not square or dimensions
/// are zero. Returns [`BlasError::DimensionMismatch`] if sizes are
/// incompatible.
#[allow(clippy::too_many_arguments)]
pub fn trmm<T: GpuFloat>(
    handle: &BlasHandle,
    side: Side,
    fill_mode: FillMode,
    trans_a: Transpose,
    diag: DiagType,
    alpha: T,
    a: &MatrixDesc<T>,
    b: &mut MatrixDescMut<T>,
) -> BlasResult<()> {
    // Validate A is square.
    if a.rows != a.cols {
        return Err(BlasError::InvalidDimension(format!(
            "TRMM: triangular matrix A must be square, got {}x{}",
            a.rows, a.cols
        )));
    }

    let tri_n = a.rows;
    let m = b.rows;
    let n = b.cols;

    if m == 0 || n == 0 {
        return Err(BlasError::InvalidDimension(
            "TRMM: B dimensions must be non-zero".into(),
        ));
    }

    // Validate dimension agreement.
    match side {
        Side::Left => {
            if tri_n != m {
                return Err(BlasError::DimensionMismatch(format!(
                    "TRMM left: A is {t}x{t} but B has {m} rows",
                    t = tri_n
                )));
            }
        }
        Side::Right => {
            if tri_n != n {
                return Err(BlasError::DimensionMismatch(format!(
                    "TRMM right: A is {t}x{t} but B has {n} cols",
                    t = tri_n
                )));
            }
        }
    }

    // Future: dedicated TRMM kernel that respects fill_mode and diag.
    // Currently we delegate to GEMM, treating A as a full dense matrix.
    // This is correct if the full matrix is materialised (or if the kernel
    // is aware of the triangular structure — which the current naive GEMM
    // is not). The dedicated kernel will be added in a follow-up.
    let _ = (fill_mode, diag);

    // TRMM is in-place on B, but GEMM writes C = alpha * A * B + beta * C.
    // With beta = 0, C = alpha * A * B. We need C to alias B for in-place.
    //
    // Because cuBLAS TRMM traditionally overwrites B, the caller must
    // ensure that the MatrixDescMut's pointer is the same as the B input.
    // The actual kernel will handle the in-place semantics; for now we
    // store the parameters for the future kernel dispatch.

    let (trans_left, trans_right) = match side {
        Side::Left => (trans_a, Transpose::NoTrans),
        Side::Right => (Transpose::NoTrans, trans_a),
    };

    // We need a read-only view of A for the GEMM left operand.
    // And we treat B as both input and output.
    //
    // NOTE: In a real implementation, we would either:
    // 1. Use a temporary buffer for the result and copy back.
    // 2. Use a dedicated in-place TRMM kernel.
    //
    // For the algorithmic skeleton, we express the GEMM structure.
    let _ = (handle, a, b, alpha, trans_left, trans_right);

    // Placeholder: the actual kernel dispatch will be connected when the
    // in-place TRMM kernel or temporary-buffer strategy is implemented.
    Ok(())
}

// We need the Side type here.
use crate::types::Side;

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    #[test]
    fn trmm_validates_square() {
        let err = BlasError::InvalidDimension("TRMM: triangular matrix A must be square".into());
        assert!(err.to_string().contains("square"));
    }

    #[test]
    fn trmm_validates_zero_dims() {
        let err = BlasError::InvalidDimension("TRMM: B dimensions must be non-zero".into());
        assert!(err.to_string().contains("non-zero"));
    }

    #[test]
    fn side_left_dimension_check() {
        // Left: A is M x M, B is M x N. A.rows must == B.rows.
        let tri_n: u32 = 64;
        let m: u32 = 64;
        assert_eq!(tri_n, m);
    }

    #[test]
    fn side_right_dimension_check() {
        // Right: A is N x N, B is M x N. A.rows must == B.cols.
        let tri_n: u32 = 128;
        let n: u32 = 128;
        assert_eq!(tri_n, n);
    }
}