How to Optimize Python 3.12 Code with Cython 3 and Rust 1.85 Bindings for 10x Speedups

import time
import sys
from typing import List, Tuple, Optional

def validate_matrix(mat: List[List[float]], name: str) -> None:
    """Validate that a matrix is non-empty and rectangular.

    Args:
        mat: Input matrix to validate
        name: Name of the matrix for error messages

    Raises:
        ValueError: If matrix is empty, not rectangular, or contains non-float values
    """
    if not mat:
        raise ValueError(f"{name} cannot be empty")
    row_len = len(mat[0])
    if row_len == 0:
        raise ValueError(f"{name} rows cannot be empty")
    for i, row in enumerate(mat):
        if len(row) != row_len:
            raise ValueError(f"{name} row {i} has length {len(row)}, expected {row_len}")
        for j, val in enumerate(row):
            if not isinstance(val, (float, int)):
                raise ValueError(f"{name} row {i} col {j} has non-numeric value {val}")

def pure_python_matmul(a: List[List[float]], b: List[List[float]]) -> List[List[float]]:
    """Multiply two matrices using pure Python nested loops.

    Args:
        a: First matrix (m x n)
        b: Second matrix (n x p)

    Returns:
        Result matrix (m x p)

    Raises:
        ValueError: If matrix dimensions are incompatible
    """
    validate_matrix(a, "Matrix A")
    validate_matrix(b, "Matrix B")

    a_rows = len(a)
    a_cols = len(a[0])
    b_rows = len(b)
    b_cols = len(b[0])

    if a_cols != b_rows:
        raise ValueError(f"Incompatible dimensions: A is {a_rows}x{a_cols}, B is {b_rows}x{b_cols}")

    # Initialize result matrix with zeros
    result = [[0.0 for _ in range(b_cols)] for _ in range(a_rows)]

    # Naive triple nested loop multiplication
    for i in range(a_rows):
        for j in range(b_cols):
            cell_sum = 0.0
            for k in range(a_cols):
                cell_sum += a[i][k] * b[k][j]
            result[i][j] = cell_sum
    return result

def benchmark_matmul(
    mat_size: int = 100, 
    iterations: int = 10
) -> Tuple[float, float]:
    """Benchmark matrix multiplication for given matrix size.

    Args:
        mat_size: Size of square matrices to multiply (mat_size x mat_size)
        iterations: Number of times to run multiplication for averaging

    Returns:
        Tuple of (average_time_ms, std_dev_ms)
    """
    # Generate random square matrices
    a = [[float((i * j) % 100) / 100.0 for j in range(mat_size)] for i in range(mat_size)]
    b = [[float((i + j) % 100) / 100.0 for j in range(mat_size)] for i in range(mat_size)]

    times = []
    for _ in range(iterations):
        start = time.perf_counter()
        try:
            pure_python_matmul(a, b)
        except Exception as e:
            print(f"Benchmark failed: {e}", file=sys.stderr)
            raise
        end = time.perf_counter()
        times.append((end - start) * 1000)  # Convert to ms

    avg_time = sum(times) / len(times)
    std_dev = (sum((t - avg_time) ** 2 for t in times) / len(times)) ** 0.5
    return avg_time, std_dev

if __name__ == "__main__":
    # Run benchmarks for increasing matrix sizes
    print("Pure Python 3.12 Matrix Multiplication Benchmark")
    print("=" * 50)
    for size in [10, 50, 100, 200, 500]:
        try:
            avg, std = benchmark_matmul(size, iterations=5)
            print(f"Size {size}x{size}: Avg {avg:.2f}ms, Std Dev {std:.2f}ms")
        except Exception as e:
            print(f"Failed to benchmark size {size}: {e}", file=sys.stderr)
            sys.exit(1)

    # Validate correctness with small matrix
    test_a = [[1.0, 2.0], [3.0, 4.0]]
    test_b = [[5.0, 6.0], [7.0, 8.0]]
    expected = [[19.0, 22.0], [43.0, 50.0]]
    try:
        result = pure_python_matmul(test_a, test_b)
        assert result == expected, f"Correctness check failed: got {result}, expected {expected}"
        print("\nCorrectness check passed for 2x2 matrices")
    except AssertionError as e:
        print(f"Correctness check failed: {e}", file=sys.stderr)
        sys.exit(1)

# cython: language_level=3, boundscheck=False, wraparound=False
import numpy as np
cimport numpy as np
from libc.math cimport fabs
from typing import List

# Declare numpy types for static typing
np.import_array()

