Remove cuDNN from CUDA dependencies, removed BenchmarkTools, added mean pooling for Metal, updated testing

.github/workflows/CI.yml

@@ -25,29 +25,24 @@ env:
 
 jobs:
   test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.group }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
-        version:
-          - '1.12'
-        os:
-          - ubuntu-latest
-          - macOS-latest   # Apple Silicon (ARM64)
-        arch:
-          - x64
-        exclude:
-          - os: macOS-latest
-            arch: x64
         include:
+          - os: ubuntu-latest
+            arch: x64
+            version: '1.12'
+            group: CUDA
           - os: macOS-latest
             arch: aarch64
             version: '1.12'
+            group: Metal
 
     steps:
       - uses: actions/checkout@v4
-      
+
       - name: Check out ContinuousWavelets.jl
         uses: actions/checkout@v4
         with:
@@ -116,9 +111,7 @@ jobs:
             using Pkg
             Pkg.add([
               PackageSpec(name = "CUDA"),
-              PackageSpec(name = "cuDNN"),
               PackageSpec(name = "cuFFT"),
-              PackageSpec(name = "CUDNN_jll", version = "9.20"),
             ])'
       - name: Add Metal (macOS only)
         if: runner.os == 'macOS'
@@ -127,12 +120,35 @@ jobs:
       - name: Precompile
         run: julia --project=ci -e 'using Pkg; Pkg.precompile()'
 
+      - name: Environment status
+        run: julia --project=ci -e 'using Pkg; Pkg.status()'
+
+      - name: Report GPU availability
+        continue-on-error: true
+        run: |
+          julia --project=ci -e '
+            if Sys.islinux()
+                using CUDA
+                @info "CUDA backend" functional = CUDA.functional()
+            elseif Sys.isapple()
+                using Metal
+                @info "Metal backend" functional = Metal.functional()
+            end'
+
       - name: Run tests
         env:
-          GROUP: ${{ runner.os == 'Linux' && 'CUDA' || runner.os == 'macOS' && 'Metal' || 'CPU' }}
+          GROUP: ${{ matrix.group }}
         run: julia --project=ci --code-coverage=user -e 'include("test/runtests.jl")'
 
       - uses: julia-actions/julia-processcoverage@v1
+
+      - uses: actions/upload-artifact@v4
+        if: always()
+        with:
+          name: lcov-${{ matrix.os }}-${{ matrix.group }}
+          path: lcov.info
+          if-no-files-found: warn
+
       - uses: codecov/codecov-action@v4
         with:
           files: lcov.info

Project.toml

@@ -9,7 +9,6 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ContinuousWavelets = "96eb917e-2868-4417-9cb6-27e7ff17528f"
-CUDNN_jll = "62b44479-cb7b-5706-934f-f13b2eb2e645"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
@@ -25,14 +24,12 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-CUDA_Runtime_Discovery = "1af6417a-86b4-443c-805f-a4643ffb695f"
-cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 cuFFT = "533571aa-0936-420e-b4be-9c66f5f626ca"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 
 [extensions]
-CUDAExt = ["CUDA", "cuDNN", "cuFFT"]
+CUDAExt = ["CUDA", "cuFFT"]
 MetalExt = "Metal"
 ScatteringPlotsExt = "Plots"
 
@@ -41,20 +38,21 @@ AbstractFFTs = "1"
 Adapt = "4"
 ContinuousWavelets = "1"
 CUDA = "6.1"
-cuDNN = "6.1"
-CUDNN_jll = "~9.20"
 cuFFT = "6.1"
 FFTW = "1"
 Flux = "0.13, 0.14, 0.15, 0.16"
 Functors = "0.5.2"
 julia = "1.12"
+Metal = "1"
 Plots = "1.41.6"
+Statistics = "1"
 Wavelets = "0.9, 0.10"
 Zygote = "0.6, 0.7"
 
 [extras]
-BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "BenchmarkTools"]
+test = ["Test", "Logging", "Random"]

ext/CUDAExt.jl

@@ -8,13 +8,11 @@ module CUDAExt
     FourierFilterFlux is loaded.
 
     This file's real job is triggering activation: `[extensions] CUDAExt =
