广播阵列和在Julia中没有临时矩阵分配的求和

Question

我需要执行以下形式的计算：

A = reshape(big_mat,m1,1,1,n1,n2,n3)
B = reshape(big_mat,1,m1,n1,1,n2,n3)
C = reshape(another_mat,m1,m1,n1,n1,1,1)
D = sum(A.*B.*C,dims=(5,6))

A.*B.*C正在创建大小为(m1，m1，n1，n1，n2，n3(的临时大矩阵。假设D的大小只有(m1，m1，n1，n1(，是否有一种更有效的过程可以在不调用循环的情况下进行求和？

Answer 1

您可以要求计算机为您编写循环。@einsum D[a,b,c,d] := mat[a,d,e,f] * mat[b,c,e,f] * another[a,b,c,d]将写入6个嵌套的for循环，对未出现在左侧的e,f求和。

@tullio也会这样做，并添加平铺内存访问(和多线程(，这应该会更快一些。

julia> using Einsum, Tullio, BenchmarkTools
julia> let
n = 25
big_mat = rand(n,n,n,n)
another_mat = rand(n,n,n,n)
D1 = @btime let n = $n
A = reshape($big_mat,     n,1,1,n, n,n)
B = reshape($big_mat,     1,n,n,1, n,n)
C = reshape($another_mat, n,n,n,n, 1,1)
D = sum(A.*B.*C,dims=(5,6))
end
D2 = @btime @einsum D[a,b,c,d] := $big_mat[a,d,e,f] * $big_mat[b,c,e,f] * $another_mat[a,b,c,d]
D3 = @btime @tullio D[a,b,c,d] := $big_mat[a,d,e,f] * $big_mat[b,c,e,f] * $another_mat[a,b,c,d]
D1 ≈ D2 ≈ D3
end
min 462.545 ms, mean 494.505 ms (20 allocations, 1.82 GiB)
min 213.126 ms, mean 214.412 ms (3 allocations, 2.98 MiB)
min 80.585 ms, mean 80.785 ms (53 allocations, 2.98 MiB)
true
julia> 2.98 * 25^2  # memory 2.98 MiB -> 1.82 GiB
1862.5
julia> @macroexpand1 @einsum D[a,b,c,d] := mat[a,d,e,f] * mat[b,c,e,f] * another[a,b,c,d]
quote  # this will print the loops