feat: progressive lowering of LU and SVD via lapack ops (#1625)

* feat: more lapack operations for LU/SVD

* fix: lower linalg ops to lapack

* chore: stub code for lowering

* feat: lowering for getrf

* feat: svd lapack op lowering

* feat: cpu lowering

* fix: missing impl in header

* fix: export names

* fix: update workspace buffer correctly

* test: update all tests

* feat: support batching

Author: avik-pal

src/enzyme_ad/jax/Dialect/EnzymeXLAAttrs.td

@@ -82,6 +82,24 @@ def EnzymeXLA_QrAlgorithmAttr : EnumAttr<EnzymeXLA_Dialect,
   let assemblyFormat = "`<` $value `>`";
 }
 
+def EnzymeXLA_SVDAlgorithm : I32EnumAttr<"SVDAlgorithm",
+    "Algorithm to use for the SVD factorization",
+    [
+      I32EnumAttrCase<"DEFAULT", 0>,
+      I32EnumAttrCase<"QRIteration", 1>,
+      I32EnumAttrCase<"DivideAndConquer", 2>,
+      I32EnumAttrCase<"Jacobi", 3>
+    ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::enzymexla";
+}
+
+def EnzymeXLA_SVDAlgorithmAttr : EnumAttr<EnzymeXLA_Dialect,
+    EnzymeXLA_SVDAlgorithm, "svd_algorithm"> {
+  let assemblyFormat = "`<` $value `>`";
+  let cppType = "::mlir::enzymexla::SVDAlgorithmAttr";
+}
+
 // Machine Learning
 
 def EnzymeXLA_GeluApproximation : I32EnumAttr<"GeluApproximation",

src/enzyme_ad/jax/Dialect/EnzymeXLAOps.td

@@ -446,15 +446,15 @@ def QRFactorizationOp: EnzymeXLA_Op<"linalg.qr", [Pure]> {
   let summary = "QR factorization operation (high-level op)";
 
   let description = [{
-    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.
   }];
 
@@ -473,28 +473,49 @@ def QRFactorizationOp: EnzymeXLA_Op<"linalg.qr", [Pure]> {
   }];
 }
 
