4D张量旋转的优化

Question

我必须在 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)

有什么优化代码的方法吗？

Answer 1

您在这里进行了大量索引，因此进行了大量边界检查。在这里节省一些时间的一种方法是使用 @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) 如果您需要循环根据输入更改迭代次数。

Answer 2

这似乎是一个适当的收缩（每个索引都出现在输出中，或者恰好在右侧出现两次），因此可以用 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 也可能会有所回报，因为您的张量很小且大小不变。我不知道它是否适用于任何 Einstein 求和包，但无论如何你都可以为收缩生成一个完全展开的函数。

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

Answer 3

我对各种 einsum 包进行了计时。 Einsum 仅通过添加 @inbounds 就更快了。对于如此小的矩阵，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

4D张量旋转的优化

Optimisation of 4D tensor rotation

arrays

math

julia

tensor