加速numpy嵌套循环

Question

我正在使用numpy和cython为python中的无线网络编写模拟,其中假设no_nodes在2d平面上随机散布了许多节点,这些节点发送一些波形及其各自的接收器,再次在2d平面上随机散射.每个发送节点产生一个I调用的波形output(每个波形可以产生不同长度的输出).

我想要做的是将每个节点的输出相加到一个大波形,这个波形将成为每个接收器的输入以进行解调等.现在有两个关键点:

• 发送器异步发送,因此每个发送节点必须维护a start_clock和a ,end_clock以便正确地对波形求和
• 在根据功能j接收i节点之前,发送节点的输出将被衰减attenuate(i,j)

所以这是代码:

#create empty 2d array (no_rx_nodes x no_samples for each waveform)
waveforms = np.zeros((no_nodes, max(end_clock))) 

for i in range(no_nodes): #calculate the waveform for each receiver
    for j in range(no_nodes): #sum the waveforms produced by each transmitter
        waveforms[i, start_clock[j]:end_clock[j]] += output[j,:] * attenuate(i,j)
return waveforms

关于上述的一些评论:

• output[j, :] 是发射器j的输出波形
• waveforms[i,:] 是接收器i接收的波形

我希望我在这里要完成的工作相当清楚.因为产生的波形非常大(大约10 ^ 6个样本),我也尝试将此代码转换为cython,但没有注意到任何特定的加速(可能是5-10倍更好但不多).我在想,如果有什么事我可以诉诸这样才能得到一个速度,因为它是一个真正的瓶颈整个模拟(它需要计算几乎一样多时间的代码的其余部分,这实际上是相当比这更复杂).

Answer 1

这是一个内存带宽限制问题,大约3GB/s的内存带宽,你可以得到的最好是内循环大约2-4ms.要达到该范围,您需要阻止内部循环以更好地利用cpu缓存(numexpr为您执行此操作):

for i in range(no_nodes):
    for j in range(no_nodes):
        # should be chosen so all operands fit in the (next-to-)last level cache
        # first level is normally too small to be usable due to python overhead
        s  = 15000 
        a = attenuation[i,j]
        o = output[j]
        w = waveforms[i]
        for k in range(0, w.size, s): 
            u = min(k + s, w.size)
            w[k:u] += o[k:u] * a
        # or: numexpr.evaluate("w + o * a", out=w)

使用float32数据而不是float64也应该是内存带宽要求的一半,性能提高一倍.

要获得更大的加速,您必须重新设计完整算法以获得更好的数据位置