Julia中有效的元素矩阵运算

Question

我需要执行(复杂)矩阵的离散卷积,并在Julia中定义以下函数:

function convolve(M::Array{Complex{Float64},2}, K::Array{Float64,2}, p::Int)
    (n,m) = size(M)
    res = zeros(Complex{Float64},n)
    for k=1:p
        for l=1:n
            res[l] += M[l,k]*K[l,end-p+k]
        end
    end
    return res
end

我这样使用它:

M=complex(rand(2000,2000))
K=rand(2000,2000)
@time convolve(M,K,2000,0)

现在这是相对快速和令人惊讶的,比矢量化版本更快(约3倍),我用它替换内部循环res += M[:,k].*K[:,end-p+k].(我认为这是由于临时阵列的很多记忆分配,我可以忍受).

但是矢量化的MATLAB代码运行速度提高了大约5倍:

function res = convolve(M, K, p)
    n = size(M,1);
    res = zeros(n,1);
    for k=1:p
        res = res + M(:,k).*K(:,end-p+k);
    end
end

我做错了什么？如何让Julia像MATLAB一样快速地执行元素乘法？这是一个索引问题吗？

注意:我已经检查过@code_warntype没有类型犹豫不决的有趣业务(没有Any或Union等),但问题可能更微妙.宏@code_llvm产生了令人惊讶的长输出,但我不是专家所以我很难看到发生了什么.

Answer 1

以下版本在我的机器上速度更快：

function convolve2(M::Array{Complex{Float64},2}, K::Array{Float64,2}, p::Int)
    (n,m) = size(M)
    res = zeros(Complex{Float64},n)
    offset = size(K,2)-p
    (p>m || offset<0) && error("parameter p ($p) out of bounds")
    @inbounds for k=1:p
        @inbounds @simd for l=1:n
            res[l] += M[l,k]*K[l,offset+k]
        end
    end
    return res
end

请注意@simd当前在许多 CPU 中使用向量指令的添加。

编辑：OP代码中的性能影响似乎源于在热循环行end的索引中使用。K重新定义Base.trailingsizewith@inline使 LLVM 内联end（在我的机器上）并使两个版本的运行速度大致相同。使用的代码：

import Base: trailingsize
@inline function Base.trailingsize(A, n)
  s = 1
  for i=n:ndims(A)
    s *= size(A,i)
  end
  return s
end

请参阅有关此问题的#19389问题的评论。