为什么在Python中使用NumPy进行矩阵乘法比使用ctypes更快?

66

我试图找到最快的矩阵乘法方法,并尝试了三种不同的方式:

  • 纯Python实现:没有什么意外。
  • Numpy实现,使用numpy.dot(a, b)
  • 使用Python中的ctypes模块与C进行交互。

以下是被转换为共享库的C代码:

#include <stdio.h>
#include <stdlib.h>

void matmult(float* a, float* b, float* c, int n) {
    int i = 0;
    int j = 0;
    int k = 0;

    /*float* c = malloc(nay * sizeof(float));*/

    for (i = 0; i < n; i++) {
        for (j = 0; j < n; j++) {
            int sub = 0;
            for (k = 0; k < n; k++) {
                sub = sub + a[i * n + k] * b[k * n + j];
            }
            c[i * n + j] = sub;
        }
    }
    return ;
}

接下来是调用它的 Python 代码:

def C_mat_mult(a, b):
    libmatmult = ctypes.CDLL("./matmult.so")

    dima = len(a) * len(a)
    dimb = len(b) * len(b)

    array_a = ctypes.c_float * dima
    array_b = ctypes.c_float * dimb
    array_c = ctypes.c_float * dima

    suma = array_a()
    sumb = array_b()
    sumc = array_c()

    inda = 0
    for i in range(0, len(a)):
        for j in range(0, len(a[i])):
            suma[inda] = a[i][j]
            inda = inda + 1
        indb = 0
    for i in range(0, len(b)):
        for j in range(0, len(b[i])):
            sumb[indb] = b[i][j]
            indb = indb + 1

    libmatmult.matmult(ctypes.byref(suma), ctypes.byref(sumb), ctypes.byref(sumc), 2);

    res = numpy.zeros([len(a), len(a)])
    indc = 0
    for i in range(0, len(sumc)):
        res[indc][i % len(a)] = sumc[i]
        if i % len(a) == len(a) - 1:
            indc = indc + 1

    return res

我原本认为使用C的版本会更快...但结果证明我错了!以下是我的基准测试,显示我可能要么执行不正确,要么numpy真的异常快:

benchmark

我想了解为什么numpy版本比ctypes版本更快,我甚至不是在讨论纯Python实现,因为那显然很慢。


9
不错的问题 - 结果表明,np.dot() 比 C 语言中一个朴素的 GPU 实现更快。 - user2398029
12
你的naive C矩阵乘法速度慢的最大原因之一是内存访问模式。内部循环(关于k)中的b[k * n + j];具有步幅为n,因此它在每次访问时都会触及不同的缓存行,并且你的循环无法使用SSE/AVX进行自动向量化。通过提前转置b来解决这个问题,这会花费O(n^2)的时间,但在从b进行O(n^3)加载时,通过减少缓存未命中可以回报其本身的代价。尽管如此,这仍然是一个naive实现,没有使用缓存块(又称循环分块)。 - Peter Cordes
1
由于您使用了一个 int sum(出于某种原因...),如果内部循环正在访问两个连续的数组,则您的循环实际上可以进行矢量化而无需 -ffast-math。FP数学不是结合律,因此编译器无法在没有 -ffast-math 的情况下重新排序操作,但整数数学是结合律的(并且比FP加法具有更低的延迟,这有助于如果您不打算使用多个累加器或其他延迟隐藏技术来优化循环)。float -> int 转换的成本与 FP add 相当(实际上在 Intel CPU 上使用 FP add ALU),因此在优化的代码中不值得。 - Peter Cordes
6个回答

40

NumPy使用高度优化、精心调整的BLAS方法进行矩阵乘法(也可参见:ATLAS)。在这种情况下,使用的是GEMM函数(用于通用矩阵乘法)。您可以通过搜索dgemm.f(它在Netlib中)来查找原始函数。

顺便提一句,这种优化超出了编译器优化的范畴。Philip上面提到的是Coppersmith–Winograd算法。如果我没记错的话,这是ATLAS中大多数矩阵乘法的算法(尽管评论者指出可能是Strassen算法)。

换句话说,您的matmult算法是简单的实现。有更快的方式来完成相同的事情。


5
顺便提一下,np.show_config()会显示它链接到的lapack / blas信息。 - denis
3
你和Philip提出了正确的观点(问题在于OP的实现速度较慢),但我猜测NumPy使用的是Strassen算法或其变种,而不是Coppersmith-Winograd算法,因为后者具有很大的常数,在实践中通常没有用。 - Huck Bennett
Numpy在内部使用BLAS库。在大多数平台上,默认情况下使用OpenBLAS(基于GotoBLAS)。据我理解,OpenBLAS不使用Strassen算法(BLIS库也是如此)。如果非常小心地实现,对于非常大的矩阵,Strassen可能会更快。话虽如此,它的数值稳定性不如标准的基于瓦片的方法。OpenBLAS利用多个线程和SIMD指令。 - Jérôme Richard

