Эффективные поэлементно операции с матрицами в Julia

Question

мне нужно выполнить discretised свертку (комплексных) матриц, и определили следующую функцию в 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]. (Я думаю, что это связано с большим количеством распределения памяти для временных массивов, я могу с этим справиться).

Но vectorised 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 

Что я делаю не так, и как я могу получить Джулии выполнять такие поэлементно умножений так быстро, как 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 дополнение, которое использует векторные инструкции в настоящее время во многих процессорах.

РЕДАКТИРОВАТЬ: Чертёж в коде 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 относительно этого.