feat: update to allow algorithms

Author: avik-pal

deps/ReactantExtra/WORKSPACE

@@ -4,7 +4,7 @@ NSYNC_COMMIT = "82b118aa7ace3132e517e2c467f8732978cf4023"
 
 NSYNC_SHA256 = ""
 
-ENZYMEXLA_COMMIT = "719107a26418020ecee54c18cd2a299df7004786"
+ENZYMEXLA_COMMIT = "634866abf929fce4cb2d1ac25e76dca9b7a33f3c"
 
 ENZYMEXLA_SHA256 = ""
 

src/Compiler.jl

@@ -1722,6 +1722,8 @@ function compile_mlir!(
 
     blas_int_width = sizeof(BlasInt) * 8
     lower_enzymexla_linalg_pass = "lower-enzymexla-linalg{backend=$backend \
+                                   blas_int_width=$blas_int_width},\
+                                   lower-enzymexla-lapack{backend=$backend \
                                    blas_int_width=$blas_int_width}"
 
     legalize_chlo_to_stablehlo =

src/Ops.jl

@@ -3316,6 +3316,7 @@ end
     x::TracedRArray{T,N},
     ::Type{iT}=Int32;
     full::Bool=false,
+    algorithm::String="DEFAULT",
     location=mlir_stacktrace("svd", @__FILE__, @__LINE__),
 ) where {T,iT,N}
     @assert N >= 2
@@ -3329,13 +3330,26 @@ end
     Vt_size = (batch_sizes..., full ? n : r, n)
     info_size = batch_sizes
 
+    if algorithm == "DEFAULT"
+        algint = 0
+    elseif algorithm == "QRIteration"
+        algint = 1
+    elseif algorithm == "DivideAndConquer"
+        algint = 2
+    elseif algorithm == "Jacobi"
+        algint = 3
+    else
+        error("Unsupported SVD algorithm: $algorithm")
+    end
+
     svd_op = enzymexla.linalg_svd(
         x.mlir_data;
         U=mlir_type(TracedRArray{T,N}, U_size),
         S=mlir_type(TracedRArray{Base.real(T),N - 1}, S_size),
         Vt=mlir_type(TracedRArray{T,N}, Vt_size),
         info=mlir_type(TracedRArray{iT,N - 2}, info_size),
         full=full,
+        algorithm=MLIR.API.enzymexlaSVDAlgorithmAttrGet(MLIR.IR.context(), algint),
         location,
     )
 

src/mlir/Dialects/EnzymeXLA.jl

