Julia 的广播速度是 Matlab 的两倍

Question

我正在尝试熟悉Julia以便从Matlab迁移，到目前为止一切顺利，直到我开始使用广播来移植一个特定的功能，该函数的执行速度或多或少是Matlab的两倍。

function features(X::Vector{Float64},M::Int,hyper::Float64,mid::Float64)
X = X.-mid
H = 4.0.*hyper.+maximum(abs.(X))
X = (X.+H)./(2.0.*H)
w = transpose(1:M)
S = (sqrt.(2.0.*pi).*hyper).*exp.((-0.5.*hyper.^2).*((pi.*w./(2.0.*H)).^2))
f = H.^(-0.5).*sin.(pi.*X.*w).*sqrt.(S)
end

任何帮助将不胜感激！

Answer 1

首先，您对广播的使用不是最佳的。你用得太多了，;)不够

其次，几乎所有的运行时(99.9%)都发生在广播sin表达式中，因此应该集中精力。

第三，你不应该指望Julia在这种情况下会胜过Matlab。这正是 Matlab 优化的目的：直接对优化的 C/Fortran 例程进行元素调用。此外，Matlab 默认是多线程的，隐式并行运行元素调用，而 Julia 要求您明确多线程。

就目前而言，2 倍的差异似乎并非不合理。

不过，让我们努力一下。先说几点意见：

X = X .- mid

您错过了就地进行操作，请使用

X .= X .- mid

相反。这样可以节省中间阵列的分配。

H = 4.0.*hyper.+maximum(abs.(X))

通过标量(hyper)进行广播是徒劳的，最坏的情况是浪费。abs.(X)会创建一个不必要的临时数组。而是使用带有函数输入的maximum版本，这样更有效：

H = 4 * hyper + maximum(abs, X)

以下是一些不必要的点：

S = (sqrt.(2.0.*pi).*hyper).*exp.((-0.5.*hyper.^2).*((pi.*w./(2.0.*H)).^2))

避免再次通过标量广播，并在大多数位置使用整数而不是浮点数：

S = (sqrt(2pi) * hyper) .* exp.((-0.5 * hyper^2 * (pi/2H)^2) .* w.^2)

请注意，x^(-0.5)比1/sqrt(x)慢得多，所以

f = H.^(-0.5).*sin.(pi.*X.*w).*sqrt.(S)

应该是

f = sin.(pi .* X .* w') .* (sqrt.(S)' ./ sqrt(H))

让我们把这个放在一起：

function features2(X::Vector{Float64},M::Int,hyper::Float64,mid::Float64)
X .= X .- mid
H = 4 * hyper + maximum(abs, X)
X .= (X .+ H) ./ (2 * H)
w = 1:M
S = (sqrt(2pi) * hyper) .* exp.((-0.5 * hyper^2 * (pi/2H)^2) .* w.^2)
f = sin.(pi .* X .* w') .* (sqrt.(S)' ./ sqrt(H))
return f
end

基准：

jl> X = rand(10000);
jl> M = 100;
jl> hyper = rand();
jl> mid = 0.4;
jl> @btime features($X, $M, $hyper, $mid);
17.339 ms (9 allocations: 7.86 MiB)
jl> @btime features2($X, $M, $hyper, $mid);
17.173 ms (4 allocations: 7.63 MiB)

这不是太大的加速。不过，分配较少。问题在于运行时在很大程度上由sin广播主导。

让我们尝试多线程。我有 8 个内核，所以我使用 8 个线程：

function features3(X::Vector{Float64},M::Int,hyper::Float64,mid::Float64)
X .= X .- mid
H = 4 * hyper + maximum(abs, X)
X .= (X .+ H) ./ (2 * H)
w = transpose(1:M)
S = (sqrt(2pi) * hyper) .* exp.((-0.5 * hyper^2 * (pi/2H)^2) .* w.^2)
f = similar(X, length(X), M)
temp = sqrt.(S) ./ sqrt(H)
Threads.@threads for j in axes(f, 2)
wj = w[j]
tempj = temp[j]
for i in axes(f, 1)
@inbounds f[i, j] = tempj * sin(pi * X[i] * w[j])
end
end
return f
end

本克马克：

jl> @btime features3($X, $M, $hyper, $mid);
1.919 ms (45 allocations: 7.63 MiB)

这要好得多，循环和显式线程的速度提高了 9 倍。

但是仍然有一些选择：例如LoopVectorization.jl。您可以安装这个惊人的软件包，但您需要新版本，可能会有一些安装问题，具体取决于您拥有的其他软件包。LoopVectorization有两个特别感兴趣的宏，@avx和@avxt，前者做了很多工作来矢量化(在simd意义上)你的代码，单线程，而后者做同样的事情，但多线程。

using LoopVectorization
function features4(X::Vector{Float64},M::Int,hyper::Float64,mid::Float64)
X .= X .- mid
H = 4 * hyper + maximum(abs, X)
X .= (X .+ H) ./ (2 * H)
w = collect(1:M)  # I have to use collect here due to some issue with LoopVectorization
S = (sqrt(2pi) * hyper) .* exp.((-0.5 * hyper^2 * (pi/2H)^2) .* w.^2)
f = @avx sin.(pi .* X .* w') .* (sqrt.(S)' ./ sqrt(H))
return f
end
function features4t(X::Vector{Float64},M::Int,hyper::Float64,mid::Float64)
X .= X .- mid
H = 4 * hyper + maximum(abs, X)
X .= (X .+ H) ./ (2 * H)
w = collect(1:M)  # I have to use collect here due to some issue with LoopVectorization
S = (sqrt(2pi) * hyper) .* exp.((-0.5 * hyper^2 * (pi/2H)^2) .* w.^2)
f = @avxt sin.(pi .* X .* w') .* (sqrt.(S)' ./ sqrt(H))
return f
end

这些功能之间的唯一区别是@avx与@avxt.

基准：

jl> @btime features4($X, $M, $hyper, $mid);
2.695 ms (5 allocations: 7.63 MiB)

对于单线程情况来说，这是一个非常好的加速。

jl> @btime features4t($X, $M, $hyper, $mid);
431.700 μs (5 allocations: 7.63 MiB)

多线程 avx 代码的速度是笔记本电脑原始代码的 40 倍。不错？