基于np.dot的multi_dot的Cython版本实现及代码优化咨询

Got it, let's break this down for you. You're working with small-scale matrix-vector chain dot products (like A·B·C·D·e) and want a Cython implementation that matches the optimized order of numpy.linalg.multi_dot but is tailored to small matrices—since the built-in multi_dot might not squeeze out all possible performance for tiny arrays. Here's a practical, step-by-step solution:

Core Idea First #

For chain products ending with a vector, the optimal order is right-associative: compute D·e first, then C·(D·e), then B·(C·(D·e)), and finally A·(B·(C·(D·e))). This avoids expensive matrix-matrix multiplications (O(n³)) and sticks to cheaper matrix-vector operations (O(n²)), which is way more efficient for small matrices.

We'll lean on BLAS's highly optimized dgemv (double-precision matrix-vector multiply) instead of writing raw Cython loops—BLAS is hand-tuned for CPU architectures and will outperform custom loops every time.

Cython Implementation #

Create a file named multi_dot_cython.pyx with this code:

import numpy as np
cimport numpy as np
from cython cimport boundscheck, wraparound
from libc.stdlib cimport malloc, free

# Declare BLAS's double-precision matrix-vector multiply function
cextern void dgemv(char* trans, int* m, int* n, double* alpha, double* a, int* lda, double* x, int* incx, double* beta, double* y, int* incy) nogil;

@boundscheck(False)  # Disable bounds checking for speed
@wraparound(False)   # Disable negative index wrapping for speed
def multi_dot_chain(list matrices, np.ndarray[np.double_t, ndim=1] vec):
    """
    Compute optimized chain dot product: A·B·C·...·vec
    Args:
        matrices: List of 2D numpy arrays (order: A, B, C, ...)
        vec: 1D numpy array (final vector in the chain)
    Returns:
        1D numpy array with the result
    """
    cdef int num_mats = len(matrices)
    cdef np.ndarray[np.double_t, ndim=1] current_vec = vec.copy()
    cdef np.ndarray[np.double_t, ndim=2] mat
    cdef int m, n, lda, incx, incy
    cdef double alpha = 1.0, beta = 0.0
    cdef char trans = 'T'  # Account for numpy's C-order vs BLAS's Fortran-order

    # Traverse matrices from right to left (optimal order for vector chain)
    for i in range(num_mats - 1, -1, -1):
        mat = matrices[i]
        # Adjust dimensions for BLAS (numpy uses C-order, BLAS uses Fortran-order)
        m = mat.shape[1]  # Rows of the Fortran-order matrix we're effectively using
        n = mat.shape[0]  # Columns of that Fortran-order matrix
        lda = mat.shape[0]  # Leading dimension for BLAS
        incx = 1
        incy = 1

        # Allocate output vector
        cdef np.ndarray[np.double_t, ndim=1] new_vec = np.empty(m, dtype=np.double)

        # Call BLAS dgemv without holding the GIL
        with nogil:
            dgemv(&trans, &m, &n, &alpha, <double*>mat.data, &lda, <double*>current_vec.data, &incx, &beta, <double*>new_vec.data, &incy)
        
        current_vec = new_vec
    
    return current_vec

Compilation & Usage #

Step 1: Setup Script #

Create a setup.py file to compile the Cython code:

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

ext_modules = [
    Extension(
        "multi_dot_cython",
        sources=["multi_dot_cython.pyx"],
        include_dirs=[np.get_include()],
        libraries=["blas"],  # Use "openblas" or "mkl" on Windows if needed
        extra_compile_args=["-O3", "-ffast-math"],  # Aggressive optimizations
    )
]

setup(
    name="multi_dot_cython",
    ext_modules=cythonize(ext_modules),
)

Step 2: Compile #

Run this command in your terminal:

python setup.py build_ext --inplace

Step 3: Test It #

import numpy as np
import multi_dot_cython

# Generate small test matrices and vector
A = np.random.rand(10, 20)
B = np.random.rand(20, 15)
C = np.random.rand(15, 5)
vec = np.random.rand(5)

# Compare with numpy's multi_dot
result_np = np.linalg.multi_dot([A, B, C, vec])
result_cython = multi_dot_cython.multi_dot_chain([A, B, C], vec)

# Verify correctness
print(np.allclose(result_np, result_cython))  # Should print True

Key Optimizations & Notes #

1. BLAS Integration: Directly calling BLAS's dgemv avoids the overhead of Python-level loops and leverages optimized assembly code.
2. GIL Release: The nogil context frees the Global Interpreter Lock during computation, letting other Python threads run (and boosting single-threaded performance too).
3. Memory Layout Handling: Numpy uses C-order (row-first) arrays, while BLAS expects Fortran-order (column-first). The trans='T' parameter fixes this mismatch without copying data.
4. Bounds Checking Disabled: Turning off bounds checks and index wrapping removes unnecessary runtime checks for speed (safe if you trust your input dimensions).

Additional Tips #

• For float32 data, switch to BLAS's sgemv function instead of dgemv for even faster computation.
• Add input validation (e.g., check matrix/vector dimensions match) if you need robustness for untrusted inputs.
• If you need to handle general multi_dot (not just chains ending with a vector), implement a dynamic programming approach to find the optimal parenthesization for matrix-matrix products, then call BLAS's dgemm for those operations.

内容的提问来源于stack exchange，提问作者Vincent D.