feat: add triangular_solve op

Author: avik-pal

src/Ops.jl

@@ -2846,6 +2846,16 @@ end
     ]
 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;
@@ -2862,30 +2872,76 @@ end
     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))),
+        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{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
-    )
+    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{unwrapped_eltype(pT)}((), MLIR.IR.result(op, 4))
+        info = TracedRNumber{pT}((), MLIR.IR.result(op, 4))
     else
-        info = TracedRArray{unwrapped_eltype(pT),ndims(x) - 2}(
-            (), MLIR.IR.result(op, 4), info_shape
-        )
+        info = TracedRArray{pT,ndims(x) - 2}((), MLIR.IR.result(op, 4), info_shape)
     end
     return (res, ipiv, perm, info)
 end
 
+@noinline function triangular_solve(
+    a::TracedRArray{T,N},
+    b::TracedRArray{T,M};
+    left_side::Bool=true,
+    lower::Bool,
+    unit_diagonal::Bool=false,
+    transpose_a::Symbol=:N,
+    location=mlir_stacktrace("triangular_solve", @__FILE__, @__LINE__),
+) where {T,N,M}
+    @assert N >= 2
+
+    if M == N - 1
+        if left_side
+            b = Ops.reshape(b, size(b)..., 1; location)
+        else
+            b = Ops.reshape(b, size(b)[1:(M - 1)]..., 1, size(b, M); location)
+        end
+    end
+
+    @assert size(a, N - 1) == size(a, N) == size(b, N - left_side)
+    @assert N == ndims(b)
+    @assert transpose_a ∈ (:N, :T, :C)
+    @assert size(a)[1:(N - 2)] == size(b)[1:(N - 2)]
+
+    transpose_attr = if transpose_a === :N
+        MLIR.API.stablehloTransposeAttrGet(MLIR.IR.context(), "NO_TRANSPOSE")
+    elseif transpose_a === :T
+        MLIR.API.stablehloTransposeAttrGet(MLIR.IR.context(), "TRANSPOSE")
+    else
+        MLIR.API.stablehloTransposeAttrGet(MLIR.IR.context(), "ADJOINT")
+    end
+
+    op = stablehlo.triangular_solve(
+        a.mlir_data,
+        b.mlir_data;
+        left_side=MLIR.IR.Attribute(left_side),
+        lower=MLIR.IR.Attribute(lower),
+        unit_diagonal=MLIR.IR.Attribute(unit_diagonal),
+        transpose_a=transpose_attr,
+        location,
+    )
+
+    result_shape = size(a)[1:(N - 2)]
+    if left_side
+        result_shape = (result_shape..., size(a, N), size(b, N))
+    else
+        result_shape = (result_shape..., size(b, N - 1), size(a, N - 1))
+    end
+
+    sol = TracedRArray{T,N}((), MLIR.IR.result(op, 1), result_shape)
+    N == M && return sol
+    return Ops.reshape(sol, Int64[size(a)[1:(N - 2)]..., size(b, N - left_side)]; location)
+end
+
 end # module Ops