Fix Test Loads (#374)

* Change how tests are run * Update action * Fix CI formatting * Fix CI * Fix CI file * Load up Test globally * Fix labelling * Fix a couple of load paths * Improve runner * Fix CI * Fix typo * Fix up CI * Do not specify project * Fix typo * Fix CI * Fix CI * Fix CI * More fixes * Finish rename and fix loading * Fix loading for main tests * Move stuff around * Fix CI * Do not load up front_matter twice * Fix LuxLib tests * Hopefully fix buildkite * Qualify bitcast uses * Qualify more globalrefs * Rule for compilerbarrier * Comment out failing test + leave issue number reference * Maybe fix buildkite * Hmmm * Correct file naming * Hopefully run codecov conditionally * Hopefully fix CI * Fix CUDA runner * Use JET in CUDA extension tests * Enable coverage from buildkite * Split up rules testing * Move coverage plugin to single command section * Update the docs * Tweak pipeline labelling * Bump patch version
compintell · Nov 18, 2024 · 2d2ddbd · 2d2ddbd · willtebbutt · Nov 18, 2024
1 parent c4fbfc8
commit 2d2ddbd
Show file tree

Hide file tree

Showing 42 changed files with 321 additions and 179 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -6,15 +6,16 @@ steps:
     plugins:
       - JuliaCI/julia#v1:
           version: "1"
-      - JuliaCI/julia-test#v1: ~
       - JuliaCI/julia-coverage#v1:
           dirs:
             - src
             - ext
+    command: julia --code-coverage=user --eval='include("test/run_extra.jl")'
     agents:
       queue: "juliagpu"
       cuda: "*"
     if: build.message !~ /\[skip tests\]/
     timeout_in_minutes: 60
     env:
-      TEST_GROUP: "gpu"
+      LABEL: cuda
+      TEST_TYPE: ext
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -20,26 +20,25 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        test_group:
-          - 'aqua'
-          - 'basic'
-          - 'rrules'
-          - 'ext/differentiation_interface'
-          - 'ext/dynamic_ppl'
-          - 'ext/luxlib'
-          - 'ext/nnlib'
-          - 'ext/special_functions'
-          - 'integration_testing/array'
-          - 'integration_testing/bijectors'
-          - 'integration_testing/diff_tests'
-          - 'integration_testing/distributions'
-          - 'integration_testing/gp'
-          - 'integration_testing/logexpfunctions'
-          - 'integration_testing/lux'
-          - 'integration_testing/misc'
-          - 'integration_testing/misc_abstract_array'
-          - 'integration_testing/temporalgps'
-          - 'integration_testing/turing'
+        test_group: [
+          'aqua',
+          'basic',
+          'rrules/avoiding_non_differentiable_code',
+          'rrules/blas',
+          'rrules/builtins',
+          'rrules/fastmath',
+          'rrules/foreigncall',
+          'rrules/functionwrappers',
+          'rrules/iddict',
+          'rrules/lapack',
+          'rrules/linear_algebra',
+          'rrules/low_level_maths',
+          'rrules/memory',
+          'rrules/misc',
+          'rrules/new',
+          'rrules/tasks',
+          'rrules/twice_precision',
+        ]
         version:
           - '1'
         arch:
@@ -48,13 +47,9 @@ jobs:
           - test_group: 'basic'
             version: '1.10'
             arch: x64
-          - test_group: 'integration_testing/turing'
-            version: '1.10'
-            arch: x64
           - test_group: 'basic'
             version: '1.10'
             arch: x86
-
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
@@ -73,6 +68,63 @@ jobs:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
           fail_ci_if_error: false
+  extra:
+    name: ${{matrix.test_group.test_type}}-${{ matrix.test_group.label }}-${{ matrix.version }}-${{ matrix.arch }}
+    runs-on: ubuntu-latest
+    if: github.event_name != 'schedule'
+    strategy:
+      fail-fast: false
+      matrix:
+        test_group: [
+          {test_type: 'ext', label: 'differentiation_interface'},
+          {test_type: 'ext', label: 'dynamic_ppl'},
+          {test_type: 'ext', label: 'luxlib'},
+          {test_type: 'ext', label: 'nnlib'},
+          {test_type: 'ext', label: 'special_functions'},
+          {test_type: 'integration_testing', label: 'array'},
+          {test_type: 'integration_testing', label: 'bijectors'},
+          {test_type: 'integration_testing', label: 'diff_tests'},
+          {test_type: 'integration_testing', label: 'distributions'},
+          {test_type: 'integration_testing', label: 'gp'},
+          {test_type: 'integration_testing', label: 'logexpfunctions'},
+          {test_type: 'integration_testing', label: 'lux'},
+          {test_type: 'integration_testing', label: 'battery_tests'},
+          {test_type: 'integration_testing', label: 'misc_abstract_array'},
+          {test_type: 'integration_testing', label: 'temporalgps'},
+          {test_type: 'integration_testing', label: 'turing'},
+        ]
+        version:
+          - '1'
+        arch:
+          - x64
+        include:
+          - test_group: {test_type: 'integration_testing', label: 'turing'}
+            version: '1.10'
+            arch: x64
+    steps:
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: ${{ matrix.version }}
+          arch: ${{ matrix.arch }}
+          include-all-prereleases: false
+      - uses: julia-actions/cache@v2
+      - uses: julia-actions/julia-buildpkg@v1
+      - run: |
+          if [ ${{ matrix.test_group.test_type }} == 'ext' ]; then
+            julia --code-coverage=user --eval 'include("test/run_extra.jl")'
+          else
+            julia --eval 'include("test/run_extra.jl")'
+          fi
+        env:
+          LABEL: ${{ matrix.test_group.label }}
+          TEST_TYPE: ${{ matrix.test_group.test_type }}
+      - uses: julia-actions/julia-processcoverage@v1
+      - uses: codecov/codecov-action@v4
+        with:
+          files: lcov.info
+          token: ${{ secrets.CODECOV_TOKEN }}
+          fail_ci_if_error: false
   perf:
     name: "Performance (${{ matrix.perf_group }})"
     runs-on: ubuntu-latest

diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Mooncake"
 uuid = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.4.43"
+version = "0.4.44"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

diff --git a/docs/src/running_tests_locally.md b/docs/src/running_tests_locally.md
@@ -1,6 +1,8 @@
 # Running Tests Locally
 
-Mooncake.jl's test suite is fairly extensive. While you can use `Pkg.test` to run the test suite in the standard manner, this is not usually optimal in Mooncake.jl. When editing some code, you typically only want to run the tests associated with it, not the entire test suite.
+Mooncake.jl's test suite is fairly extensive. While you can use `Pkg.test` to run the test suite in the standard manner, this is not usually optimal in Mooncake.jl, and will not run all of the tests. When editing some code, you typically only want to run the tests associated with it, not the entire test suite.
+
+## Main Testing Functionality
 
 Mooncake's tests are organised as follows:
 1. Things that are required for most / all test suites are loaded up in `test/front_matter.jl`.
@@ -21,3 +23,12 @@ The purpose of this approach is to:
 
 If you find that this strategy leaves you running more of the test suite than you would like, consider copy + pasting specific tests into the REPL, or commenting out a chunk of tests in the file that you are editing during development (try not to commit this).
 I find this is rather crude strategy effective in practice.
+
+## Extension and Integration Testing
+
+Mooncake now has quite a lot of package extensions, and a large number of integration tests.
+Unfortunately, these come with a lot of additional dependencies.
+To avoid these dependencies causing CI to take much longer to run, we locate all tests for extensions and integration testing in their own environments. These can be found in the `test/ext` and `test/integration_testing` directories respectively.
+
+These directories comprise a single `.jl` file, and a `Project.toml`.
+You should run these tests by simply `include`ing the `.jl` file. Doing so will activate the environemnt, ensure that the correct version of Mooncake is used, and run the tests.
diff --git a/src/Mooncake.jl b/src/Mooncake.jl
@@ -24,7 +24,7 @@ using Base.Experimental: @opaque
 using Base.Iterators: product
 using Core:
     Intrinsics, bitcast, SimpleVector, svec, ReturnNode, GotoNode, GotoIfNot, PhiNode,
-    PiNode, SSAValue, Argument, OpaqueClosure
+    PiNode, SSAValue, Argument, OpaqueClosure, compilerbarrier
 using Core.Compiler: IRCode, NewInstruction
 using Core.Intrinsics: pointerref, pointerset
 using LinearAlgebra.BLAS: @blasfunc, BlasInt, trsm!

diff --git a/src/rrules/builtins.jl b/src/rrules/builtins.jl
@@ -494,7 +494,11 @@ function rrule!!(f::CoDual{typeof(Core.apply_type)}, args...)
     return CoDual{_typeof(T), NoFData}(T, NoFData()), NoPullback(f, args...)
 end
 
-# Core.compilerbarrier
+function rrule!!(::CoDual{typeof(compilerbarrier)}, setting::CoDual{Symbol}, val::CoDual)
+    compilerbarrier_pb(dout) = NoRData(), NoRData(), dout
+    return compilerbarrier(setting.x, val), compilerbarrier_pb
+end
+
 # Core.donotdelete
 # Core.finalizer
 # Core.get_binding_type
@@ -801,7 +805,7 @@ function generate_hand_written_rrule!!_test_cases(rng_ctor, ::Val{:builtins})
         (true, :stability, nothing, Core._typevar, :T, Union{}, Any),
         (false, :none, _range, Core.apply_type, Vector, Float64),
         (false, :none, _range, Core.apply_type, Array, Float64, 2),
-        # Core.compilerbarrier -- NEEDS IMPLEMENTING AND TESTING
+        (false, :none, nothing, compilerbarrier, :type, 5.0),
         # Core.const_arrayref -- NEEDS IMPLEMENTING AND TESTING
         # Core.donotdelete -- NEEDS IMPLEMENTING AND TESTING
         # Core.finalizer -- NEEDS IMPLEMENTING AND TESTING

diff --git a/test/ext/cuda/Project.toml b/test/ext/cuda/Project.toml
@@ -1,4 +1,6 @@
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/ext/cuda/cuda.jl b/test/ext/cuda/cuda.jl
@@ -2,19 +2,20 @@ using Pkg
 Pkg.activate(@__DIR__)
 Pkg.develop(; path=joinpath(@__DIR__, "..", "..", ".."))
 
-using CUDA, Mooncake, Test
+using CUDA, JET, Mooncake, StableRNGs, Test
+using Mooncake.TestUtils: test_tangent, test_rule
 
 @testset "cuda" begin
 
     # Check we can operate on CuArrays.
     test_tangent(
-        Xoshiro(123456), CuArray{Float32, 2, CUDA.DeviceMemory}(undef, 8, 8);
+        StableRNG(123456), CuArray{Float32, 2, CUDA.DeviceMemory}(undef, 8, 8);
         interface_only=false,
     )
 
     # Check we can instantiate a CuArray.
     test_rule(
-        sr(123456), CuArray{Float32, 1, CUDA.DeviceMemory}, undef, 256;
+        StableRNG(123456), CuArray{Float32, 1, CUDA.DeviceMemory}, undef, 256;
         interface_only=true, is_primitive=true, debug_mode=true,
     )
 end
diff --git a/test/ext/differentiation_interface/di.jl → ...on_interface/differentiation_interface.jl b/test/ext/differentiation_interface/di.jl → ...on_interface/differentiation_interface.jl
diff --git a/test/ext/dynamic_ppl/Project.toml b/test/ext/dynamic_ppl/Project.toml
@@ -1,4 +1,5 @@
 [deps]
 DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/ext/dynamic_ppl/dynamic_ppl.jl b/test/ext/dynamic_ppl/dynamic_ppl.jl
@@ -2,8 +2,10 @@ using Pkg
 Pkg.activate(@__DIR__)
 Pkg.develop(; path=joinpath(@__DIR__, "..", "..", ".."))
 
-using DynamicPPL, Mooncake, Test
+using DynamicPPL, Mooncake, StableRNGs, Test
+using DynamicPPL: istrans, VarInfo
+using Mooncake.TestUtils: test_rule
 
 @testset "DynamicPPLMooncakeExt" begin
-    test_rule(sr(123456), DynamicPPL.istrans, DynamicPPL.VarInfo(); unsafe_perturb=true)
+    test_rule(StableRNG(123456), istrans, VarInfo(); unsafe_perturb=true)
 end
diff --git a/test/ext/luxlib/Project.toml b/test/ext/luxlib/Project.toml
@@ -1,8 +1,10 @@
 [deps]
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]

