如何在Julia中加速多个广播

Question

这个Julia函数似乎效率很低(比同等的Pythran/C++代码慢一个数量级,即使在Julia热身之后)......

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).有没有办法在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)不再独立于维度.

看来,我们有哪些简单的Python/numpy的代码可以Pythran比我们与朱莉娅得到(除非有在朱莉娅加速了这种方式？和未来的朱莉娅版本可以解决这个问题)更快获得加速的情况下.有趣...

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就地修改或创建一个新数组,所以我展示了两种方法).