如何在Julia中加速多个广播

Question

这个Julia函数似乎效率很低（即使在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)

我猜只是因为我写的不正确...朱莉娅（Julia）如何加快诸如“多播”的速度？我猜/希望我不需要扩展循环...

在第一个答案之后进行编辑

谢谢！通过我的设置，Pythran解决方案（就地和就地）仍然快1.5到2倍（没有OpenMP）。有没有办法在Julia中激活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）不再与尺寸无关。

似乎我们有一种情况，使用Pythran可以比使用Julia更快地简化简单的Python / Numpy代码（除非在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)

当我们获得大约10倍的加速时，它与Pythran / C ++相比如何？

最后请注意，如果您可以通过编写以下内容来对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还是创建一个新数组，所以我将展示两种方法。）