diff --git a/test/ext/luxlib/luxlib.jl b/test/ext/luxlib/luxlib.jl
@@ -2,8 +2,9 @@ using Pkg
 Pkg.activate(@__DIR__)
 Pkg.develop(; path=joinpath(@__DIR__, "..", "..", ".."))
 
-using Lux, LuxLib, Mooncake, NNlib, Test
+using JET, Lux, LuxLib, Mooncake, NNlib, StableRNGs, Test
 using LuxLib.Impl: SLEEFActivations
+using Mooncake.TestUtils: test_rule
 
 @testset "luxlib" begin
     @testset "$(typeof(fargs))" for (interface_only, perf_flag, is_primitive, fargs...) in vcat(
@@ -65,6 +66,6 @@ using LuxLib.Impl: SLEEFActivations
             )
         end),
     )
-        test_rule(sr(1), fargs...; perf_flag, is_primitive, interface_only)
+        test_rule(StableRNG(123), fargs...; perf_flag, is_primitive, interface_only)
     end
 end
diff --git a/test/ext/nnlib/Project.toml b/test/ext/nnlib/Project.toml
@@ -1,4 +1,6 @@
 [deps]
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/ext/nnlib/nnlib.jl b/test/ext/nnlib/nnlib.jl
@@ -2,7 +2,8 @@ using Pkg
 Pkg.activate(@__DIR__)
 Pkg.develop(; path=joinpath(@__DIR__, "..", "..", ".."))
 
