Incremental work on Enzyme support

Project.toml

@@ -23,6 +23,8 @@ VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
@@ -31,6 +33,8 @@ cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 TensorKitAMDGPUExt = "AMDGPU"
 TensorKitCUDAExt = ["CUDA", "cuTENSOR"]
 TensorKitChainRulesCoreExt = "ChainRulesCore"
+TensorKitEnzymeExt = "Enzyme"
+TensorKitEnzymeTestUtilsExt = "EnzymeTestUtils"
 TensorKitFiniteDifferencesExt = "FiniteDifferences"
 TensorKitMooncakeExt = "Mooncake"
 
@@ -43,18 +47,20 @@ AMDGPU = "2"
 CUDA = "6"
 ChainRulesCore = "1"
 Dictionaries = "0.4"
+Enzyme = "0.13.146"
+EnzymeTestUtils = "0.2.7"
 FiniteDifferences = "0.12"
 LRUCache = "1.0.2"
 LinearAlgebra = "1"
-MatrixAlgebraKit = "0.6.7"
+MatrixAlgebraKit = "0.6.8"
 Mooncake = "0.5.27"
 OhMyThreads = "0.8.0"
 Printf = "1"
 Random = "1"
 ScopedValues = "1.3.0"
 Strided = "2"
 TensorKitSectors = "0.3.7"
-TensorOperations = "5.5"
+TensorOperations = "5.5.2"
 TupleTools = "1.5"
 VectorInterface = "0.4.8, 0.5, 0.6"
 cuTENSOR = "6"

ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl

@@ -0,0 +1,20 @@
+module TensorKitEnzymeExt
+
+using Enzyme
+using TensorKit
+import TensorKit as TK
+using VectorInterface
+using TensorOperations: TensorOperations, IndexTuple, Index2Tuple, linearize
+import TensorOperations as TO
+using MatrixAlgebraKit
+using TupleTools
+using Random: AbstractRNG
+
+include("utility.jl")
+include("linalg.jl")
+include("tensoroperations.jl")
+include("factorizations.jl")
+include("indexmanipulations.jl")
+#include("planaroperations.jl")
+
+end

ext/TensorKitEnzymeExt/factorizations.jl

