两个数组的朱莉娅乘法

Question

有没有办法加快/更优雅地编写此数组乘法（在 numpy 数组中，我会写为 A*B）？

A = rand(8,15,10)
B = rand(10,5)
C = zeros(8,15,5)
for i in 1:8
    for j in 1:15
        for k in 1:10
            for l in 1:5
                C[i,j,l] = A[i,j,:]⋅B[:,l]
            end
        end
    end
end

Answer 1

有一堆 Julia 包可以让你在一行中写下你的收缩。这里有几个例子基于 Einsum.jl, OMEinsum.jl, and TensorOperations.jl:

using OMEinsum
f_omeinsum(A,B) = ein"ijk,km->ijm"(A,B)

using Einsum
f_einsum(A,B) = @einsum C[i,j,l] := A[i,j,k] * B[k,l]

using TensorOperations
f_tensor(A,B) = @tensor C[i,j,l] := A[i,j,k] * B[k,l]

除了这些优雅（且快速，见下文）的版本之外，您还可以大大改进循环代码。这是你的代码，包装在一个函数中，以及一个带有注释的改进版本：

function f(A,B)
    C = zeros(8,15,5)
    for i in 1:8
        for j in 1:15
            for k in 1:10
                for l in 1:5
                    C[i,j,l] = A[i,j,:]⋅B[:,l]
                end
            end
        end
    end
    return C
end

function f_fast(A,B)
    # check bounds
    n1,n2,n3 = size(A)
    m1, m2 = size(B)
    @assert m1 == n3
    C = zeros(n1,n2,m2)

    # * @inbounds to skip boundchecks inside the loop
    # * different order of the loops to account for Julia's column major order
    # * written out the k-loop (dot product) explicitly to avoid temporary allocations
    @inbounds for l in 1:m2
                for k in 1:m1
                        for j in 1:n2
                            for i in 1:n1
                            C[i,j,l] += A[i,j,k]*B[k,l]
                        end
                    end
                end
            end
    return C
end

让我们比较一下所有的方法。首先我们检查正确性：

using Test
@test f(A,B) ≈ f_omeinsum(A,B) # Test passed
@test f(A,B) ≈ f_einsum(A,B) # Test passed
@test f(A,B) ≈ f_tensor(A,B) # Test passed
@test f(A,B) ≈ f_fast(A,B) # Test passed

现在，让我们使用 BenchmarkTools.jl 进行基准测试。我把时间放在我的机器上作为评论。

using BenchmarkTools
@btime f($A,$B); # 663.500 μs (12001 allocations: 1.84 MiB)
@btime f_omeinsum($A,$B); # 33.799 μs (242 allocations: 20.20 KiB)
@btime f_einsum($A,$B); # 4.200 μs (1 allocation: 4.81 KiB)
@btime f_tensor($A,$B); # 2.367 μs (3 allocations: 4.94 KiB)
@btime f_fast($A,$B); # 7.375 μs (1 allocation: 4.81 KiB)

正如我们所看到的，所有基于 einsum/tensor 符号的方法都比您原来的循环实现快得多 - 而且只有一行！我们的 f_fast 的表现大致相同，但仍落后于最快的 f_tensor。

最后，让我们全力以赴，因为我们可以。利用 LoopVectorization.jl 中的魔法，我们将 f_fast 中的 @inbounds 替换为 @avx（我们在下面称这个新版本为 f_avx）并自动获得另外 2 倍的加速相对于上面的 f_tensor 表现：

@test f(A,B) ≈ f_avx(A,B) # Test passed
@btime f_avx($A,$B); # 930.769 ns (1 allocation: 4.81 KiB)

但是，由于它的简单性，我仍然更喜欢 f_tensor 除非每一微秒在您的应用程序中都很重要。

两个数组的朱莉娅乘法

julia multiplication of two arrays

linear-algebra

julia