优化此现有3D numpy矩阵乘法

Question

我有一些刚刚完成的代码。它按预期工作。我选择在Numpy中使用dot，因为根据我的有限经验，如果在系统上安装了BLAS，则它比通常的矩阵乘法形式快。但是，您会注意到比我不得不转置了很多东西。我注意到这一点实际上超过了使用dot的好处。

我提供了论文中发现的数学方程。请注意是递归

这是我的代码实现：

我首先提供必要的组件及其尺寸

P = array([[0.73105858, 0.26894142],
           [0.26894142, 0.73105858]])  # shape (K,K)

B = array([[6.07061629e-09, 0.00000000e+00],
           [0.00000000e+00, 2.57640371e-10]])  # shape (K,K)

dP = array([[[ 0.19661193, -0.19661193],
             [ 0.        ,  0.        ]],

           [[ 0.        ,  0.        ],
            [ 0.19661193, -0.19661193]],

           [[ 0.        ,  0.        ],
            [ 0.        ,  0.        ]],

           [[ 0.        ,  0.        ],
            [ 0.        ,  0.        ]]])  # shape (L,K,K)

dB = array([[[ 0.00000000e+00,  0.00000000e+00],
             [ 0.00000000e+00,  0.00000000e+00]],

            [[ 0.00000000e+00,  0.00000000e+00],
             [ 0.00000000e+00,  0.00000000e+00]],

            [[-1.16721049e-09,  0.00000000e+00],
             [ 0.00000000e+00,  0.00000000e+00]],

            [[ 0.00000000e+00,  0.00000000e+00],
             [ 0.00000000e+00, -1.27824683e-09]]])  # shape (L,K,K)

a = array([[[ 7.60485178e-08,  5.73923956e-07]],

           [[-5.54100398e-09, -8.75213012e-08]],

           [[-1.25878643e-08, -1.48361081e-07]],

           [[-2.73494035e-08, -1.74585971e-07]]])  # shape (L,1,K)

alpha = array([[0.11594542, 0.88405458]])  # shape (1,K)

c = 1  # scalar

这里是实际的Numpy计算。注意所有转置用法。

term1 = (a/c).dot(P).dot(B)
term2 = dP.transpose((0,2,1)).dot(alpha.T).transpose((0,2,1)).dot(B)
term3 = dB.dot(  P.T.dot(alpha.T) ).transpose((0,2,1))
a = term1 + term2 + term3

然后应该得到：

>>> a
array([[[ 1.38388584e-10, -5.87312190e-12]],

       [[ 1.05516813e-09, -4.47819530e-11]],

       [[-3.76451117e-10, -2.88160549e-17]],

       [[-4.06412069e-16, -8.65984406e-10]]])

请注意，我已经选择了alpha和a的形状。如果发现它们可以提供卓越的性能，则可以更改这些设置。

我想指出的是，我认为现有代码是快速的。实际上，非常快。但是，我一直在想是否可以做得更好。请放手一搏。我已经剖析了我的代码（其中包含大量的Numpy广播和矢量化），这不一定是瓶颈，因为在我的旧机器上进行评估需要23微秒的时间。但是，这是递归的单个步骤。这意味着将按顺序评估N次。因此，即使是最小的收益，对于大的序列也将产生很大的差异。

谢谢您的时间。

Answer 1

Q：…如果可以做得更好？

您的原样代码在我的_{（似乎更旧）}的计算机上执行，而不是在发布的~ 23 [us]中执行，而是在首次调用时使用~ 45 [ms]，并享受到iCACHE和[ C0]层次结构介于[dCACHE：

之间

~77..1xx [us]

有趣的是，多次重新运行代码，将处理结果重新分配回

>>> from zmq import Stopwatch; aClk = Stopwatch()
>>> import numpy as np
>>>
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
44679
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
149
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
113
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
128
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
82
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
100
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
77
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
97
>>> a
array([[[ 1.38387304e-10, -5.87323502e-12]],
       [[ 1.05516829e-09, -4.47819355e-11]],
       [[-3.76450816e-10, -2.60843400e-20]],
       [[-1.41384088e-18, -8.65984377e-10]]])

实际上不改变 a中的值：

a

这意味着代码按原样进行了大量工作，最终产生了不变值

>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
97
>>> a
array([[[ 1.38387304e-10, -5.87323502e-12]],    
       [[ 1.05516829e-09, -4.47819355e-11]],
       [[-3.76450816e-10, -2.60843400e-20]],
       [[-1.41384088e-18, -8.65984377e-10]]])

（复制的身份，但花费〜XY [我们]这样做-您是唯一决定的人，无论您的目标应用程序是否可以使用）

---

关于需要改善的地方的评论：

好吧，考虑到[a N和~ 1E(3..6)和K ~ 10，通过重新解决（ ]的身份结果）希望提高性能。

寻求改进的目标处理将依次重复大于L ~ 100…小于a，这意味着：

RAM绑定问题不是主要问题，因为对静态部件的缓存影响（所有大小只有几个~1,000x，将以尽可能短的延迟从缓存中重用]

与CPU相关的问题已在~ 1,000,000x-工具的设计和工程中预先解决（在可行的情况下，利用CPU的SIMD资源和矢量对齐的跨步技巧-因此，对用户级编码的期望不高） ）

最后但并非最不重要的是，人们可以评论转置的

[[[costs

”-

[MB]

不需要做任何其他事情，因为必须转置矩阵，但是索引顺序-没有别的。如果这可能会产生某些效果，则可以通过查看数据单元的基础存储到物理环境中的numpy类型的排序或numpy语言类型的排序来预期RAM，但最大规模为[FORTRAN〜C，这使得该方面可以忽略不计，并且被in-cache性质掩盖了具有零回写或其他与性能相关的影响的处理。结果：鉴于上述所有事实和进一步的观察，实际上要获得此处定义的计算的更高吞吐量的最便宜，最合理的步骤是，移至此处的线性工作自由度-更高的1E1 x 1E1 CPU-芯片，性能越好（此处线性增长），并且具有尽可能多的AVX-512寄存器和尽可能多的L1i + L1d高速缓存（对于任何其他O / S噪声，亲和映射避免的策略对于HPC级性能目标），并且依赖于已经很智能的[[1E2 x 1E1 x 1E1工具，针对用于矩阵处理的CPU资源的混合进行了微调（如果需要超越[GHz] IEEE-764，表示，另一个故事开始）。

---

[不要期望用户级别的代码做得比这更好，

numpy

-原始的SIMD对齐处理可以并且将实现。