-using Mooncake, NNlib, Test
+using JET, Mooncake, NNlib, StableRNGs, Test
+using Mooncake.TestUtils: test_rule
 using NNlib: dropout
 
 @testset "nnlib" begin
@@ -27,9 +28,18 @@ using NNlib: dropout
         (false, :none, true, batched_mul, randn(3, 2, 3), randn(2, 5, 3)),
 
         # dropout
-        (true, :none, false, (x, p) -> dropout(sr(1), x, p; dims=1), randn(2, 2), 0.5),
-        (true, :none, false, (x, p) -> dropout(sr(1), x, p; dims=2), randn(2, 2), 0.1),
-        (true, :none, false, (x, p) -> dropout(sr(1), x, p; dims=(1, 2)), randn(2, 2), 0.4),
+        (
+            true, :none, false,
+            (x, p) -> dropout(StableRNG(1), x, p; dims=1), randn(2, 2), 0.5,
+        ),
+        (
+            true, :none, false,
+            (x, p) -> dropout(StableRNG(1), x, p; dims=2), randn(2, 2), 0.1,
+        ),
+        (
+            true, :none, false,
+            (x, p) -> dropout(StableRNG(1), x, p; dims=(1, 2)), randn(2, 2), 0.4,
+        ),
 
         # softmax
         (false, :stability, true, softmax, randn(2)),
