make NNlibCUDA an extension (#492)

* make NNlibCUDA an extension

* cuDNN extension

* update workflows

* uncomment when 1.10 is out

* cleanup

* cleanup

* test Project.toml

* cleanup

* rethink extensions

* batchnorm

* cleanup

Author: CarloLucibello

.buildkite/pipeline.yml

@@ -1,44 +1,53 @@
 steps:
-  - label: "GPU julia v1.6"
+  - label: ":julia: Julia {{matrix.julia}} + CUDA GPU"
     plugins:
       - JuliaCI/julia#v1:
-          version: "1.6"
-      - JuliaCI/julia-test#v1: ~
+          version: "{{matrix.julia}}"
+      - JuliaCI/julia-test#v1:
+          test_args: "--quickfail"
       - JuliaCI/julia-coverage#v1:
           codecov: true
           dirs:
             - src
-#     commands:
-#       - julia --project=test -e """
-#         Pkg.develop(url = \"https://github.com/FluxML/NNlibCUDA.jl\")
-#         Pkg.instantiate()
-#         Pkg.build()
-#         Pkg.status()
-#         Pkg.test()
-#         Pkg.test(\"NNlibCUDA\")
-#         """
+            - ext
     agents:
       queue: "juliagpu"
       cuda: "*"
     env:
       NNLIB_TEST_CUDA: true
+    if: build.message !~ /\[skip tests\]/
     timeout_in_minutes: 60
+    matrix:
+      setup:
+        julia:
+          # - "1.9" # uncomment when 1.10 is out
+          - "1"
+          - "nightly"
+      adjustments:
+        - with:
+            julia: "nightly"
+          soft_fail: true
 
-  - label: "GPU julia v1"
+  - label: ":julia: Julia 1 + AMD GPU"
     plugins:
       - JuliaCI/julia#v1:
           version: "1"
-      - JuliaCI/julia-test#v1: ~
+      - JuliaCI/julia-test#v1:
       - JuliaCI/julia-coverage#v1:
           codecov: true
           dirs:
             - src
+            - ext
     agents:
       queue: "juliagpu"
-      cuda: "*"
-    env:
-      NNLIB_TEST_CUDA: true
+      rocm: "*"
+      rocmgpu: "*"
     timeout_in_minutes: 60
+    env:
+      JULIA_AMDGPU_CORE_MUST_LOAD: "1"
+      JULIA_AMDGPU_HIP_MUST_LOAD: "1"
+      NNLIB_TEST_AMDGPU: true
+
 
   - label: "Benchmarks"
     plugins:
@@ -55,38 +64,5 @@ steps:
     if: build.pull_request.labels includes "benchmark"
     timeout_in_minutes: 30
 
-  - label: "AMDGPU - Julia 1.9"
-    plugins:
-      - JuliaCI/julia#v1:
-          version: 1.9-nightly
-      - JuliaCI/julia-test#v1:
-      - JuliaCI/julia-coverage#v1:
-          codecov: true
-          dirs:
-            - src
-            - ext
-    agents:
-      queue: "juliagpu"
-      rocm: "*"
-      rocmgpu: "*"
-    timeout_in_minutes: 60
-    env:
-      JULIA_AMDGPU_CORE_MUST_LOAD: "1"
-      JULIA_AMDGPU_HIP_MUST_LOAD: "1"
-      NNLIB_TEST_AMDGPU: true
-
-  # - label: "GPU julia nightly"
-  #   plugins:
-  #     - JuliaCI/julia#v1:
-  #         version: "nightly"
-  #     - JuliaCI/julia-test#v1: ~
-  #     - JuliaCI/julia-coverage#v1:
-  #         codecov: true
-  #         dirs:
-  #           - src
-  #   agents:
-  #     queue: "juliagpu"
-  #     cuda: "*"
-  #   timeout_in_minutes: 60
 env:
   SECRET_CODECOV_TOKEN: "IlEMvDI6RciJQr5eX7qBBpHYFAe8+Svf3lNJh9gZi0MeJZQvMZWzHfW/lVncA9d9K+gDBBTv/zwqF86xOaIFLuACNdcGZiGgHS+NGeXN5CEppjqLnqKuaeHmLgJ43jygxRwgF88LhwTGcHG7pmESIp1Bn3Jd23UUv4t8hJLBDF+KJLZMefzCXnEVzfwJYxhJktnKJPA4dOv59w33Vj1x5uCYZbQlLP54IJPBm8UGdXS+JrUX8Z7lhxbkJUi6c+R6cvVBw27uRjF0pUJY26mt1frx8MzTGTOweXTpi+Kc5JhzlokMlan17j6T/b7qMC13IuKopfqu1GhkSBQD3ZhQqA==;U2FsdGVkX19l7JMB48k4oJHLoaqC7/MmvQWmaiBxRN472ZC6AcQ0uCBRy6Fw8tI0YcjIxKDScaBnJ2v/deOfhg=="

.github/workflows/CompatHelper.yml

@@ -10,6 +10,16 @@ jobs:
   CompatHelper:
     runs-on: ubuntu-latest
     steps:
+      - name: Check if Julia is already available in the PATH
+        id: julia_in_path
+        run: which julia
+        continue-on-error: true
+      - name: Install Julia, but only if it is not already available in the PATH
+        uses: julia-actions/setup-julia@v1
+        with:
+          version: '1'
+          arch: ${{ runner.arch }}
+        if: steps.julia_in_path.outcome != 'success'
       - name: "Add the General registry via Git"
         run: |
           import Pkg
@@ -32,4 +42,4 @@ jobs:
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           COMPATHELPER_PRIV: ${{ secrets.DOCUMENTER_KEY }}
-          # COMPATHELPER_PRIV: ${{ secrets.COMPATHELPER_PRIV }}
+          # COMPATHELPER_PRIV: ${{ secrets.COMPATHELPER_PRIV }}

.github/workflows/TagBot.yml

@@ -1,13 +1,33 @@
 name: TagBot
 on:
-  issue_comment:  # https://discourse.julialang.org/t/ann-required-updates-to-tagbot-yml/49249
+  issue_comment:
     types:
       - created
+  workflow_dispatch:
+    inputs:
+      lookback:
+        default: 3
+permissions:
+  actions: read
+  checks: read
+  contents: write
+  deployments: read
+  issues: read
+  discussions: read
+  packages: read
+  pages: read
+  pull-requests: read
+  repository-projects: read
+  security-events: read
+  statuses: read
 jobs:
   TagBot:
+    if: github.event_name == 'workflow_dispatch' || github.actor == 'JuliaTagBot'
     runs-on: ubuntu-latest
     steps:
       - uses: JuliaRegistries/TagBot@v1
         with:
           token: ${{ secrets.GITHUB_TOKEN }}
+          # Edit the following line to reflect the actual name of the GitHub Secret containing your private key
           ssh: ${{ secrets.DOCUMENTER_KEY }}
+          # ssh: ${{ secrets.NAME_OF_MY_SSH_PRIVATE_KEY_SECRET }}

.github/workflows/ci.yml

@@ -15,15 +15,15 @@ defaults:
 
 jobs:
   test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.julia-threads }} thread(s) - ${{ github.event_name }} 
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.julia-threads }} thread(s) 
     runs-on: ${{ matrix.os }}
     env:
       JULIA_NUM_THREADS: ${{ matrix.julia-threads }}
     strategy:
       fail-fast: false
       matrix:
         version:
