GPU Gradient Kernels

Technical reference for the KernelAbstractions-based gradient kernels used by the batched zone driver. For the rationale behind the batched design, see GPU Backends and When They Help.

Kernel

Sentinel uses KernelAbstractions.jl to write portable GPU kernels. The same mu_gradient_kernel! source runs on CPU, Metal, and CUDA; the backend is selected at call time.

@kernel function mu_gradient_kernel!(grad, elem_data, disp_fwd, disp_adj, @Const(lookup))
    i = @index(Global)
    # gradient accumulation over 27x27 shape function pairs per Gauss point
end

The kernel computes $\partial \mathrm{\Phi }/\partial \theta$ at each Gauss point by contracting forward and adjoint displacement fields through the element stiffness sensitivity. At each Gauss point it performs an O(27²) double loop over the 27 nodes of the Hex27 element, accumulating the product of shape function derivatives, forward displacements, and adjoint displacements. The resulting per-Gauss-point contribution is scattered to the material mesh via the GP2MTR mapping.

Only the $\mu$ gradient runs on the GPU. Other property gradients ($\kappa$, $\rho$) are computed on the CPU, where they are cheap.

Backend selection

Backend selection is handled by KAGradientBackend:

using Metal
backend = KAGradientBackend(MetalBackend())  # Metal GPU (Float32)

using CUDA
backend = KAGradientBackend(CUDABackend())   # NVIDIA GPU (Float64)

backend = KAGradientBackend(CPU())           # CPU (Float64, default)

The backend is passed via the gradient_backend keyword argument to the batched driver. When no backend is specified, the gradient runs on CPU.

Backend	Float type	Extension module
CPU()	Float64	(built-in)
MetalBackend()	Float32	SentinelMetalExt
CUDABackend()	Float64	SentinelCUDAExt

Data structures

Two data layouts support per-zone and batched execution:

ElementGradientData (per-zone) holds shape function derivatives, Jacobian determinants, and material indices for a single zone. It is used in the CPU gradient path and as the building block for batched data.

BatchedElementData (cross-zone) packs data from all zones into flat arrays with O(1) lookup tables:

• work_zone[i] — which zone work item i belongs to

• work_el[i] — which element within that zone

• work_q[i] — which Gauss point within that element

This flat mapping allows a single GPU kernel launch to cover all Gauss points across all zones without nested indexing or zone-boundary logic. concatenate_element_data builds a BatchedElementData from the per-zone data; compute_batched_mu_gradient! launches the batched kernel over it. Per-zone and cross-zone displacement/adjoint arrays live in GradientKAWorkspace and BatchedWorkspace respectively.

Extension hooks

Each GPU backend extends two functions that integrate it into the dispatch system:

# Float type used on the device
Sentinel._gpu_float_type(::MetalBackend) = Float32
Sentinel._gpu_float_type(::CUDABackend)  = Float64

# Array transfer: convert CPU arrays to device arrays
Sentinel._to_gpu_array(::MetalBackend, x::AbstractArray) = MtlArray(x)
Sentinel._to_gpu_array(::CUDABackend,  x::AbstractArray) = CuArray(x)

These are called internally by the batched driver and gradient kernel infrastructure; user code selects a backend only through KAGradientBackend.

API

Sentinel.AbstractGradientBackend Type

AbstractGradientBackend

Type hierarchy for gradient computation backends.

source

Sentinel.CPUGradientBackend Type

CPUGradientBackend()

Default CPU gradient backend — uses the existing sequential code in gradient.jl.

source

Sentinel.KAGradientBackend Type

KAGradientBackend(device)

KernelAbstractions gradient backend. device is a KA backend:

• CPU() for multithreaded CPU

• CUDABackend() for NVIDIA GPU (requires CUDA.jl)

• MetalBackend() for Apple GPU (requires Metal.jl)

source

Sentinel.ElementGradientData Type

