4D张量旋转的优化

Question

我必须在斯托克斯求解器中每个时间步执行 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

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