Support ArrayInterface.can_avx and speed up tests on Github Actions

Author: chriselrod

src/costs.jl

@@ -458,13 +458,16 @@ const FUNCTIONSYMBOLS = IdDict{Type{<:Function},Instruction}(
     typeof(<<) => :<<,
     typeof(>>) => :>>,
     typeof(>>>) => :>>>,
+    typeof(%) => :(%),
+    typeof(÷) => :(÷),
     typeof(Base.ifelse) => :ifelse,
     typeof(ifelse) => :ifelse,
     typeof(identity) => :identity,
     typeof(conj) => :conj
 )
 
 # implement whitelist for avx_support that package authors may use to conservatively guard `@avx` application
-# for f ∈ keys(FUNCTIONSYMBOLS)
-#     @eval ArrayInterface.can_avx(::$(typeof(f))) = true
-# end
+for f ∈ keys(FUNCTIONSYMBOLS)
+    @eval ArrayInterface.can_avx(::$f) = true
+end
+

test/broadcast.jl

@@ -83,19 +83,20 @@
         D1 = C .+ A * B;
         D2 = @avx C .+ A .*ˡ B;
         @test D1 ≈ D2
-        fill!(D2, -999999); D2 = @avx C .+ At' *ˡ B;
-        @test D1 ≈ D2
-        fill!(D2, -999999); @test A * B ≈ (@avx @. D2 = A *ˡ B)
-        D1 .= view(C, 1, :)' .+ A * B;
-        fill!(D2, -999999);
-        @avx D2 .= view(C, 1, :)' .+ A .*ˡ B;
-        @test D1 ≈ D2
-        C3d = rand(R,3,M,N);
-        D1 .= view(C3d, 1, :, :) .+ A * B;
-        fill!(D2, -999999);
-        @avx D2 .= view(C3d, 1, :, :) .+ A .*ˡ B;
-        @test D1 ≈ D2
-
+        if RUN_SLOW_TESTS
+            fill!(D2, -999999); D2 = @avx C .+ At' *ˡ B;
+            @test D1 ≈ D2
+            fill!(D2, -999999); @test A * B ≈ (@avx @. D2 = A *ˡ B)
+            D1 .= view(C, 1, :)' .+ A * B;
+            fill!(D2, -999999);
+            @avx D2 .= view(C, 1, :)' .+ A .*ˡ B;
+            @test D1 ≈ D2
+            C3d = rand(R,3,M,N);
+            D1 .= view(C3d, 1, :, :) .+ A * B;
+            fill!(D2, -999999);
+            @avx D2 .= view(C3d, 1, :, :) .+ A .*ˡ B;
+            @test D1 ≈ D2
+        end
         D1 .= 9999;
         @avx D2 .= 9999;
         @test D1 == D2

test/gemm.jl

@@ -710,14 +710,21 @@
                 @test C ≈ C2
                 fill!(C, 9999.999); mulCAtB_2x2blockavx_noinline!(C, A', B);
                 @test C ≈ C2
-                fill!(C, 9999.999); gemm_accurate!(C, A, B);
-                @test C ≈ C2
-                fill!(C, 9999.999); gemm_accurate!(C, At', B);
-                @test C ≈ C2
-                fill!(C, 9999.999); gemm_accurate!(C, A, Bt');
-                @test C ≈ C2
-                fill!(C, 9999.999); gemm_accurate!(C, At', Bt');
-                @test C ≈ C2
+                if RUN_SLOW_TESTS
+                    fill!(C, 9999.999); gemm_accurate!(C, A, B);
+                    @test C ≈ C2
+                    fill!(C, 9999.999); gemm_accurate!(C, At', B);
+                    @test C ≈ C2
+                    fill!(C, 9999.999); gemm_accurate!(C, A, Bt');
+                    @test C ≈ C2
+                    fill!(C, 9999.999); gemm_accurate!(C, At', Bt');
+                    @test C ≈ C2
+                    Ab = zeros(eltype(C), size(A)); Bb = zeros(eltype(C), size(B)); Cb = zero(C);
+                    threegemms!(Ab, Bb, Cb, A, B, C)
+                    @test Ab ≈ C * B'
+                    @test Bb ≈ A' * C
+                    @test Cb ≈ A * B
+                end
                 if iszero(size(A,1) % 8)
                     Abit = A .> 0.5;
                     fill!(C, 9999.999); AmulBavx1!(C, Abit, B);
@@ -728,11 +735,6 @@
                     fill!(C, 9999.999); AmulBavx1!(C, A, Bbit);
                     @test C ≈ A * Bbit
                 end
-                Ab = zeros(eltype(C), size(A)); Bb = zeros(eltype(C), size(B)); Cb = zero(C);
-                threegemms!(Ab, Bb, Cb, A, B, C)
-                @test Ab ≈ C * B'
-                @test Bb ≈ A' * C
-                @test Cb ≈ A * B
             end
             # exceeds_time_limit() && break
             @time @testset "_avx $T dynamic gemm" begin

test/runtests.jl

@@ -2,6 +2,10 @@ using Test
 using LoopVectorization
 using LinearAlgebra
 
+import InteractiveUtils
+
+InteractiveUtils.versioninfo(stdout; verbose = true)
+
 # const START_TIME = time()
 # exceeds_time_limit() = (time() - START_TIME) > 35 * 60
 
@@ -30,12 +34,17 @@ Base.IndexStyle(::Type{<:FallbackArrayWrapper}) = IndexLinear()
 
 @show LoopVectorization.REGISTER_COUNT
 
+const RUN_SLOW_TESTS = LoopVectorization.REGISTER_COUNT ≤ 16 || !parse(Bool, get(ENV, "GITHUB_ACTIONS", "false"))
+@show RUN_SLOW_TESTS
+
 @time @testset "LoopVectorization.jl" begin
 
     @test isempty(detect_unbound_args(LoopVectorization))
 
     @time include("printmethods.jl")
 
+    @time include("can_avx.jl")
+
     @time include("fallback.jl")
 
     @time include("utils.jl")
@@ -44,9 +53,7 @@ Base.IndexStyle(::Type{<:FallbackArrayWrapper}) = IndexLinear()
 
     @time include("check_empty.jl")
 
-    if isnothing(get(ENV, "TRAVIS_BRANCH", nothing)) || LoopVectorization.REGISTER_COUNT ≠ 32 || VERSION ≥ v"1.4"
-        @time include("offsetarrays.jl")
-    end
+    @time include("offsetarrays.jl")
 
     @time include("tensors.jl")
 
@@ -70,8 +77,5 @@ Base.IndexStyle(::Type{<:FallbackArrayWrapper}) = IndexLinear()
 
     @time include("broadcast.jl")
 
-    # I test  locally on master; times out on Travis.
-    if isnothing(get(ENV, "TRAVIS_BRANCH", nothing)) || LoopVectorization.REGISTER_COUNT ≠ 32 || VERSION ≥ v"1.4"
-        @time include("gemm.jl")
-    end
+    @time include("gemm.jl")
 end