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)

现在这相对较快，而且令人惊讶的是，比我用 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 产生出奇的长输出，但我不是专家所以我很难看出发生了什么。

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 代码中的性能下降似乎源于在热循环行中 K 的索引中使用 end。用 @inline 重新定义 Base.trailingsize 使 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 的问题的评论。

Julia 中高效的逐元素矩阵运算

Efficient element-wise matrix operations in Julia

matrix-multiplication

julia