31

我对Numpy不是很熟悉,但源代码在Github上。点积的一部分实现在https://github.com/numpy/numpy/blob/master/numpy/core/src/multiarray/arraytypes.c.src中,我假设它被翻译成每种数据类型的具体C实现。例如:

/**begin repeat
 *
 * #name = BYTE, UBYTE, SHORT, USHORT, INT, UINT,
 * LONG, ULONG, LONGLONG, ULONGLONG,
 * FLOAT, DOUBLE, LONGDOUBLE,
 * DATETIME, TIMEDELTA#
 * #type = npy_byte, npy_ubyte, npy_short, npy_ushort, npy_int, npy_uint,
 * npy_long, npy_ulong, npy_longlong, npy_ulonglong,
 * npy_float, npy_double, npy_longdouble,
 * npy_datetime, npy_timedelta#
 * #out = npy_long, npy_ulong, npy_long, npy_ulong, npy_long, npy_ulong,
 * npy_long, npy_ulong, npy_longlong, npy_ulonglong,
 * npy_float, npy_double, npy_longdouble,
 * npy_datetime, npy_timedelta#
 */
static void
@name@_dot(char *ip1, npy_intp is1, char *ip2, npy_intp is2, char *op, npy_intp n,
           void *NPY_UNUSED(ignore))
{
    @out@ tmp = (@out@)0;
    npy_intp i;

    for (i = 0; i < n; i++, ip1 += is1, ip2 += is2) {
        tmp += (@out@)(*((@type@ *)ip1)) *
               (@out@)(*((@type@ *)ip2));
    }
    *((@type@ *)op) = (@type@) tmp;
}
/**end repeat**/

这似乎是计算一维点积,即向量。在我浏览Github的几分钟中,我无法找到矩阵的源代码,但它可能对每个结果矩阵中的元素使用一次调用。这意味着此函数中的循环对应于您最内层的循环。

它们之间的一个区别是“步幅” - 输入中连续元素之间的差异 - 在调用函数之前明确计算。在您的情况下,没有步幅,并且每个输入的偏移量每次都会被计算,例如a[i * n + k]。我本来希望优秀的编译器能够将其优化为类似于Numpy步幅的内容,但也许它无法证明步长是常数(或者它没有被优化)。

Numpy还可以在调用此函数的更高级别代码中执行一些智能缓存效果操作。一个常见的技巧是考虑每行是否连续,或每列是否连续 - 并尝试首先迭代每个连续部分。对于每个点积,似乎很难完全达到最佳状态,因为一个输入矩阵必须通过行进行遍历,而另一个输入矩阵必须通过列进行遍历(除非它们被存储在不同的主要顺序中)。但它至少可以对结果元素执行此操作。

Numpy还包含从不同基本实现中选择某些操作(包括“点”)的代码。例如,它可以使用BLAS库。从上面的讨论中可以看出,使用了CBLAS。这是从Fortran翻译成C的。我认为您测试中使用的实现是在此处找到的:http://www.netlib.org/clapack/cblas/sdot.c

请注意,此程序是由一台机器编写的,供另一台机器阅读。但是您可以在底部看到它使用展开的循环每次处理5个元素:

for (i = mp1; i <= *n; i += 5) {
stemp = stemp + SX(i) * SY(i) + SX(i + 1) * SY(i + 1) + SX(i + 2) * 
    SY(i + 2) + SX(i + 3) * SY(i + 3) + SX(i + 4) * SY(i + 4);
}

这个展开因子很可能是在对多个方案进行剖析后选择的。但其理论优势之一是,在每个分支点之间会执行更多的算术操作,编译器和CPU可以更好地安排它们的调度,以尽可能地实现指令流水线。

3
我再次错了,看起来在Numpy的/linalg/blas_lite.c中调用了例行程序。第一个daxpy_是针对浮点数点积的展开内部循环,并基于很久以前的代码。请查看那里的注释:"常数乘以向量再加上向量。对于增量等于一,使用展开循环。Jack Dongarra,Linpack,3/11/78。修改于12/3/93,array(1)声明更改为array(*)" - John Lyon
5
我猜这两个算法都不会用于浮点数、双精度浮点数、单精度复数或双精度复数。NumPy需要ATLAS,它有自己的“daxpy”和“dgemm”的版本。有针对浮点数和复数的版本;对于整数等类型,NumPy可能会退回到您链接的C模板。 - Translunar

