Fix API

JuliaSmoothOptimizers · Nov 30, 2023 · 0ccd647 · 0ccd647
1 parent 04678e6
commit 0ccd647
Show file tree

Hide file tree

Showing 9 changed files with 21 additions and 17 deletions.
diff --git a/ext/AMDGPU/blockjacobi.jl b/ext/AMDGPU/blockjacobi.jl
@@ -33,7 +33,7 @@ function _update_gpu(p, j_rowptr, j_colval, j_nzval, device::ROCBackend)
 end
 
 """
-    function update(J::ROCSparseMatrixCSR, p)
+    function update!(J::ROCSparseMatrixCSR, p)
 
 Update the preconditioner `p` from the sparse Jacobian `J` in CSR format for ROCm
 
@@ -42,6 +42,6 @@ Update the preconditioner `p` from the sparse Jacobian `J` in CSR format for ROC
 3) Extract the preconditioner matrix `p.P` from the dense blocks `cuJs`
 
 """
-function KP.update(p::BlockJacobiPreconditioner, J::rocSPARSE.ROCSparseMatrixCSR, device::ROCBackend)
-    _update_gpu(p, J.rowPtr, J.colVal, J.nzVal, device)
+function KP.update!(p::BlockJacobiPreconditioner, J::rocSPARSE.ROCSparseMatrixCSR)
+    _update_gpu(p, J.rowPtr, J.colVal, J.nzVal, p.device)
 end
diff --git a/ext/AMDGPU/ic0.jl b/ext/AMDGPU/ic0.jl
@@ -45,6 +45,6 @@ for ArrayType in (:(ROCVector{T}), :(ROCMatrix{T}))
   end
 end
 
-function KP.update(p::AMD_ICU0, A, device::CUDABackend)
-    p.P = CUSPARSE.ic02(A, 'O')
+function KP.update!(p::AMD_IC0, A)
+    p.P = rocSPARSE.ic0(A, 'O')
 end
diff --git a/ext/AMDGPU/ilu0.jl b/ext/AMDGPU/ilu0.jl
@@ -45,6 +45,6 @@ for ArrayType in (:(ROCVector{T}), :(ROCMatrix{T}))
   end
 end
 
-function KP.update!(p::AMD_ILU0, A, device::CUDABackend)
+function KP.update!(p::AMD_ILU0, A)
     p.P = rocSPARSE.ilu0(A, 'O')
 end
diff --git a/ext/CUDA/blockjacobi.jl b/ext/CUDA/blockjacobi.jl
@@ -33,7 +33,7 @@ function _update_gpu(p, j_rowptr, j_colval, j_nzval, device::CUDABackend)
 end
 
 """
-    function update(J::CuSparseMatrixCSR, p)
+    function update!(J::CuSparseMatrixCSR, p)
 
 Update the preconditioner `p` from the sparse Jacobian `J` in CSR format for CUDA
 
@@ -42,6 +42,6 @@ Update the preconditioner `p` from the sparse Jacobian `J` in CSR format for CUD
 3) Extract the preconditioner matrix `p.P` from the dense blocks `cuJs`
 
 """
-function KP.update(p::BlockJacobiPreconditioner, J::CUSPARSE.CuSparseMatrixCSR, device::CUDABackend)
-    _update_gpu(p, J.rowPtr, J.colVal, J.nzVal, device)
+function KP.update!(p::BlockJacobiPreconditioner, J::CUSPARSE.CuSparseMatrixCSR)
+    _update_gpu(p, J.rowPtr, J.colVal, J.nzVal, p.device)
 end
diff --git a/ext/CUDA/ic0.jl b/ext/CUDA/ic0.jl
@@ -46,7 +46,7 @@ for ArrayType in (:(CuVector{T}), :(CuMatrix{T}))
   end
 end
 
-function KP.update(p::NVIDIA_ICU0, A, device::CUDABackend)
+function KP.update!(p::NVIDIA_IC0, A)
     p.P = CUSPARSE.ic02(A)
 end
 