@@ -346,14 +346,13 @@ end
 """
 `lapack_geqrf`
 
-This operation computes the QR factorization of a matrix using Householder 
-reflections. Mathematically, it decomposes A into the product of an 
+This operation computes the QR factorization of a matrix using Householder
+reflections. Mathematically, it decomposes A into the product of an
 orthogonal matri x Q and an upper triangular matrix R,
-  such that A = QR.
+such that A = QR.
 
-                This operation is modeled after
-                    LAPACK\'s *GEQRF routines, which returns the  result in
-                        the QR packed format.
+This operation is modeled after LAPACK\'s *GEQRF routines, which returns the
+result in the QR packed format.
 """
 function lapack_geqrf(
     input::Value; output::IR.Type, tau::IR.Type, info::IR.Type, location=Location()
@@ -379,11 +378,11 @@ end
 """
 `lapack_geqrt`
 
-This operation computes the QR factorization of a matrix using Householder 
-reflections. Mathematically, it decomposes A into the product of an 
+This operation computes the QR factorization of a matrix using Householder
+reflections. Mathematically, it decomposes A into the product of an
 orthogonal matrix Q and an upper triangular matrix R, such that A = QR.
 
-This operation is modeled after LAPACK\'s *GEQRT routines, which returns the 
+This operation is modeled after LAPACK\'s *GEQRT routines, which returns the
 result in the QR CompactWY format.
 """
 function lapack_geqrt(
@@ -413,6 +412,90 @@ function lapack_geqrt(
     )
 end
 
+function lapack_gesdd(
+    input::Value;
+    U::IR.Type,
+    S::IR.Type,
+    Vt::IR.Type,
+    info::IR.Type,
+    full=nothing,
+    location=Location(),
+)
+    op_ty_results = IR.Type[U, S, Vt, info]
+    operands = Value[input,]
+    owned_regions = Region[]
+    successors = Block[]
+    attributes = NamedAttribute[]
+    !isnothing(full) && push!(attributes, namedattribute("full", full))
+
+    return create_operation(
+        "enzymexla.lapack.gesdd",
+        location;
+        operands,
+        owned_regions,
+        successors,
+        attributes,
+        results=op_ty_results,
+        result_inference=false,
+    )
+end
+
+function lapack_gesvd(
+    input::Value;
+    U::IR.Type,
+    S::IR.Type,
+    Vt::IR.Type,
+    info::IR.Type,
+    full=nothing,
+    location=Location(),
+)
+    op_ty_results = IR.Type[U, S, Vt, info]
+    operands = Value[input,]
+    owned_regions = Region[]
+    successors = Block[]
+    attributes = NamedAttribute[]
+    !isnothing(full) && push!(attributes, namedattribute("full", full))
+
+    return create_operation(
+        "enzymexla.lapack.gesvd",
+        location;
+        operands,
+        owned_regions,
+        successors,
+        attributes,
+        results=op_ty_results,
+        result_inference=false,
+    )
+end
+
+function lapack_gesvj(
+    input::Value;
+    U::IR.Type,
+    S::IR.Type,
+    Vt::IR.Type,
+    info::IR.Type,
+    full=nothing,
+    location=Location(),
+)
+    op_ty_results = IR.Type[U, S, Vt, info]
+    operands = Value[input,]
+    owned_regions = Region[]
+    successors = Block[]
+    attributes = NamedAttribute[]
+    !isnothing(full) && push!(attributes, namedattribute("full", full))
+
+    return create_operation(
+        "enzymexla.lapack.gesvj",
+        location;
+        operands,
+        owned_regions,
+        successors,
+        attributes,
+        results=op_ty_results,
+        result_inference=false,
+    )
+end
+
 function get_stream(; result::IR.Type, location=Location())
     op_ty_results = IR.Type[result,]
     operands = Value[]
@@ -432,6 +515,51 @@ function get_stream(; result::IR.Type, location=Location())
     )
 end
 
+function lapack_getrf(
+    input::Value;
+    output::IR.Type,
+    pivots::IR.Type,
+    permutation::IR.Type,
+    info::IR.Type,
+    location=Location(),
+)
+    op_ty_results = IR.Type[output, pivots, permutation, info]
+    operands = Value[input,]
+    owned_regions = Region[]
+    successors = Block[]
+    attributes = NamedAttribute[]
+
+    return create_operation(
+        "enzymexla.lapack.getrf",
+        location;
+        operands,
+        owned_regions,
+        successors,
+        attributes,
+        results=op_ty_results,
+        result_inference=false,
+    )
+end
+
+function lapack_getri(input::Value, ipiv::Value; output::IR.Type, location=Location())
+    op_ty_results = IR.Type[output,]
+    operands = Value[input, ipiv]
+    owned_regions = Region[]
+    successors = Block[]
+    attributes = NamedAttribute[]
+
+    return create_operation(
+        "enzymexla.lapack.getri",
+        location;
+        operands,
+        owned_regions,
+        successors,
+        attributes,
+        results=op_ty_results,
+        result_inference=false,
+    )
+end
+
 function jit_call(
     inputs::Vector{Value};
     result_0::Vector{IR.Type},
@@ -754,15 +882,15 @@ end
 """
 `linalg_qr`
 
-This operation computes the QR factorization of a matrix using Householder 
-reflections. Mathematically, it decomposes A into the product of an 
-orthogonal (unitary if complex) matrix Q and an upper triangular matrix R, 
+This operation computes the QR factorization of a matrix using Householder
+reflections. Mathematically, it decomposes A into the product of an
+orthogonal (unitary if complex) matrix Q and an upper triangular matrix R,
 such that A = QR.
 
 If A has size m x n and m > n, Q is an m x n isometric matrix. If m < n, R
 will be a m x n trapezoidal matrix.
 
-This operation is modeled after the mathematical formulation of the QR 
+This operation is modeled after the mathematical formulation of the QR
 factorization, and not after LAPACK\'s compact formats.
 """
 function linalg_qr(
@@ -842,6 +970,7 @@ function linalg_svd(
     Vt::IR.Type,
     info::IR.Type,
     full=nothing,
+    algorithm=nothing,
     location=Location(),
 )
     op_ty_results = IR.Type[U, S, Vt, info]
@@ -850,6 +979,7 @@ function linalg_svd(
     successors = Block[]
     attributes = NamedAttribute[]
     !isnothing(full) && push!(attributes, namedattribute("full", full))
+    !isnothing(algorithm) && push!(attributes, namedattribute("algorithm", algorithm))
 
     return create_operation(
         "enzymexla.linalg.svd",

src/mlir/libMLIR_h.jl

@@ -11655,6 +11655,10 @@ function enzymexlaQRAlgorithmAttrGet(ctx, mode)
     @ccall mlir_c.enzymexlaQRAlgorithmAttrGet(ctx::MlirContext, mode::Int32)::MlirAttribute
 end
 
+function enzymexlaSVDAlgorithmAttrGet(ctx, mode)
+    @ccall mlir_c.enzymexlaSVDAlgorithmAttrGet(ctx::MlirContext, mode::Int32)::MlirAttribute
+end
+
 function enzymexlaGeluApproximationAttrGet(ctx, mode)
     @ccall mlir_c.enzymexlaGeluApproximationAttrGet(
         ctx::MlirContext, mode::Int32