update 2

.buildkite/pipeline.yml

@@ -1,5 +1,7 @@
 steps:
   - label: ":julia: Julia {{matrix.julia}} + CUDA GPU"
+    command:
+      - echo 'CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"' >> test/Project.toml
     plugins:
       - JuliaCI/julia#v1:
           version: "{{matrix.julia}}"
@@ -14,15 +16,16 @@ steps:
       queue: "juliagpu"
       cuda: "*"
     env:
+      JULIA_NUM_THREADS: 4
       NNLIB_TEST_CUDA: "true"
-      NNLIB_TEST_CPU: "false"
+      NNLIB_TEST_CPU: "true" # Could be useful to uncover multithreading related issues
+                             # Buildkite workers have more threads.
     if: build.message !~ /\[skip tests\]/
     timeout_in_minutes: 180
     matrix:
       setup:
         julia:
-          # - "1.9" # uncomment when 1.10 is out
-          - "1"
+          - "1.10"
           - "nightly"
       adjustments:
         - with:
@@ -32,8 +35,9 @@ steps:
   - label: ":julia: Julia 1 + AMD GPU"
     plugins:
       - JuliaCI/julia#v1:
-          version: "1"
+          version: "1.10"
       - JuliaCI/julia-test#v1:
+          test_args: "--quickfail"
       - JuliaCI/julia-coverage#v1:
           codecov: true
           dirs:
@@ -49,8 +53,7 @@ steps:
       JULIA_AMDGPU_HIP_MUST_LOAD: "1"
       JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
       NNLIB_TEST_AMDGPU: "true"
-      NNLIB_TEST_CPU: "true" # Could be useful to uncover multithreading related issues
-                             # Buildkite workers have more threads. 
+      NNLIB_TEST_CPU: "false"
       JULIA_NUM_THREADS: 4
 
   - label: "Benchmarks"

.github/workflows/Downstream.yml

@@ -5,6 +5,11 @@ on:
     tags: [v*]
   pull_request:
 
+# needed to allow julia-actions/cache to delete old caches that it has created
+permissions:
+  actions: write
+  contents: read
+
 jobs:
   test:
     name: ${{ matrix.package.repo }}/${{ matrix.package.group }}
@@ -19,24 +24,14 @@ jobs:
         package:
           - {user: FluxML, repo: Flux.jl, group: All}
           - {user: FluxML, repo: Tracker.jl, group: All}
-          - {user: denizyuret, repo: Knet.jl, group: All}
-          - {user: dfdx, repo: Avalon.jl, group: All}
-          - {user: JuliaOptimalTransport, repo: OptimalTransport.jl, group: All}
-          - {user: avik-pal, repo: Lux.jl, group: All}
+          - {user: LuxDL, repo: Lux.jl, group: All}
     steps:
       - uses: actions/checkout@v3
-      # for OptimalTransport.jl
-      - name: Install python
-        uses: actions/setup-python@v4
-        with:
-          python-version: '3.9'
-          architecture: ${{ matrix.arch }}
-      # for OptimalTransport.jl
-      - run: python -m pip install pot
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.julia-version }}
           arch: x64
+      - uses: julia-actions/cache@v1
       - uses: julia-actions/julia-buildpkg@latest
       - name: Clone Downstream
         uses: actions/checkout@v3
@@ -60,3 +55,7 @@ jobs:
             @info "Not compatible with this release. No problem." exception=err
             exit(0)  # Exit immediately, as a success
           end
+        env:
+          RETESTITEMS_NWORKERS: 4
+          BACKEND_GROUP: CPU  # for Lux.jl
+

.github/workflows/ci.yml

@@ -9,6 +9,11 @@ on:
     tags: '*'
   pull_request:
 
+# needed to allow julia-actions/cache to delete old caches that it has created
+permissions:
+  actions: write
+  contents: read
+
 defaults:
   run:
     shell: bash
@@ -24,8 +29,9 @@ jobs:
       matrix:
         version:
           # - '1.9' # uncomment when julia 1.10 is out
+          - 'lts'
           - '1'   # automatically expands to the latest stable 1.x release of Julia
-          - 'nightly'
+          - 'pre'
         os:
           - ubuntu-latest
           # - macOS-latest