+def SVDFactorizationOp : EnzymeXLA_Op<"linalg.svd", [Pure]> {
+  let summary = "Singular Value Decomposition (SVD) factorization operation.";
+
+  let arguments = (ins
+    HLO_Tensor:$input,
+    DefaultValuedAttr<BoolAttr, "false">:$full,
+    DefaultValuedAttr<EnzymeXLA_SVDAlgorithmAttr, "::mlir::enzymexla::SVDAlgorithm::DEFAULT">:$algorithm
+  );
+
+  let results = (outs
+    HLO_Tensor:$U,
+    HLO_Tensor:$S,
+    HLO_Tensor:$Vt,
+    HLO_Tensor:$info
+  );
+
+  let assemblyFormat = [{
+    $input attr-dict `:` functional-type($input, results)
+  }];
+}
+
 def GeqrfOp : EnzymeXLA_Op<"lapack.geqrf", [Pure]> {
   let summary = "QR factorization operation (low-level op)";
 
   let description = [{
-    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.
+  }];
 
-let arguments = (ins HLO_Tensor : $input);
+  let arguments = (ins HLO_Tensor : $input);
 
-let results = (outs HLO_Tensor
-               : $output, HLO_Tensor
-               : $tau, HLO_Tensor
-               : $info);
+  let results = (outs
+    HLO_Tensor:$output,
+    HLO_Tensor:$tau,
+    HLO_Tensor:$info
+  );
 
-let assemblyFormat = [{
+  let assemblyFormat = [{
     $input attr-dict `:` functional-type($input, results)
   }];
 }
@@ -503,11 +524,11 @@ def GeqrtOp : EnzymeXLA_Op<"lapack.geqrt", [Pure]> {
   let summary = "QR factorization operation (low-level op)";
 
   let description = [{
-    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.
   }];
 
@@ -605,7 +626,81 @@ def GemqrtOp : EnzymeXLA_Op<"lapack.gemqrt", [Pure]> {
   }];
 }
 
-def SVDFactorizationOp : EnzymeXLA_Op<"linalg.svd", [Pure]> {
+def GetrfOp : EnzymeXLA_Op<"lapack.getrf", [Pure]> {
+  let summary = "LU factorization operation with row-major pivoting.";
+
+  let arguments = (ins HLO_Tensor:$input);
+
+  let results = (outs
+    HLO_Tensor:$output,
+    HLO_Tensor:$pivots,
+    HLO_Tensor:$permutation,
+    HLO_Tensor:$info
+  );
+
+  let assemblyFormat = [{
+    $input attr-dict `:` functional-type($input, results)
+  }];
+}
+
+def GetriOp : EnzymeXLA_Op<"lapack.getri", [Pure]> {
+  let summary = "Computes the inverse of a matrix using the LU factorization.";
+
+  let arguments = (ins
+    HLO_Tensor:$input,
+    HLO_Tensor:$ipiv
+  );
+
+  let results = (outs
+    HLO_Tensor:$output
+  );
+
+  let assemblyFormat = [{
+    $input `,` $ipiv attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def GesddOp : EnzymeXLA_Op<"lapack.gesdd", [Pure]> {
+  let summary = "Singular Value Decomposition (SVD) factorization operation.";
+
+  let arguments = (ins
+    HLO_Tensor:$input,
+    DefaultValuedAttr<BoolAttr, "false">:$full
+  );
+
+  let results = (outs
+    HLO_Tensor:$U,
+    HLO_Tensor:$S,
+    HLO_Tensor:$Vt,
+    HLO_Tensor:$info
+  );
+
+  let assemblyFormat = [{
+    $input attr-dict `:` functional-type($input, results)
+  }];
+}
+
+def GesvdOp : EnzymeXLA_Op<"lapack.gesvd", [Pure]> {
+  let summary = "Singular Value Decomposition (SVD) factorization operation.";
+
+  let arguments = (ins
+    HLO_Tensor:$input,
+    DefaultValuedAttr<BoolAttr, "false">:$full
+  );
+
+  let results = (outs
+    HLO_Tensor:$U,
+    HLO_Tensor:$S,
+    HLO_Tensor:$Vt,
+    HLO_Tensor:$info
+  );
+
+  let assemblyFormat = [{
+    $input attr-dict `:` functional-type($input, results)
+  }];
+}
+
+def GesvjOp : EnzymeXLA_Op<"lapack.gesvj", [Pure]> {
   let summary = "Singular Value Decomposition (SVD) factorization operation.";
 
   let arguments = (ins

src/enzyme_ad/jax/Integrations/c/EnzymeXLA.cpp

@@ -64,6 +64,25 @@ MlirAttribute enzymexlaQRAlgorithmAttrGet(MlirContext ctx, int32_t mode) {
   return wrap(mlir::enzymexla::QrAlgorithmAttr::get(unwrap(ctx), algorithm));
 }
 
+MlirAttribute enzymexlaSVDAlgorithmAttrGet(MlirContext ctx, int32_t mode) {
+  mlir::enzymexla::SVDAlgorithm algorithm;
+  switch (mode) {
+  case 0:
+    algorithm = mlir::enzymexla::SVDAlgorithm::DEFAULT;
+    break;
+  case 1:
+    algorithm = mlir::enzymexla::SVDAlgorithm::QRIteration;
+    break;
+  case 2:
+    algorithm = mlir::enzymexla::SVDAlgorithm::DivideAndConquer;
+    break;
+  case 3:
+    algorithm = mlir::enzymexla::SVDAlgorithm::Jacobi;
+    break;
+  }
+  return wrap(mlir::enzymexla::SVDAlgorithmAttr::get(unwrap(ctx), algorithm));
+}
+
 MlirAttribute enzymexlaGeluApproximationAttrGet(MlirContext ctx, int32_t mode) {
   mlir::enzymexla::GeluApproximation approximation;
   switch (mode) {

src/enzyme_ad/jax/Integrations/c/EnzymeXLA.h

@@ -33,6 +33,9 @@ MLIR_CAPI_EXPORTED MlirAttribute enzymexlaLapackUploAttrGet(MlirContext ctx,
 MLIR_CAPI_EXPORTED MlirAttribute enzymexlaQRAlgorithmAttrGet(MlirContext ctx,
                                                              int32_t mode);
 
+MLIR_CAPI_EXPORTED MlirAttribute enzymexlaSVDAlgorithmAttrGet(MlirContext ctx,
+                                                              int32_t mode);
+
 //===----------------------------------------------------------------------===//
 // Machine Learning Ops
 //===----------------------------------------------------------------------===//