fix: lower linalg ops to lapack

avik-pal · avik-pal · commit 3782e4e5a7e6 · 2025-11-16T20:05:31.000-06:00
diff --git a/src/enzyme_ad/jax/Dialect/EnzymeXLAAttrs.td b/src/enzyme_ad/jax/Dialect/EnzymeXLAAttrs.td
@@ -87,7 +87,8 @@ def EnzymeXLA_SVDAlgorithm : I32EnumAttr<"SVDAlgorithm",
     [
       I32EnumAttrCase<"DEFAULT", 0>,
       I32EnumAttrCase<"QRIteration", 1>,
-      I32EnumAttrCase<"DivideAndConquer", 2>
+      I32EnumAttrCase<"DivideAndConquer", 2>,
+      I32EnumAttrCase<"Jacobi", 3>
     ]> {
   let genSpecializedAttr = 0;
   let cppNamespace = "::mlir::enzymexla";
diff --git a/src/enzyme_ad/jax/Integrations/c/EnzymeXLA.cpp b/src/enzyme_ad/jax/Integrations/c/EnzymeXLA.cpp
@@ -75,6 +75,10 @@ MlirAttribute enzymexlaSVDAlgorithmAttrGet(MlirContext ctx, int32_t mode) {
     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));
 }
diff --git a/src/enzyme_ad/jax/Passes/LowerEnzymeXLALinalg.cpp b/src/enzyme_ad/jax/Passes/LowerEnzymeXLALinalg.cpp

Original file line number	Diff line number	Diff line change
`@@ -75,6 +75,10 @@ MlirAttribute enzymexlaSVDAlgorithmAttrGet(MlirContext ctx, int32_t mode) {`
`75`	`75`	`break;`
`76`	`76`	`case 2:`
`77`	`77`	`algorithm = mlir::enzymexla::SVDAlgorithm::DivideAndConquer;`
	`78`	`+ break;`
	`79`	`+ case 3:`
	`80`	`+ algorithm = mlir::enzymexla::SVDAlgorithm::Jacobi;`
	`81`	`+ break;`
`78`	`82`	`}`
`79`	`83`	`return wrap(mlir::enzymexla::SVDAlgorithmAttr::get(unwrap(ctx), algorithm));`
`80`	`84`	`}`