优化 4D 张量旋转答案

【问题标题】：Optimisation of 4D tensor rotation优化 4D 张量旋转
【发布时间】：2020-06-07 13:16:25
【问题描述】：

我必须在 Stokes 求解器中每个时间步执行 3x3x3x3 4D 张量 +100k 次的旋转，其中旋转的 4D 张量是 Crot[i,j,k,l] = Crot[i,j,k, l] + Q[m,i] * Q[n,j] * Q[o,k] * Q[p,l] * C[m,n,o,p]，所有索引从 1 到 3。

到目前为止，我已经在 Julia 中天真地编写了以下代码：

Q    = rand(3,3)
C    = rand(3,3,3,3)
Crot = Array{Float64}(undef,3,3,3,3)
function rotation_4d!(Crot::Array{Float64,4},Q::Array{Float64,2},C::Array{Float64,4})
aux = 0.0
for i = 1:3
    for j = 1:3
        for k = 1:3
            for l = 1:3

                for m = 1:3
                    for n = 1:3
                        for o = 1:3
                            for p = 1:3                                     
                                aux += Q[m,i] * Q[n,j] * Q[o,k] * Q[p,l] * C[m,n,o,p];
                            end
                        end
                    end
                end

                Crot[i,j,k,l] += aux

            end
        end
    end
end

end

与：

@btime rotation_4d(Crot,Q,C)
14.255 μs (0 allocations: 0 bytes)

有没有办法优化代码？

【问题讨论】：

搜索爱因斯坦求和；似乎有Einsum 以及更有效的在线实现。

标签： arrays math julia tensor

【解决方案1】：

我对各种 einsum 包进行了计时。仅通过添加@inbounds，Einsum 就更快了。对于这样的小矩阵，TensorOperations 比较慢。 LoopVectorization 在这里编译需要一些时间，但最终结果更快。

（我假设您的意思是每个元素将 aux 归零一次，for l = 1:3; aux = 0.0; for m = 1:3，我设置了 Crot .= 0 以免堆积在垃圾上。）

@btime rotation_4d!($Crot,$Q,$C)  # 14.556 μs (0 allocations: 0 bytes)
Crot .= 0; # surely!
rotation_4d!(Crot,Q,C)
res = copy(Crot);

using Einsum # just adds @inbounds really
rot_ei!(Crot,Q,C) = @einsum Crot[i,j,k,l] += Q[m,i] * Q[n,j] * Q[o,k] * Q[p,l] * C[m,n,o,p]
Crot .= 0;
rot_ei!(Crot,Q,C) ≈ res # true
@btime rot_ei!($Crot,$Q,$C);      # 7.445 μs (0 allocations: 0 bytes)

using TensorOperations # sends to BLAS
rot_to!(Crot,Q,C) = @tensor Crot[i,j,k,l] += Q[m,i] * Q[n,j] * Q[o,k] * Q[p,l] * C[m,n,o,p]
Crot .= 0; 
rot_to!(Crot,Q,C) ≈ res # true
@btime rot_to!($Crot,$Q,$C);      # 22.810 μs (106 allocations: 11.16 KiB)

using Tullio, LoopVectorization
rot_lv!(Crot,Q,C) = @tullio Crot[i,j,k,l] += Q[m,i] * Q[n,j] * Q[o,k] * Q[p,l] * C[m,n,o,p]  tensor=false
Crot .= 0; 
@time rot_lv!(Crot,Q,C) ≈ res # 50 seconds!
@btime rot_lv!($Crot,$Q,$C);      # 2.662 μs (8 allocations: 256 bytes)

但是，这仍然是一个糟糕的算法。这只是 4 次小矩阵乘法，但每次都完成了很多次。连续执行它们要快得多——9*4 * 27 乘法，而不是 [更正！] 4 * 9^4 用于上面的简单嵌套。

function rot2_ein!(Crot, Q, C)
    @einsum mid[m,n,k,l] := Q[o,k] * Q[p,l] * C[m,n,o,p]
    @einsum Crot[i,j,k,l] += Q[m,i] * Q[n,j] * mid[m,n,k,l]
end
Crot .= 0; rot2_ein!(Crot,Q,C) ≈ res # true 
@btime rot2_ein!($Crot, $Q, $C);  # 1.585 μs (2 allocations: 784 bytes)

function rot4_ein!(Crot, Q, C) # overwrites Crot without addition
    @einsum Crot[m,n,o,l] = Q[p,l] * C[m,n,o,p]
    @einsum Crot[m,n,k,l] = Q[o,k] * Crot[m,n,o,l]
    @einsum Crot[m,j,k,l] = Q[n,j] * Crot[m,n,k,l]
    @einsum Crot[i,j,k,l] = Q[m,i] * Crot[m,j,k,l]
end
rot4_ein!(Crot,Q,C) ≈ res # true
@btime rot4_ein!($Crot, $Q, $C);  # 1.006 μs

【讨论】：

通过优化最后一个算法（4 个矩阵乘法），大约还有一个 5 的因子，在 this LoopVectorization.jl issue 中计算出来。 @btime rotation_nexpr!($Crot, $Q, $C); 给我 183.101 ns（在与上述相同的计算机上）。

【解决方案2】：

您在这里做了很多索引，因此需要进行很多边界检查。在这里减少一些时间的一种方法是使用@inbounds 宏，它会关闭边界检查。将代码重写为：

function rotation_4d!(Crot::Array{Float64,4},Q::Array{Float64,2},C::Array{Float64,4})
    aux = 0.0
    @inbounds for i = 1:3, j = 1:3, k = 1:3, l = 1:3
        for m = 1:3, n = 1:3, o = 1:3, p = 1:3                                     
            aux += Q[m,i] * Q[n,j] * Q[o,k] * Q[p,l] * C[m,n,o,p];
        end
    Crot[i,j,k,l] += aux

    end
end

给了我大约 3 倍的加速（在我的系统上为 6μs 对 18μs）。

您可以在手册here 中阅读有关此内容的信息。但是请注意，您需要确保所有维度的大小都正确，这使得在函数中使用硬编码范围变得很棘手 - 如果需要，请考虑使用 Julia 的一些内置迭代语法（如eachindex）或使用size(Q, 1)您的循环根据输入更改迭代次数。

【讨论】：

【解决方案3】：

这似乎是一个适当的收缩（每个索引都出现在输出中，或者恰好出现在右侧两次），因此可以使用TensorOperations.jl：

@tensor Crot[i,j,k,l] = Crot[i,j,k,l] + Q[m,i] * Q[n,j] * Q[o,k] * Q[p,l] * C[m,n,o,p]

或OMEinsum.jl。

使用StaticArrays.jl 也可能会有所回报，因为您的张量很小且大小恒定。我不知道它是否适用于任何爱因斯坦求和包，但无论如何你都可以为收缩生成一个完全展开的函数。

（注意：我实际上并没有针对这种情况测试它们中的任何一个。如果不是正确的收缩，TensorOperations 会在（我认为）编译时抱怨。）

【讨论】：