@@ -106,6 +116,6 @@ using NNlib: dropout
         (false, :none, false, x -> pad_constant(x, 1, 2.0; dims=:), x),
     ]
         @info "$(typeof(fargs))"
-        test_rule(sr(1), fargs...; perf_flag, is_primitive, interface_only)
+        test_rule(StableRNG(123), fargs...; perf_flag, is_primitive, interface_only)
     end
 end
diff --git a/test/ext/special_functions/Project.toml b/test/ext/special_functions/Project.toml
@@ -1,4 +1,7 @@
 [deps]
+AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/ext/special_functions/special_functions.jl b/test/ext/special_functions/special_functions.jl
@@ -2,7 +2,8 @@ using Pkg
 Pkg.activate(@__DIR__)
 Pkg.develop(; path=joinpath(@__DIR__, "..", "..", ".."))
 
-using Mooncake, SpecialFunctions, Test
+using AllocCheck, JET, Mooncake, SpecialFunctions, StableRNGs, Test
+using Mooncake.TestUtils: test_rule
 
 # Rules in this file are only lightly tester, because they are all just @from_rrule rules.
 @testset "special_functions" begin
@@ -44,7 +45,7 @@ using Mooncake, SpecialFunctions, Test
         (:stability_and_allocs, ellipe, 0.3),
         (:stability_and_allocs, logfactorial, 3),
     ]
-        test_rule(Xoshiro(123456), f, x...; perf_flag)
+        test_rule(StableRNG(123456), f, x...; perf_flag)
     end
     @testset for (perf_flag, f, x...) in [
         (:allocs, logerf, 0.3, 0.5), # first branch
@@ -61,6 +62,6 @@ using Mooncake, SpecialFunctions, Test
         (:allocs, SpecialFunctions.loggamma1p, -0.3),
         (:none, SpecialFunctions.lambdaeta, 5.0),
     ]
-        test_rule(Xoshiro(123456), f, x...; perf_flag, is_primitive=false)
+        test_rule(StableRNG(123456), f, x...; perf_flag, is_primitive=false)
     end
 end
diff --git a/test/integration_testing/array/Project.toml b/test/integration_testing/array/Project.toml
@@ -0,0 +1,5 @@
+[deps]
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"