Unreasonable slow just in time compiling of `permutedims` when the array dimension is high

[GiggleLiu]

Permuting the dimensions of high dimensional tensors is frequently used in tensor networks.

julia> using Random, Profile

julia> (n=24; t = randn(fill(2, n)...); @time permutedims(t, randperm(n)));
 39.144106 seconds (1.21 M allocations: 184.261 MiB, 0.09% gc time, 99.45% compilation time)

julia> (n=24; t = randn(fill(2, n)...); @time permutedims(t, randperm(n)));
  0.194273 seconds (9 allocations: 128.001 MiB, 25.46% gc time)

julia> (n=28; t = randn(fill(2, n)...); @time permutedims(t, randperm(n)));
657.982409 seconds (1.04 M allocations: 2.044 GiB, 0.00% gc time, 99.20% compilation time)

julia> (n=28; t = randn(fill(2, n)...); @time permutedims(t, randperm(n)));
  6.854463 seconds (3.84 k allocations: 2.000 GiB, 1.76% gc time, 0.07% compilation time)

This time is much longer than TensorOperations.tensorcopy (~6s).
The crazy fact (from my observation) is: the compiling time of permutedims scales exponentially with the tensor rank!

Related code block:

julia/base/multidimensional.jl

Line 1532 in 93d375c

