Merge pull request #483 from pxl-th/conditional-miopen

[NNlibAMDGPUExt] Load MIOpen module only if it is available

Author: ToucheSir

.buildkite/pipeline.yml

@@ -55,7 +55,7 @@ steps:
     if: build.pull_request.labels includes "benchmark"
     timeout_in_minutes: 30
 
-  - label: "AMDGPU - Julia 1.9 - No Artifacts"
+  - label: "AMDGPU - Julia 1.9"
     plugins:
       - JuliaCI/julia#v1:
           version: 1.9-nightly
@@ -69,25 +69,17 @@ steps:
       queue: "juliagpu"
       rocm: "*"
       rocmgpu: "*"
-    command:
-      - julia -e """
-        using TOML;
-        conf = TOML.parse(read(\"Project.toml\", String));
-        push!(conf[\"targets\"][\"test\"], \"AMDGPU\");
-        open(io -> TOML.print(io, conf), \"Project.toml\", \"w\");
-        """
-    timeout_in_minutes: 30
+    timeout_in_minutes: 60
     env:
       JULIA_AMDGPU_CORE_MUST_LOAD: "1"
       JULIA_AMDGPU_HIP_MUST_LOAD: "1"
-      JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
       NNLIB_TEST_AMDGPU: true
 
   # - label: "GPU julia nightly"
   #   plugins:
   #     - JuliaCI/julia#v1:
   #         version: "nightly"
-  #     - JuliaCI/julia-test#v1: ~  
+  #     - JuliaCI/julia-test#v1: ~
   #     - JuliaCI/julia-coverage#v1:
   #         codecov: true
   #         dirs:

Project.toml

@@ -26,18 +26,3 @@ julia = "1.6"
 
 [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"
-NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
-Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
-
-[targets]
-test = ["ChainRulesTestUtils", "CUDA", "Documenter", "FiniteDifferences", "ForwardDiff", "Logging", "NNlibCUDA", "Random", "StableRNGs", "Test", "UnicodePlots", "Zygote"]

ext/NNlibAMDGPUExt/NNlibAMDGPUExt.jl

@@ -2,7 +2,6 @@ module NNlibAMDGPUExt
 
 using Adapt
 using AMDGPU
-using AMDGPU.MIOpen
 using ChainRulesCore
 using NNlib
 using NNlib: BatchedAdjoint, BatchedTranspose, BatchedAdjOrTrans
@@ -53,9 +52,18 @@ function nnlib_padding(dims)
     pd[1:2:end]
 end
 
-include("conv.jl")
-include("pool.jl")
-include("softmax.jl")
-include("activations.jl")
+@static if AMDGPU.functional(:MIOpen)
+    using AMDGPU.MIOpen
+
+    include("conv.jl")
+    include("pool.jl")
+    include("softmax.jl")
+    include("activations.jl")
+else
+    @warn """
+    ROCm MIOpen is not available for AMDGPU.
+    NNlib has limited functionality for AMDGPU.
+    """
+end
 
 end

test/Project.toml

@@ -0,0 +1,17 @@
+[deps]
+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"
+NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

test/runtests.jl

@@ -6,32 +6,44 @@ import ForwardDiff
 import Zygote
 using Zygote: gradient
 using StableRNGs
-using CUDA
 using Documenter
 DocMeta.setdocmeta!(NNlib, :DocTestSetup, :(using NNlib, UnicodePlots); recursive=true)
 
 const rng = StableRNG(123)
 include("test_utils.jl")
 
 @testset verbose=true "NNlib.jl" begin
-    if CUDA.functional()
-        if get(ENV, "NNLIB_TEST_CUDA", "false") == "true"
+    if get(ENV, "NNLIB_TEST_CUDA", "false") == "true"
+        using CUDA
+        if CUDA.functional()
             import Pkg
             using NNlibCUDA
             @testset "CUDA" begin
                 Pkg.test("NNlibCUDA")
             end
         else
-            @info "Skipping CUDA tests, set NNLIB_TEST_CUDA=true to run them"
+            @info "Insufficient version or CUDA not found; Skipping CUDA tests"
         end
     else
-        @info "Insufficient version or CUDA not found; Skipping CUDA tests"
+        @info "Skipping CUDA tests, set NNLIB_TEST_CUDA=true to run them"
     end
 
     if get(ENV, "NNLIB_TEST_AMDGPU", "false") == "true"
+        import Pkg
+        test_info = Pkg.project()
+
+        # Add MIOpen_jll to AMDGPU.
+        Pkg.develop("AMDGPU")
+        Pkg.activate(joinpath(Pkg.devdir(), "AMDGPU"))
+        Pkg.add("MIOpen_jll")
+        Pkg.update()
+        # Update test project.
+        Pkg.activate(test_info.path)
+        Pkg.update()
+
         using AMDGPU
+        AMDGPU.versioninfo()
         if AMDGPU.functional() && AMDGPU.functional(:MIOpen)
-            AMDGPU.versioninfo()
             @show AMDGPU.MIOpen.version()
             @testset "AMDGPU" begin
                 include("ext_amdgpu/runtests.jl")