Add rule for triangular solves (#1264)

Author: simsurace

src/internal_rules.jl

@@ -447,6 +447,65 @@ function EnzymeRules.reverse(config, func::Const{typeof(\)}, ::Type{RT}, cache,
     return (nothing,nothing)
 end
 
+const EnzymeTriangulars = Union{
+    UpperTriangular,
+    LowerTriangular,
+    UnitUpperTriangular,
+    UnitLowerTriangular
+}
+
+function EnzymeRules.augmented_primal(
+    config,
+    func::Const{typeof(ldiv!)},
+    ::Type{RT},
+    Y::Annotation{YT},
+    A::Annotation{AT},
+    B::Annotation{BT}
+) where {RT, YT <: Array, AT <: EnzymeTriangulars, BT <: Array}
+    cache_Y = EnzymeRules.overwritten(config)[1] ? copy(Y.val) : Y.val
+    cache_A = EnzymeRules.overwritten(config)[2] ? copy(A.val) : A.val
+    cache_A = compute_lu_cache(cache_A, B.val)
+    cache_B = EnzymeRules.overwritten(config)[3] ? copy(B.val) : nothing
+    primal = EnzymeRules.needs_primal(config) ? Y.val : nothing
+    shadow = EnzymeRules.needs_shadow(config) ? Y.dval : nothing
+    func.val(Y.val, A.val, B.val)
+    return EnzymeRules.AugmentedReturn{typeof(primal), typeof(shadow), Any}(
+        primal, shadow, (cache_Y, cache_A, cache_B))
+end
+
+function EnzymeRules.reverse(
+    config,
+    func::Const{typeof(ldiv!)},
+    ::Type{RT},
+    cache,
+    Y::Annotation{YT},
+    A::Annotation{AT},
+    B::Annotation{BT}
+) where {YT <: Array, RT, AT <: EnzymeTriangulars, BT <: Array}
+    if !isa(Y, Const)
+        (cache_Yout, cache_A, cache_B) = cache
+        for b in 1:EnzymeRules.width(config)
+            dY = EnzymeRules.width(config) == 1 ? Y.dval : Y.dval[b]
+            z = adjoint(cache_A) \ dY
+            if !isa(B, Const)
+                dB = EnzymeRules.width(config) == 1 ? B.dval : B.dval[b]
+                dB .+= z
+            end
+            if !isa(A, Const)
+                dA = EnzymeRules.width(config) == 1 ? A.dval : A.dval[b]
+                dA.data .-= _zero_unused_elements!(AT(z * adjoint(cache_Yout)))
+            end
+            dY .= zero(eltype(dY))
+        end
+    end
+    return (nothing, nothing, nothing)
+end
+
+_zero_unused_elements!(A::UpperTriangular) = triu!(A.data)
+_zero_unused_elements!(A::LowerTriangular) = tril!(A.data)
+_zero_unused_elements!(A::UnitUpperTriangular) = triu!(A.data, 1)
+_zero_unused_elements!(A::UnitLowerTriangular) = tril!(A.data, -1)
+
 @static if VERSION >= v"1.7-"
 # Force a rule around hvcat_fill as it is type unstable if the tuple is not of the same type (e.g., int, float, int, float)
 function EnzymeRules.augmented_primal(config, func::Const{typeof(Base.hvcat_fill!)}, ::Type{RT}, out::Annotation{AT}, inp::Annotation{BT}) where {RT, AT <: Array, BT <: Tuple}

test/Project.toml

@@ -19,3 +19,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 Aqua = "0.8"
+EnzymeTestUtils = "0.1.4"

test/internal_rules.jl

@@ -2,6 +2,7 @@ module InternalRules
 
 using Enzyme
 using Enzyme.EnzymeRules
+using EnzymeTestUtils
 using FiniteDifferences
 using LinearAlgebra
 using SparseArrays
@@ -386,5 +387,29 @@ end
         @test isapprox(dA, dA_sym)
     end
 end
+
+@testset "Linear solve for triangular matrices" begin
+    @testset for T in (UpperTriangular, LowerTriangular, UnitUpperTriangular, UnitLowerTriangular),
+        TE in (Float64, ComplexF64), sizeB in ((3,), (3, 3))
+        n = sizeB[1]
+        M = rand(TE, n, n)
+        B = rand(TE, sizeB...)
+        Y = zeros(TE, sizeB...)
+        A = T(M)
+        @testset "test against EnzymeTestUtils through constructor" begin
+            _A = T(A)
+            function f!(Y, A, B, ::T) where T
+                ldiv!(Y, T(A), B)
+                return nothing
+            end
+            for TY in (Const, Duplicated, BatchDuplicated),
+                TM in (Const, Duplicated, BatchDuplicated),
+                TB in (Const, Duplicated, BatchDuplicated)
+                are_activities_compatible(Const, TY, TM, TB) || continue
+                test_reverse(f!, Const, (Y, TY), (M, TM), (B, TB), (_A, Const))
+            end
+        end
+    end
+end
 end
 end # InternalRules