如何加快朱莉娅的多播

Question

这个Julia函数看起来效率很低(比等价的Pythran / C++代码慢一个数量级，即使在朱莉娅热身之后也是如此)。

function my_multi_broadcast(a)
    10 * (2*a.^2 + 4*a.^3) + 2 ./ a
end

arr = ones(1000, 1000)
my_multi_broadcast(arr)

我想只是我写得不对.一个人怎么能加快朱莉娅的“多播”呢？我想/希望我不需要扩大循环.

在第一个答案之后编辑

谢谢!通过我的设置，Pythran解决方案(就位和不合适)仍然要快1.5到2倍(没有OpenMP)。有办法激活朱莉娅的SIMD指令吗？或者用另一种方法来加速这样的CPU计算？

Python代码：

from transonic import jit

@jit
def broadcast(a):
    return 10 * (2*a**2 + 4*a**3) + 2 / a

@jit
def broadcast_inplace(a):
    a[:] = 10 * (2*a**2 + 4*a**3) + 2 / a

在@simd建议之后进行编辑

看起来，@simd并不是开箱即用的，只是在行的开头添加它。

ERROR: LoadError: LoadError: Base.SimdLoop.SimdError("for loop expected")
Stacktrace:
 [1] compile(::Expr, ::Bool) at ./simdloop.jl:54
 [2] @simd(::LineNumberNode, ::Module, ::Any) at ./simdloop.jl:126
 [3] include at ./boot.jl:317 [inlined]
 [4] include_relative(::Module, ::String) at ./loading.jl:1044
 [5] include(::Module, ::String) at ./sysimg.jl:29
 [6] exec_options(::Base.JLOptions) at ./client.jl:231
 [7] _start() at ./client.jl:425

我想我们必须扩展for循环，但是代码(i)变得更难读，并且(ii)不再独立于维度。

似乎我们有一种情况，简单的Python/Numpy代码可以用Pythran比Julia更快地加速(除了在Julia中有加速的方法吗？)未来的朱莉娅版本可能会解决这个问题)。有趣的是。

Answer 1

广播所有这样的操作：

julia> function my_multi_broadcast2(a)
           @. 10 * (2*a^2 + 4*a^3) + 2 / a
       end
my_multi_broadcast2 (generic function with 1 method)

区别在于，在10 * (2*a.^2 + 4*a.^3) + 2 ./ a中，实际上没有利用广播融合，因为*和两个+没有广播。

编写@. 10 * (2*a^2 + 4*a^3) + 2 / a等同于10 .* (2 .* a.^2 .+ 4 .* a.^3) .+ 2 ./ a。

下面是性能的比较

julia> @btime my_multi_broadcast($arr);
  58.146 ms (18 allocations: 61.04 MiB)

julia> @btime my_multi_broadcast2($arr);
  5.982 ms (4 allocations: 7.63 MiB)

与Pythran / C++相比，我们得到了大约10倍的加速比？

最后，请注意，如果您可以通过以下方式改变arr的位置：

julia> function my_multi_broadcast3(a)
           @. a = 10 * (2*a^2 + 4*a^3) + 2 / a
       end
my_multi_broadcast3 (generic function with 1 method)

julia> @btime my_multi_broadcast3($arr);
  1.840 ms (0 allocations: 0 bytes)

这样做更快，并且可以实现零分配(我不知道您是想修改arr，还是要创建一个新的数组，以便展示这两种方法)。