def cython_matmul(
    np.ndarray[np.float64_t, ndim=2] a,
    np.ndarray[np.float64_t, ndim=2] b
) -> np.ndarray[np.float64_t, ndim=2]:
    """Multiply two matrices using Cython 3 with static typing and bounds checking disabled.

    Args:
        a: First matrix (m x n) as float64 numpy array
        b: Second matrix (n x p) as float64 numpy array

    Returns:
        Result matrix (m x p) as float64 numpy array

    Raises:
        ValueError: If matrix dimensions are incompatible
    """
    # Validate input dimensions
    cdef int a_rows = a.shape[0]
    cdef int a_cols = a.shape[1]
    cdef int b_rows = b.shape[0]
    cdef int b_cols = b.shape[1]

    if a_cols != b_rows:
        raise ValueError(f"Incompatible dimensions: A is {a_rows}x{a_cols}, B is {b_rows}x{b_cols}")

    # Initialize result array with zeros
    cdef np.ndarray[np.float64_t, ndim=2] result = np.zeros((a_rows, b_cols), dtype=np.float64)
    cdef int i, j, k
    cdef double cell_sum

    # Triple nested loop with C-level variables
    for i in range(a_rows):
        for j in range(b_cols):
            cell_sum = 0.0
            for k in range(a_cols):
                cell_sum += a[i, k] * b[k, j]
            result[i, j] = cell_sum

    return result

def validate_numpy_matrix(
    mat: np.ndarray,
    name: str
) -> None:
    """Validate that a numpy matrix is 2D float64.

    Args:
        mat: Input matrix to validate
        name: Name of the matrix for error messages

    Raises:
        ValueError: If matrix is not 2D float64
    """
    if mat.ndim != 2:
        raise ValueError(f"{name} must be 2D, got {mat.ndim}D")
    if mat.dtype != np.float64:
        raise ValueError(f"{name} must be float64, got {mat.dtype}")
    if mat.size == 0:
        raise ValueError(f"{name} cannot be empty")

def benchmark_cython_matmul(
    mat_size: int = 100,
    iterations: int = 10
) -> tuple:
    """Benchmark Cython matrix multiplication.

    Args:
        mat_size: Size of square matrices to multiply
        iterations: Number of iterations for averaging

    Returns:
        Tuple of (average_time_ms, std_dev_ms)
    """
    # Generate random float64 matrices
    a = np.random.rand(mat_size, mat_size).astype(np.float64)
    b = np.random.rand(mat_size, mat_size).astype(np.float64)

    # Validate inputs
    validate_numpy_matrix(a, "Matrix A")
    validate_numpy_matrix(b, "Matrix B")

    import time
    times = []
    for _ in range(iterations):
        start = time.perf_counter()
        try:
            cython_matmul(a, b)
        except Exception as e:
            print(f"Benchmark failed: {e}")
            raise
        end = time.perf_counter()
        times.append((end - start) * 1000)

    avg_time = sum(times) / len(times)
    std_dev = (sum((t - avg_time) ** 2 for t in times) / len(times)) ** 0.5
    return avg_time, std_dev

if __name__ == "__main__":
    import sys
    print("Cython 3 Matrix Multiplication Benchmark")
    print("=" * 50)
    for size in [10, 50, 100, 200, 500, 1000]:
        try:
            avg, std = benchmark_cython_matmul(size, iterations=5)
            print(f"Size {size}x{size}: Avg {avg:.2f}ms, Std Dev {std:.2f}ms")
        except Exception as e:
            print(f"Failed to benchmark size {size}: {e}", file=sys.stderr)
            sys.exit(1)

    # Correctness check
    test_a = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float64)
    test_b = np.array([[5.0, 6.0], [7.0, 8.0]], dtype=np.float64)
    expected = np.array([[19.0, 22.0], [43.0, 50.0]], dtype=np.float64)
    try:
        result = cython_matmul(test_a, test_b)
        assert np.allclose(result, expected), f"Correctness check failed: got {result}, expected {expected}"
        print("\nCorrectness check passed for 2x2 matrices")
    except AssertionError as e:
        print(f"Correctness check failed: {e}", file=sys.stderr)
        sys.exit(1)

from setuptools import setup, Extension
import numpy as np
from Cython.Build import cythonize

# Define the Cython extension
matmul_ext = Extension(
    name="matmul",
    sources=["matmul.pyx"],
    include_dirs=[np.get_include()],
    extra_compile_args=["-O3", "-march=native"],  # Maximize optimization
    extra_link_args=["-O3"]
)

setup(
    name="cython-matmul",
    version="0.1.0",
    ext_modules=cythonize(matmul_ext, language_level=3),
    zip_safe=False
)

