GPU: integration test kernel abstractions

In #1291 ([comment](https://github.com/Ferrite-FEM/Ferrite.jl/pull/1291#discussion_r3707868202)), the following integration test (`test/integration/test_kernel_abstractions.jl`) was not included. It needs to be updated (especially the comments) and should ideally follow the how-to (gpu assembly) closer. Adding it here to make it easy to pick up and fix when someone has time.

```julia
using Ferrite
import Adapt: adapt
import KernelAbstractions: @kernel, @index
import KernelAbstractions as KA
using SparseArrays
using Test
using GPUArrays
import LinearAlgebra

@testset "KernelAbstractions.jl integration" begin

# We start with some to be used in the following for simple convenience.
NUM_THREADS = Threads.nthreads()
NUM_TASKS_PER_THREAD = 4

function compute_threads_and_blocks(n)
## Let's assign, arbitrarily, two element assembly tasks per GPU thread.
tasks_per_thread = min(NUM_TASKS_PER_THREAD, n)
## To do so, let us first compute how many element groups we have to assemble.
n_effective = cld(n, tasks_per_thread)
## This potentially limits the number of usable threads, e.g. when a color just has a small
## number of elements.
threads = min(NUM_THREADS, n_effective)
## Furthermore, for CPU computing we typically group the tasks into blocks of worker threads.
blocks = cld(n, tasks_per_thread * threads)

return threads, blocks
end

# In this how-to we want to use an existing assembly routine on the GPU with Ferrite.
# We implicitly assume that nothing dynamic happens inside the routine, i.e. the routine
# is type stable, does not allocate and also does not have any dynamic dispatches.
function assemble_element!(Ke::AbstractMatrix, fe::AbstractVector, cv::CellValues)
n_basefuncs = getnbasefunctions(cv)

for q_point in 1:getnquadpoints(cv)
dΩ = getdetJdV(cv, q_point)
for i in 1:n_basefuncs
∇δuᵢ = shape_gradient(cv, q_point, i)
δuᵢ = shape_value(cv, q_point, i)
fe[i] += δuᵢ * dΩ
for j in 1:n_basefuncs
∇δuⱼ = shape_gradient(cv, q_point, j)
Ke[i, j] += (∇δuᵢ ⋅ ∇δuⱼ) * dΩ
end
end
end
return nothing
end

# We also have a simple cell assembly wrapping the element in two variants.
# In the first variant we assemble using an assembler. In the second variant we
# only fill Ke, e.g. as part of element-assembly techniques.
function assemble_cell!(Ke, fe, cell, cv, assembler)
reinit!(cv, nothing, cell.coords)
fill!(Ke, 0)
fill!(fe, 0)
assemble_element!(Ke, fe, cv)
assemble!(assembler, celldofs(cell), Ke, fe)
return nothing
end

# Now to the actual assembly kernel. To ensure portability we show how to use KernelAbstractions.jl
# as the kernel language, although we will also show how to use CUDA directly below. In this kernel
# we use a grid-stride loop, which has several benefits in terms of performance and debuggability.
# For more details please consult https://developer.nvidia.com/blog/cuda-pro-tip-write-flexible-kernels-grid-stride-loops/ .
@kernel function ka_assembly_kernel(assemblers, @Const(color), cc, cv, Kes, fes)
## This is the classical grid-stride-loop
task_index = @index(Global, Linear)
stride = prod(KA.@ndrange())

## Query the local evaluation buffer of the GPU worker.
## As explained later this is the secret sauce.
cv_i = cv[task_index]
cc_i = cc[task_index]
assembler = assemblers[task_index]
for i in task_index:stride:length(color)
## Work item index
cellid = color[i]

## Query work item cell cache
reinit!(cc_i, cellid)

## Query assembly buffer.
Ke = view(Kes, i, :, :)
fe = view(fes, i, :)

## Actual assembly routine.
assemble_cell!(Ke, fe, cc_i, cv_i, assembler)
end
end
function assemble_global_ka!(backend, cv::Ferrite.SoAContainer{<:CellValues}, K, f, cc, colors::Vector, Ke, fe, n_workers)
assembler = if K === nothing
nothing
else
Ferrite.distribute_to_workers(backend, start_assemble(K, f; fillzero = false), n_workers)
end
for color in colors
## We divide the work into blocks and fire up the kernel.
n = length(color)
threads, blocks = compute_threads_and_blocks(n)
## Now, we can build and execute the Kernel.
ka_kernel = ka_assembly_kernel(backend, threads)
ka_kernel(assembler, color, cc, cv, Ke, fe, ndrange = threads * blocks)
## Since the kernel launches asynchronously we need to add a synchronization
## point before proceeding here. Otherwise we will start assembling the next color,
## while there are still threads working on the current color, therefore potentially
## causing race conditions.
KA.synchronize(backend)
end
return nothing
end

# Reference for internal testing #hide
function assemble_global!(cv::CellValues, K::SparseMatrixCSC, f, dh::DofHandler) #hide
n_basefuncs = getnbasefunctions(cv) #hide
Ke = zeros(Float32, n_basefuncs, n_basefuncs) #hide
fe = zeros(Float32, n_basefuncs) #hide
assembler = start_assemble(K, f) #hide
for cell in CellIterator(dh) #hide
assemble_cell!(Ke, fe, cell, cv, assembler) #hide
end #hide
return nothing #hide
end #hide

function assemble_global!(cv::CellValues, K::SparseMatrixCSC, f, dh::SubDofHandler) #hide
n_basefuncs = getnbasefunctions(cv) #hide
Ke = zeros(Float32, n_basefuncs, n_basefuncs) #hide
fe = zeros(Float32, n_basefuncs) #hide
assembler = start_assemble(K, f; fillzero = false) #hide
for cell in CellIterator(dh) #hide
assemble_cell!(Ke, fe, cell, cv, assembler) #hide
end #hide
return nothing #hide
end #hide

# Now we first setup the problem almost as usual on the host (CPU).
# The only major difference here is that we instantiate everything
# using Float32 and Int32 whenever it makes sense to lower memory
# pressure on the GPU, and because Float32 is on most GPUs quite
# a bit faster than using Float64 -- outside of high-end server GPUs.
# Please note that GPU kernels have a launch overhead. Therefore out problem
# must be sufficiently large to see any benefits of utilizing the GPU.
# The small number of elements here is just for demonstration purposes.
num_elements = 10
# We generate a Float32 coordinate grid by passing in Float32 corner coordinates.
grid = generate_grid(Hexahedron, (num_elements, num_elements, num_elements), Vec{3}((-1.0f0, -1.0f0, -1.0f0)), Vec{3}((1.0f0, 1.0f0, 1.0f0)))
ip = Lagrange{RefHexahedron, 1}()
qr = QuadratureRule{RefHexahedron}(Float32, 2)
cv = CellValues(Float32, qr, ip)
dh = DofHandler(grid)
add!(dh, :u, ip)
close!(dh)

# If we assemble into a matrix, then we still need to color the grid as usual.
# See also the threading how-to. Note that we still leave some of the integers
# 64 bit to still enable the indexing of large problems.
colors = create_coloring(grid)

# Now to the GPU side. Here we use Adapt.jl to generate GPU counterparts of
# all relevant objects.
backend = KA.CPU()
colors_device = [adapt(backend, c) for c in colors]
dh_device = adapt(backend, dh)
K_device = allocate_matrix(SparseMatrixCSC{Float32, Int32}, dh)
f_device = KA.zeros(backend, Float32, (ndofs(dh),))

# Furthermore, the individual GPU workers need local buffers.
# Ferrite comes with a little helper to transform common buffers
# into a suitable GPU format.
n_workers = prod(compute_threads_and_blocks(maximum(length.(colors))))
cv_device = Ferrite.distribute_to_workers(backend, cv, n_workers)
cc = CellCache(dh_device)
cc_device = Ferrite.distribute_to_workers(backend, cc, n_workers)
# Technically we can also just get one Ke or fe per worker, but for demonstration
# purposes we allocate the full block here for element-assembly style matrix-free GPU
# usage.
Kes = KA.zeros(backend, Float32, getncells(grid), getnbasefunctions(cv), getnbasefunctions(cv))
fes = KA.zeros(backend, Float32, getncells(grid), getnbasefunctions(cv))

# Now everything is set to launch the assembly via KernelAbstractions.
assemble_global_ka!(backend, cv_device, K_device, f_device, cc_device, colors_device, Kes, fes, n_workers)

# Finally, we can apply the Dirichlet constraints and solve our linear system.
ch = ConstraintHandler(Float32, Int32, dh)
∂Ω = union(
getfacetset(grid, "left"), getfacetset(grid, "right"),
getfacetset(grid, "top"), getfacetset(grid, "bottom")
)
add!(ch, Dirichlet(:u, ∂Ω, (x, t) -> 1.0))
close!(ch)

ch_device = adapt(backend, ch)
apply!(K_device, f_device, ch_device)
# Julia 1.10 cannot handle \ on Float32, Int32 CSC matrices
u_device = Float32.(LinearAlgebra.lu(K_device) \ f_device)

K = allocate_matrix(SparseMatrixCSC{Float32, Int32}, dh)
f = zeros(Float32, (ndofs(dh),))
assemble_global!(cv, K, f, dh)
apply!(K, f, ch)
# Julia 1.10 cannot handle \ on Float32, Int32 CSC matrices
uref = Float32.(LinearAlgebra.lu(K) \ f)
@test u_device ≈ uref

end
```