9
实现某个功能所使用的语言本身并不是衡量性能的好方法。通常,使用更合适的算法是决定性因素。
在您的情况下,您正在使用学校教授的矩阵乘法的朴素方法,其时间复杂度为O(n^3)。然而,对于某些类型的矩阵,例如方阵、稀疏矩阵等,您可以做得更好。
请参考Coppersmith–Winograd algorithm(O(n^2.3737)的方阵乘法)作为快速矩阵乘法的良好起点。还可以查看“参考文献”部分,其中列出了一些指向更快方法的指针。
对于更具惊人性能提升的更贴近实际的例子,请尝试编写快速的strlen()并将其与glibc实现进行比较。如果您无法打败它,请阅读glibc的strlen()源代码,它有相当好的注释。

1
+1 针对使用大O符号和分析(我总是记得朴素方法n^3与Strassen算法约为n^2.8)。再次强调,检查算法速度的好方法是使用大O符号,而不是语言。 - Juan Antonio Gomez Moriano
1
在这种情况下,OP的天真C matmul没有进行缓存块处理,甚至没有转置其中一个输入。它在一个矩阵的行和另一个矩阵的列上循环,当它们都按行主序排列时,所以它会出现大量的缓存未命中。(转置是O(n^2)的工作,用于使行*列向量点积进行顺序访问,这也使它们可以使用SSE/AVX/等自动向量化,如果您使用-ffast-math)。 - Peter Cordes
1
使用Coppersmith-Winograd算法并不是一个好主意,因为它有一个巨大的隐藏常数因子。事实上,据我所知,这就是为什么没有主流高度优化的BLAS库使用它的原因。更不用说它实现复杂,不适合现代处理器架构。它只对真正巨大的矩阵(对于大多数实际问题来说太大了)有用。然而,Strassen算法实际上在一些BLAS库中使用。但是,它只用于相对较大的矩阵,如>512x512。在实践中,从使用Strassen中获得的显着加速只对>4096x4096这样的矩阵可见。 - Jérôme Richard

5

编写NumPy的人显然知道他们在做什么。

优化矩阵乘法有很多方法。例如,遍历矩阵的顺序会影响内存访问模式,从而影响性能。
良好使用SSE是另一种优化方式,NumPy可能采用了这种方式。
可能还有更多方式,这些都是NumPy开发人员所知道的,而我不知道的。

顺便问一下,你是否使用优化编译了你的C代码?

你可以尝试以下C的优化方法。它可以并行工作,我想NumPy也是沿着同样的路线做的。
注意:仅适用于偶数大小。通过额外的工作,您可以消除此限制并保持性能提升。

for (i = 0; i < n; i++) {
        for (j = 0; j < n; j+=2) {
            int sub1 = 0, sub2 = 0;
            for (k = 0; k < n; k++) {
                sub1 = sub1 + a[i * n + k] * b[k * n + j];
                sub1 = sub1 + a[i * n + k] * b[k * n + j + 1];
            }
            c[i * n + j]     = sub;
            c[i * n + j + 1] = sub;
        }
    }
}

是的,我尝试了不同级别的编译优化,但与numpy相比,这并没有改变结果太多。 - Charles Menguy
一个好的乘法实现会胜过任何优化级别。我猜测,没有任何优化会显著更差。 - ugoren
3
这个答案对Numpy的功能作出了许多假设。但实际上,Numpy几乎不会直接完成这些任务,而是在有BLAS库可用时将工作卸载给它。矩阵乘法的性能在很大程度上取决于BLAS的实现。 - Fred Foo

5

Numpy是高度优化的代码。在书籍Beautiful Code中,有一篇关于它的文章。

Ctypes需要进行从C到Python再返回的动态翻译,这会增加一些开销。而在Numpy中,大多数矩阵运算都是完全内部完成的。


4
NumPy 本身并非经过优化的代码。它利用了经过优化的代码,例如 ATLAS。 - Translunar

3
Fortran在数值计算中速度优势最常见的原因,据我所知,是因为该语言更容易检测到别名 - 编译器可以确定要相乘的矩阵不共享同一内存,这有助于改善缓存(无需确保结果立即写回“共享”内存)。这就是为什么C99引入了restrict的原因。
然而,在这种情况下,我想知道numpy代码是否也能够使用一些特殊指令,而C代码则不行(因为差异似乎特别大)。

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接