Fuse the GPU boundary-condition updates into one kernel per color

src/PotentialCalculation/PotentialCalculationSetup/BoundaryConditions/BoundaryConditions.jl

@@ -1,2 +1,3 @@
-include("BoundaryConditionsCylindrical.jl")
-include("BoundaryConditionsCartesian.jl")
+include("BoundaryConditionsCylindrical.jl")
+include("BoundaryConditionsCartesian.jl")
+include("BoundaryConditionsGPU.jl")

src/PotentialCalculation/PotentialCalculationSetup/BoundaryConditions/BoundaryConditionsCylindrical.jl

@@ -30,8 +30,60 @@ function apply_boundary_conditions_on_r_axis!(  rbpot::AbstractArray{T, 4}, rbi:
     nothing
 end
 
+# One work-item per z-row: sum the row's φ columns of matching parity with the
+# precomputed weights and broadcast the result back to those columns. Replaces
+# a getindex + gemv + broadcast sequence (a temporary and several kernel
+# launches per SOR iteration) on GPU backends.
+@kernel function r0_bc_kernel!(rbpot, rbi::Int, nnz::Int, gw, even_points::Bool)
+    irbz = @index(Global, Linear)
+    rbz = irbz + 1
+    inφ_ub = size(rbpot, 2) - 1
+    @inbounds begin
+        last_odd = nnz - (iseven(nnz) ? 1 : 0)   # odd real-z rows
+        if rbz <= div(last_odd, 2) + mod(last_odd, 2) + 1
+            gwrow = even_points ? 6 : 5
+            φ0 = even_points ? 2 : 1
+            s = zero(eltype(rbpot))
+            j = 0
+            for inφ in φ0:2:inφ_ub
+                j += 1
+                s += rbpot[rbz, inφ + 1, 2, rbi] * gw[gwrow, j]
+            end
+            for inφ in φ0:2:inφ_ub
+                rbpot[rbz, inφ + 1, 2, rbi] = s
+            end
+        end
+        last_even = nnz - (isodd(nnz) ? 1 : 0)   # even real-z rows
+        if rbz <= div(last_even, 2) + mod(last_even, 2) + 1
+            gwrow = even_points ? 5 : 6
+            φ0 = even_points ? 1 : 2
+            s = zero(eltype(rbpot))
+            j = 0
+            for inφ in φ0:2:inφ_ub
+                j += 1
+                s += rbpot[rbz, inφ + 1, 2, rbi] * gw[gwrow, j]
+            end
+            for inφ in φ0:2:inφ_ub
+                rbpot[rbz, inφ + 1, 2, rbi] = s
+            end
+        end
+    end
+end
+
+function apply_boundary_conditions_at_r0!(
+    rbpot::GPUArrays.AbstractGPUArray{T, 4},
+    rbi::Int,
+    nnz::Int,
+    gw::AbstractArray{T},
+    update_even_points::Val{even_points}
+    ) where {T, even_points}
+    kern = r0_bc_kernel!(_ka_get_backend(rbpot))
+    kern(rbpot, rbi, nnz, gw, even_points, ndrange = rbidx(nnz) - 1)
+    nothing
+end
+
 function apply_boundary_conditions_at_r0!(
-    rbpot::AbstractArray{T, 4}, 
+    rbpot::AbstractArray{T, 4},
     rbi::Int,
     nnz::Int,
     gw::AbstractArray{T},

src/PotentialCalculation/PotentialCalculationSetup/BoundaryConditions/BoundaryConditionsGPU.jl

@@ -0,0 +1,123 @@
+# GPU boundary-condition path: the CPU implementation applies the ghost-face
+# updates as ~10 view broadcasts per red-black color, whose per-call host
+# overhead dominates the GPU iteration loop. Here each face update is encoded
+# as (active, source index, factor) with ghost = factor * pot[source] and all
+# six faces are applied by one kernel launch per color.
+#
+# The op tables below are 1:1 transcriptions of the broadcast methods in
+# BoundaryConditionsCylindrical.jl / BoundaryConditionsCartesian.jl.
+# Ghost entries with more than one ghost coordinate are never read by the SOR
+# stencil, so the faces may be applied concurrently.
+
+const _GhostFaceOp{T} = Tuple{Bool, Int, T}
+_no_ghost_op(::Type{T}) where {T} = (false, 1, one(T))
+
+# faces of the compressed (red-black) dimension: neighbors sit at rows 2 / end-1
+_ghost_ops_compressed(::DiscreteAxis{T, :fixed, :fixed}, n::Int, gbf) where {T} =
+    (_no_ghost_op(T), _no_ghost_op(T))
+_ghost_ops_compressed(::DiscreteAxis{T, :infinite, :infinite}, n::Int, gbf) where {T} =
+    ((true, 2, gbf[1]), (true, n - 1, gbf[2]))
+_ghost_ops_compressed(::DiscreteAxis{T, :fixed, :infinite}, n::Int, gbf) where {T} =
+    (_no_ghost_op(T), (true, n - 1, gbf[2]))
+_ghost_ops_compressed(::DiscreteAxis{T, :infinite, :fixed}, n::Int, gbf) where {T} =
+    ((true, 2, gbf[1]), _no_ghost_op(T))
+_ghost_ops_compressed(::DiscreteAxis{T, :infinite, :reflecting}, n::Int, gbf) where {T} =
+    ((true, 2, gbf[1]), (true, n - 1, one(T)))
+_ghost_ops_compressed(::DiscreteAxis{T, :reflecting, :infinite}, n::Int, gbf) where {T} =
+    ((true, 2, one(T)), (true, n - 1, gbf[2]))
+_ghost_ops_compressed(::DiscreteAxis{T, :reflecting, :reflecting}, n::Int, gbf) where {T} =
+    ((true, 2, one(T)), (true, n - 1, one(T)))
+_ghost_ops_compressed(::DiscreteAxis{T, :fixed, :reflecting}, n::Int, gbf) where {T} =
+    (_no_ghost_op(T), (true, n - 1, one(T)))
+_ghost_ops_compressed(::DiscreteAxis{T, :reflecting, :fixed}, n::Int, gbf) where {T} =
+    ((true, 2, one(T)), _no_ghost_op(T))
+_ghost_ops_compressed(::DiscreteAxis{T, :periodic, :periodic}, n::Int, gbf) where {T} =
+    ((true, n - 1, one(T)), (true, 2, one(T)))
+
+# faces of the non-compressed dimensions: reflecting neighbors sit at rows 3 / end-2
+_ghost_ops(::DiscreteAxis{T, :fixed, :fixed}, n::Int, gbf) where {T} =
+    (_no_ghost_op(T), _no_ghost_op(T))
+_ghost_ops(::DiscreteAxis{T, :infinite, :infinite}, n::Int, gbf) where {T} =
+    ((true, 2, gbf[1]), (true, n - 1, gbf[2]))
+_ghost_ops(::DiscreteAxis{T, :fixed, :infinite}, n::Int, gbf) where {T} =
+    (_no_ghost_op(T), (true, n - 1, gbf[2]))
+_ghost_ops(::DiscreteAxis{T, :infinite, :fixed}, n::Int, gbf) where {T} =
+    ((true, 2, gbf[1]), _no_ghost_op(T))
+_ghost_ops(::DiscreteAxis{T, :infinite, :reflecting}, n::Int, gbf) where {T} =
+    ((true, 2, gbf[1]), (true, n - 2, one(T)))
+_ghost_ops(::DiscreteAxis{T, :reflecting, :infinite}, n::Int, gbf) where {T} =
+    ((true, 3, one(T)), (true, n - 1, gbf[2]))
+_ghost_ops(::DiscreteAxis{T, :reflecting, :reflecting}, n::Int, gbf) where {T} =
+    ((true, 3, one(T)), (true, n - 2, one(T)))
+_ghost_ops(::DiscreteAxis{T, :fixed, :reflecting}, n::Int, gbf) where {T} =
+    (_no_ghost_op(T), (true, n - 2, one(T)))
+_ghost_ops(::DiscreteAxis{T, :reflecting, :fixed}, n::Int, gbf) where {T} =
+    ((true, 3, one(T)), _no_ghost_op(T))
+_ghost_ops(::DiscreteAxis{T, :periodic, :periodic}, n::Int, gbf) where {T} =
+    ((true, n - 1, one(T)), (true, 2, one(T)))
+# the radial axis: no left face (r = 0 is handled separately), and the
+# infinite/reflecting right face reads row end-2 (as in the broadcast methods)
+_ghost_ops(::DiscreteAxis{T, :r0, :infinite}, n::Int, gbf) where {T} =
+    (_no_ghost_op(T), (true, n - 2, gbf[2]))
+_ghost_ops(::DiscreteAxis{T, :r0, :reflecting}, n::Int, gbf) where {T} =
+    (_no_ghost_op(T), (true, n - 2, one(T)))
+_ghost_ops(::DiscreteAxis{T, :r0, :fixed}, n::Int, gbf) where {T} =
+    (_no_ghost_op(T), _no_ghost_op(T))
+
+@kernel function faces_bc_kernel!(rbpot, rbi::Int,
+        a1l::Bool, s1l::Int, f1l, a1r::Bool, s1r::Int, f1r,
+        a2l::Bool, s2l::Int, f2l, a2r::Bool, s2r::Int, f2r,
+        a3l::Bool, s3l::Int, f3l, a3r::Bool, s3r::Int, f3r)
+    i, j = @index(Global, NTuple)
+    n1 = size(rbpot, 1); n2 = size(rbpot, 2); n3 = size(rbpot, 3)
+    @inbounds begin
+        if i <= n1 && j <= n2
+            if a3l rbpot[i, j,  1, rbi] = f3l * rbpot[i, j, s3l, rbi] end
+            if a3r rbpot[i, j, n3, rbi] = f3r * rbpot[i, j, s3r, rbi] end
+        end
+        if i <= n1 && j <= n3
+            if a2l rbpot[i,  1, j, rbi] = f2l * rbpot[i, s2l, j, rbi] end
+            if a2r rbpot[i, n2, j, rbi] = f2r * rbpot[i, s2r, j, rbi] end
+        end
+        if i <= n2 && j <= n3
+            if a1l rbpot[ 1, i, j, rbi] = f1l * rbpot[s1l, i, j, rbi] end
+            if a1r rbpot[n1, i, j, rbi] = f1r * rbpot[s1r, i, j, rbi] end
+        end
+    end
+end
+
+function _apply_faces_bc!(rbpot::GPUArrays.AbstractGPUArray{T, 4}, rbi::Int,
+        ops1::NTuple{2, _GhostFaceOp{T}}, ops2::NTuple{2, _GhostFaceOp{T}}, ops3::NTuple{2, _GhostFaceOp{T}}) where {T}
+    kern = faces_bc_kernel!(_ka_get_backend(rbpot))
+    n1, n2, n3 = size(rbpot, 1), size(rbpot, 2), size(rbpot, 3)
+    kern(rbpot, rbi,
+        ops1[1]..., ops1[2]..., ops2[1]..., ops2[2]..., ops3[1]..., ops3[2]...,
+        ndrange = (max(n1, n2), max(n2, n3)))
+    nothing
+end
+
+function apply_boundary_conditions!(pcs::PotentialCalculationSetup{T, Cylindrical, 3, <:GPUArrays.AbstractGPUArray},
+        update_even_points::Val{even_points}, only2d::Val{only_2d}) where {T, even_points, only_2d}
+    rbi::Int = even_points ? rb_even::Int : rb_odd::Int
+    rbpot = pcs.potential
+    if !only_2d
+        apply_boundary_conditions_at_r0!(rbpot, rbi, size(pcs.ϵ_r, 3) - 1, pcs.geom_weights[2], update_even_points)
+    end
+    # red-black dimension order is (z, φ, r)
+    ops_z = _ghost_ops_compressed(pcs.grid.axes[3], size(rbpot, 1), pcs.grid_boundary_factors[3])
+    ops_φ = only_2d ? (_no_ghost_op(T), _no_ghost_op(T)) : _ghost_ops(pcs.grid.axes[2], size(rbpot, 2), pcs.grid_boundary_factors[2])
+    ops_r = _ghost_ops(pcs.grid.axes[1], size(rbpot, 3), pcs.grid_boundary_factors[1])
+    _apply_faces_bc!(rbpot, rbi, ops_z, ops_φ, ops_r)
+    nothing
+end
+
+function apply_boundary_conditions!(pcs::PotentialCalculationSetup{T, Cartesian, 3, <:GPUArrays.AbstractGPUArray},
+        update_even_points::Val{even_points}, only2d::Val{only_2d}) where {T, even_points, only_2d}
+    rbi::Int = even_points ? rb_even::Int : rb_odd::Int
+    rbpot = pcs.potential
+    ops_x = _ghost_ops_compressed(pcs.grid.axes[1], size(rbpot, 1), pcs.grid_boundary_factors[1])
+    ops_y = _ghost_ops(pcs.grid.axes[2], size(rbpot, 2), pcs.grid_boundary_factors[2])
+    ops_z = _ghost_ops(pcs.grid.axes[3], size(rbpot, 3), pcs.grid_boundary_factors[3])
+    _apply_faces_bc!(rbpot, rbi, ops_x, ops_y, ops_z)
+    nothing
+end

src/PotentialCalculation/SuccessiveOverRelaxation/GPU_kernel.jl

@@ -164,6 +164,11 @@ Maybe this will change in the future.
 """
 get_sor_kernel(::Type{Cylindrical}, backend, CPU_via_KernelAbstractions, args...) = sor_cyl_gpu!(backend, args...)
 get_sor_kernel(::Type{Cartesian},   backend, CPU_via_KernelAbstractions, args...) = sor_car_gpu!(backend, args...)
+
+# Workgroup shapes tuned on an A100; both are 2.4-2.6x faster than the
+# KernelAbstractions default across 2D/3D cylindrical and Cartesian grids.
+# Dimension 1 is the memory-contiguous red-black dimension.
+_sor_kernel_workgroupsize(ndrange) = ndrange[2] == 1 ? (128, 1, 1) : (32, 4, 1)
 get_sor_kernel(::Type{Cylindrical}, ::CPU, CPU_via_KernelAbstractions::Val{false}) = nothing
 get_sor_kernel(::Type{Cartesian},   ::CPU, CPU_via_KernelAbstractions::Val{false}) = nothing
 get_sor_kernel(::Type{Cylindrical}, backend::CPU, CPU_via_KernelAbstractions::Val{true}, args...) = sor_cyl_gpu!(backend, args...)

src/PotentialCalculation/SuccessiveOverRelaxation/GPU_update.jl

@@ -8,14 +8,14 @@
     only2d::Val{only_2d} = Val{false}()
 )::Nothing where {T, only_2d, depletion_handling_enabled, _is_weighting_potential}
     update_even_points = true
-    _ka_synchronize(kernel, kernel( 
+    _ka_synchronize(kernel, kernel(
         pcs.potential, pcs.imp_scale, pcs.point_types, pcs.volume_weights, pcs.q_eff_imp, pcs.q_eff_fix, pcs.ϵ_r,
-        pcs.geom_weights, pcs.sor_const, update_even_points, depletion_handling_enabled, _is_weighting_potential, only_2d, 
+        pcs.geom_weights, pcs.sor_const, update_even_points, depletion_handling_enabled, _is_weighting_potential, only_2d,
         ndrange = ndrange
     ))
     apply_boundary_conditions!(pcs, Val(update_even_points), only2d)
     update_even_points = false
-    _ka_synchronize(kernel, kernel( 
+    _ka_synchronize(kernel, kernel(
         pcs.potential, pcs.imp_scale, pcs.point_types, pcs.volume_weights, pcs.q_eff_imp, pcs.q_eff_fix, pcs.ϵ_r,
         pcs.geom_weights, pcs.sor_const, update_even_points, depletion_handling_enabled, _is_weighting_potential, only_2d,
         ndrange = ndrange

src/PotentialCalculation/SuccessiveOverRelaxation/convergence.jl

@@ -11,7 +11,11 @@ function _update_till_convergence!( pcs::PotentialCalculationSetup{T, S, 3},
                                 ) where {T, S, depletion_handling_enabled, only_2d, _is_weighting_potential} #  <: GPUArrays.AbstractGPUArray
     backend = _ka_get_backend(pcs.potential)
     ndrange = size(pcs.potential)[1:3] .- 2
-    kernel = get_sor_kernel(S, backend, Val(via_KernelAbstractions))
+    kernel = if backend isa _KA_Backend && !(backend isa KernelAbstractions.CPU)
+        get_sor_kernel(S, backend, Val(via_KernelAbstractions), _sor_kernel_workgroupsize(ndrange))
+    else
+        get_sor_kernel(S, backend, Val(via_KernelAbstractions))
+    end
     c_limit = _is_weighting_potential ? convergence_limit : abs(convergence_limit * (iszero(pcs.bias_voltage) ? maximum(abs.(pcs.potential)) : pcs.bias_voltage))
     c = (one(c_limit) + c_limit) * 10 # Has to be larger than c_limit at the beginning
     n_performed_iterations = 0

src/ka_compat.jl

@@ -4,7 +4,11 @@
 @static if isdefined(KernelAbstractions, :Backend)
     import KernelAbstractions.Backend as _KA_Backend
     import KernelAbstractions.get_backend as _ka_get_backend
-    _ka_synchronize(kernel, kernel_ret) = KernelAbstractions.synchronize(kernel.backend)
+    # GPU backends are stream-ordered: subsequent operations on the same queue see
+    # the kernel's results without a host-side synchronization. Only the CPU
+    # backend (task-based) needs an explicit wait.
+    _ka_synchronize(kernel, kernel_ret) = kernel.backend isa KernelAbstractions.CPU ?
+        KernelAbstractions.synchronize(kernel.backend) : nothing
 else
     # KernelAbstractions < v0.9:
     import KernelAbstractions.Device as _KA_Backend

test/Project.toml

@@ -3,6 +3,7 @@ ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018"
 Geant4 = "559df036-b7a0-42fd-85df-7d5dd9d70f44"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LegendDataTypes = "99e09c13-5545-5ee2-bfa2-77f358fb75d8"
 LegendHDF5IO = "c9265ca6-b027-5446-b1a4-febfa8dd10b0"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -22,5 +23,6 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [compat]
 Geant4 = "0.3.1"
+JLArrays = "0.2, 0.3"
 LegendDataTypes = "0.1.3"
 LegendHDF5IO = "0.1"

test/runtests.jl

@@ -15,6 +15,10 @@ end
     include("test_gpu_backend.jl")
 end
 
+@timed_testset "SOR boundary conditions" begin
+    include("test_boundary_conditions.jl")
+end
+
 @timed_testset "ConstructiveSolidGeometry" begin
     include("ConstructiveSolidGeometry/CSG_IO.jl")
     include("ConstructiveSolidGeometry/CSG_coordinates.jl")