feat: factorizations (#1234)

* fix: mlir regeneration

* feat: factorizations

* chore: fmt

* fix: symname

* fix: flags

* feat: lapack integration working 🎉

* feat: Ops.lu

* feat: add triangular_solve op

* chore: rename

* test: add tests for lu fact

* chore: fmt

* test: fix

* chore: bump jll

Author: avik-pal

Project.toml

@@ -87,7 +87,7 @@ PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.9"
-Reactant_jll = "0.0.182"
+Reactant_jll = "0.0.183"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

src/Compiler.jl

@@ -2,6 +2,7 @@ module Compiler
 
 using Reactant_jll
 using Libdl: dlsym
+using LinearAlgebra: BLAS
 
 import ..Reactant:
     Reactant,
@@ -1273,6 +1274,10 @@ function compile_mlir!(
         "canonicalize"
     end
 
+    blas_int_width = sizeof(BLAS.BlasInt) * 8
+    lower_enzymexla_linalg_pass = "lower-enzymexla-linalg{backend=$backend \
+                                   blas_int_width=$blas_int_width}"
+
     if optimize === :all
         run_pass_pipeline!(
             mod,
@@ -1291,6 +1296,7 @@ function compile_mlir!(
                         "remove-unnecessary-enzyme-ops",
                         "enzyme-simplify-math",
                         opt_passes2,
+                        lower_enzymexla_linalg_pass,
                         jit,
                     ]
                 else
@@ -1307,6 +1313,7 @@ function compile_mlir!(
                         opt_passes2,
                         kern,
                         raise_passes,
+                        lower_enzymexla_linalg_pass,
                         jit,
                     ]
                 end,
@@ -1453,6 +1460,7 @@ function compile_mlir!(
                         "remove-unnecessary-enzyme-ops",
                         "enzyme-simplify-math",
                         opt_passes2,
+                        lower_enzymexla_linalg_pass,
                         jit,
                     ]
                 else
@@ -1466,6 +1474,7 @@ function compile_mlir!(
                         opt_passes2,
                         kern,
                         raise_passes,
+                        lower_enzymexla_linalg_pass,
                         jit,
                     ]
                 end,
@@ -1487,6 +1496,7 @@ function compile_mlir!(
                         opt_passes2,
                         enzyme_pass,
                         "canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math",
+                        lower_enzymexla_linalg_pass,
                         jit,
                     ]
                 else
@@ -1499,6 +1509,7 @@ function compile_mlir!(
                         "canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math",
                         kern,
                         raise_passes,
+                        lower_enzymexla_linalg_pass,
                         jit,
                     ]
                 end,

src/Ops.jl

@@ -2930,4 +2930,48 @@ function triangular_solve(
     return TracedRArray{T,N}((), res, size(res))
 end
 
+"""
+    lu(
+        x::TracedRArray{T},
+        ::Type{pT}=Int32;
+        location=mlir_stacktrace("lu", @__FILE__, @__LINE__)
+    ) where {T,pT}
+
+Compute the row maximum pivoted LU factorization of `x` and return the factors `LU`,
+`ipiv`, `permutation` tensor, and `info`.
+"""
+@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(pT)),
+        permutation=MLIR.IR.TensorType(permutation_shape, MLIR.IR.Type(pT)),
+        info=MLIR.IR.TensorType(info_shape, MLIR.IR.Type(pT)),
+        location,
+    )
+
+    res = TracedRArray{T,ndims(x)}((), MLIR.IR.result(op, 1), size(x))
+    ipiv = TracedRArray{pT,ndims(x) - 1}((), MLIR.IR.result(op, 2), pivots_shape)
+    perm = TracedRArray{pT,ndims(x) - 1}((), MLIR.IR.result(op, 3), permutation_shape)
+
+    if ndims(x) == 2
+        info = TracedRNumber{pT}((), MLIR.IR.result(op, 4))
+    else
+        info = TracedRArray{pT,ndims(x) - 2}((), MLIR.IR.result(op, 4), info_shape)
+    end
+    return (res, ipiv, perm, info)
+end
+
 end # module Ops

src/stdlibs/LinearAlgebra.jl

@@ -18,6 +18,25 @@ using ReactantCore: materialize_traced_array
 using ..TracedUtils: TracedUtils, get_mlir_data, set_mlir_data!
 
 using LinearAlgebra
+using Libdl: Libdl
+
+function __init__()
+    libblastrampoline_handle = Libdl.dlopen(LinearAlgebra.BLAS.libblas)
+
+    for (cname, enzymexla_name) in [
+        (LinearAlgebra.BLAS.@blasfunc(sgetrf_), :enzymexla_lapack_sgetrf_),
+        (LinearAlgebra.BLAS.@blasfunc(dgetrf_), :enzymexla_lapack_dgetrf_),
+        (LinearAlgebra.BLAS.@blasfunc(cgetrf_), :enzymexla_lapack_cgetrf_),
+        (LinearAlgebra.BLAS.@blasfunc(zgetrf_), :enzymexla_lapack_zgetrf_),
+    ]
+        sym = Libdl.dlsym(libblastrampoline_handle, cname)
+        @ccall MLIR.API.mlir_c.EnzymeJaXMapSymbol(
+            enzymexla_name::Cstring, sym::Ptr{Cvoid}
+        )::Cvoid
+    end
+
+    return nothing
+end
 
 # Various Wrapper Arrays defined in LinearAlgebra
 function ReactantCore.materialize_traced_array(

test/ops.jl

@@ -1180,3 +1180,43 @@ end
         @test Array(y_ra) == ones(Float32, 2, 3)
     end
 end
+
+function recon_from_lu(lu_res::AbstractArray{T,4}) where {T}
+    y = similar(lu_res)
+    for i in 1:size(lu_res, 1), j in 1:size(lu_res, 2)
+        y[i, j, :, :] .= recon_from_lu(lu_res[i, j, :, :])
+    end
+    return y
+end
+
+function apply_permutation(x::AbstractArray{T,4}, perm) where {T}
+    y = similar(x)
+    for i in 1:size(x, 1), j in 1:size(x, 2)
+        y[i, j, :, :] .= x[i, j, perm[i, j, :], :]
+    end
+    return y
+end
+
+function recon_from_lu(lu_res::AbstractMatrix)
+    return UnitLowerTriangular(lu_res) * UpperTriangular(lu_res)
+end
+
+@testset "lu factorization" begin
+    @testset "unbatched" begin
+        x_ra = Reactant.to_rarray(randn(6, 6))
+        lu_ra, ipiv, perm, info = @jit Ops.lu(x_ra)
+
+        @test @jit(recon_from_lu(lu_ra)) ≈ @jit(getindex(x_ra, perm, :))
+    end
+
+    @testset "batched" begin
+        x_ra = Reactant.to_rarray(randn(4, 3, 6, 6))
+        lu_ra, ipiv, perm, info = @jit Ops.lu(x_ra)
+        @test size(lu_ra) == (4, 3, 6, 6)
+        @test size(ipiv) == (4, 3, 6)
+        @test size(perm) == (4, 3, 6)
+        @test size(info) == (4, 3)
+
+        @test @jit(recon_from_lu(lu_ra)) ≈ @jit(apply_permutation(x_ra, perm))
+    end
+end