@nloops($N, i, P,

Julia version: 1.7.0-beta4

The problem might be related to the compiling of multi-loops.

@kshyatt

[inkydragon]

Confirmed on 1.7.0-rc1.

julia> using Random, Profile

julia> for n in 1:24
           print(n)
           t = randn(fill(2, n)...);
           @time permutedims(t, randperm(n));
           print("  ")
           @time permutedims(t, randperm(n));
       end
1  0.071749 seconds (176.90 k allocations: 9.236 MiB, 99.94% compilation time)
    0.000004 seconds (6 allocations: 336 bytes)
2  0.049844 seconds (176.38 k allocations: 9.112 MiB, 99.91% compilation time)
    0.000005 seconds (6 allocations: 368 bytes)
3  0.061044 seconds (201.82 k allocations: 10.444 MiB, 99.93% compilation time)
    0.000004 seconds (6 allocations: 400 bytes)
4  0.075158 seconds (220.40 k allocations: 11.334 MiB, 6.67% gc time, 99.94% compilation time)
    0.000004 seconds (7 allocations: 544 bytes)
5  0.076597 seconds (235.45 k allocations: 12.050 MiB, 99.94% compilation time)
    0.000004 seconds (7 allocations: 672 bytes)
6  0.087803 seconds (266.48 k allocations: 13.564 MiB, 99.95% compilation time)
    0.000004 seconds (7 allocations: 992 bytes)
7  0.104104 seconds (287.17 k allocations: 14.520 MiB, 5.24% gc time, 99.95% compilation time)
    0.000005 seconds (7 allocations: 1.500 KiB)
8  0.113047 seconds (294.48 k allocations: 14.789 MiB, 99.95% compilation time)
    0.000007 seconds (7 allocations: 2.578 KiB)
9  0.138857 seconds (331.71 k allocations: 16.611 MiB, 3.40% gc time, 99.97% compilation time)
    0.000008 seconds (7 allocations: 4.578 KiB)
10  0.167496 seconds (355.46 k allocations: 17.703 MiB, 99.97% compilation time)
    0.000010 seconds (7 allocations: 8.609 KiB)
11  0.160835 seconds (374.31 k allocations: 18.609 MiB, 99.96% compilation time)
    0.000015 seconds (8 allocations: 16.750 KiB)
12  0.180515 seconds (399.03 k allocations: 19.721 MiB, 2.36% gc time, 99.96% compilation time)
    0.000021 seconds (9 allocations: 32.734 KiB)
13  0.193186 seconds (426.11 k allocations: 21.038 MiB, 99.96% compilation time)
    0.000033 seconds (9 allocations: 64.750 KiB)
14  0.226863 seconds (453.70 k allocations: 22.192 MiB, 1.94% gc time, 99.95% compilation time)
    0.000062 seconds (9 allocations: 128.797 KiB)
15  0.265839 seconds (482.30 k allocations: 23.470 MiB, 2.11% gc time, 99.91% compilation time)
    0.000207 seconds (9 allocations: 256.812 KiB)
16  0.316385 seconds (511.99 k allocations: 25.052 MiB, 99.88% compilation time)
    0.000346 seconds (9 allocations: 512.859 KiB)
17  0.444436 seconds (542.79 k allocations: 26.792 MiB, 1.19% gc time, 99.85% compilation time)
    0.000447 seconds (9 allocations: 1.001 MiB)
18  0.637155 seconds (574.68 k allocations: 29.155 MiB, 1.00% gc time, 99.82% compilation time)
    0.001078 seconds (9 allocations: 2.001 MiB)
19  1.036371 seconds (607.65 k allocations: 32.400 MiB, 0.46% gc time, 99.30% compilation time)
    0.001833 seconds (9 allocations: 4.001 MiB)
20  1.772452 seconds (641.40 k allocations: 37.857 MiB, 0.78% gc time, 98.89% compilation time)
    0.005123 seconds (9 allocations: 8.001 MiB)
21  3.309725 seconds (676.79 k allocations: 47.263 MiB, 1.50% gc time, 99.72% compilation time)
    0.049899 seconds (9 allocations: 16.001 MiB, 79.28% gc time)
22  6.159004 seconds (713.72 k allocations: 64.715 MiB, 0.14% gc time, 99.45% compilation time)
    0.031260 seconds (9 allocations: 32.001 MiB)
23 11.872545 seconds (751.27 k allocations: 98.139 MiB, 0.03% gc time, 99.51% compilation time)
    0.077926 seconds (9 allocations: 64.001 MiB, 19.33% gc time)
24 23.564871 seconds (790.00 k allocations: 163.708 MiB, 0.20% gc time, 99.04% compilation time)
    0.131734 seconds (9 allocations: 128.001 MiB, 8.39% gc time)

julia> versioninfo()
Julia Version 1.7.0-rc1
Commit 9eade6195e (2021-09-12 06:45 UTC)
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: Intel(R) Core(TM) i5-9400F CPU @ 2.90GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-12.0.1 (ORCJIT, skylake)

---

Why does it keep slowing down on 1.6.3, as if it hadn't been compiled

julia> using Random, Profile

julia> (n=24; t = randn(fill(2, n)...); @time permutedims(t, randperm(n)));
 20.145386 seconds (470.57 M allocations: 14.418 GiB, 8.69% gc time, 2.94% compilation time)

julia> (n=24; t = randn(fill(2, n)...); @time permutedims(t, randperm(n)));
 20.329427 seconds (469.76 M allocations: 14.375 GiB, 10.13% gc time)

julia> (n=24; t = randn(fill(2, n)...); @time permutedims(t, randperm(n)));
 21.299640 seconds (469.76 M allocations: 14.375 GiB, 10.19% gc time)

julia> versioninfo()
Julia Version 1.6.3
Commit ae8452a9e0 (2021-09-23 17:34 UTC)
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: Intel(R) Core(TM) i5-9400F CPU @ 2.90GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-11.0.1 (ORCJIT, skylake)

[GiggleLiu]

Thanks for confirming it.
The 1.6 issue is about the dynamic allocation in tuple splatting, you can see it from the abnormal allocations. It solved by @kshyatt in PR #40468 so that you do not see this issue in v1.7+.

[GiggleLiu]

Update: This compiling overhead is related to the change of tuple splatting (changed from 16 to 32 in Julia 1.7).

(base) ➜  ~ julia-1.6
               _
   _       _ _(_)_     |  Documentation: https://docs.julialang.org
  (_)     | (_) (_)    |
   _ _   _| |_  __ _   |  Type "?" for help, "]?" for Pkg help.
  | | | | | | |/ _` |  |
  | | |_| | | | (_| |  |  Version 1.6.1 (2021-04-23)
 _/ |\__'_|_|_|\__'_|  |  Official https://julialang.org/ release
|__/                   |

julia> using Random

julia> (n=28; t = randn(fill(1, n)...); @time permutedims(t, randperm(n)));
  0.655323 seconds (908.77 k allocations: 48.766 MiB, 1.91% gc time, 99.98% compilation time)

(base) ➜  ~ julia-1.7 
               _
   _       _ _(_)_     |  Documentation: https://docs.julialang.org
  (_)     | (_) (_)    |
   _ _   _| |_  __ _   |  Type "?" for help, "]?" for Pkg help.
  | | | | | | |/ _` |  |
  | | |_| | | | (_| |  |  Version 1.7.0-beta4 (2021-08-24)
 _/ |\__'_|_|_|\__'_|  |  Official https://julialang.org/ release
|__/                   |

julia> using Random

julia> (n=28; t = randn(fill(1, n)...); @time permutedims(t, randperm(n)));
319.175964 seconds (1.46 M allocations: 66.736 MiB, 0.00% gc time, 100.00% compilation time)

julia> (n=33; t = randn(fill(1, n)...); @time permutedims(t, randperm(n)));
  0.763701 seconds (1.04 M allocations: 49.455 MiB, 0.98% gc time, 99.98% compilation time)

The following MWE may help debugging.

using Base.Cartesian
using Base: size_to_strides, checkdims_perm
using Random

function myperm!(P, B, perm)
    checkdims_perm(P, B, perm)
    native_strides = size_to_strides(1, size(B)...)
    strides_1 = 0
    begin
        strides_2 = native_strides[perm[1]]
        strides_3 = native_strides[perm[2]]
        strides_4 = native_strides[perm[3]]
        strides_5 = native_strides[perm[4]]
        strides_6 = native_strides[perm[5]]
        strides_7 = native_strides[perm[6]]
        strides_8 = native_strides[perm[7]]
        strides_9 = native_strides[perm[8]]
        strides_10 = native_strides[perm[9]]
        strides_11 = native_strides[perm[10]]
        strides_12 = native_strides[perm[11]]
        strides_13 = native_strides[perm[12]]
        strides_14 = native_strides[perm[13]]
        strides_15 = native_strides[perm[14]]
        strides_16 = native_strides[perm[15]]
        strides_17 = native_strides[perm[16]]
        strides_18 = native_strides[perm[17]]
        strides_19 = native_strides[perm[18]]
        strides_20 = native_strides[perm[19]]
        strides_21 = native_strides[perm[20]]
        strides_22 = native_strides[perm[21]]
        strides_23 = native_strides[perm[22]]
        strides_24 = native_strides[perm[23]]
        strides_25 = native_strides[perm[24]]
        strides_26 = native_strides[perm[25]]
    end
    offset = 1 - sum((strides_2, strides_3, strides_4, strides_5, strides_6, strides_7, strides_8, strides_9, strides_10, strides_11, strides_12, strides_13, strides_14, strides_15, strides_16, strides_17, strides_18, strides_19, strides_20, strides_21, strides_22, strides_23, strides_24, strides_25, strides_26))
    ind = 1
    begin
        counts_26 = strides_26
    end
    begin
        for i_25 = Base.axes(P, 25)
            counts_25 = strides_25
            begin
                for i_24 = Base.axes(P, 24)
                    counts_24 = strides_24
                    begin
                        for i_23 = Base.axes(P, 23)
                            counts_23 = strides_23
                            begin
                                for i_22 = Base.axes(P, 22)
                                    counts_22 = strides_22
                                    begin
                                        for i_21 = Base.axes(P, 21)
                                            counts_21 = strides_21
                                            begin
                                                for i_20 = Base.axes(P, 20)
                                                    counts_20 = strides_20
                                                    begin
                                                        for i_19 = Base.axes(P, 19)
                                                            counts_19 = strides_19
                                                            begin
                                                                for i_18 = Base.axes(P, 18)
                                                                    counts_18 = strides_18
                                                                    begin
                                                                        for i_17 = Base.axes(P, 17)
                                                                            counts_17 = strides_17
                                                                            begin
                                                                                for i_16 = Base.axes(P, 16)
                                                                                    counts_16 = strides_16
                                                                                    begin
                                                                                        for i_15 = Base.axes(P, 15)
                                                                                            counts_15 = strides_15
                                                                                            begin
                                                                                                for i_14 = Base.axes(P, 14)
                                                                                                    counts_14 = strides_14
                                                                                                    begin
                                                                                                        for i_13 = Base.axes(P, 13)
                                                                                                            counts_13 = strides_13
                                                                                                            begin
                                                                                                                for i_12 = Base.axes(P, 12)
                                                                                                                    counts_12 = strides_12
                                                                                                                    begin
                                                                                                                        for i_11 = Base.axes(P, 11)
                                                                                                                            counts_11 = strides_11
                                                                                                                            begin
                                                                                                                                for i_10 = Base.axes(P, 10)
                                                                                                                                    counts_10 = strides_10
                                                                                                                                    begin
                                                                                                                                        for i_9 = Base.axes(P, 9)
                                                                                                                                            counts_9 = strides_9
                                                                                                                                            begin
                                                                                                                                                for i_8 = Base.axes(P, 8)
                                                                                                                                                    counts_8 = strides_8
                                                                                                                                                    begin
                                                                                                                                                        for i_7 = Base.axes(P, 7)
                                                                                                                                                            counts_7 = strides_7
                                                                                                                                                            begin
                                                                                                                                                                for i_6 = Base.axes(P, 6)
                                                                                                                                                                    counts_6 = strides_6
                                                                                                                                                                    begin
                                                                                                                                                                        for i_5 = Base.axes(P, 5)
                                                                                                                                                                            counts_5 = strides_5
                                                                                                                                                                            begin
                                                                                                                                                                                for i_4 = Base.axes(P, 4)
                                                                                                                                                                                    counts_4 = strides_4
                                                                                                                                                                                    begin
                                                                                                                                                                                        for i_3 = Base.axes(P, 3)
                                                                                                                                                                                            counts_3 = strides_3
                                                                                                                                                                                            begin
                                                                                                                                                                                                for i_2 = Base.axes(P, 2)
                                                                                                                                                                                                    counts_2 = strides_2
                                                                                                                                                                                                    begin
                                                                                                                                                                                                        for i_1 = Base.axes(P, 1)
                                                                                                                                                                                                            counts_1 = strides_1
                                                                                                                                                                                                            begin
                                                                                                                                                                                                                sumc = sum((counts_2, counts_3, counts_4, counts_5, counts_6, counts_7, counts_8, counts_9, counts_10, counts_11, counts_12, counts_13, counts_14, counts_15, counts_16, counts_17, counts_18, counts_19, counts_20, counts_21, counts_22, counts_23, counts_24, counts_25, counts_26))
                                                                                                                                                                                                                @inbounds P[ind] = B[sumc + offset]
                                                                                                                                                                                                                ind += 1
                                                                                                                                                                                                            end
                                                                                                                                                                                                            counts_2 += strides_2
                                                                                                                                                                                                        end
                                                                                                                                                                                                    end
                                                                                                                                                                                                    counts_3 += strides_3
                                                                                                                                                                                                end
                                                                                                                                                                                            end
                                                                                                                                                                                            counts_4 += strides_4
                                                                                                                                                                                        end
                                                                                                                                                                                    end
                                                                                                                                                                                    counts_5 += strides_5
                                                                                                                                                                                end
                                                                                                                                                                            end
                                                                                                                                                                            counts_6 += strides_6
                                                                                                                                                                        end
                                                                                                                                                                    end
                                                                                                                                                                    counts_7 += strides_7
                                                                                                                                                                end
                                                                                                                                                            end
                                                                                                                                                            counts_8 += strides_8
                                                                                                                                                        end
                                                                                                                                                    end
                                                                                                                                                    counts_9 += strides_9
                                                                                                                                                end
                                                                                                                                            end
                                                                                                                                            counts_10 += strides_10
                                                                                                                                        end
                                                                                                                                    end
                                                                                                                                    counts_11 += strides_11
                                                                                                                                end
                                                                                                                            end
                                                                                                                            counts_12 += strides_12
                                                                                                                        end
                                                                                                                    end
                                                                                                                    counts_13 += strides_13
                                                                                                                end
                                                                                                            end
                                                                                                            counts_14 += strides_14
                                                                                                        end
                                                                                                    end
                                                                                                    counts_15 += strides_15
                                                                                                end
                                                                                            end
                                                                                            counts_16 += strides_16
                                                                                        end
                                                                                    end
                                                                                    counts_17 += strides_17
                                                                                end
                                                                            end
                                                                            counts_18 += strides_18
                                                                        end
                                                                    end
                                                                    counts_19 += strides_19
                                                                end
                                                            end
                                                            counts_20 += strides_20
                                                        end
                                                    end
                                                    counts_21 += strides_21
                                                end
                                            end
                                            counts_22 += strides_22
                                        end
                                    end
                                    counts_23 += strides_23
                                end
                            end
                            counts_24 += strides_24
                        end
                    end
                    counts_25 += strides_25
                end
            end
            counts_26 += strides_26
        end
    end
    return P
end

(n=25; p=randn(fill(1, n)...); t = randn(fill(1, n)...); @time myperm!(p, t, randperm(n)));

The sum statement in the middle matters! I tried to use Vector instead of Tuple, it compiles fast.

[inkydragon]

Related issue #19158

[timholy]

It's pretty much all LLVM. I suspect it might be good to cut off the @generated specialization of permutedims! at some upper dimensionality limit and provide a fallback.

[GiggleLiu]

Thanks for clarifying, let me file an issue in LLVM's git repo later. I just submitted a PR to fix the permutedims issue.
(I do not know much about llvm, so I just try and error, hope someone can give it a review)