@@ -0,0 +1,275 @@
+# need these due to Enzyme choking on blocks
+
+for f in (:project_hermitian, :project_antihermitian)
+    f! = Symbol(f, :!)
+    @eval begin
+        function EnzymeRules.augmented_primal(
+                config::EnzymeRules.RevConfigWidth{1},
+                func::Const{typeof($f!)},
+                ::Type{RT},
+                A::Annotation{<:AbstractTensorMap},
+                arg::Annotation{<:AbstractTensorMap},
+                alg::Const,
+            ) where {RT}
+            $f!(A.val, arg.val, alg.val)
+            primal = EnzymeRules.needs_primal(config) ? arg.val : nothing
+            shadow = EnzymeRules.needs_shadow(config) ? arg.dval : nothing
+            cache = nothing
+            return EnzymeRules.AugmentedReturn(primal, shadow, cache)
+        end
+        function EnzymeRules.reverse(
+                config::EnzymeRules.RevConfigWidth{1},
+                func::Const{typeof($f!)},
+                ::Type{RT},
+                cache,
+                A::Annotation{<:AbstractTensorMap},
+                arg::Annotation{<:AbstractTensorMap},
+                alg::Const,
+            ) where {RT}
+            if !isa(A, Const) && !isa(arg, Const)
+                $f!(arg.dval, arg.dval, alg.val)
+                if A.dval !== arg.dval
+                    A.dval .+= arg.dval
+                    make_zero!(arg.dval)
+                end
+            end
+            return (nothing, nothing, nothing)
+        end
+        function EnzymeRules.augmented_primal(
+                config::EnzymeRules.RevConfigWidth{1},
+                func::Const{typeof($f)},
+                ::Type{RT},
+                A::Annotation{<:AbstractTensorMap},
+                alg::Const,
+            ) where {RT}
+            ret = $f(A.val, alg.val)
+            dret = make_zero(ret)
+            primal = EnzymeRules.needs_primal(config) ? ret : nothing
+            shadow = EnzymeRules.needs_shadow(config) ? dret : nothing
+            cache = dret
+            return EnzymeRules.AugmentedReturn(primal, shadow, cache)
+        end
+        function EnzymeRules.reverse(
+                config::EnzymeRules.RevConfigWidth{1},
+                func::Const{typeof($f)},
+                ::Type{RT},
+                cache,
+                A::Annotation{<:AbstractTensorMap},
+                alg::Const,
+            ) where {RT}
+            dret = cache
+            if !isa(A, Const)
+                $f!(dret, dret, alg.val)
+                add!(A.dval, dret)
+            end
+            make_zero!(dret)
+            return (nothing, nothing)
+        end
+    end
+end
+
+for (f, pb) in (
+        (:eig_full, :(MatrixAlgebraKit.eig_pullback!)),
+        (:eigh_full, :(MatrixAlgebraKit.eigh_pullback!)),
+        (:lq_compact, :(MatrixAlgebraKit.lq_pullback!)),
+        (:qr_compact, :(MatrixAlgebraKit.qr_pullback!)),
+        (:lq_full, :(MatrixAlgebraKit.lq_pullback!)),
+        (:qr_full, :(MatrixAlgebraKit.qr_pullback!)),
+        (:lq_null, :(MatrixAlgebraKit.lq_null_pullback!)),
+        (:qr_null, :(MatrixAlgebraKit.qr_null_pullback!)),
+    )
+    @eval begin
+        function EnzymeRules.augmented_primal(
+                config::EnzymeRules.RevConfigWidth{1},
+                func::Const{typeof($f)},
+                ::Type{RT},
+                A::Annotation{<:AbstractTensorMap},
+                alg::Const,
+            ) where {RT}
+            ret = $f(A.val, alg.val)
+            primal = EnzymeRules.needs_primal(config) ? ret : nothing
+            shadow = EnzymeRules.needs_shadow(config) ? make_zero(ret) : nothing
+            cache = (ret, shadow)
+            return EnzymeRules.AugmentedReturn(primal, shadow, cache)
+        end
+        function EnzymeRules.reverse(
+                config::EnzymeRules.RevConfigWidth{1},
+                func::Const{typeof($f)},
+                ::Type{RT},
+                cache,
+                A::Annotation{<:AbstractTensorMap},
+                alg::Const,
+            ) where {RT}
+            !isa(A, Const) && $pb(A.dval, A.val, cache...)
+            return (nothing, nothing)
+        end
+    end
+end
+
+for (f, f_full, pb) in (
+        (:eig_vals, :eig_full, :(MatrixAlgebraKit.eig_vals_pullback!)),
+        (:eigh_vals, :eigh_full, :(MatrixAlgebraKit.eigh_vals_pullback!)),
+    )
+    @eval begin
+        function EnzymeRules.augmented_primal(
+                config::EnzymeRules.RevConfigWidth{1},
+                func::Const{typeof($f)},
+                ::Type{RT},
+                A::Annotation{<:AbstractTensorMap},
+                alg::Const,
+            ) where {RT}
+            ret_full = $f_full(A.val, alg.val)
+            ret = diagview(ret_full[1])
+            primal = EnzymeRules.needs_primal(config) ? ret : nothing
+            shadow = EnzymeRules.needs_shadow(config) ? make_zero(ret) : nothing
+            cache = (ret, shadow, ret_full[2])
+            return EnzymeRules.AugmentedReturn(primal, shadow, cache)
+        end
+        function EnzymeRules.reverse(
+                config::EnzymeRules.RevConfigWidth{1},
+                func::Const{typeof($f)},
+                ::Type{RT},
+                cache,
+                A::Annotation{<:AbstractTensorMap},
+                alg::Const,
+            ) where {RT}
+            D, dD, V = cache
+            !isa(A, Const) && $pb(A.dval, A.val, (DiagonalTensorMap(D), V), dD)
+            return (nothing, nothing)
+        end
+    end
+end
+
+for f in (:svd_compact, :svd_full)
+    @eval begin
+        function EnzymeRules.augmented_primal(
+                config::EnzymeRules.RevConfigWidth{1},
+                func::Const{typeof($f)},
+                ::Type{RT},
+                A::Annotation{<:AbstractTensorMap},
+                alg::Const,
+            ) where {RT}
+            USVᴴ = $f(A.val, alg.val)
+            primal = EnzymeRules.needs_primal(config) ? USVᴴ : nothing
+            shadow = EnzymeRules.needs_shadow(config) ? make_zero(USVᴴ) : nothing
+            cache = (USVᴴ, shadow)
+            return EnzymeRules.AugmentedReturn(primal, shadow, cache)
+        end
+        function EnzymeRules.reverse(
+                config::EnzymeRules.RevConfigWidth{1},
+                func::Const{typeof($f)},
+                ::Type{RT},
+                cache,
+                A::Annotation{<:AbstractTensorMap},
+                alg::Const,
+            ) where {RT}
+            !isa(A, Const) && MatrixAlgebraKit.svd_pullback!(A.dval, A.val, cache...)
+            return (nothing, nothing)
+        end
+    end
+
+    # mutating version is not guaranteed to actually mutate
+    # so we can simply use the non-mutating version instead
+    f! = Symbol(f, :!)
+    #=@eval begin
+        function EnzymeRules.augmented_primal(
+            config::EnzymeRules.RevConfigWidth{1},
+            func::Const{typeof($f!)},
+            ::Type{RT},
+            A::Annotation{<:AbstractTensorMap},
+            USVᴴ::Annotation,
+            alg::Const,
+        ) where {RT}
+            EnzymeRules.augmented_primal(func, RT, A, alg)
+        end
+        function EnzymeRules.reverse(
+            config::EnzymeRules.RevConfigWidth{1},
+            func::Const{typeof($f!)},
+            ::Type{RT},
+            cache,
+            A::Annotation{<:AbstractTensorMap},
+            USVᴴ::Annotation,
+            alg::Const,
+        ) where {RT}
+            EnzymeRules.reverse(func, RT, A, alg)
+        end
+    end=# #hmmmm
+end
+
+function EnzymeRules.augmented_primal(
+        config::EnzymeRules.RevConfigWidth{1},
+        func::Const{typeof(svd_trunc_no_error)},
+        ::Type{RT},
+        A::Annotation{<:AbstractTensorMap},
+        alg::Const,
+    ) where {RT}
+    USVᴴ = svd_compact(A.val, alg.val.alg)
+    USVᴴtrunc, ind = MatrixAlgebraKit.truncate(svd_trunc!, USVᴴ, alg.val.trunc)
+    dUSVᴴtrunc = make_zero(USVᴴtrunc)
+    cache = (USVᴴ, USVᴴtrunc, dUSVᴴtrunc, ind)
+    return EnzymeRules.AugmentedReturn(USVᴴtrunc, dUSVᴴtrunc, cache)
+end
+function EnzymeRules.reverse(
+        config::EnzymeRules.RevConfigWidth{1},
+        func::Const{typeof(svd_trunc_no_error)},
+        ::Type{RT},
+        cache,
+        A::Annotation{<:AbstractTensorMap},
+        alg::Const,
+    ) where {RT}
+    USVᴴ, USVᴴtrunc, dUSVᴴtrunc, ind = cache
+    MatrixAlgebraKit.svd_pullback!(A.dval, A.val, USVᴴ, dUSVᴴtrunc, ind)
+    return (nothing, nothing)
+end
+
+function EnzymeRules.augmented_primal(
+        config::EnzymeRules.RevConfigWidth{1},
+        func::Const{typeof(eig_trunc_no_error)},
+        ::Type{RT},
+        A::Annotation{<:AbstractTensorMap},
+        alg::Const,
+    ) where {RT}
+    DV = eig_full(A.val, alg.val.alg)
+    DVtrunc, ind = MatrixAlgebraKit.truncate(eig_trunc!, DV, alg.val.trunc)
+    dDVtrunc = make_zero(DVtrunc)
+    cache = (DV, DVtrunc, dDVtrunc, ind)
+    return EnzymeRules.AugmentedReturn(DVtrunc, dDVtrunc, cache)
+end
+function EnzymeRules.reverse(
+        config::EnzymeRules.RevConfigWidth{1},
+        func::Const{typeof(eig_trunc_no_error)},
+        ::Type{RT},
+        cache,
+        A::Annotation{<:AbstractTensorMap},
+        alg::Const,
+    ) where {RT}
+    DV, DVtrunc, dDVtrunc, ind = cache
+    MatrixAlgebraKit.eig_pullback!(A.dval, A.val, DV, dDVtrunc, ind)
+    return (nothing, nothing)
+end
+
+function EnzymeRules.augmented_primal(
+        config::EnzymeRules.RevConfigWidth{1},
+        func::Const{typeof(eigh_trunc_no_error)},
+        ::Type{RT},
+        A::Annotation{<:AbstractTensorMap},
+        alg::Const,
+    ) where {RT}
+    DV = eigh_full(A.val, alg.val.alg)
+    DVtrunc, ind = MatrixAlgebraKit.truncate(eigh_trunc!, DV, alg.val.trunc)
+    dDVtrunc = make_zero(DVtrunc)
+    cache = (DV, DVtrunc, dDVtrunc, ind)
+    return EnzymeRules.AugmentedReturn(DVtrunc, dDVtrunc, cache)
+end
+function EnzymeRules.reverse(
+        config::EnzymeRules.RevConfigWidth{1},
+        func::Const{typeof(eigh_trunc_no_error)},
+        ::Type{RT},
+        cache,
+        A::Annotation{<:AbstractTensorMap},
+        alg::Const,
+    ) where {RT}
+    DV, DVtrunc, dDVtrunc, ind = cache
+    MatrixAlgebraKit.eigh_pullback!(A.dval, A.val, DV, dDVtrunc, ind)
+    return (nothing, nothing)
+end