-    ["CUDA", "cuDNN", "cuFFT"]` in Project.toml means all three must be
-    loaded before this extension (and therefore this package's CUDA support)
-    activates at all, so this `using` line is what actually makes that
-    happen, not a formality. It also guarantees FourierFilterFlux's and
-    ContinuousWavelets' own CUDA extensions are active by the time anything
-    in ScatteringTransform runs on the GPU, since their trigger sets
-    (CUDA+cuFFT, and CUDA+cuDNN+cuFFT, respectively) are subsets of this one. =#
-using ScatteringTransform, CUDA, cuDNN, cuFFT
+    ["CUDA", "cuFFT"]` in Project.toml means both must be loaded before this 
+    extension (and therefore this package's CUDA support) activates at all. 
+    It also guarantees FourierFilterFlux's and ContinuousWavelets' own CUDA 
+    extensions are active by the time anything in ScatteringTransform runs on 
+    the GPU, since their trigger sets are the same as this one. =#
+using ScatteringTransform, CUDA, cuFFT
 
 end

ext/MetalExt.jl

@@ -2,4 +2,94 @@ module MetalExt
 
 using ScatteringTransform, Metal
 
+#= NNlib ships no Metal pooling kernel, so meanpool! falls through to 
+meanpool_direct!, which scalar-indexes and is therefore rejected on MtlArray.
+cuDNN provides the equivalent on the CUDA side, so this only shows up on Metal.
+
+    NNlib normalises 3-D and 4-D pooling to the 5-D method, so implementing 
+the 5-D case intercepts every call. Mean pooling is prod(kernel) strided reads
+and an accumulate: for a fixed window offset the source indices are strided and 
+non-overlapping, so this is race-free without atomics.
+
+    This is type piracy in the strict sense, as neither NNlib.meanpool! nor 
+MtlArray is ours, but it is the same shape as what NNlibCUDA does for CUDA, and 
+it lives in a Metal-triggered extension. =#
+
+using NNlib
+using NNlib: PoolDims, kernel_size, stride, dilation, padding, output_size
+
+function _mtlZeroPad(x::MtlArray{T,5}, p::NTuple{6,Int}) where {T}
+    all(iszero, p) && return x
+    s = size(x)
+    xp = similar(x, T, (s[1] + p[1] + p[2], s[2] + p[3] + p[4],
+                        s[3] + p[5] + p[6], s[4], s[5]))
+    fill!(xp, zero(T))
+    @views xp[p[1].+(1:s[1]), p[3].+(1:s[2]), p[5].+(1:s[3]), :, :] .= x
+    return xp
+end
+
+# indices of the window-offset `c` (0-based) along one spatial axis
+_poolIdx(c, st, dil, nOut) = (1 + c * dil):st:(1 + c * dil + st * (nOut - 1))
+
+function NNlib.meanpool!(y::MtlArray{T,5}, x::MtlArray{T,5}, pdims::PoolDims;
+                         alpha = true, beta = false, kwargs...) where {T}
+    kern = kernel_size(pdims)
+    st = stride(pdims)
+    dil = dilation(pdims)
+    o1, o2, o3 = output_size(pdims)
+    xp = _mtlZeroPad(x, padding(pdims))
+
+    acc = similar(y)
+    fill!(acc, zero(T))
+    for c3 = 0:kern[3]-1, c2 = 0:kern[2]-1, c1 = 0:kern[1]-1
+        i1 = _poolIdx(c1, st[1], dil[1], o1)
+        i2 = _poolIdx(c2, st[2], dil[2], o2)
+        i3 = _poolIdx(c3, st[3], dil[3], o3)
+        @views acc .+= xp[i1, i2, i3, :, :]
+    end
+    acc ./= T(prod(kern))
+
+    if iszero(beta)
+        y .= T(alpha) .* acc
+    else
+        y .= T(alpha) .* acc .+ T(beta) .* y
+    end
+    return y
+end
+
+function NNlib.∇meanpool!(dx::MtlArray{T,5}, dy::MtlArray{T,5}, y::MtlArray{T,5},
+                          x::MtlArray{T,5}, pdims::PoolDims;
+                          alpha = true, beta = false, kwargs...) where {T}
+    kern = kernel_size(pdims)
+    st = stride(pdims)
+    dil = dilation(pdims)
+    p = padding(pdims)
+    o1, o2, o3 = output_size(pdims)
+
+    scaled = dy ./ T(prod(kern))
+    padded = !all(iszero, p)
+    s = size(dx)
+    dxp = padded ?
+          fill!(similar(dx, T, (s[1] + p[1] + p[2], s[2] + p[3] + p[4],
+                                s[3] + p[5] + p[6], s[4], s[5])), zero(T)) :
+          fill!(dx, zero(T))
+
+    # each offset writes a strided, non-overlapping set of destinations, so the
+    # accumulate is safe; the overlap between offsets is serialised by the loop. 
+    for c3 = 0:kern[3]-1, c2 = 0:kern[2]-1, c1 = 0:kern[1]-1
+        i1 = _poolIdx(c1, st[1], dil[1], o1)
+        i2 = _poolIdx(c2, st[2], dil[2], o2)
+        i3 = _poolIdx(c3, st[3], dil[3], o3)
+        @views dxp[i1, i2, i3, :, :] .+= scaled
+    end
+
+    if padded
+        @views dx .= dxp[p[1].+(1:s[1]), p[3].+(1:s[2]), p[5].+(1:s[3]), :, :]
+    end
+    if !iszero(alpha) && alpha != true
+        dx .*= T(alpha)
+    end
+    return dx
+end
+
 end

test/CUDATests.jl

@@ -1,5 +1,3 @@
-using BenchmarkTools
-
 const cuda_available = @isdefined(CUDA) && CUDA.functional()
 
 @testset "GPU Tests (CUDA)" begin
@@ -148,9 +146,20 @@ const cuda_available = @isdefined(CUDA) && CUDA.functional()
             end
         end
 
+        @testset "CUDA FFT plan round-trip" begin
+            x = randn(Float32, 16, 2, 3)
+            for (P, inp) in ((plan_rfft(x, 1:1), x), (plan_fft(complex(x), 1:1), complex(x)))
+                Q = adapt(Array, adapt(CuArray, P))
+                @test typeof(Q) == typeof(P)
+                @test Q.region == P.region
+                @test size(Q) == size(P)
+                @test Q * inp ≈ P * inp
+            end
+        end
+
         @testset "CPU/GPU timing" begin
             sizes = [256, 2048, 16384, 131072]
-            cpu_max_size = 16384
+            nSamples = 5 # Number of runs we test for CPU and GPU timing.
 
             for sz in sizes
                 GC.gc()
@@ -161,25 +170,17 @@ const cuda_available = @isdefined(CUDA) && CUDA.functional()
                 sstGPU = gpu(sst)
                 initGPU = gpu(init)
 
-                if sz > cpu_max_size
-                    CUDA.@sync sstGPU(initGPU)  # warmup
-                    GC.gc(); CUDA.reclaim()
-                    tGPU = @elapsed (CUDA.@sync sstGPU(initGPU))
-                else
-                    tGPU = @belapsed (CUDA.@sync $sstGPU($initGPU))
-                end
+                sst(init)
+                CUDA.@sync sstGPU(initGPU)
+                GC.gc(); CUDA.reclaim()
 
-                tCPU = @belapsed $sst($init)
-                speedup = tCPU / tGPU
-                @info "size=$sz" tCPU tGPU speedup
-                if sz >= 512
-                    @test tGPU < tCPU
-                end
+                tCPU = minimum(@elapsed(sst(init)) for _ = 1:nSamples)
+                tGPU = minimum(@elapsed(CUDA.@sync sstGPU(initGPU)) for _ = 1:nSamples)
+                @info "size=$sz" tCPU tGPU speedup = tCPU / tGPU
 
                 sstGPU = nothing
                 initGPU = nothing
-                GC.gc()
-                CUDA.reclaim()
+                GC.gc(); CUDA.reclaim()
             end
         end
     end

test/MetalTests.jl

@@ -1,5 +1,3 @@
-using BenchmarkTools
-
 const metal_available = @isdefined(Metal) && Metal.functional()
 
 @testset "GPU Tests (Metal)" begin
@@ -148,9 +146,20 @@ const metal_available = @isdefined(Metal) && Metal.functional()
             end
         end
 
+        @testset "Metal FFT plan round-trip" begin
+            x = randn(Float32, 16, 2, 3)
+            for (P, inp) in ((plan_rfft(x, 1:1), x), (plan_fft(complex(x), 1:1), complex(x)))
+                Q = adapt(Array, Metal.mtl(P))
+                @test typeof(Q) == typeof(P)
+                @test Q.region == P.region
+                @test size(Q) == size(P)
+                @test Q * inp ≈ P * inp
+            end
+        end
+
         @testset "CPU/GPU timing" begin
             sizes = [256, 2048, 16384, 131072]
-            cpu_max_size = 16384
+            nSamples = 5 # Number of runs we test for CPU and GPU timing. 
 
             for sz in sizes
                 GC.gc()
@@ -160,20 +169,13 @@ const metal_available = @isdefined(Metal) && Metal.functional()
                 sstGPU = gpu(sst)
                 initGPU = gpu(init)
 
-                if sz > cpu_max_size
-                    Metal.@sync sstGPU(initGPU)  # warmup
-                    GC.gc()
-                    tGPU = @elapsed (Metal.@sync sstGPU(initGPU))
-                else
-                    tGPU = @belapsed (Metal.@sync $sstGPU($initGPU))
-                end
+                sst(init)
+                Metal.@sync sstGPU(initGPU)
+                GC.gc()
 
-                tCPU = @belapsed $sst($init)
-                speedup = tCPU / tGPU
-                @info "size=$sz" tCPU tGPU speedup
-                if sz >= 512
-                    @test tGPU < tCPU
-                end
+                tCPU = minimum(@elapsed(sst(init)) for _ = 1:nSamples)
+                tGPU = minimum(@elapsed(Metal.@sync sstGPU(initGPU)) for _ = 1:nSamples)
+                @info "size=$sz" tCPU tGPU speedup = tCPU / tGPU
 
                 sstGPU = nothing
                 initGPU = nothing