Update IC(0) and ILU(0) for CUDA GPUs

JuliaSmoothOptimizers · Nov 30, 2023 · 63f14a1 · 63f14a1
1 parent 9745f5f
commit 63f14a1
Show file tree

Hide file tree

Showing 4 changed files with 130 additions and 33 deletions.
diff --git a/ext/AMDGPU/ic0.jl b/ext/AMDGPU/ic0.jl
@@ -8,7 +8,11 @@ for (SparseMatrixType, BlasType) in ((:(ROCSparseMatrixCSR{T,Cint}), :BlasFloat)
   @eval begin
     function KP.kp_ic0(A::$SparseMatrixType) where T <: $BlasType
       P = rocSPARSE.ic0(A, 'O')
-      return AMD_IC0(P,0.0)
+      return AMD_IC0(P, 0.0)
+    end
+
+    function KP.update!(p::AMD_IC0{$SparseMatrixType}, A::$SparseMatrixType) where T <: $BlasType
+      p.P = rocSPARSE.ic0(A, 'O')
     end
   end
 end
@@ -44,7 +48,3 @@ for ArrayType in (:(ROCVector{T}), :(ROCMatrix{T}))
     end
   end
 end
-
-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
@@ -8,7 +8,11 @@ for (SparseMatrixType, BlasType) in ((:(ROCSparseMatrixCSR{T,Cint}), :BlasFloat)
   @eval begin
     function KP.kp_ilu0(A::$SparseMatrixType) where T <: $BlasType
       P = rocSPARSE.ilu0(A, 'O')
-      return AMD_ILU0(P,0.0)
+      return AMD_ILU0(P, 0.0)
+    end
+
+    function KP.update!(p::AMD_ILU0{$SparseMatrixType}, A::$SparseMatrixType) where T <: $BlasType
+      p.P = rocSPARSE.ilu0(A, 'O')
     end
   end
 end
@@ -44,7 +48,3 @@ for ArrayType in (:(ROCVector{T}), :(ROCMatrix{T}))
     end
   end
 end
-
-function KP.update!(p::AMD_ILU0, A)
-    p.P = rocSPARSE.ilu0(A, 'O')
-end
diff --git a/ext/CUDA/ic0.jl b/ext/CUDA/ic0.jl
@@ -1,15 +1,68 @@
+mutable struct IC0Info
+    info::CUSPARSE.csric02Info_t
+
+    function IC0Info()
+        info_ref = Ref{CUSPARSE.csric02Info_t}()
+        CUSPARSE.cusparseCreateCsric02Info(info_ref)
+        obj = new(info_ref[])
+        finalizer(CUSPARSE.cusparseDestroyCsric02Info, obj)
+        obj
+    end
+end
+
+Base.unsafe_convert(::Type{CUSPARSE.csric02Info_t}, info::IC0Info) = info.info
+
 mutable struct NVIDIA_IC0{SM} <: AbstractKrylovPreconditioner
-  P::SM
+  desc::CUSPARSE.cusparseMatDescr_t
+  buffer::CuVector{UInt8}
+  info::IC0Info
   timer_update::Float64
+  P::SM
 end
 
-for (SparseMatrixType, BlasType) in ((:(CuSparseMatrixCSR{T,Cint}), :BlasFloat),
-                                     (:(CuSparseMatrixCSC{T,Cint}), :BlasReal))
+for (bname, aname, sname, T) in ((:cusparseScsric02_bufferSize, :cusparseScsric02_analysis, :cusparseScsric02, :Float32),
+                                 (:cusparseDcsric02_bufferSize, :cusparseDcsric02_analysis, :cusparseDcsric02, :Float64),
+                                 (:cusparseCcsric02_bufferSize, :cusparseCcsric02_analysis, :cusparseCcsric02, :ComplexF32),
+                                 (:cusparseZcsric02_bufferSize, :cusparseZcsric02_analysis, :cusparseZcsric02, :ComplexF64))
   @eval begin
-    function KP.kp_ic0(A::$SparseMatrixType) where T <: $BlasType
-      P = CUSPARSE.ic02(A)
-      n = checksquare(A)
-      return NVIDIA_IC0(P,0.0)
+    function KP.kp_ic0(A::CuSparseMatrixCSR{$T,Cint})
+      P = copy(A)
+      n = checksquare(P)
+      desc = CUSPARSE.CuMatrixDescriptor('G', 'L', 'N', 'O')
+      info = ILU0Info()
+      buffer_size = Ref{Cint}()
+      CUSPARSE.$bname(CUSPARSE.handle(), n, nnz(P), desc, P.nzVal, P.rowPtr, P.colVal, info, buffer_size)
+      buffer = CuVector{UInt8}(undef, buffer_size[])
+      CUSPARSE.$aname(CUSPARSE.handle(), n, nnz(P), desc, P.nzVal, P.rowPtr, P.colVal, info, CUSPARSE.CUSPARSE_SOLVE_POLICY_USE_LEVEL, buffer)
+      posit = Ref{Cint}(1)
+      CUSPARSE.cusparseXcsric02_zeroPivot(CUSPARSE.handle(), info, posit)
+      (posit[] ≥ 0) && error("Structural/numerical zero in A at ($(posit[]),$(posit[])))")
+      CUSPARSE.$sname(CUSPARSE.handle(), n, nnz(P), desc, P.nzVal, P.rowPtr, P.colVal, info, CUSPARSE.CUSPARSE_SOLVE_POLICY_USE_LEVEL, buffer)
+      return NVIDIA_IC0(desc, buffer, info, 0.0, P)
+    end
+
+    function KP.update!(p::NVIDIA_IC0, A::CuSparseMatrixCSR{$T,Cint})
+      p.P = CUSPARSE.ic02(A)
+    end
+
+    function KP.kp_ic0(A::CuSparseMatrixCSC{$T,Cint})
+      P = copy(A)
+      n = checksquare(P)
+      desc = CUSPARSE.CuMatrixDescriptor('G', 'L', 'N', 'O')
+      info = ILU0Info()
+      buffer_size = Ref{Cint}()
+      CUSPARSE.$bname(CUSPARSE.handle(), n, nnz(P), desc, P.nzVal, P.colPtr, P.rowVal, info, buffer_size)
+      buffer = CuVector{UInt8}(undef, buffer_size[])
+      CUSPARSE.$aname(CUSPARSE.handle(), n, nnz(P), desc, P.nzVal, P.colPtr, P.rowVal, info, CUSPARSE.CUSPARSE_SOLVE_POLICY_USE_LEVEL, buffer)
+      posit = Ref{Cint}(1)
+      CUSPARSE.cusparseXcsric02_zeroPivot(CUSPARSE.handle(), info, posit)
+      (posit[] ≥ 0) && error("Structural/numerical zero in A at ($(posit[]),$(posit[])))")
+      CUSPARSE.$sname(CUSPARSE.handle(), n, nnz(P), desc, P.nzVal, P.colPtr, P.rowVal, info, CUSPARSE.CUSPARSE_SOLVE_POLICY_USE_LEVEL, buffer)
+      return NVIDIA_IC0(desc, buffer, info, 0.0, P)
+    end
+
+    function KP.update!(p::NVIDIA_IC0, A::CuSparseMatrixCSC{$T,Cint})
+      p.P = CUSPARSE.ic02(A)
     end
   end
 end
@@ -45,8 +98,3 @@ for ArrayType in (:(CuVector{T}), :(CuMatrix{T}))
     end
   end
 end
-
-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
@@ -1,15 +1,68 @@
+mutable struct ILU0Info
+    info::CUSPARSE.csrilu02Info_t
+
+    function ILU0Info()
+        info_ref = Ref{CUSPARSE.csrilu02Info_t}()
+        CUSPARSE.cusparseCreateCsrilu02Info(info_ref)
+        obj = new(info_ref[])
+        finalizer(CUSPARSE.cusparseDestroyCsrilu02Info, obj)
+        obj
+    end
+end
+
+Base.unsafe_convert(::Type{CUSPARSE.csrilu02Info_t}, info::ILU0Info) = info.info
+
 mutable struct NVIDIA_ILU0{SM} <: AbstractKrylovPreconditioner
-  P::SM
+  desc::CUSPARSE.cusparseMatDescr_t
+  buffer::CuVector{UInt8}
+  info::ILU0Info
   timer_update::Float64
+  P::SM
 end
 
-for (SparseMatrixType, BlasType) in ((:(CuSparseMatrixCSR{T,Cint}), :BlasFloat),
-                                     (:(CuSparseMatrixCSC{T,Cint}), :BlasReal))
+for (bname, aname, sname, T) in ((:cusparseScsrilu02_bufferSize, :cusparseScsrilu02_analysis, :cusparseScsrilu02, :Float32),
+                                 (:cusparseDcsrilu02_bufferSize, :cusparseDcsrilu02_analysis, :cusparseDcsrilu02, :Float64),
+                                 (:cusparseCcsrilu02_bufferSize, :cusparseCcsrilu02_analysis, :cusparseCcsrilu02, :ComplexF32),
+                                 (:cusparseZcsrilu02_bufferSize, :cusparseZcsrilu02_analysis, :cusparseZcsrilu02, :ComplexF64))
   @eval begin
-    function KP.kp_ilu0(A::$SparseMatrixType) where T <: $BlasType
-      P = CUSPARSE.ilu02(A)
-      n = checksquare(A)
-      return NVIDIA_ILU0(P,0.0)
+    function KP.kp_ilu0(A::CuSparseMatrixCSR{$T,Cint})
+      P = copy(A)
+      n = checksquare(P)
+      desc = CUSPARSE.CuMatrixDescriptor('G', 'L', 'N', 'O')
+      info = ILU0Info()
+      buffer_size = Ref{Cint}()
+      CUSPARSE.$bname(CUSPARSE.handle(), n, nnz(P), desc, P.nzVal, P.rowPtr, P.colVal, info, buffer_size)
+      buffer = CuVector{UInt8}(undef, buffer_size[])
+      CUSPARSE.$aname(CUSPARSE.handle(), n, nnz(P), desc, P.nzVal, P.rowPtr, P.colVal, info, CUSPARSE.CUSPARSE_SOLVE_POLICY_USE_LEVEL, buffer)
+      posit = Ref{Cint}(1)
+      CUSPARSE.cusparseXcsric02_zeroPivot(CUSPARSE.handle(), info, posit)
+      (posit[] ≥ 0) && error("Structural/numerical zero in A at ($(posit[]),$(posit[])))")
+      CUSPARSE.$sname(CUSPARSE.handle(), n, nnz(P), desc, P.nzVal, P.rowPtr, P.colVal, info, CUSPARSE.CUSPARSE_SOLVE_POLICY_USE_LEVEL, buffer)
+      return NVIDIA_ILU0(desc, buffer, info, 0.0, P)
+    end
+
+    function KP.update!(p::NVIDIA_ILU0{CuSparseMatrixCSC{$T,Cint}}, A::CuSparseMatrixCSR{$T,Cint})
+      p.P = CUSPARSE.ilu02(A)
+    end
+
+    function KP.kp_ilu0(A::CuSparseMatrixCSC{$T,Cint})
+      P = copy(A)
+      n = checksquare(P)
+      desc = CUSPARSE.CuMatrixDescriptor('G', 'L', 'N', 'O')
+      info = ILU0Info()
+      buffer_size = Ref{Cint}()
+      CUSPARSE.$bname(CUSPARSE.handle(), n, nnz(P), desc, P.nzVal, P.colPtr, P.rowVal, info, buffer_size)
+      buffer = CuVector{UInt8}(undef, buffer_size[])
+      CUSPARSE.$aname(CUSPARSE.handle(), n, nnz(P), desc, P.nzVal, P.colPtr, P.rowVal, info, CUSPARSE.CUSPARSE_SOLVE_POLICY_USE_LEVEL, buffer)
+      posit = Ref{Cint}(1)
+      CUSPARSE.cusparseXcsric02_zeroPivot(CUSPARSE.handle(), info, posit)
+      (posit[] ≥ 0) && error("Structural/numerical zero in A at ($(posit[]),$(posit[])))")
+      CUSPARSE.$sname(CUSPARSE.handle(), n, nnz(P), desc, P.nzVal, P.colPtr, P.rowVal, info, CUSPARSE.CUSPARSE_SOLVE_POLICY_USE_LEVEL, buffer)
+      return NVIDIA_ILU0(desc, buffer, info, 0.0, P)
+    end
+
+    function KP.update!(p::NVIDIA_ILU0{CuSparseMatrixCSC{$T,Cint}}, A::CuSparseMatrixCSC{$T,Cint})
+      p.P = CUSPARSE.ilu02(A)
     end
   end
 end
@@ -45,7 +98,3 @@ for ArrayType in (:(CuVector{T}), :(CuMatrix{T}))
     end
   end
 end
-
-function KP.update!(p::NVIDIA_ILU0, A)
-    p.P = CUSPARSE.ilu02(A)
-end