时间戳

首页 > 解决方案 > 优化 4D 张量旋转

arrays - 优化 4D 张量旋转

问题描述

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

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

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)

有没有办法优化代码？

标签： arraysmathjuliatensor

解决方案

我对各种 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

推荐阅读