-          - '1.6'
+          # - '1.9' # uncomment when julia 1.10 is out
           - '1'   # automatically expands to the latest stable 1.x release of Julia
           - 'nightly'
         os:
@@ -69,17 +69,17 @@ jobs:
 
       - name: "Run test without coverage"
         uses: julia-actions/julia-runtest@v1
-        if: ${{ !contains(fromJson('["1", "1.6"]'), matrix.version) || matrix.os != 'ubuntu-latest' }}
+        if: ${{ !contains(fromJson('["1", "1.9"]'), matrix.version) || matrix.os != 'ubuntu-latest' }}
         with:
           coverage: false
 
       - name: "Run test with coverage"
         uses: julia-actions/julia-runtest@v1
-        if: contains(fromJson('["1", "1.6"]'), matrix.version) && matrix.os == 'ubuntu-latest'
+        if: contains(fromJson('["1", "1.9"]'), matrix.version) && matrix.os == 'ubuntu-latest'
       - uses: julia-actions/julia-processcoverage@v1
-        if: contains(fromJson('["1", "1.6"]'), matrix.version) && matrix.os == 'ubuntu-latest'
+        if: contains(fromJson('["1", "1.9"]'), matrix.version) && matrix.os == 'ubuntu-latest'
       - uses: codecov/codecov-action@v3