ElementGradientData

Pre-extracted per-(element, quadrature point) data needed by the gradient kernel. Built once per compute_gradient! call by iterating over Ferrite CellValues on CPU, then transferred to GPU as flat arrays.

source

Sentinel.BatchedElementData Type

BatchedElementData

Concatenated element gradient data for multiple zones, stored as flat arrays with precomputed lookup tables for O(1) work-item-to-zone mapping.

Memory Layout

All per-zone arrays are concatenated along their first dimension:

grad_x = [zone1_grad_x; zone2_grad_x; ...; zoneN_grad_x]
          ├─ ne1×nqp ─┤├─ ne2×nqp ─┤     ├─ neN×nqp ─┤

Work Item Mapping

Flat lookup tables provide O(1) mapping with zero branch divergence:

• work_zone[I] → which zone (1-based)

• work_el[I] → local element index within that zone (1-based)

• work_q[I] → local quadrature point (1-based)

Memory cost: ~7.6 MB for 630K work items × 3 × Int32.

Node Indexing

Node IDs in node_ids and mtrl_conn are offset-adjusted to index into concatenated displacement/gradient arrays directly: global_node = zone_node_offsets[zi] + local_node_id

source

Sentinel.GradientKAWorkspace Type

GradientKAWorkspace

Pre-flattened workspace for GPU gradient computation. Converts Sentinel's Displacement struct-of-arrays into contiguous vectors.

All arrays are indexed by DOF node id (1:n_dof_nodes).

source

Sentinel.BatchedWorkspace Type

BatchedWorkspace

Concatenated displacement and adjoint vectors across all zones. Each zone's DOF nodes occupy a contiguous slice at zone_node_offsets[zi]+1 through zone_node_offsets[zi+1].

Adjoint values are stored conjugated (as needed by the gradient formula).

source

Sentinel.concatenate_element_data Function

concatenate_element_data(edatas::AbstractVector) -> BatchedElementData

Concatenate per-zone ElementGradientData structs into a single BatchedElementData with flat lookup tables and offset-adjusted indices.

Arguments

• edatas: vector of per-zone ElementGradientData (from build_element_gradient_data)

Returns

CPU-resident BatchedElementData. Call adapt_to_backend to transfer to GPU.

Example

edatas = [build_element_gradient_data(zone.dh, zone.cv, ...) for zone in zones]
batched = concatenate_element_data(edatas)
# batched.total_eq ≈ 630K for 290 zones × ~80 elements × 27 GPs

source

Sentinel.compute_batched_mu_gradient! Function

compute_batched_mu_gradient!(output, batched, bws, zone_materials, backend;
                              adjind_bases_real, adjind_bases_imag)

Compute μ gradient for all zones in a single batched kernel launch.

Arguments

• output: Float64 vector, length = total_nodes (concatenated gradient output)

• batched: BatchedElementData with concatenated element data

• bws: BatchedWorkspace with concatenated displacement/adjoint data

• zone_materials: vector of per-zone Material (for scalar weights)

• backend: KAGradientBackend (CPU() or MetalBackend() etc.)

• adjind_bases_real/imag: per-zone gradient parameter base indices

Returns nothing; accumulates into output via atomic add.

source

Sentinel.adapt_to_backend Function

adapt_to_backend(edata, ws, grad_rvalue, backend) -> (edata_gpu, ws_gpu, rvalue_gpu)

Convert CPU element data and workspace arrays to the appropriate array type for the given KA backend. Returns new structs wrapping GPU arrays.

For CPU() backend, returns the inputs unchanged. For MetalBackend(), converts to MtlArray{Float32}. For CUDABackend(), converts to CuArray{Float64}.

source

adapt_to_backend(batched_edata, batched_ws, backend) -> (edata_gpu, ws_gpu)

Convert batched data structures to GPU arrays for the given backend. For CPU backend, returns inputs unchanged.

source