Runtime dispatch overhead for small tensor contractions using the `StridedBLAS` backend

[leburgel]

When profiling some code involving many contractions with small tensors I noticed that there is a lot of overhead due to runtime dispatch and resulting allocations and garbage collection when using the StridedBLAS backend. I've figured out some of it but I thought it would be good to report here to see if something more can be done.

I ran a profile for a dummy example contraction of small complex tensors which can sort of reproduce the issue:

using TensorOperations

T = ComplexF64
L = randn(T, 8, 8, 7)
R = randn(T, 8, 8, 7)
O = randn(T, 4, 4, 7, 7)

function local_update!(psi::Array{T,3})::Array{T,3} where {T}
    @tensor begin
        psi[-1, -2, -3] =
            psi[1, 3, 5] *
            L[-1, 1, 2] *
            O[-2, 3, 2, 4] *
            R[-3, 5, 4]
    end
    return psi
end

psi0 = randn(T, 8, 4, 8)
@profview begin
    for _ in 1:100000
        local_update!(psi0)
    end
end

which gives:

So a lot of runtime dispatch overhead in TensorOperations.tensoradd! and Strided._mapreduce_order!. The effect becomes negligible for large tensor dimensions, but it turns out to be a real pain if there are a lot of these small contractions being performed.

For TensorOperations.tensoradd!, I managed to track the problem to an ambiguity caused by patterns like flag2op(conjA)(A), where the return type at compile time can be a Union of two StridedView concrete types with typeof(identity) and typeof(conj) as their op field types respectively. This leads to an issue here:
https://github.com/Jutho/TensorOperations.jl/blob/c1e37eca9c2d1ab468ddd1beb2ffb9eb0993bb4b/src/implementation/strided.jl#L99-L103
where the last argument in the call to Strided._mapreducedim! is a Tuple with mixed concrete and abstract (the union described above) types in its type parameters, which seems to mess things up.

At that level I managed to fix things by just splitting into two branches to de-confuse the compiler

opA = flag2op(conjA)
if opA isa typeof(identity)
    A′ = permutedims(identity(A), linearize(pA))
    op1 = Base.Fix2(scale, α)
    op2 = Base.Fix2(scale, β)
    Strided._mapreducedim!(op1, +, op2, size(C), (C, A′))
    return C
elseif opA isa typeof(conj)
    A′ = permutedims(conj(A), linearize(pA))
    op1 = Base.Fix2(scale, α)
    op2 = Base.Fix2(scale, β)
    Strided._mapreducedim!(op1, +, op2, size(C), (C, A′))
    return C
end

which gets rid of the runtime dispatch in tensoradd! and already makes a big difference for small contractions.

I haven't descended all the way down into Strided._mapreduce_dim!, so I don't know what the issue is there.

So in the end my questions are:

• Is there a cleaner way to deal with the conj flags and their effect on StridedViews that could avoid these kinds of type ambiguities?
• Do you have any idea if the runtime dispatch in Strided._mapreduce_dim! could be prevented by changing the StridedBLAS backend implementation here, and if not, if it could be avoided altogether by a reasonable change to Strided.jl?

[Jutho]

So I keep getting sucked into the lowest levels of the code, never having time to properly study PEPSKit.jl 😄 .

[Jutho]

Certainly, for conjA, the hope was that constant propagation would resolve the ambiguity, but apparently it does not. Maybe we need some more aggressive constant propagation.

[Jutho]

I assume the actual coding pattern does not include a function local_update! that captures global variables, right?

[Jutho]

Ok, it seems adding @constprop :aggressive before every tensorcontract! and tensoradd! definition solves the problem, but it is of course not particularly nice.

An alternative would be to explicitly encode the conjugation flag in the type domain, e.g. by having a parametric singleton struct ConjugationFlag{true}() and ConjugationFlag{false}(). I am interested in hearing the opinion of @lkdvos .

With Cthulhu.jl, I couldn't spot an additional ambiguity/dynamic dispatch in mapreduce_order!, so I am not sure if those two are related or if that is yet another issue.

[lkdvos]

We are already doing this for the istemp flag with a Val{true}, so in principle I have nothing against this change. It's a bit unfortunate that this has to be breaking, so we might want to investigate a bit more...
Do you know if we can do const propagation for single arguments?

[Jutho]

No that doesn't work. I also tried only annotating the flag2op function, so that also doesn't work. Maybe there is some in between balance where you only need to annotate some of the tensor operation definitions, but that can only be found by trial and error I think.

