feat: lower svd

Author: avik-pal

src/Ops.jl

@@ -3312,6 +3312,39 @@ Compute the row maximum pivoted LU factorization of `x` and return the factors `
     return (res, ipiv, perm, info)
 end
 
+@noinline function svd(
+    x::TracedRArray{T,N};
+    full::Bool=false,
+    location=mlir_stacktrace("svd", @__FILE__, @__LINE__),
+) where {T,N}
+    @assert N >= 2
+
+    batch_sizes = size(x)[1:(end - 2)]
+    m, n = size(x)[(end - 1):end]
+    r = min(m, n)
+
+    U_size = (batch_sizes..., m, full ? m : r)
+    S_size = (batch_sizes..., r)
+    Vt_size = (batch_sizes..., full ? n : r, n)
+    info_size = batch_sizes
+
+    svd_op = enzymexla.linalg_svd(
+        x.mlir_data;
+        U=mlir_type(TracedRArray{T,N}, U_size),
+        S=mlir_type(TracedRArray{T,N - 1}, S_size),
+        Vt=mlir_type(TracedRArray{T,N}, Vt_size),
+        info=mlir_type(TracedRArray{Int32,N - 2}, info_size),
+        full=full,
+        location,
+    )
+
+    U = TracedRArray{T,N}((), MLIR.IR.result(svd_op, 1), U_size)
+    S = TracedRArray{T,N - 1}((), MLIR.IR.result(svd_op, 2), S_size)
+    Vt = TracedRArray{T,N}((), MLIR.IR.result(svd_op, 3), Vt_size)
+
+    return U, S, Vt
+end
+
 @noinline function reduce_window(
     f::F,
     inputs::Vector{TracedRArray{T,N}},

src/stdlibs/LinearAlgebra.jl

@@ -15,7 +15,7 @@ using LinearAlgebra: LinearAlgebra, BLAS
 using LinearAlgebra: Adjoint, Transpose, Factorization, RowMaximum, NoPivot
 using LinearAlgebra: SymTridiagonal, Symmetric, Bidiagonal, Diagonal, Tridiagonal
 using LinearAlgebra: LowerTriangular, UnitLowerTriangular, UpperTriangular
-using LinearAlgebra: diag, diagm, ldiv!, inv, rmul!
+using LinearAlgebra: diag, diagm, ldiv!, rmul!, normalize, norm
 using Libdl: Libdl
 using GPUArraysCore: @allowscalar
 
@@ -330,7 +330,7 @@ end
 # LinearAlgebra defines norm with some conditionals which cannot be traced directly
 function LinearAlgebra.norm(x::TracedRArray{T,N}, p::Real=2) where {T,N}
     isinf(p) && return maximum(abs, x)
-    return mapreduce(Base.Fix2(^, p), +, x)^(1 / p)
+    return mapreduce(Base.Fix2(^, p), +, x)^(T(1 / p))
 end
 
 function LinearAlgebra._diagm(shape, kv::Pair{<:Integer,<:AnyTracedRVector}...)

src/stdlibs/factorization/SVD.jl

@@ -1,4 +1,4 @@
-struct GeneralizedSVD{T,Tr,M<:AbstractArray{T},C<:AbstractArray{T}} <: Factorization{T}
+struct GeneralizedSVD{T,Tr,M<:AbstractArray,C<:AbstractArray} <: Factorization{T}
     U::M
     S::C
     Vt::M
@@ -9,6 +9,12 @@ struct GeneralizedSVD{T,Tr,M<:AbstractArray{T},C<:AbstractArray{T}} <: Factoriza
     end
 end
 
+function GeneralizedSVD(
+    U::AbstractArray{T}, S::AbstractArray{Tr}, Vt::AbstractArray{T}
+) where {T,Tr}
+    return GeneralizedSVD{T,Tr,typeof(U),typeof(S)}(U, S, Vt)
+end
+
 function overloaded_svd(A::AbstractArray; kwargs...)
     return overloaded_svd(Reactant.promote_to(TracedRArray, A); kwargs...)
 end
@@ -17,17 +23,55 @@ function overloaded_svd(
     A::AnyTracedRArray{T,N}; full::Bool=false, algorithm=nothing
 ) where {T,N}
     # TODO: don't ignore the algorithm kwarg
-    return error("TODO: Not implemented yet")
+    U, S, Vt = @opcall svd(A; full)
+    return GeneralizedSVD(U, S, Vt)
 end
 
 function overloaded_svd(
     A::AnyTracedRVector{T}; full::Bool=false, algorithm=nothing
 ) where {T}
     # TODO: don't ignore the algorithm kwarg
-    m = length(A)
-    normA = LinearAlgebra.norm(A)
+    normA = Reactant.call_with_reactant(LinearAlgebra.norm, A)
+    U, S, Vt = if full
+        ReactantCore.traced_if(
+            iszero(normA), zeronorm_vector_svd_full, vector_svd_full, (A, normA)
+        )
+    else
+        ReactantCore.traced_if(iszero(normA), zeronorm_vector_svd, vector_svd, (A, normA))
+    end
+    return GeneralizedSVD(U, S, Vt)
+end
+
+function zeronorm_vector_svd(A::AbstractVector{T}, normA) where {T}
+    return zeronorm_vector_svd(A, false, normA)
+end
+function zeronorm_vector_svd_full(A::AbstractVector{T}, normA) where {T}
+    return zeronorm_vector_svd(A, true, normA)
+end
+
+function zeronorm_vector_svd(A::AbstractVector{T}, full::Bool, normA) where {T}
+    U = Reactant.promote_to(
+        TracedRArray,
+        Matrix{Reactant.unwrapped_eltype(T)}(
+            LinearAlgebra.I, length(A), full ? length(A) : 1
+        ),
+    )
+    return U, fill(normA, 1), ones(T, 1, 1)
+end
 
-    return error("TODO: Not implemented yet")
+vector_svd(A::AbstractVector{T}, normA) where {T} = vector_svd(A, false, normA)
+function vector_svd_full(A::AbstractVector{T}, normA) where {T}
+    return vector_svd(A, true, normA)
+end
+
+function vector_svd(A::AbstractVector{T}, full::Bool, normA) where {T}
+    if !full
+        U = materialize_traced_array(reshape(normalize(A), length(A), 1))
+        return U, fill(normA, 1), ones(T, 1, 1)
+    end
+    return @opcall svd(
+        materialize_traced_array(reshape(normalize(A), length(A), 1)); full
+    )
 end
 
 # TODO: compute svdvals without computing the full svd. In principle we should