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μsvs18μs）。

您可以在手册here中阅读有关此内容的信息。但是请注意，您需要确保所有尺寸的大小都正确，这使得处理像函数中那样的硬编码范围非常棘手-考虑使用Julia的一些内置迭代语法（例如eachindex）或使用{{1} }，如果您需要循环根据输入更改迭代次数。

Answer 2

我定时了各种einsum软件包。仅通过添加@inbounds即可使Einsum更快。对于如此小的矩阵，TensorOperations的速度较慢。 LoopVectorization在这里编译需要一定的时间，但最终结果更快。

（我想您是想将每个元素aux的{{1}}设为零，并且我将for l = 1:3; aux = 0.0; for m = 1:3设置为不堆积在垃圾之上。）

Crot .= 0

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

@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)

Answer 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将在（我认为）编译时抱怨。）

优化4D张量旋转

3 个答案: