diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index c296bb0..2320147 100644 --- a/.github/workflows/CI.yml +++ b/.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 diff --git a/Project.toml b/Project.toml index eadf0f9..deb57c1 100644 --- a/Project.toml +++ b/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"] diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl index c930723..2cca59c 100644 --- a/ext/CUDAExt.jl +++ b/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 \ No newline at end of file diff --git a/ext/MetalExt.jl b/ext/MetalExt.jl index 42a1bb6..e9c484b 100644 --- a/ext/MetalExt.jl +++ b/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 \ No newline at end of file diff --git a/test/CUDATests.jl b/test/CUDATests.jl index 62283f4..8970bf3 100644 --- a/test/CUDATests.jl +++ b/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 diff --git a/test/MetalTests.jl b/test/MetalTests.jl index c1acef3..5d78943 100644 --- a/test/MetalTests.jl +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index 453abb8..0587261 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -2,7 +2,7 @@ using ScatteringTransform using ContinuousWavelets using AbstractFFTs, FFTW -using Test, LinearAlgebra, Statistics +using Test, LinearAlgebra, Statistics, Adapt using Flux, FourierFilterFlux, MonogenicFilterFlux using Zygote @@ -19,20 +19,28 @@ const GROUP = get(ENV, "GROUP", "All") include("2DTests.jl") if GROUP in ("All", "CUDA") - try - using CUDA, cuDNN, cuFFT - include("CUDATests.jl") + haveCUDA = try + @eval using CUDA, cuFFT + true catch e - @info "CUDA/cuDNN/cuFFT not available in this environment -- skipping CUDATests.jl" exception=e + @info "CUDA/cuFFT not available in this environment -- skipping CUDATests.jl" exception=e + false + end + if haveCUDA + include("CUDATests.jl") end end if GROUP in ("All", "Metal") - try - using Metal - include("MetalTests.jl") + haveMetal = try + @eval using Metal + true catch e - @info "Metal not available in this environment -- skipping MetalTests.jl" exception=e + @info "Metal.jl not installed -- skipping Metal tests" exception = e + false + end + if haveMetal + include("MetalTests.jl") end end end \ No newline at end of file