Refactor Enzyme testsuite

ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl

@@ -15,15 +15,17 @@ using LinearAlgebra
 
 @inline EnzymeRules.inactive_type(::Type{Alg}) where {Alg <: MatrixAlgebraKit.AbstractAlgorithm} = true
 @inline EnzymeRules.inactive_type(::Type{TS}) where {TS <: MatrixAlgebraKit.TruncationStrategy} = true
-@inline EnzymeRules.inactive(f::typeof(MatrixAlgebraKit.select_algorithm), func::F, A::AbstractMatrix, alg::Alg) where {F, Alg} = true
-@inline EnzymeRules.inactive(f::typeof(MatrixAlgebraKit.default_algorithm), func::F, A::AbstractMatrix) where {F} = true
-@inline EnzymeRules.inactive(f::typeof(MatrixAlgebraKit.check_input), func::F, A::AbstractMatrix, alg::Alg) where {F, Alg} = true
+@inline EnzymeRules.inactive(::typeof(MatrixAlgebraKit.select_algorithm), func::F, A::AbstractMatrix, alg::Alg) where {F, Alg} = true
+@inline EnzymeRules.inactive(::typeof(MatrixAlgebraKit.default_algorithm), func::F, A::AbstractMatrix) where {F} = true
+@inline EnzymeRules.inactive(::typeof(MatrixAlgebraKit.check_input), func::F, A::AbstractMatrix, alg::Alg) where {F, Alg} = true
+@inline EnzymeRules.inactive(::typeof(MatrixAlgebraKit.check_input), func::F, A::AbstractMatrix, arg::Any, alg::Alg) where {F, Alg} = true
+@inline EnzymeRules.inactive(::typeof(MatrixAlgebraKit.check_hermitian), A::AbstractMatrix, alg::Alg) where {Alg} = true
 @inline EnzymeRules.inactive(::typeof(MatrixAlgebraKit.defaulttol), ::Any) = true
 @inline EnzymeRules.inactive(::typeof(MatrixAlgebraKit.default_pullback_gauge_atol), ::Any) = true
 @inline EnzymeRules.inactive(::typeof(MatrixAlgebraKit.default_pullback_gauge_atol), ::Any, ::Any...) = true
 @inline EnzymeRules.inactive(::typeof(MatrixAlgebraKit.default_pullback_degeneracy_atol), ::Any) = true
 @inline EnzymeRules.inactive(::typeof(MatrixAlgebraKit.default_pullback_rank_atol), ::Any) = true
-@inline EnzymeRules.inactive(::typeof(MatrixAlgebraKit.default_hermitian_tol), ::Any) = true
+@inline EnzymeRules.inactive(::typeof(MatrixAlgebraKit.default_hermitian_tol), ::AbstractMatrix) = true
 
 #----------- NOTE about derivatives ---------
 # Each Enzyme augmented_return + reverse pair

test/enzyme/eig.jl

