feat: triangular_solve lowering (#1291)

* feat: triangular_solve lowering

* feat: more overloads

* feat: overload generic_trimatdiv!

* feat: add rdiv! support

* test: add tests

Author: avik-pal

src/Ops.jl

@@ -2846,4 +2846,66 @@ end
     ]
 end
 
+function triangular_solve(
+    a::TracedRArray{T,N},
+    b::TracedRArray{T,M};
+    left_side::Bool,
+    location=mlir_stacktrace("triangular_solve", @__FILE__, @__LINE__),
+    kwargs...,
+) where {T,N,M}
+    @assert M == N - 1
+
+    if left_side
+        b = reshape(b, size(b)..., 1)
+    else
+        b = reshape(b, size(b)[1:(M - 1)]..., 1, size(b, M))
+    end
+
+    return dropdims(
+        triangular_solve(a, b; location, left_side, kwargs...); dims=(N - 1 + left_side)
+    )
+end
+
+function triangular_solve(
+    a::TracedRArray{T,N},
+    b::TracedRArray{T,N};
+    left_side::Bool,
+    lower::Bool,
+    transpose_a::Char,
+    unit_diagonal::Bool,
+    location=mlir_stacktrace("triangular_solve", @__FILE__, @__LINE__),
+) where {T,N}
+    @assert N >= 2
+    @assert size(a, N - 1) == size(a, N) == size(b, N - left_side)
+    @assert size(a)[1:(N - 2)] == size(b)[1:(N - 2)] "a and b must have the same leading \
+                                                      dimensions"
+
+    @assert transpose_a in ('N', 'T', 'C') "transpose_a must be one of 'N', 'T', or 'C'"
+    transpose_attr = MLIR.API.stablehloTransposeAttrGet(
+        MLIR.IR.context(),
+        if transpose_a == 'N'
+            "NO_TRANSPOSE"
+        elseif transpose_a == 'T'
+            "TRANSPOSE"
+        else
+            "ADJOINT"
+        end,
+    )
+
+    res = MLIR.IR.result(
+        MLIR.Dialects.stablehlo.triangular_solve(
+            a.mlir_data,
+            b.mlir_data;
+            left_side=left_side,
+            lower=lower,
+            transpose_a=transpose_attr,
+            unit_diagonal=unit_diagonal,
+            location,
+        ),
+        1,
+    )
+
+    return TracedRArray{T,N}((), res, size(res))
+end
+
 end # module Ops

src/stdlibs/LinearAlgebra.jl

@@ -469,4 +469,55 @@ function LinearAlgebra.dot(x::AnyTracedRVector, y::AnyTracedRVector)
     return TracedRNumber{unwrapped_eltype(res)}((), res.mlir_data)
 end
 
+# ldiv & rdiv interfaces
+tfun_to_char(::typeof(identity)) = 'N'
+tfun_to_char(::typeof(transpose)) = 'T'
+tfun_to_char(::typeof(adjoint)) = 'C'
+
+function LinearAlgebra.generic_trimatdiv!(
+    C::AbstractVecOrMat{TracedRNumber{T}},
+    uploc,
+    isunitc,
+    tfun::Function,
+    A::AbstractMatrix,
+    B::AbstractVecOrMat,
+) where {T}
+    @assert uploc in ('L', 'U')
+    @assert isunitc in ('N', 'U')
+
+    res = Ops.triangular_solve(
+        TracedUtils.promote_to(TracedRArray{T,2}, materialize_traced_array(A)),
+        TracedUtils.promote_to(TracedRArray{T,ndims(B)}, materialize_traced_array(B));
+        left_side=true,
+        lower=(uploc == 'L'),
+        transpose_a=tfun_to_char(tfun),
+        unit_diagonal=(isunitc == 'U'),
+    )
+    set_mlir_data!(C, get_mlir_data(res))
+    return C
+end
+
+function LinearAlgebra.generic_mattridiv!(
+    C::AbstractMatrix{TracedRNumber{T}},
+    uploc,
+    isunitc,
+    tfun::Function,
+    A::AbstractMatrix,
+    B::AbstractMatrix,
+) where {T}
+    @assert uploc in ('L', 'U')
+    @assert isunitc in ('N', 'U')
+
+    res = Ops.triangular_solve(
+        TracedUtils.promote_to(TracedRArray{T,2}, materialize_traced_array(B)),
+        TracedUtils.promote_to(TracedRArray{T,2}, materialize_traced_array(A));
+        left_side=false,
+        lower=(uploc == 'L'),
+        transpose_a=tfun_to_char(tfun),
+        unit_diagonal=(isunitc == 'U'),
+    )
+    set_mlir_data!(C, get_mlir_data(res))
+    return C
+end
+
 end

test/integration/linear_algebra.jl

@@ -279,3 +279,43 @@ end
 
     @test @jit(dot(x_ra, y_ra)) ≈ dot(x, y)
 end
+
+@testset "Triangular ldiv and rdiv" begin
+    fn1(A, b) = A \ b
+    fn2(A, b) = A' \ b
+    fn3(A, b) = transpose(A) \ b
+
+    fn4(A, B) = B / A
+    fn5(A, B) = B / A'
+    fn6(A, B) = B / transpose(A)
+
+    @testset for T in (Float32, Float64, ComplexF32, ComplexF64)
+        A = rand(T, 6, 6)
+        B = rand(T, 6, 6)
+        b = rand(T, 6)
+        b_ra = Reactant.to_rarray(b)
+        B_ra = Reactant.to_rarray(B)
+
+        @testset for wT in (
+            UnitLowerTriangular, UnitUpperTriangular, LowerTriangular, UpperTriangular
+        )
+            A_wrapped = wT(A)
+            A_ra = Reactant.to_rarray(A_wrapped)
+
+            @testset "no_tranpose" begin
+                @test @jit(fn1(A_ra, b_ra)) ≈ fn1(A_wrapped, b)
+                @test @jit(fn4(A_ra, B_ra)) ≈ fn4(A_wrapped, B)
+            end
+
+            @testset "adjoint" begin
+                @test @jit(fn2(A_ra, b_ra)) ≈ fn2(A_wrapped, b)
+                @test @jit(fn5(A_ra, B_ra)) ≈ fn5(A_wrapped, B)
+            end
+
+            @testset "transpose" begin
+                @test @jit(fn3(A_ra, b_ra)) ≈ fn3(A_wrapped, b)
+                @test @jit(fn6(A_ra, B_ra)) ≈ fn6(A_wrapped, B)
+            end
+        end
+    end
+end