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 倍。还不错？

Julia 的广播速度是 Matlab 的两倍

Julia broadcasting twice as slow as Matlab

performance

matlab

overhead

broadcasting

julia