diff --git a/ext/CUDA/ilu0.jl b/ext/CUDA/ilu0.jl
@@ -46,6 +46,6 @@ for ArrayType in (:(CuVector{T}), :(CuMatrix{T}))
   end
 end
 
-function KP.update(p::NVIDIA_ILU0, A, device::CUDABackend)
+function KP.update!(p::NVIDIA_ILU0, A)
     p.P = CUSPARSE.ilu02(A)
 end
diff --git a/src/KrylovPreconditioners.jl b/src/KrylovPreconditioners.jl
@@ -11,8 +11,8 @@ import LinearAlgebra: ldiv!
 
 abstract type AbstractKrylovPreconditioner end
 export AbstractKrylovPreconditioner
-update(p::AbstractKrylovPreconditioner, A, device::KA.Backend) = error("Update for $(typeof(p)) not implemented")
-export update, get_timer, reset_timer!
+update!(p::AbstractKrylovPreconditioner, A) = error("update!() for $(typeof(p)) not implemented")
+export update!, get_timer, reset_timer!
 
 function get_timer(p::AbstractKrylovPreconditioner)
     return p.timer_update

diff --git a/src/blockjacobi.jl b/src/blockjacobi.jl
@@ -60,6 +60,7 @@ mutable struct BlockJacobiPreconditioner{AT,GAT,VI,GVI,GMT,MI,GMI} <: AbstractKr
     id::GMT
     blocklist::Vector{GMT}
     timer_update::Float64
+    device::KA.Backend
 end
 
 function create_blocklist(blocks::Array, npart)
@@ -135,7 +136,8 @@ function BlockJacobiPreconditioner(J, npart::Int64, device=CPU(), olevel=0) wher
         curest_size, blocks,
         cublocks, map,
         cumap, part,
-        cupart, id, blocklist, 0.0
+        cupart, id, blocklist, 0.0,
+        device
     )
 end
 
@@ -311,14 +313,14 @@ Fill the dense blocks of the preconditioner from the sparse CSR matrix arrays
 end
 
 """
-    function update(J::SparseMatrixCSC, p)
+    function update!(J::SparseMatrixCSC, p)
 
 Update the preconditioner `p` from the sparse Jacobian `J` in CSC format for the CPU
 
 Note that this implements the same algorithm as for the GPU and becomes very slow on CPU with growing number of blocks.
 
 """
-function update(p, J::SparseMatrixCSC, device)
+function update!(p, J::SparseMatrixCSC)
     # TODO: Enabling threading leads to a crash here
     for b in 1:p.nblocks
         p.blocks[:,:,b] = p.id[:,:]

diff --git a/test/gpu/gpu.jl b/test/gpu/gpu.jl
@@ -14,6 +14,7 @@ function test_ic0(FC, V, M)
   A_gpu = M(A_cpu)
   b_gpu = V(b_cpu)
   P = kp_ic0(A_gpu)
+  update!(P, A_gpu)
 
   x_gpu, stats = cg(A_gpu, b_gpu, M=P, ldiv=true)
   r_gpu = b_gpu - A_gpu * x_gpu
@@ -35,6 +36,7 @@ function test_ilu0(FC, V, M)
   A_gpu = M(A_cpu)
   b_gpu = V(b_cpu)
   P = kp_ilu0(A_gpu)
+  update!(P, A_gpu)
 
   x_gpu, stats = gmres(A_gpu, b_gpu, N=P, ldiv=true)
   r_gpu = b_gpu - A_gpu * x_gpu
@@ -132,7 +134,7 @@ function test_block_jacobi(device, AT, SMT)
     x = similar(b); r = similar(b)
     nblocks = 2
     precond = BlockJacobiPreconditioner(A, nblocks, device)
-    KrylovPreconditioners.update(precond, A, device)
+    update!(precond, A)
 
     S = _get_type(A)
     linear_solver = Krylov.BicgstabSolver(n, m, S)