cuBLAS矩阵乘法结果与NumPy不一致问题排查

问题 #

尝试使用cuBLAS执行矩阵乘法，并将结果与NumPy对比，但二者结果存在差异。

C++ cuBLAS代码 #

#include <iostream>
#include <vector>
#include <cublas_v2.h>
#include <cuda_runtime.h>

void printMatrix(const std::vector<double>& matrix, int rows, int cols) {
    for (int i = 0; i < rows; ++i) {
        for (int j = 0; j < cols; ++j) {
            std::cout << matrix[i * cols + j] << " ";
        }
        std::cout << std::endl;
    }
}

int main() {
    // Initialize cuBLAS
    cublasHandle_t handle;
    cublasCreate(&handle);

    // Matrix dimensions
    int m = 3;
    int n = 2;
    int k = 4;

    // Matrix A: shape (3, 4)
    std::vector<double> A = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0};

    // Matrix B: shape (4, 2)
    std::vector<double> B = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0};

    // Matrix C: shape (3, 2)
    std::vector<double> C(m * n, 0.0);

    // Allocate device memory
    double* d_A;
    double* d_B;
    double* d_C;
    cudaMalloc(&d_A, A.size() * sizeof(double));
    cudaMalloc(&d_B, B.size() * sizeof(double));
    cudaMalloc(&d_C, C.size() * sizeof(double));

    // Copy data from host to device
    cublasSetVector(A.size(), sizeof(double), A.data(), 1, d_A, 1);
    cublasSetVector(B.size(), sizeof(double), B.data(), 1, d_B, 1);

    // Perform matrix multiplication using cuBLAS
    double alpha = 1.0;
    double beta = 0.0;
    cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, d_A, m, d_B, k, &beta, d_C, m);

    // Copy result from device to host
    cublasGetVector(C.size(), sizeof(double), d_C, 1, C.data(), 1);

    // Print the result
    std::cout << "Result matrix C:" << std::endl;
    printMatrix(C, m, n);

    // Verify the result
    std::vector<double> expected_C = {50.0, 60.0, 114.0, 140.0, 178.0, 220.0};
    bool success = true;
    for (int i = 0; i < m * n; ++i) {
        if (std::abs(C[i] - expected_C[i]) > 1e-6) {
            success = false;
            break;
        }
    }

    if (success) {
        std::cout << "Test passed!" << std::endl;
    } else {
        std::cout << "Test failed!" << std::endl;
    }

    // Free device memory
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);

    // Destroy cuBLAS handle
    cublasDestroy(handle);

    return 0;
}

NumPy计算结果 #

>>> import numpy as np
>>> a = np.array(list(range(1, 13))).reshape(3, 4)
>>> b = np.array(list(range(1, 9))).reshape(4, 2)
>>> np.matmul(a, b)
array([[ 50,  60],
       [114, 140],
       [178, 220]])

cuBLAS输出结果 #

Result matrix C:
70 80 
90 158 
184 210 
Test failed!

问题原因与修复方案 #

核心原因 #

cuBLAS默认采用**列优先（Fortran风格）存储矩阵，而你的代码中矩阵A、B是按行优先（C风格）**初始化的，同时cublasDgemm的leading dimension（lda、ldb、ldc）参数设置错误，导致cuBLAS读取矩阵的方式与预期不符，计算出错误结果。

在cublasDgemm中，lda表示矩阵A的leading dimension：

• 对于列优先存储的M×K矩阵，lda等于矩阵的行数M（每列有M个元素，下一列第一个元素与当前列第一个元素的间隔是M）。
• 但你的矩阵A是行优先存储的3×4矩阵，此时对应的leading dimension应该是矩阵的列数K=4（每行有4个元素，下一行第一个元素与当前行第一个元素的间隔是4）。

同理，矩阵B是行优先的4×2矩阵，ldb应设为2；结果矩阵C是行优先的3×2矩阵，ldc应设为2。

修复代码 #

修改cublasDgemm的调用参数，调整lda、ldb、ldc的值：

// 修正后的矩阵乘法调用
cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, d_A, k, d_B, n, &beta, d_C, n);

另一种修复思路 #

如果你想保持leading dimension参数为矩阵行数，也可以将矩阵转置后传递给cuBLAS，并设置转置操作参数：

// 转置矩阵操作，此时A和B按列优先读取对应原行优先矩阵的转置
cublasDgemm(handle, CUBLAS_OP_T, CUBLAS_OP_T, m, n, k, &alpha, d_A, m, d_B, k, &beta, d_C, m);

验证结果 #

修复后，cuBLAS的输出结果将与NumPy完全一致：

Result matrix C:
50 60 
114 140 
178 220 
Test passed!

内容的提问来源于stack exchange，提问作者musako