feat: Ops.lu

Author: avik-pal

src/Compiler.jl

@@ -1273,7 +1273,7 @@ function compile_mlir!(
     end
 
     blas_int_width = sizeof(BLAS.BlasInt) * 8
-    lower_factorization_pass = "lower-factorization{backend=$backend \
+    lower_factorization_pass = "lower-enzymexla-linalg{backend=$backend \
                                 blas_int_width=$blas_int_width}"
 
     if optimize === :all

src/Ops.jl

@@ -2846,4 +2846,46 @@ end
     ]
 end
 
+@noinline function lu(
+    x::TracedRArray{T},
+    ::Type{pT}=Int32;
+    location=mlir_stacktrace("lu", @__FILE__, @__LINE__),
+) where {T,pT}
+    @assert ndims(x) >= 2
+
+    output_shape = collect(Int64, size(x))
+    batch_shape = output_shape[1:(end - 2)]
+    pivots_shape = vcat(batch_shape, min(size(x, ndims(x) - 1), size(x, ndims(x))))
+    permutation_shape = vcat(batch_shape, size(x, ndims(x) - 1))
+    info_shape = batch_shape
+
+    op = MLIR.Dialects.enzymexla.linalg_lu(
+        x.mlir_data;
+        output=MLIR.IR.TensorType(output_shape, MLIR.IR.Type(unwrapped_eltype(T))),
+        pivots=MLIR.IR.TensorType(pivots_shape, MLIR.IR.Type(unwrapped_eltype(pT))),
+        permutation=MLIR.IR.TensorType(
+            permutation_shape, MLIR.IR.Type(unwrapped_eltype(pT))
+        ),
+        info=MLIR.IR.TensorType(info_shape, MLIR.IR.Type(unwrapped_eltype(pT))),
+        location,
+    )
+
+    res = TracedRArray{unwrapped_eltype(T),ndims(x)}((), MLIR.IR.result(op, 1), size(x))
+    ipiv = TracedRArray{unwrapped_eltype(pT),ndims(x) - 1}(
+        (), MLIR.IR.result(op, 2), pivots_shape
+    )
+    perm = TracedRArray{unwrapped_eltype(pT),ndims(x) - 1}(
+        (), MLIR.IR.result(op, 3), permutation_shape
+    )
+
+    if ndims(x) == 2
+        info = TracedRNumber{unwrapped_eltype(pT)}((), MLIR.IR.result(op, 4))
+    else
+        info = TracedRArray{unwrapped_eltype(pT),ndims(x) - 2}(
+            (), MLIR.IR.result(op, 4), info_shape
+        )
+    end
+    return (res, ipiv, perm, info)
+end
+
 end # module Ops

src/mlir/Dialects/EnzymeXLA.jl

@@ -189,9 +189,14 @@ function kernel_call(
 end
 
 function linalg_lu(
-    input::Value; output::IR.Type, pivots::IR.Type, info::IR.Type, location=Location()
+    input::Value;
+    output::IR.Type,
+    pivots::IR.Type,
+    permutation::IR.Type,
+    info::IR.Type,
+    location=Location(),
 )
-    op_ty_results = IR.Type[output, pivots, info]
+    op_ty_results = IR.Type[output, pivots, permutation, info]
     operands = Value[input,]
     owned_regions = Region[]
     successors = Block[]

src/stdlibs/LinearAlgebra.jl

@@ -488,31 +488,4 @@ function LinearAlgebra.dot(x::AnyTracedRVector, y::AnyTracedRVector)
     return TracedRNumber{unwrapped_eltype(res)}((), res.mlir_data)
 end
 
-# Factorizations
-function LinearAlgebra.lu!(
-    A::TracedRArray{T,2},
-    ::LinearAlgebra.RowMaximum;
-    check::Bool=true,
-    allowsingular::Bool=false,
-) where {T}
-    # TODO: stop ignoring `check` and `allowsingular`
-    m, n = size(A)
-
-    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[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), (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 (UnitLowerTriangular(A), UpperTriangular(A), pivots, info)
-end
-
 end