Speed up block-Jacobi update

JuliaSmoothOptimizers · Jan 30, 2024 · f3142b4 · f3142b4
1 parent e85608e
commit f3142b4
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 7 deletions.
diff --git a/ext/AMDGPU/blockjacobi.jl b/ext/AMDGPU/blockjacobi.jl
@@ -10,14 +10,15 @@ end
 
 function _update_gpu(p, j_rowptr, j_colval, j_nzval, device::ROCBackend)
     nblocks = p.nblocks
+    blocksize = p.blocksize
     fillblock_gpu_kernel! = KP._fillblock_gpu!(device)
     # Fill Block Jacobi" begin
     fillblock_gpu_kernel!(
         p.cublocks, size(p.id,1),
         p.cupartitions, p.cumap,
         j_rowptr, j_colval, j_nzval,
         p.cupart, p.culpartitions, p.id,
-        ndrange=nblocks,
+        ndrange=(nblocks, blocksize),
     )
     KA.synchronize(device)
     # Invert blocks begin

diff --git a/ext/CUDA/blockjacobi.jl b/ext/CUDA/blockjacobi.jl
@@ -10,14 +10,15 @@ end
 
 function _update_gpu(p, j_rowptr, j_colval, j_nzval, device::CUDABackend)
     nblocks = p.nblocks
+    blocksize = p.blocksize
     fillblock_gpu_kernel! = KP._fillblock_gpu!(device)
     # Fill Block Jacobi" begin
     fillblock_gpu_kernel!(
         p.cublocks, size(p.id,1),
         p.cupartitions, p.cumap,
         j_rowptr, j_colval, j_nzval,
         p.cupart, p.culpartitions, p.id,
-        ndrange=nblocks,
+        ndrange=(nblocks, blocksize),
     )
     KA.synchronize(device)
     # Invert blocks begin

diff --git a/src/blockjacobi.jl b/src/blockjacobi.jl
@@ -288,14 +288,14 @@ Fill the dense blocks of the preconditioner from the sparse CSR matrix arrays
 
 """
 @kernel function _fillblock_gpu!(blocks, blocksize, partition, map, rowPtr, colVal, nzVal, part, lpartitions, id)
-    b = @index(Global, Linear)
+    b,k = @index(Global, NTuple)
     for i in 1:blocksize
-        for j in 1:blocksize
-            blocks[i,j,b] = id[i,j]
-        end
+        blocks[k,i,b] = id[k,i]
     end
 
-    @inbounds for k in 1:lpartitions[b]
+    @synchronize
+
+    @inbounds if k <= lpartitions[b]
         # select row
         i = partition[k, b]
         # iterate matrix