如何用C函数对非方阵n×m的NumPy数组执行计算？

编辑说明 #

我重新整理了一个更清晰实用的同类问题，核心需求是：用C函数从Python中复制二维数组（矩阵），执行密集计算后返回二维数组。

原始问题 #

我需要在Python中调用C函数，对大型非方阵（尺寸n×m，n≠m）的所有元素执行计算。参考的教程代码只适用于方阵，方阵情况下运行正常，但把n和m设为不同值时就会出现段错误。我猜测需要为n和m单独分配内存，但C语言知识不足，求支持m=3000、n=2000这类非方阵的可行代码。

现有编译与运行步骤 #

编译C脚本：

$ cc -fPIC -shared -o c_sum.so c_sum.c

运行Python脚本：

$ python main.py

现有代码 #

c_sum.c

#include <stdlib.h>

double * c_sum(const double * matrix, int n, int m){
    double * results = (double *)malloc(sizeof(double) * n);
    int index = 0;
    for(int i=0; i< n*m; i+=n){
        results[index] = 0;
        for(int j=0; j<m; j++){
            results[index] += matrix[i+j];
        }
        index += 1;
    }
    return results;
}

main.py

# 参考教程代码
from ctypes import c_void_p, c_double, c_int, cdll
from numpy.ctypeslib import ndpointer
import numpy as np
import time

def py_sum(matrix: np.array, n: int, m: int) -> np.array:
    result = np.zeros(n)
    for i in range(0, n):
        for j in range(0, m):
            result[i] += matrix[i][j]
    return result

n = 3000
m = 3000
matrix = np.random.randn(n, m)

time1 = time.time()
py_result = py_sum(matrix, n, m)
time2 = time.time() - time1
print("py running time in seconds:", time2)
py_time = time2

lib = cdll.LoadLibrary("c_sum.so")
c_sum = lib.c_sum
c_sum.restype = ndpointer(dtype=c_double,
                          shape=(n,))

time1 = time.time()
result = c_sum(c_void_p(matrix.ctypes.data),
            c_int(n),
            c_int(m))
time2 = time.time() - time1
print("c  running time in seconds:", time2)

c_time = time2
print("speedup:", py_time/c_time)

问题分析与解决方案 #

错误原因 #

原C代码的循环逻辑是为方阵设计的：for(int i=0; i< n*m; i+=n)，这里每次跳n步，只在方阵（n=m）时能正确定位每一行的起始位置。当n≠m时，这个步长会导致越界访问内存，触发段错误。正确的行起始索引应该是i * m（因为每行有m个元素），而非用i+=n的方式累加。

修改后的代码 #

修正后的c_sum.c

#include <stdlib.h>

double * c_sum(const double * matrix, int n, int m){
    // 为结果数组分配内存，共n个元素（每行一个求和结果）
    double * results = (double *)malloc(sizeof(double) * n);
    if (results == NULL) {
        return NULL; // 内存分配失败时返回空指针
    }
    
    for(int i=0; i < n; i++){
        results[i] = 0.0;
        // 计算第i行的起始位置：i * m
        const double *row_start = matrix + i * m;
        for(int j=0; j < m; j++){
            results[i] += row_start[j];
        }
    }
    return results;
}

适配非方阵的main.py

from ctypes import c_void_p, c_double, c_int, cdll
from numpy.ctypeslib import ndpointer
import numpy as np
import time

def py_sum(matrix: np.array, n: int, m: int) -> np.array:
    result = np.zeros(n)
    for i in range(n):
        result[i] = np.sum(matrix[i])
    return result

# 设置非方阵尺寸
n = 2000
m = 3000
matrix = np.random.randn(n, m)

# Python实现计时
time1 = time.time()
py_result = py_sum(matrix, n, m)
py_time = time.time() - time1
print(f"Python 运行时间（秒）: {py_time:.4f}")

# 加载C库并设置类型
lib = cdll.LoadLibrary("c_sum.so")
c_sum = lib.c_sum
# 指定返回值类型：n个double元素的数组
c_sum.restype = ndpointer(dtype=c_double, shape=(n,))
# 指定参数类型
c_sum.argtypes = [ndpointer(dtype=c_double, flags="C_CONTIGUOUS"), c_int, c_int]

# C实现计时
time1 = time.time()
c_result = c_sum(matrix, c_int(n), c_int(m))
c_time = time.time() - time1
print(f"C 运行时间（秒）: {c_time:.4f}")

print(f"加速比: {py_time / c_time:.2f}")

# 验证结果一致性（浮点数误差范围内）
assert np.allclose(py_result, c_result), "结果不一致！"
print("结果验证通过！")

关键改进点 #

• C代码循环逻辑修正：用i * m定位每行起始位置，适配任意n和m的矩阵。
• 添加内存分配检查：避免内存分配失败导致的未知错误。
• Python端参数类型优化：使用argtypes明确指定参数类型，替代c_void_p，提升安全性和兼容性。
• 结果验证：添加np.allclose验证C和Python实现的结果一致性，确保正确性。
• Python求和优化：用np.sum(matrix[i])替代嵌套循环，提升Python端基准测试的效率。

编译与运行 #

重新编译C代码：

$ cc -fPIC -shared -o c_sum.so c_sum.c

运行修改后的Python脚本：

$ python main.py

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