Refactor Enzyme testsuite

Author: kshyatt

test/enzyme.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...)
+    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/eig.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...)
+    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/eigh.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/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_orthnull(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_polar(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_qr(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+    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!(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/enzyme/svd.jl

@@ -26,11 +26,10 @@ if filter_tests!(testsuite, args)
     else
         is_apple_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")
+        (Sys.iswindows() || is_apple_ci) && filter!(p -> !startswith(first(p), "enzyme/"), testsuite)
     end
 end
 

test/runtests.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,16 @@ 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")
+
 end