[documentation] AMDGPU.jl supports rocSPARSE

JuliaSmoothOptimizers · Sep 7, 2023 · 04757ec · 04757ec
1 parent 78f3a35
commit 04757ec
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 22 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -19,25 +19,27 @@ steps:
         include("test/gpu/nvidia.jl")'
     timeout_in_minutes: 30
 
-  # - label: "AMD GPUs -- AMDGPU.jl"
-  #   plugins:
-  #     - JuliaCI/julia#v1:
-  #         version: 1.9
-  #   agents:
-  #     queue: "juliagpu"
-  #     rocm: "*"
-  #     rocmgpu: "gfx1031"
-  #   env:
-  #     JULIA_AMDGPU_CORE_MUST_LOAD: "1"
-  #     JULIA_AMDGPU_HIP_MUST_LOAD: "1"
-  #     JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
-  #   command: |
-  #     julia --color=yes --project -e '
-  #     using Pkg
-  #       Pkg.add("AMDGPU")
-  #       Pkg.instantiate()
-  #       include("test/gpu/amd.jl")'
-  #   timeout_in_minutes: 30
+  - label: "AMD GPUs -- AMDGPU.jl"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: 1.9
+    agents:
+      queue: "juliagpu"
+      rocm: "*"
+      rocmgpu: "gfx1031"
+    env:
+      JULIA_NUM_THREADS: 4
+      JULIA_AMDGPU_CORE_MUST_LOAD: "1"
+      JULIA_AMDGPU_HIP_MUST_LOAD: "1"
+      JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
+    command: |
+      julia --color=yes --project -e '
+        using Pkg
+        # Pkg.add("AMDGPU")
+        Pkg.add(url="https://github.com/JuliaGPU/AMDGPU.jl", rev="master")
+        Pkg.instantiate()
+        include("test/gpu/amd.jl")'
+    timeout_in_minutes: 30
 
   - label: "Intel GPUs -- oneAPI.jl"
     plugins:

diff --git a/docs/src/gpu.md b/docs/src/gpu.md
@@ -171,8 +171,29 @@ if AMDGPU.functional()
 end
 ```
 
-!!! info
-    The library `rocSPARSE` is not interfaced yet in AMDGPU.jl and only dense linear systems are supported.
+Sparse matrices have a specific storage on AMD GPUs (`ROCSparseMatrixCSC`, `ROCSparseMatrixCSR` or `ROCSparseMatrixCOO`):
+
+```julia
+using AMDGPU, Krylov
+using AMDGPU.rocSPARSE, SparseArrays
+
+if AMDGPU.functional()
+  # CPU Arrays
+  A_cpu = sprand(100, 200, 0.3)
+  b_cpu = rand(100)
+
+  # GPU Arrays
+  A_csc_gpu = ROCSparseMatrixCSC(A_cpu)
+  A_csr_gpu = ROCSparseMatrixCSR(A_cpu)
+  A_coo_gpu = ROCSparseMatrixCOO(A_cpu)
+  b_gpu = CuVector(b_cpu)
+
+  # Solve a rectangular and sparse system on an AMD GPU
+  x_csc, y_csc, stats_csc = lnlq(A_csc_gpu, b_gpu)
+  x_csr, y_csr, stats_csr = craig(A_csr_gpu, b_gpu)
+  x_coo, y_coo, stats_coo = craigmr(A_coo_gpu, b_gpu)
+end
+```
 
 ## Intel GPUs
 

diff --git a/test/gpu/amd.jl b/test/gpu/amd.jl
@@ -1,4 +1,4 @@
-using AMDGPU
+using AMDGPU, AMDGPU.rocSPARSE
 
 include("gpu.jl")
 
@@ -14,6 +14,25 @@ include("gpu.jl")
     A_gpu = ROCMatrix(A_cpu)
     b_gpu = ROCVector(b_cpu)
     x, stats = minres(A_gpu, b_gpu)
+    r_gpu = b_gpu - A_gpu * x
+    @test norm(r_gpu) ≤ 1e-4
+
+    A_cpu = sprand(100, 200, 0.3)
+    b_cpu = rand(100)
+    A_csc_gpu = ROCSparseMatrixCSC(A_cpu)
+    A_csr_gpu = ROCSparseMatrixCSR(A_cpu)
+    A_coo_gpu = ROCSparseMatrixCOO(A_cpu)
+    b_gpu = ROXVector(b_cpu)
+    b_gpu = ROCVector(b_cpu)
+    x_csc, y_csc, stats_csc = lnlq(A_csc_gpu, b_gpu)
+    x_csr, y_csr, stats_csr = craig(A_csr_gpu, b_gpu)
+    x_coo, y_coo, stats_coo = craigmr(A_coo_gpu, b_gpu)
+    r_csc = b_gpu - A_csc_gpu * x_csc
+    r_csr = b_gpu - A_csr_gpu * x_csr
+    r_coo = b_gpu - A_coo_gpu * x_coo
+    @test norm(r_csc) ≤ 1e-4
+    @test norm(r_csr) ≤ 1e-4
+    @test norm(r_coo) ≤ 1e-4
   end
 
   for FC in (Float32, Float64, ComplexF32, ComplexF64)