@@ -51,20 +57,11 @@ jobs:
 
     steps:
       - uses: actions/checkout@v3
-      - uses: julia-actions/setup-julia@v1
+      - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
-      - uses: actions/cache@v3
-        env:
-          cache-name: cache-artifacts
-        with:
-          path: ~/.julia/artifacts
-          key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
-          restore-keys: |
-            ${{ runner.os }}-test-${{ env.cache-name }}-
-            ${{ runner.os }}-test-
-            ${{ runner.os }}-
+      - uses: julia-actions/cache@v1
       - uses: julia-actions/julia-buildpkg@v1
 
       - name: "Run test without coverage"
@@ -88,22 +85,15 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
-      - uses: julia-actions/setup-julia@v1
+      - uses: julia-actions/setup-julia@v2
         with:
           version: '1.9'
+      - uses: julia-actions/cache@v1
       - run: |
           julia --project=docs -e '
             using Pkg
             Pkg.develop(PackageSpec(path=pwd()))
             Pkg.instantiate()'
-      - run: |
-          julia --color=yes --project=docs/ -e '
-            using NNlib
-            # using Pkg; Pkg.activate("docs")
-            using Documenter
-            using Documenter: doctest
-            DocMeta.setdocmeta!(NNlib, :DocTestSetup, :(using NNlib); recursive=true)
-            doctest(NNlib)'
       - run: julia --project=docs docs/make.jl
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

Project.toml

@@ -1,6 +1,6 @@
 name = "NNlib"
 uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
-version = "0.9.6"
+version = "0.9.30"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -9,49 +9,43 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
+ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
-cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [extensions]
 NNlibAMDGPUExt = "AMDGPU"
-NNlibCUDAExt = "CUDA"
 NNlibCUDACUDNNExt = ["CUDA", "cuDNN"]
+NNlibCUDAExt = "CUDA"
+NNlibEnzymeCoreExt = "EnzymeCore"
+NNlibFFTWExt = "FFTW"
+NNlibForwardDiffExt = "ForwardDiff"
+NNlibSpecialFunctionsExt = "SpecialFunctions"
 
 [compat]
-AMDGPU = "0.5, 0.6"
-Adapt = "3.2"
-Atomix = "0.1"
-ChainRulesCore = "1.13"
+AMDGPU = "1"
+Adapt = "3.2, 4"
+Atomix = "0.1, 1"
 CUDA = "4, 5"
-cuDNN = "1"
-GPUArraysCore = "0.1"
+ChainRulesCore = "1.25"
+EnzymeCore = "0.5, 0.6, 0.7, 0.8"
+FFTW = "1.8.0"
+ForwardDiff = "0.10.36, 1"
+GPUArraysCore = "0.1, 0.2"
 KernelAbstractions = "0.9.2"
-Requires = "1.0"
+LinearAlgebra = "<0.0.1, 1"
+Random = "<0.0.1, 1"
+ScopedValues = "1.3.0"
+SpecialFunctions = "2"
+Statistics = "1"
+cuDNN = "1"
 julia = "1.9"
-
-[extras]
-AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
-Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
-FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
-ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
-ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
-StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
-cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
-
-[targets]
-test = ["AMDGPU", "CUDA", "ChainRulesTestUtils", "Documenter", 
-        "FiniteDifferences", "ForwardDiff", "Logging", "ReverseDiff", 
-        "StableRNGs", "Test", "UnicodePlots", "Zygote", "cuDNN"]

README.md

