Performance issue in Zygote AD for TensorOperations.@tensor on CUDA: costly CPU copies in rrule

[XingyuZhang2018]

I am benchmarking reverse-mode AD with Zygote on CUDA tensors and observe that contractions written using TensorOperations.@tensor are significantly slower than equivalent implementations using OMEinsum or explicit reshape + matmul.

The slowdown seems to originate from the AD rule (_rrule_tensorcontract!) in TensorOperationsChainRulesCoreExt, which appears to introduce expensive CPU copies during the backward pass in the section.

Here is the minimal working example:

using TensorOperations
using cuTENSOR
using OMEinsum
using Zygote
using Test
using LinearAlgebra
using BenchmarkTools

@testset "ad" begin
    D = 2^5
    A = [CUDA.rand(ComplexF64, D,D,D) for _ in 1:10]
    B = [CUDA.rand(ComplexF64, D,D) for _ in 1:10]

    function foo1(A)
        C = Zygote.Buffer(A)
        for i in 1:length(A)
            @tensor C[i][1,2,4] := A[i][1,2,3] * B[i][3,4]
        end
        return real(dot(C, C))
    end

    function foo2(A)
        C = Zygote.Buffer(A)
        for i in 1:length(A)
            C[i] = ein"abc,cd->abd"(A[i], B[i])
        end
        return real(dot(C, C))
    end

    function foo3(A)
        C = Zygote.Buffer(A)
        for i in 1:length(A)
            C[i] =  reshape(reshape(A[i], D^2, D) * B[i], D, D, D)
        end
        return real(dot(C, C))
    end

    @btime CUDA.@sync Zygote.gradient($foo1, $A)
    @btime CUDA.@sync Zygote.gradient($foo2, $A)
    @btime CUDA.@sync Zygote.gradient($foo3, $A)
end


foo1 (@tensor):     370.627 ms (41219 allocations: 947.14 KiB)
foo2 (OMEinsum):     91.993 ms (25300 allocations: 749.11 KiB)
foo3 (reshape+*):  107.602 ms (14940 allocations: 400.62 KiB)

[Jutho]

Thanks for figuring this out, this is a great catch. Is it the tensorscalar call that is generating the data transfer to the CPU? That should only be a single scalar, and furthermore it is in a @thunk closure, so this shouldn't be called in the first place. Very weird.

[lkdvos]

I'm not sure what the current state is of the support for actually using the @thunk functionality in Zygote, I know there was a recent(-ish) update that started adding this in, but in general this system never really got fully worked out and most thunks were never really used

[XingyuZhang2018]

From my profiling results, even transferring a 0-dimensional CuArray to a scalar incurs significant latency. In the profiler output above (highlighted by arrows), these costs all come from calls to synchronize.

# wait for the current owner of memory to finish processing
function synchronize(managed::Managed)
  synchronize(managed.stream)
  managed.dirty = false
end

In contrast, OMEinsum.jl never converts 0-dimensional CuArrays to scalars during the computation; instead, this conversion can be done manually at the very end. This avoids repeated synchronization and significantly reduces overhead.

[lkdvos]

One somewhat easy way around this issue might be to specialize when Alfa and beta are One and Zero, or booleans, since we then know ZeroTangent is the correct tangent, and we can avoid computing this altogether. It might be worth investigating separately what to do with the thunk not being used, but this could cover quite a few cases already. If I find some time I'll try to put something together this evening