Another nonbreaking solution is to manually split the code based on the flag values, but there are quite a few instances where flag2op is called, and in contract there are two of those, so then one has 4 different cases.

[lkdvos]

To be fair, if just splitting the flag2op into if statements works, that might be the cleanest. I do see how that function is inherently type unstable, so either this information needs to be type domain or we do a manual Union-splitting type construction, which I had hoped the compiler would have figured out but apparently it doesn't.

[Jutho]

Ok, I'll take a stab at it asap.

[leburgel]

I assume the actual coding pattern does not include a function local_update! that captures global variables, right?

No, in the realistic setting everything is passed through locally and all of the concrete array types are properly inferred in the call to tensorcontract!. I wasn't really paying attention with to the globals in the example, but it runs into the same problem lower down.

[leburgel]

With Cthulhu.jl, I couldn't spot an additional ambiguity/dynamic dispatch in mapreduce_order!, so I am not sure if those two are related or if that is yet another issue.

After the branch splitting with an if statement the ambiguity in tensoradd! is completely removed, but the runtime dispatch in mapreduce_order! remains. So that might be another issue, but I couldn't really figure it out.

[Jutho]

I think the dynamic dispatch in the mapreduce lowering chain comes from the fact that in _mapreduce_fuse! and _mapreduce_block!, the compiler doesn't specify on the function arguments f, op and initop. I think this is a known behaviour of the Julia compiler to reduce compile times: if you want to specify function arguments on the specific function, you should actually make it a type parameter in your method. In this case, the non-specialisation is in principle fine, as this has no effect on any of the types of the variables in those function bodies, so everything is perfectly type stable.

But then in finally selecting _mapreduce_kernel, that being a generated function, it does again specialise on f etc, and so there is some dynamic dispatch going on. Not sure what the impact of that is.

[Jutho]

Ok, so I am not sure how reliable this is, this is from simply timing the first call (so timing compilation) and then benchmarking the call on my computer, both including the fix from the PR, but the first one with the vanilla version of Strided.jl:

julia> @time local_update!(psi0);
 19.850534 seconds (34.72 M allocations: 2.236 GiB, 4.29% gc time, 100.00% compilation time)

julia> @benchmark local_update!($psi0)
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):  30.459 μs … 123.357 ms  ┊ GC (min … max):  0.00% … 99.86%
 Time  (median):     41.032 μs               ┊ GC (median):     0.00%
 Time  (mean ± σ):   71.495 μs ±   1.255 ms  ┊ GC (mean ± σ):  32.50% ±  4.91%

  ▄▇▄▂▂▅█▅▁▂▂▄▃                     ▁▂▂▂▃▂▂▂▂▂▁                ▂
  ██████████████▇▆▇▆▇▇▆▆▇▆▄▄▆▅▄▆▅▆▇██████████████▇▇▇▆▆▆▅▄▄▅▁▃▅ █
  30.5 μs       Histogram: log(frequency) by time       130 μs <

 Memory estimate: 135.19 KiB, allocs estimate: 59.

and then with a modified Strided.jl where f::F1, op::F2 and initop::F3 are now explicit type parameters of the relevant methods:

julia> @time local_update!(psi0);
 36.712724 seconds (27.60 M allocations: 1.740 GiB, 1.54% gc time, 100.00% compilation time)

julia> @benchmark local_update!($psi0)
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):  26.941 μs … 122.038 ms  ┊ GC (min … max):  0.00% … 99.86%
 Time  (median):     38.208 μs               ┊ GC (median):     0.00%
 Time  (mean ± σ):   69.025 μs ±   1.243 ms  ┊ GC (mean ± σ):  33.64% ±  4.82%

  ▃▆▇▃▂▄█▇▃▁▂▄▂                     ▁▁▂▂▂▂▂▂▂▁▁▁               ▂
  ██████████████▇▇█▇▅▆▆▇▆▆▅▅▄▃▂▄▄▅▆█████████████████▇▇▆▅▆▄▅▅▄▄ █
  26.9 μs       Histogram: log(frequency) by time       127 μs <

 Memory estimate: 131.98 KiB, allocs estimate: 15.

So allocations and runtime are a little bit lower, but compilation time is up by a factor of almost two.

[leburgel]

Thanks for figuring out the mapreduce issue! It is a small effect and a lot more compilation time, so I'm not sure if you think it's worth actually making the change in Strided.jl?

I can anyway make the change locally now that I understand what's going on, so as far as I'm concerned my problems are solved :)

[Jutho]

Ok, let us know if you encounter cases where specialising the methods in Strided do actually lead to a noticeable speedup. I will close this and merge the PR.