-        if: contains(fromJson('["1", "1.6"]'), matrix.version) && matrix.os == 'ubuntu-latest'
+        if: contains(fromJson('["1", "1.9"]'), matrix.version) && matrix.os == 'ubuntu-latest'
         with:
           file: lcov.info
 
@@ -90,7 +90,7 @@ jobs:
       - uses: actions/checkout@v3
       - uses: julia-actions/setup-julia@v1
         with:
-          version: '1.6'
+          version: '1.9'
       - run: |
           julia --project=docs -e '
             using Pkg

.gitignore

@@ -12,7 +12,7 @@ deps.jl
 *.log
 .vscode/
 /Manifest.toml
-lib/NNlibCUDA/Manifest.toml
+test/Manifest.toml
 benchmark/Manifest.toml
 benchmark/*.json
 benchmark/report.md

Project.toml

@@ -1,6 +1,6 @@
 name = "NNlib"
 uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
-version = "0.8.20"
+version = "0.9.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -16,19 +16,22 @@ 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"
 
 [extensions]
 NNlibAMDGPUExt = "AMDGPU"
+NNlibCUDAExt = "CUDA"
+NNlibCUDACUDNNExt = ["CUDA", "cuDNN"]
 
 [compat]
 AMDGPU = "0.4.8"
-Adapt = "2, 3.2"
+Adapt = "3.2"
 Atomix = "0.1"
 ChainRulesCore = "1.13"
+CUDA = "4"
+cuDNN = "1"
 GPUArraysCore = "0.1"
 KernelAbstractions = "0.9.2"
-Requires = "0.5, 1.0"
-julia = "1.6"
-
-[extras]
-AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+Requires = "1.0"
+julia = "1.9"

README.md

@@ -16,4 +16,13 @@ This package provides a library of functions useful for neural networks, such as
 
 For use with automatic differentiation, this package defines gradients using [ChainRules.jl](https://github.com/JuliaDiff/ChainRules.jl). These will be seen by various packages including [Zygote.jl](https://github.com/FluxML/Zygote.jl).
 
-To use these functions with [CUDA.jl](https://github.com/JuliaGPU/CUDA.jl) you will need [NNlibCUDA.jl](https://github.com/FluxML/NNlibCUDA.jl) as well.
+GPU support is provided as package extensions (see the `ext/` folder). In order to load the extensions, use the imports
+```julia
+using NNlib, CUDA, cuDNN
+```
+for CUDA support, or
+```julia
+using NNlib, AMDGPU
+```
+for AMDGPU support.
+

docs/src/index.md

@@ -4,5 +4,12 @@
 
 For use with automatic differentiation, this package defines gradients using [ChainRules.jl](https://github.com/JuliaDiff/ChainRules.jl). These will be seen by various packages including [Zygote.jl](https://github.com/FluxML/Zygote.jl).
 
-To use these functions with [CUDA.jl](https://github.com/JuliaGPU/CUDA.jl) you will need [NNlibCUDA.jl](https://github.com/FluxML/NNlibCUDA.jl) as well.
-For [AMDGPU.jl](https://github.com/JuliaGPU/AMDGPU.jl) you will need to load it and NNlib in the same Julia session.
+GPU support is provided as package extensions. In order to load the extensions, use the imports
+```julia
+using NNlib, CUDA, cuDNN
+```
+for CUDA support, or
+```julia
+using NNlib, AMDGPU
+```
+for AMDGPU support.