我需要执行(复杂)矩阵的离散卷积,并在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)
现在这比使用res += M[:,k].*K[:,end-p+k]
替换内循环的矢量化版本相对快速且令人惊讶,更快(约3倍)。 (我认为这是由于临时阵列的大量记忆分配,我可以忍受它。)
但矢量化的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
产生了令人惊讶的长输出,但我不是专家所以我很难看到发生了什么。
答案 0 :(得分:3)
我的机器上的以下版本更快:
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.trailingsize
重新定义@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的问题的评论。