// rust_matmul/src/lib.rs
use pyo3::prelude::*;
use numpy::{PyArray2, IntoPyArray, ToPyArray};
use numpy::ndarray::{Array2, s};
use std::time::Instant;

/// Multiply two 2D float64 matrices using Rust with pyo3 and numpy bindings.
/// 
/// Args:
///     a: First matrix (m x n) as PyArray2
///     b: Second matrix (n x p) as PyArray2
/// 
/// Returns:
///     Result matrix (m x p) as PyArray2
/// 
/// Raises:
///     PyErr: If matrix dimensions are incompatible or inputs are invalid
#[pyfunction]
fn rust_matmul(
    py: Python<'_>,
    a: &PyArray2,
    b: &PyArray2
) -> PyResult>> {
    // Convert PyArray to Rust ndarray for processing
    let a_arr = a.as_array();
    let b_arr = b.as_array();

    // Validate dimensions
    let a_rows = a_arr.shape()[0];
    let a_cols = a_arr.shape()[1];
    let b_rows = b_arr.shape()[0];
    let b_cols = b_arr.shape()[1];

    if a_cols != b_rows {
        return Err(PyErr::new::(
            format!("Incompatible dimensions: A is {}x{}, B is {}x{}", a_rows, a_cols, b_rows, b_cols)
        ));
    }

    // Initialize result array with zeros
    let mut result = Array2::::zeros((a_rows, b_cols));

    // Triple nested loop with Rust-level optimizations
    for i in 0..a_rows {
        for j in 0..b_cols {
            let mut cell_sum = 0.0;
            for k in 0..a_cols {
                cell_sum += a_arr[[i, k]] * b_arr[[k, j]];
            }
            result[[i, j]] = cell_sum;
        }
    }

    // Convert back to PyArray and return
    Ok(result.into_pyarray(py).to_owned())
}

/// Validate that a PyArray2 is 2D float64 and non-empty.
#[pyfunction]
fn validate_matrix(py: Python<'_>, mat: &PyArray2) -> PyResult<()> {
    let arr = mat.as_array();
    if arr.ndim() != 2 {
        return Err(PyErr::new::(
            format!("Matrix must be 2D, got {}D", arr.ndim())
        ));
    }
    if arr.shape()[0] == 0 || arr.shape()[1] == 0 {
        return Err(PyErr::new::(
            "Matrix cannot be empty".to_string()
        ));
    }
    Ok(())
}

/// Benchmark Rust matrix multiplication for given matrix size.
#[pyfunction]
fn benchmark_rust_matmul(
    py: Python<'_>,
    mat_size: usize,
    iterations: usize
) -> PyResult<(f64, f64)> {
    // Generate random matrices using numpy (to match other benchmarks)
    let np = py.import("numpy")?;
    let a: Py> = np.call_method1("random", ("rand", mat_size, mat_size))?.extract()?;
    let a = a.bind(py);
    let a = a.to_owned();
    let b: Py> = np.call_method1("random", ("rand", mat_size, mat_size))?.extract()?;
    let b = b.bind(py);
    let b = b.to_owned();

    // Validate inputs
    validate_matrix(py, &a)?;
    validate_matrix(py, &b)?;

    let mut times = Vec::with_capacity(iterations);
    for _ in 0..iterations {
        let start = Instant::now();
        rust_matmul(py, &a, &b)?;
        let duration = start.elapsed();
        times.push(duration.as_secs_f64() * 1000.0); // Convert to ms
    }

    let avg_time = times.iter().sum::() / times.len() as f64;
    let variance = times.iter().map(|t| (t - avg_time).powi(2)).sum::() / times.len() as f64;
    let std_dev = variance.sqrt();

    Ok((avg_time, std_dev))
}

/// Define the Python module for Rust bindings
#[pymodule]
fn rust_matmul(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(rust_matmul, m)?)?;
    m.add_function(wrap_pyfunction!(validate_matrix, m)?)?;
    m.add_function(wrap_pyfunction!(benchmark_rust_matmul, m)?)?;
    Ok(())
}

[package]
name = "rust-matmul"
version = "0.1.0"
edition = "2021"

[lib]
crate-type = ["cdylib"]

[dependencies]
pyo3 = { version = "0.20", features = ["extension-module", "numpy"] }
numpy = { version = "0.10", features = ["pyo3"] }

python-cython-rust-optimization/
├── pure_python/
│   └── matmul.py
├── cython/
│   ├── matmul.pyx
│   ├── setup.py
│   └── benchmark.py
├── rust/
│   ├── Cargo.toml
│   └── src/
│       └── lib.rs
├── benchmarks/
│   ├── compare.py
│   └── results.csv
├── case-study/
│   └── recommendation-engine.md
└── README.md