Julia 的广播速度是 Matlab 的两倍

如何解决Julia 的广播速度是 Matlab 的两倍

为了从 Matlab 迁移，我正在尝试熟悉 Julia，到目前为止一切顺利，直到我开始使用广播来移植特定函数，该函数的执行速度或多或少是 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

任何帮助将不胜感激！

解决方法

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

其次，几乎所有的运行时间 (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,$mid);
  17.173 ms (4 allocations: 7.63 MiB)

这并不是什么加速。不过，分配较少。问题是运行时在很大程度上受 sin 广播的支配。

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

function features3(X::Vector{Float64},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,$mid);
  1.919 ms (45 allocations: 7.63 MiB)

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

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

using LoopVectorization

function features4(X::Vector{Float64},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},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,$mid);
  2.695 ms (5 allocations: 7.63 MiB)

对于单线程情况非常好的加速。

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

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