@@ -0,0 +1,21 @@
+using MatrixAlgebraKit
+using Test
+using LinearAlgebra: Diagonal
+using CUDA, AMDGPU
+
+# infinity-norm doesn't play nicely with Float32, Enzyme, and 1.12
+# see https://github.com/EnzymeAD/Enzyme.jl/issues/2985
+BLASFloats = (Float64, ComplexF64) # full suite is too expensive on CI
+GenericFloats = ()
+@isdefined(TestSuite) || include("../testsuite/TestSuite.jl")
+using .TestSuite
+
+is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
+
+m = 19
+for T in (BLASFloats..., GenericFloats...)
+    TestSuite.seed_rng!(123)
+    if !is_buildkite
+        TestSuite.test_enzyme_eig(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
+    end
+end

test/enzyme/eigh.jl

@@ -0,0 +1,21 @@
+using MatrixAlgebraKit
+using Test
+using LinearAlgebra: Diagonal
+using CUDA, AMDGPU
+
+# infinity-norm doesn't play nicely with Float32, Enzyme, and 1.12
+# see https://github.com/EnzymeAD/Enzyme.jl/issues/2985
+BLASFloats = (Float64, ComplexF64) # full suite is too expensive on CI
+GenericFloats = ()
+@isdefined(TestSuite) || include("../testsuite/TestSuite.jl")
+using .TestSuite
+
+is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
+
+m = 19
+for T in (BLASFloats..., GenericFloats...)
+    TestSuite.seed_rng!(123)
+    if !is_buildkite
+        TestSuite.test_enzyme_eigh(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
+    end
+end

test/enzyme/lq.jl

@@ -0,0 +1,19 @@
+using MatrixAlgebraKit
+using Test
+using LinearAlgebra: Diagonal
+using CUDA, AMDGPU
+
+BLASFloats = (Float32, ComplexF64) # full suite is too expensive on CI
+GenericFloats = ()
+@isdefined(TestSuite) || include("../testsuite/TestSuite.jl")
+using .TestSuite
+
+is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
+
+m = 19
+for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
+    TestSuite.seed_rng!(123)
+    if !is_buildkite
+        TestSuite.test_enzyme_lq(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+    end
+end

test/enzyme/orthnull.jl

@@ -0,0 +1,19 @@
+using MatrixAlgebraKit
+using Test
+using LinearAlgebra: Diagonal
+using CUDA, AMDGPU
+
+BLASFloats = (Float32, ComplexF64) # full suite is too expensive on CI
+GenericFloats = ()
+@isdefined(TestSuite) || include("../testsuite/TestSuite.jl")
+using .TestSuite
+
+is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
+
+m = 19
+for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
+    TestSuite.seed_rng!(123)
+    if !is_buildkite
+        TestSuite.test_enzyme_orthnull(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+    end
+end

test/enzyme/polar.jl

@@ -0,0 +1,19 @@
+using MatrixAlgebraKit
+using Test
+using LinearAlgebra: Diagonal
+using CUDA, AMDGPU
+
+BLASFloats = (Float32, ComplexF64) # full suite is too expensive on CI
+GenericFloats = ()
+@isdefined(TestSuite) || include("../testsuite/TestSuite.jl")
+using .TestSuite
+
+is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
+
+m = 19
+for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
+    TestSuite.seed_rng!(123)
+    if !is_buildkite
+        TestSuite.test_enzyme_polar(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+    end
+end

test/enzyme/projections.jl

@@ -0,0 +1,21 @@
+using MatrixAlgebraKit
+using Test
+using LinearAlgebra: Diagonal
+using CUDA, AMDGPU
+
+BLASFloats = (Float32, ComplexF64) # full suite is too expensive on CI
+GenericFloats = ()
+@isdefined(TestSuite) || include("../testsuite/TestSuite.jl")
+using .TestSuite
+
+is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
+
+m = 19
+for T in (BLASFloats..., GenericFloats...)
+    TestSuite.seed_rng!(123)
+    atol = rtol = m * m * TestSuite.precision(T)
+    if !is_buildkite
+        TestSuite.test_enzyme_projections(T, (m, m); atol, rtol)
+        TestSuite.test_enzyme_projections(Diagonal{T, Vector{T}}, (m, m); atol, rtol)
+    end
+end

test/enzyme/qr.jl

@@ -0,0 +1,19 @@
+using MatrixAlgebraKit
+using Test
+using LinearAlgebra: Diagonal
+using CUDA, AMDGPU
+
+BLASFloats = (Float32, ComplexF64) # full suite is too expensive on CI
+GenericFloats = ()
+@isdefined(TestSuite) || include("../testsuite/TestSuite.jl")
+using .TestSuite
+
+is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
+
+m = 19
+for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
+    TestSuite.seed_rng!(123)
+    if !is_buildkite
+        TestSuite.test_enzyme_qr(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+    end
+end

test/enzyme/svd.jl

@@ -0,0 +1,21 @@
+using MatrixAlgebraKit
+using Test
+using LinearAlgebra: Diagonal
+using CUDA, AMDGPU
+
+# infinity-norm doesn't play nicely with Float32, Enzyme, and 1.12
+# see https://github.com/EnzymeAD/Enzyme.jl/issues/2985
+BLASFloats = (Float64, ComplexF64) # full suite is too expensive on CI
+GenericFloats = ()
+@isdefined(TestSuite) || include("../testsuite/TestSuite.jl")
+using .TestSuite
+
+is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
+
+m = 19
+for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
+    TestSuite.seed_rng!(1234)
+    if !is_buildkite
+        TestSuite.test_enzyme_svd(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+    end
+end

test/runtests.jl

@@ -25,12 +25,12 @@ if filter_tests!(testsuite, args)
         delete!(testsuite, "codequality")
     else
         is_apple_ci = Sys.isapple() && get(ENV, "CI", "false") == "true"
+        is_windows_ci = Sys.isapple() && get(ENV, "CI", "false") == "true"
         if is_apple_ci
-            delete!(testsuite, "enzyme")
             filter!(p -> !startswith(first(p), "mooncake/"), testsuite)
             delete!(testsuite, "chainrules")
         end
-        Sys.iswindows() && delete!(testsuite, "enzyme")
+        (is_windows_ci || is_apple_ci) && filter!(p -> !startswith(first(p), "enzyme/"), testsuite)
     end
 end
 

test/testsuite/TestSuite.jl

@@ -15,6 +15,7 @@ using LinearAlgebra: Diagonal, norm, istriu, istril, I
 using Random, StableRNGs
 using Mooncake
 using AMDGPU, CUDA
+using Enzyme, EnzymeTestUtils
 
 const tests = Dict()
 
@@ -118,7 +119,17 @@ include("mooncake/polar.jl")
 include("mooncake/orthnull.jl")
 include("mooncake/projections.jl")
 
-include("enzyme.jl")
 include("chainrules.jl")
 
+# Enzyme
+# ------
+include("enzyme/eig.jl")
+include("enzyme/eigh.jl")
+include("enzyme/qr.jl")
+include("enzyme/lq.jl")
+include("enzyme/svd.jl")
+include("enzyme/polar.jl")
+include("enzyme/orthnull.jl")
+include("enzyme/projections.jl")
+
 end

test/testsuite/ad_utils.jl

@@ -163,6 +163,28 @@ function call_and_zero!(f!, A, alg)
     return F′
 end
 
+is_cpu(A) = typeof(parent(A)) <: Array
+
+"""
+    project_hermitian_inplace!(A, alg)
+
+Wrapper for `project_hermitian!(A, A, alg)`, invoked this way
+to avoid Enzyme's finite differences comparison getting confused.
+"""
+project_hermitian_inplace!(A, alg) = project_hermitian!(A, A, alg)
+
+"""
+    project_antihermitian_inplace!(A, alg)
+
+Wrapper for `project_hermitian!(A, A, alg)`, invoked this way
+to avoid Enzyme's finite differences comparison getting confused.
+"""
+project_antihermitian_inplace!(A, alg) = project_antihermitian!(A, A, alg)
+
+
+enzyme_fdm(T) = eltype(T) <: Union{Float32, ComplexF32} ? EnzymeTestUtils.FiniteDifferences.central_fdm(5, 1, max_range = 1.0e-2) : EnzymeTestUtils.FiniteDifferences.central_fdm(5, 1)
+
+
 """
     eigh_wrapper(f, A, alg)
 
@@ -246,7 +268,6 @@ function ad_qr_full_setup(A)
     remove_qr_gauge_dependence!(ΔQR..., A, QR...)
     return QR, ΔQR
 end
-
 ad_qr_full_setup(A::Diagonal) = ad_qr_compact_setup(A)
 
 function ad_qr_rank_deficient_compact_setup(A)
@@ -516,8 +537,8 @@ end
 function ad_left_null_setup(A)
     m, n = size(A)
     T = eltype(A)
-    N = left_orth(A; alg = :qr)[1] * randn!(similar(A, T, min(m, n), m - min(m, n)))
-    ΔN = left_orth(A; alg = :qr)[1] * randn!(similar(A, T, min(m, n), m - min(m, n)))
+    N = left_orth(A)[1] * randn!(similar(A, T, min(m, n), m - min(m, n)))
+    ΔN = left_orth(A)[1] * randn!(similar(A, T, min(m, n), m - min(m, n)))
     return N, ΔN
 end
 
@@ -533,7 +554,7 @@ ad_right_orth_setup(A::Diagonal) = ad_left_orth_setup(A)
 function ad_right_null_setup(A)
     m, n = size(A)
     T = eltype(A)
-    Nᴴ = randn!(similar(A, T, n - min(m, n), min(m, n))) * right_orth(A; alg = :lq)[2]
-    ΔNᴴ = randn!(similar(A, T, n - min(m, n), min(m, n))) * right_orth(A; alg = :lq)[2]
+    Nᴴ = randn!(similar(A, T, n - min(m, n), min(m, n))) * right_orth(A)[2]
+    ΔNᴴ = randn!(similar(A, T, n - min(m, n), min(m, n))) * right_orth(A)[2]
     return Nᴴ, ΔNᴴ
 end