feat: lapack integration working 🎉

Author: avik-pal

deps/ReactantExtra/WORKSPACE

@@ -9,7 +9,7 @@ http_archive(
     urls = ["https://github.com/wsmoses/nsync/archive/{commit}.tar.gz".format(commit = NSYNC_COMMIT)],
 )
 
-ENZYMEXLA_COMMIT = "657c2f62e4cc3dd2b52ff0587f644b64858d7c0b"
+ENZYMEXLA_COMMIT = "abd48236a7466bd0db7d2c8ffa56d1d724463c91"
 ENZYMEXLA_SHA256 = ""
 
 http_archive(

src/Compiler.jl

@@ -2,6 +2,7 @@ module Compiler
 
 using Reactant_jll
 using Libdl: dlsym
+using LinearAlgebra: BLAS
 
 import ..Reactant:
     Reactant,
@@ -1274,6 +1275,10 @@ function compile_mlir!(
         "canonicalize"
     end
 
+    blas_int_width = sizeof(BLAS.BlasInt) * 8
+    lower_factorization_pass = "lower-factorization{backend=$backend \
+                                blas_int_width=$blas_int_width}"
+
     if optimize === :all
         run_pass_pipeline!(
             mod,
@@ -1292,7 +1297,7 @@ function compile_mlir!(
                         "remove-unnecessary-enzyme-ops",
                         "enzyme-simplify-math",
                         opt_passes2,
-                        "lower-factorization{backend=$backend}",
+                        lower_factorization_pass,
                         jit,
                     ]
                 else
@@ -1309,7 +1314,7 @@ function compile_mlir!(
                         opt_passes2,
                         kern,
                         raise_passes,
-                        "lower-factorization{backend=$backend}",
+                        lower_factorization_pass,
                         jit,
                     ]
                 end,
@@ -1456,7 +1461,7 @@ function compile_mlir!(
                         "remove-unnecessary-enzyme-ops",
                         "enzyme-simplify-math",
                         opt_passes2,
-                        "lower-factorization{backend=$backend}",
+                        lower_factorization_pass,
                         jit,
                     ]
                 else
@@ -1470,7 +1475,7 @@ function compile_mlir!(
                         opt_passes2,
                         kern,
                         raise_passes,
-                        "lower-factorization{backend=$backend}",
+                        lower_factorization_pass,
                         jit,
                     ]
                 end,
@@ -1492,7 +1497,7 @@ function compile_mlir!(
                         opt_passes2,
                         enzyme_pass,
                         "canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math",
-                        "lower-factorization{backend=$backend}",
+                        lower_factorization_pass,
                         jit,
                     ]
                 else
@@ -1505,7 +1510,7 @@ function compile_mlir!(
                         "canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math",
                         kern,
                         raise_passes,
-                        "lower-factorization{backend=$backend}",
+                        lower_factorization_pass,
                         jit,
                     ]
                 end,

src/mlir/Dialects/EnzymeXLA.jl

@@ -188,7 +188,7 @@ function kernel_call(
     )
 end
 
-function lu_factorization(
+function linalg_lu(
     input::Value; output::IR.Type, pivots::IR.Type, info::IR.Type, location=Location()
 )
     op_ty_results = IR.Type[output, pivots, info]
@@ -198,7 +198,7 @@ function lu_factorization(
     attributes = NamedAttribute[]
 
     return create_operation(
-        "enzymexla.lu_factorization",
+        "enzymexla.linalg.lu",
         location;
         operands,
         owned_regions,

src/stdlibs/LinearAlgebra.jl

@@ -498,21 +498,21 @@ function LinearAlgebra.lu!(
     # TODO: stop ignoring `check` and `allowsingular`
     m, n = size(A)
 
-    lu_op = MLIR.Dialects.enzymexla.lu_factorization(
+    lu_op = MLIR.Dialects.enzymexla.linalg_lu(
         materialize_traced_array(A).mlir_data;
         output=MLIR.IR.TensorType(
             collect(Int64, size(A)), MLIR.IR.Type(unwrapped_eltype(A))
         ),
-        pivots=MLIR.IR.TensorType(Int64[max(m, n)], MLIR.IR.Type(Int32)),
+        pivots=MLIR.IR.TensorType(Int64[min(m, n)], MLIR.IR.Type(Int32)),
         info=MLIR.IR.TensorType(Int64[], MLIR.IR.Type(Int32)),
     )
     set_mlir_data!(A, MLIR.IR.result(lu_op, 1))
-    pivots = TracedRArray{Int32,1}((), MLIR.IR.result(lu_op, 2), (max(m, n),))
+    pivots = TracedRArray{Int32,1}((), MLIR.IR.result(lu_op, 2), (min(m, n),))
     info = TracedRNumber{Int32}((), MLIR.IR.result(lu_op, 3))
 
     # XXX: `info` needs to be a BLASInt for LU
     # return LinearAlgebra.LU{T, typeof(A), typeof(pivots)}(A, pivots, -1)
-    return (A, pivots, info)
+    return (UnitLowerTriangular(A), UpperTriangular(A), pivots, info)
 end
 
 end