@@ -25,4 +25,3 @@ for CUDA support, or
 using NNlib, AMDGPU
 ```
 for AMDGPU support.
-

benchmark/perf_report.jl

@@ -37,10 +37,6 @@ for rank in (2,),
             (NNlib.depthwiseconv_im2col!, NNlib.∇depthwiseconv_data_im2col!, NNlib.∇depthwiseconv_filter_im2col!, DepthwiseConvDims, "im2col"),
     ]
 
-    if NNlib.is_nnpack_available()
-        push!(benchmark_items, (NNlib.conv_nnpack!, NNlib.∇conv_data_nnpack!, NNlib.∇conv_filter_nnpack!, DenseConvDims, "nnpack"))
-    end
-
     for (conv!, ∇conv_data!, ∇conv_filter!, cT, backend) in benchmark_items
 
         x = zeros(Float32, repeat([N], rank)..., C_in, 1)
@@ -105,15 +101,4 @@ for rank in (2,),
         @show(pdims)
         @save "results.jld2" results
     end
-
-    if NNlib.is_nnpack_available()
-        if NNlib.nnpack_supported_operation(pdims)
-            t_fwd  = @benchmark NNlib.maxpool_nnpack!($y, $x, $pdims)
-
-            add_result(t_fwd, "maxpool2d", "nnpack", pdims)
-
-            @show(pdims)
-            @save "results.jld2" results
-        end
-    end
 end

docs/.gitignore

@@ -1,2 +1,3 @@
 build/
 site/
+Manifest.toml

docs/Project.toml

@@ -1,5 +1,9 @@
 [deps]
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
-
-[compat]
-Documenter = "0.27"
+FLAC = "abae9e3b-a9a0-4778-b5c6-ca109b507d99"
+FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
+Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
+FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"

docs/make.jl

@@ -1,18 +1,21 @@
 using Documenter, NNlib
 
-DocMeta.setdocmeta!(NNlib, :DocTestSetup, :(using NNlib); recursive = true)
+DocMeta.setdocmeta!(NNlib, :DocTestSetup,
+    :(using FFTW, NNlib, UnicodePlots); recursive = true)
 
 makedocs(modules = [NNlib],
-         sitename = "NNlib.jl",
-         doctest = false,
-         pages = ["Home" => "index.md",
-                  "Reference" => "reference.md"],
-         format = Documenter.HTML(
-              canonical = "https://fluxml.ai/NNlib.jl/stable/",
-            #   analytics = "UA-36890222-9",
-              assets = ["assets/flux.css"],
-              prettyurls = get(ENV, "CI", nothing) == "true"),
-        )
+    sitename = "NNlib.jl",
+    doctest = true,
+    pages = ["Home" => "index.md",
+             "Reference" => "reference.md",
+             "Audio" => "audio.md"],
+    format = Documenter.HTML(
+        canonical = "https://fluxml.ai/NNlib.jl/stable/",
+        # analytics = "UA-36890222-9",
+        assets = ["assets/flux.css"],
+        prettyurls = get(ENV, "CI", nothing) == "true"),
+    warnonly=[:missing_docs,]
+)
 
 deploydocs(repo = "github.com/FluxML/NNlib.jl.git",
            target = "build",

docs/src/audio.md

@@ -0,0 +1,61 @@
+# Reference
+
+!!! note
+    Spectral functions require importing `FFTW` package to enable them.
+
+## Window functions
+
+```@docs
+hann_window
+hamming_window
+```
+
+## Spectral
+
+```@docs
+stft
+istft
+NNlib.power_to_db
+NNlib.db_to_power
+```
+
+## Spectrogram
+
+```@docs
+melscale_filterbanks
+spectrogram
+```
+
+Example:
+
+```@example 1
+using FFTW # <- required for STFT support.
+using NNlib
+using FileIO
+using Makie, CairoMakie
+CairoMakie.activate!()
+
+waveform, sampling_rate = load("./assets/jfk.flac")
+fig = lines(reshape(waveform, :))
+save("waveform.png", fig)
+
+# Spectrogram.
+
+n_fft = 1024
+spec = spectrogram(waveform; n_fft, hop_length=n_fft ÷ 4, window=hann_window(n_fft))
+fig = heatmap(transpose(NNlib.power_to_db(spec)[:, :, 1]))
+save("spectrogram.png", fig)
+
+# Mel-scale spectrogram.
+
+n_freqs = n_fft ÷ 2 + 1
+fb = melscale_filterbanks(; n_freqs, n_mels=128, sample_rate=Int(sampling_rate))
+mel_spec = permutedims(spec, (2, 1, 3)) ⊠ fb # (time, n_mels)
+fig = heatmap(NNlib.power_to_db(mel_spec)[:, :, 1])
+save("mel-spectrogram.png", fig)
+nothing # hide
+```
+
+|Waveform|Spectrogram|Mel Spectrogram|
+|:---:|:---:|:---:|
+|![](waveform.png)|![](spectrogram.png)|![](mel-spectrogram.png)|