diff --git a/ext/JuliaFEMCUDAExt.jl b/ext/JuliaFEMCUDAExt.jl new file mode 100644 index 0000000..cb2ff8d --- /dev/null +++ b/ext/JuliaFEMCUDAExt.jl @@ -0,0 +1,934 @@ +# This file is a part of JuliaFEM. +# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md + +""" +JuliaFEMCUDAExt - GPU Backend Extension + +Package extension that provides GPU acceleration for JuliaFEM. +Automatically loaded when user does 'using CUDA'. + +Pure GPU implementation with: +- All assembly in CUDA kernels +- Zero CPU-GPU transfer during solve +- Matrix-free conjugate gradient +""" +module JuliaFEMCUDAExt + +using CUDA +using Tensors +using LinearAlgebra + +# Import parent package +using JuliaFEM +using JuliaFEM: Element, FieldProblem, get_connectivity +using JuliaFEM: Physics, ElasticityPhysicsType, DirichletBC, NeumannBC +using JuliaFEM: AbstractElasticityData, initialize_backend, solve_backend!, GPU + +""" +ElasticityDataGPU <: AbstractElasticityData + +GPU-resident data structure (all arrays on device). + +Note: Internal type, not exported to users! +""" +mutable struct ElasticityDataGPU <: AbstractElasticityData + # Geometry (nodes + connectivity) + nodes::CuArray{Float64,2} # 3 × n_nodes + elements::CuArray{Int32,2} # 4 × n_elements (Tet4) + n_nodes::Int + n_elements::Int + + # Material properties (per element) + E::CuArray{Float64,1} # Young's modulus per element + ν::CuArray{Float64,1} # Poisson's ratio per element + + # Boundary conditions (device) + is_fixed::CuArray{Bool,1} # n_dofs (true = constrained) + prescribed::CuArray{Float64,1} # n_dofs (value if constrained) + + # Surface loads (Neumann) + surface_nodes::CuArray{Int32,2} # 3 × n_surface_tri (Tri3 connectivity) + surface_traction::CuArray{Float64,2} # 3 × n_surface_tri (traction per element) + + # Node-to-elements map (CSR format) + node_to_elem_ptr::CuArray{Int32,1} # n_nodes+1 (CSR pointers) + node_to_elem_data::CuArray{Int32,1} # (element indices for each node) + + # Solution and working arrays + u::CuArray{Float64,1} # Solution vector (n_dofs) + f_ext::CuArray{Float64,1} # External forces (n_dofs) +end + +""" +Physics{Elasticity} - GPU-resident elasticity problem + +NOTE: Physics struct is now defined in physics_api.jl (backend-agnostic) +This module implements the GPU backend. +""" + +# Alias for backward compatibility with internal code +const GPUElasticityData = ElasticityDataGPU + +# ============================================================================ +# Helper functions +# ============================================================================ + +""" +Initialize GPU data from Physics{Elasticity} + +Transfer all CPU data to GPU: +1. Extract coordinates from elements +2. Build connectivity arrays +3. Extract material properties +4. Build BC flag arrays +5. Construct node-to-elements map +""" +function initialize_gpu_data!(physics::Physics{ElasticityPhysicsType}, time::Float64=0.0) + @info "Initializing GPU data for $(physics.name)..." + + # 1. Extract geometry from elements + n_elements = length(physics.body_elements) + @assert n_elements > 0 "No body elements in physics!" + + # Assume Tet4 for now (can generalize later) + first_el = physics.body_elements[1] + conn = get_connectivity(first_el) + nnodes_per_elem = length(conn) + @assert nnodes_per_elem == 4 "Only Tet4 supported currently" + + # Build node set and renumber + node_set = Set{Int}() + for el in physics.body_elements + for node in get_connectivity(el) + push!(node_set, node) + end + end + n_nodes = length(node_set) + node_list = sort(collect(node_set)) + node_map = Dict(node => i for (i, node) in enumerate(node_list)) + + @info " Nodes: $n_nodes, Elements: $n_elements" + + # 2. Extract coordinates (assume first element has all nodes for now - FIXME) + # Better: iterate all elements, extract unique nodes + nodes_cpu = zeros(3, n_nodes) + for el in physics.body_elements + # Direct field access for immutable API (geometry is constant, not time-dependent) + X = el.fields.geometry # 3×4 matrix for Tet4 + conn = get_connectivity(el) + for (local_idx, global_node) in enumerate(conn) + renumbered = node_map[global_node] + nodes_cpu[:, renumbered] = X[:, local_idx] + end + end + + # 3. Build connectivity (renumbered) + elements_cpu = zeros(Int32, 4, n_elements) + for (i, el) in enumerate(physics.body_elements) + conn = get_connectivity(el) + for (j, node) in enumerate(conn) + elements_cpu[j, i] = node_map[node] + end + end + + # 4. Extract material properties + E_cpu = zeros(Float64, n_elements) + ν_cpu = zeros(Float64, n_elements) + for (i, el) in enumerate(physics.body_elements) + # Direct field access for immutable API (material props are constant) + E_cpu[i] = el.fields.youngs_modulus + ν_cpu[i] = el.fields.poissons_ratio + end + + # 5. Build BC flag arrays + n_dofs = 3 * n_nodes + is_fixed_cpu = fill(false, n_dofs) + prescribed_cpu = zeros(n_dofs) + + bc = physics.bc_dirichlet + for (i, node_id) in enumerate(bc.node_ids) + if haskey(node_map, node_id) + renumbered = node_map[node_id] + for (comp_idx, comp) in enumerate(bc.components[i]) + dof = 3 * (renumbered - 1) + comp + is_fixed_cpu[dof] = true + prescribed_cpu[dof] = bc.values[i][comp_idx] + end + end + end + + # 6. Build surface load data + n_surface = length(physics.bc_neumann.surface_elements) + surface_nodes_cpu = zeros(Int32, 3, n_surface) + surface_traction_cpu = zeros(Float64, 3, n_surface) + + for (i, surf_el) in enumerate(physics.bc_neumann.surface_elements) + conn = get_connectivity(surf_el) + @assert length(conn) == 3 "Only Tri3 surface elements supported" + for (j, node) in enumerate(conn) + if haskey(node_map, node) + surface_nodes_cpu[j, i] = node_map[node] + else + error("Surface element references node $node not in body mesh") + end + end + traction = physics.bc_neumann.traction[i] + surface_traction_cpu[:, i] = [traction[1], traction[2], traction[3]] + end + + # 7. Build node-to-elements map (CSR) + node_to_elems = [Int32[] for _ in 1:n_nodes] + for (el_idx, el) in enumerate(physics.body_elements) + for node in get_connectivity(el) + renumbered = node_map[node] + push!(node_to_elems[renumbered], el_idx) + end + end + + ptr_cpu = zeros(Int32, n_nodes + 1) + ptr_cpu[1] = 1 + for i in 1:n_nodes + ptr_cpu[i+1] = ptr_cpu[i] + length(node_to_elems[i]) + end + data_cpu = vcat(node_to_elems...) + + # 8. Upload to GPU + @info " Uploading to GPU..." + gpu_data = ElasticityDataGPU( + CuArray(nodes_cpu), + CuArray(elements_cpu), + n_nodes, + n_elements, + CuArray(E_cpu), + CuArray(ν_cpu), + CuArray(is_fixed_cpu), + CuArray(prescribed_cpu), + CuArray(surface_nodes_cpu), + CuArray(surface_traction_cpu), + CuArray(ptr_cpu), + CuArray(data_cpu), + CUDA.zeros(Float64, n_dofs), + CUDA.zeros(Float64, n_dofs) + ) + + @info " GPU initialization complete!" + + return gpu_data # Return instead of storing in physics object! +end + +""" +CUDA KERNEL: Compute element stresses at integration points +""" +function compute_element_stresses_kernel!( + σ_gp::CuDeviceArray{SymmetricTensor{2,3,Float64,6},1}, + u::CuDeviceArray{Float64,1}, + nodes::CuDeviceArray{Float64,2}, + elements::CuDeviceArray{Int32,2}, + E_vec::CuDeviceArray{Float64,1}, + ν_vec::CuDeviceArray{Float64,1} +) + gp_idx = (blockIdx().x - 1) * blockDim().x + threadIdx().x + n_elements = size(elements, 2) + n_gp_per_elem = 4 # Tet4 has 4 Gauss points + + if gp_idx <= n_elements * n_gp_per_elem + elem_idx = div(gp_idx - 1, n_gp_per_elem) + 1 + + # Material + E = E_vec[elem_idx] + ν = ν_vec[elem_idx] + + # Extract element nodes + n1 = elements[1, elem_idx] + n2 = elements[2, elem_idx] + n3 = elements[3, elem_idx] + n4 = elements[4, elem_idx] + + # Node coordinates + X1 = Vec{3}((nodes[1, n1], nodes[2, n1], nodes[3, n1])) + X2 = Vec{3}((nodes[1, n2], nodes[2, n2], nodes[3, n2])) + X3 = Vec{3}((nodes[1, n3], nodes[2, n3], nodes[3, n3])) + X4 = Vec{3}((nodes[1, n4], nodes[2, n4], nodes[3, n4])) + + # Displacements + u1 = Vec{3}((u[3*n1-2], u[3*n1-1], u[3*n1])) + u2 = Vec{3}((u[3*n2-2], u[3*n2-1], u[3*n2])) + u3 = Vec{3}((u[3*n3-2], u[3*n3-1], u[3*n3])) + u4 = Vec{3}((u[3*n4-2], u[3*n4-1], u[3*n4])) + + # Shape derivatives (constant for Tet4) + dN1_dxi = Vec{3}((-1.0, -1.0, -1.0)) + dN2_dxi = Vec{3}((1.0, 0.0, 0.0)) + dN3_dxi = Vec{3}((0.0, 1.0, 0.0)) + dN4_dxi = Vec{3}((0.0, 0.0, 1.0)) + + # Jacobian + J = dN1_dxi ⊗ X1 + dN2_dxi ⊗ X2 + dN3_dxi ⊗ X3 + dN4_dxi ⊗ X4 + invJ = inv(J) + + # Physical derivatives + dN1_dx = invJ ⋅ dN1_dxi + dN2_dx = invJ ⋅ dN2_dxi + dN3_dx = invJ ⋅ dN3_dxi + dN4_dx = invJ ⋅ dN4_dxi + + # Strain + ε = symmetric(dN1_dx ⊗ u1 + dN2_dx ⊗ u2 + dN3_dx ⊗ u3 + dN4_dx ⊗ u4) + + # Stress (Hooke's law) + λ = E * ν / ((1 + ν) * (1 - 2ν)) + μ = E / (2(1 + ν)) + I = one(ε) + σ = λ * tr(ε) * I + 2μ * ε + + σ_gp[gp_idx] = σ + end + + return nothing +end + +""" +CUDA KERNEL: Nodal assembly (internal forces from stresses) +""" +function nodal_assembly_kernel!( + r::CuDeviceArray{Float64,1}, + σ_gp::CuDeviceArray{SymmetricTensor{2,3,Float64,6},1}, + nodes::CuDeviceArray{Float64,2}, + elements::CuDeviceArray{Int32,2}, + node_to_elems_ptr::CuDeviceArray{Int32,1}, + node_to_elems_data::CuDeviceArray{Int32,1}, + E_vec::CuDeviceArray{Float64,1}, + ν_vec::CuDeviceArray{Float64,1} +) + node_idx = (blockIdx().x - 1) * blockDim().x + threadIdx().x + n_nodes = size(nodes, 2) + + if node_idx <= n_nodes + f_node = zero(Vec{3,Float64}) + gauss_weight = 1.0 / 24.0 + + dN_dxi = ( + Vec{3}((-1.0, -1.0, -1.0)), + Vec{3}((1.0, 0.0, 0.0)), + Vec{3}((0.0, 1.0, 0.0)), + Vec{3}((0.0, 0.0, 1.0)) + ) + + elem_start = node_to_elems_ptr[node_idx] + elem_end = node_to_elems_ptr[node_idx+1] - 1 + + for elem_offset in elem_start:elem_end + elem_idx = node_to_elems_data[elem_offset] + + n1 = elements[1, elem_idx] + n2 = elements[2, elem_idx] + n3 = elements[3, elem_idx] + n4 = elements[4, elem_idx] + + local_node = 1 + if node_idx == n2 + local_node = 2 + elseif node_idx == n3 + local_node = 3 + elseif node_idx == n4 + local_node = 4 + end + + X1 = Vec{3}((nodes[1, n1], nodes[2, n1], nodes[3, n1])) + X2 = Vec{3}((nodes[1, n2], nodes[2, n2], nodes[3, n2])) + X3 = Vec{3}((nodes[1, n3], nodes[2, n3], nodes[3, n3])) + X4 = Vec{3}((nodes[1, n4], nodes[2, n4], nodes[3, n4])) + + J = dN_dxi[1] ⊗ X1 + dN_dxi[2] ⊗ X2 + dN_dxi[3] ⊗ X3 + dN_dxi[4] ⊗ X4 + detJ = det(J) + invJ = inv(J) + + dN_dx = invJ ⋅ dN_dxi[local_node] + + for local_gp in 1:4 + gp_idx = (elem_idx - 1) * 4 + local_gp + σ = σ_gp[gp_idx] + f_node += (dN_dx ⋅ σ) * (gauss_weight * detJ) + end + end + + r[3*node_idx-2] = f_node[1] + r[3*node_idx-1] = f_node[2] + r[3*node_idx] = f_node[3] + end + + return nothing +end + +""" +CUDA KERNEL: Apply surface traction (Neumann BC) + +Integrate traction over surface elements, add to external force. +""" +function apply_surface_traction_kernel!( + f_ext::CuDeviceArray{Float64,1}, + surface_nodes::CuDeviceArray{Int32,2}, + surface_traction::CuDeviceArray{Float64,2}, + nodes::CuDeviceArray{Float64,2} +) + surf_idx = (blockIdx().x - 1) * blockDim().x + threadIdx().x + n_surface = size(surface_nodes, 2) + + if surf_idx <= n_surface + # Get surface element nodes (Tri3) + n1 = surface_nodes[1, surf_idx] + n2 = surface_nodes[2, surf_idx] + n3 = surface_nodes[3, surf_idx] + + # Coordinates + X1 = Vec{3}((nodes[1, n1], nodes[2, n1], nodes[3, n1])) + X2 = Vec{3}((nodes[1, n2], nodes[2, n2], nodes[3, n2])) + X3 = Vec{3}((nodes[1, n3], nodes[2, n3], nodes[3, n3])) + + # Surface area (for Tri3) + v1 = X2 - X1 + v2 = X3 - X1 + area = 0.5 * norm(v1 × v2) + + # Traction vector + t = Vec{3}((surface_traction[1, surf_idx], + surface_traction[2, surf_idx], + surface_traction[3, surf_idx])) + + # Distribute equally to nodes (lumped load) + force_per_node = (area / 3.0) * t + + # Atomic add to force vector (allows parallel writes) + CUDA.@atomic f_ext[3*n1-2] += force_per_node[1] + CUDA.@atomic f_ext[3*n1-1] += force_per_node[2] + CUDA.@atomic f_ext[3*n1] += force_per_node[3] + + CUDA.@atomic f_ext[3*n2-2] += force_per_node[1] + CUDA.@atomic f_ext[3*n2-1] += force_per_node[2] + CUDA.@atomic f_ext[3*n2] += force_per_node[3] + + CUDA.@atomic f_ext[3*n3-2] += force_per_node[1] + CUDA.@atomic f_ext[3*n3-1] += force_per_node[2] + CUDA.@atomic f_ext[3*n3] += force_per_node[3] + end + + return nothing +end + +""" +CUDA KERNEL: Apply Dirichlet BC to residual + +Zero out fixed DOFs in residual vector. +""" +function apply_dirichlet_kernel!( + r::CuDeviceArray{Float64,1}, + is_fixed::CuDeviceArray{Bool,1} +) + dof = (blockIdx().x - 1) * blockDim().x + threadIdx().x + n_dofs = length(r) + + if dof <= n_dofs && is_fixed[dof] + r[dof] = 0.0 + end + + return nothing +end + +""" +Compute residual: R = f_int(u) - f_ext + +All on GPU, returns CuArray. +""" +function compute_residual_gpu!( + gpu_data::GPUElasticityData, + u::CuArray{Float64,1} +) + n_gp = gpu_data.n_elements * 4 + n_nodes = gpu_data.n_nodes + + # Phase 1: Compute stresses + σ_gp = CuArray{SymmetricTensor{2,3,Float64,6}}(undef, n_gp) + threads = 256 + blocks = cld(n_gp, threads) + @cuda threads = threads blocks = blocks compute_element_stresses_kernel!( + σ_gp, u, gpu_data.nodes, gpu_data.elements, gpu_data.E, gpu_data.ν + ) + + # Phase 2: Nodal assembly (internal forces) + f_int = CUDA.zeros(Float64, 3 * n_nodes) + threads = 256 + blocks = cld(n_nodes, threads) + @cuda threads = threads blocks = blocks nodal_assembly_kernel!( + f_int, σ_gp, gpu_data.nodes, gpu_data.elements, + gpu_data.node_to_elem_ptr, gpu_data.node_to_elem_data, + gpu_data.E, gpu_data.ν + ) + + # Residual: R = f_int - f_ext + r = f_int - gpu_data.f_ext + + # Apply Dirichlet BC (zero out fixed DOFs) + threads = 256 + blocks = cld(length(r), threads) + @cuda threads = threads blocks = blocks apply_dirichlet_kernel!( + r, gpu_data.is_fixed + ) + + return r +end + +""" +Matrix-free operator: K * u (for CG solver) + +Computes stiffness operator action without forming K matrix. + +For linear elasticity: K*u = f_int(u) +For nonlinear: K(u_current)*v = ∂f_int/∂u|_{u_current} * v (tangent operator) +""" +function stiffness_operator_gpu( + gpu_data::GPUElasticityData, + u::CuArray{Float64,1} +) + # For linear elasticity: K*u = f_int(u) + return compute_residual_gpu!(gpu_data, u) +end + +""" +Matrix-free tangent operator for Newton-Krylov: K(u_current) * v + +Computes action of tangent stiffness at u_current on vector v. +For linear elasticity, K is constant so K(u_current)*v = K*v = f_int(v). +For nonlinear (future), this would linearize around u_current. +""" +function tangent_operator_gpu( + gpu_data::GPUElasticityData, + u_current::CuArray{Float64,1}, + v::CuArray{Float64,1} +) + # For linear elasticity: K independent of u_current + # Just compute K*v directly + return compute_residual_gpu!(gpu_data, v) +end + +""" + cg_solve_matfree_gpu!(Δu, rhs, gpu_data, u_current; tol, max_iter) + +Matrix-free CG solver for Newton-Krylov framework. + +Solves K(u_current)*Δu = rhs for the Newton update Δu. + +# Arguments +- `Δu`: Solution vector (modified in-place, should be initialized to zeros) +- `rhs`: Right-hand side (-R for Newton) +- `gpu_data`: GPU data structure +- `u_current`: Current solution state (for tangent linearization) +- `tol`: Convergence tolerance for CG +- `max_iter`: Maximum CG iterations + +# Returns +- `(iterations, residual)`: Number of CG iterations and final residual norm +""" +function cg_solve_matfree_gpu!( + Δu::CuArray{Float64,1}, + rhs::CuArray{Float64,1}, + gpu_data::GPUElasticityData, + u_current::CuArray{Float64,1}; + tol=1e-6, + max_iter=1000 +) + n_dofs = length(Δu) + + # Tangent operator: K(u_current) * v + K_op(v) = tangent_operator_gpu(gpu_data, u_current, v) + + # Initial residual: r = rhs - K*Δu (Δu = 0 initially) + r = rhs - K_op(Δu) + + # Apply BC to r + threads = 256 + blocks = cld(n_dofs, threads) + @cuda threads = threads blocks = blocks apply_dirichlet_kernel!(r, gpu_data.is_fixed) + + p = copy(r) + r_dot_r = dot(r, r) + + for iter in 1:max_iter + Ap = K_op(p) + + # Apply BC to Ap + @cuda threads = threads blocks = blocks apply_dirichlet_kernel!(Ap, gpu_data.is_fixed) + + alpha = r_dot_r / dot(p, Ap) + Δu .+= alpha .* p + r .-= alpha .* Ap + + r_dot_r_new = dot(r, r) + + if sqrt(r_dot_r_new) < tol + return iter, sqrt(r_dot_r_new) + end + + beta = r_dot_r_new / r_dot_r + p .= r .+ beta .* p + r_dot_r = r_dot_r_new + end + + return max_iter, sqrt(r_dot_r) +end + +""" +GPU Conjugate Gradient solver (matrix-free) - LEGACY + +Old interface for linear problems. Use `solve_newton_krylov_gpu!` for new code. +""" +function cg_solve_gpu!( + gpu_data::GPUElasticityData; + tol=1e-6, + max_iter=1000 +) + n_dofs = length(gpu_data.f_ext) + u = gpu_data.u + b = gpu_data.f_ext + + # Initial residual + r = b - stiffness_operator_gpu(gpu_data, u) + + # Apply BC to r + threads = 256 + blocks = cld(n_dofs, threads) + @cuda threads = threads blocks = blocks apply_dirichlet_kernel!( + r, gpu_data.is_fixed + ) + + p = copy(r) + r_dot_r = dot(r, r) + + for iter in 1:max_iter + Ap = stiffness_operator_gpu(gpu_data, p) + + # Apply BC to Ap + @cuda threads = threads blocks = blocks apply_dirichlet_kernel!( + Ap, gpu_data.is_fixed + ) + + alpha = r_dot_r / dot(p, Ap) + u .+= alpha .* p + r .-= alpha .* Ap + + r_dot_r_new = dot(r, r) + + if sqrt(r_dot_r_new) < tol + @info " CG converged in $iter iterations (residual: $(sqrt(r_dot_r_new)))" + return iter, sqrt(r_dot_r_new) + end + + beta = r_dot_r_new / r_dot_r + p .= r .+ beta .* p + r_dot_r = r_dot_r_new + end + + @warn "CG did not converge in $max_iter iterations" + return max_iter, sqrt(r_dot_r) +end + +# ============================================================================ +# Inexact Newton-Krylov Solver (Simultaneous Newton + CG) +# ============================================================================ + +""" + solve_newton_krylov_gpu!(gpu_data, physics; kwargs...) + +Inexact Newton-Krylov solver with adaptive forcing terms (Eisenstat-Walker). + +Solves nonlinear elasticity problem using Newton iterations with inexact +linear solves. The key insight: don't fully converge the linear system! +Use adaptive tolerance that couples Newton and Krylov iterations. + +For LINEAR problems: Converges in 1 Newton iteration (becomes pure CG). +For NONLINEAR problems: Adaptively adjusts CG tolerance per Newton step. + +# Arguments +- `gpu_data`: GPU data structure +- `physics`: Physics problem definition + +# Keyword Arguments +- `newton_tol=1e-6`: Newton convergence tolerance (residual norm) +- `max_newton=20`: Maximum Newton iterations +- `max_cg_per_newton=50`: Maximum CG iterations per Newton step +- `forcing_power=0.5`: Eisenstat-Walker parameter (η = ||R||^α, α ∈ [0.5, 1]) +- `forcing_max=0.9`: Maximum forcing term (prevents too loose solves) +- `linear_solver=:cg`: Linear solver type (:cg only for now) + +# Returns +- `(u, newton_iters, total_cg_iters, residual, history)` + - `u`: Solution vector (CuArray) + - `newton_iters`: Number of Newton iterations + - `total_cg_iters`: Total CG iterations across all Newton steps + - `residual`: Final residual norm + - `history`: Vector of (cg_iters, R_norm, η) tuples for each Newton step + +# Example +```julia +result = solve_newton_krylov_gpu!( + gpu_data, physics; + newton_tol=1e-6, + max_newton=20, + max_cg_per_newton=50 +) + +# result.history shows the simultaneous solving: +# [(5, 0.12, 0.35), (8, 0.034, 0.18), (12, 0.009, 0.09), ...] +# ↑ ↑ ↑ +# CG |R| forcing term η +``` +""" +function solve_newton_krylov_gpu!( + gpu_data::GPUElasticityData, + physics::Physics{ElasticityPhysicsType}; + newton_tol=1e-6, + max_newton=20, + max_cg_per_newton=50, + forcing_power=0.5, + forcing_max=0.9, + linear_solver=:cg +) + n_dofs = length(gpu_data.u) + u = gpu_data.u # Current solution (modified in-place) + + total_cg_iters = 0 + history = Tuple{Int,Float64,Float64}[] # (cg_iters, R_norm, η) + + @info "Starting inexact Newton-Krylov solve..." + @info " Newton tolerance: $newton_tol" + @info " Max Newton iterations: $max_newton" + @info " Max CG per Newton: $max_cg_per_newton" + @info " Forcing power: $forcing_power (Eisenstat-Walker)" + + # For LINEAR problems: Directly solve K*u = f_ext (single CG solve) + # Check if this is linear by seeing if initial residual is just f_ext + R_initial = copy(gpu_data.f_ext) + threads = 256 + blocks = cld(n_dofs, threads) + @cuda threads = threads blocks = blocks apply_dirichlet_kernel!(R_initial, gpu_data.is_fixed) + R_initial_norm = sqrt(dot(R_initial, R_initial)) + + @info "Initial residual norm: $R_initial_norm (u=0)" + @info "Detecting problem type..." + + # If initial guess is zero and we're solving linear elasticity, + # just do one CG solve: K*u = f_ext + if true # Always treat as potentially linear for now + @info "Solving K*u = f_ext (linear solve)..." + + η_k = min(forcing_max, R_initial_norm^forcing_power) + linear_tol = η_k * R_initial_norm + + @info " CG tolerance: $linear_tol (η=$η_k)" + + # Solve K*u = f_ext + cg_iters, cg_residual = cg_solve_matfree_gpu!( + u, # Solution (starts at zero) + gpu_data.f_ext, # RHS + gpu_data, + CUDA.zeros(Float64, n_dofs); # Linearization point (doesn't matter for linear) + tol=linear_tol, + max_iter=max_cg_per_newton + ) + + total_cg_iters += cg_iters + push!(history, (cg_iters, R_initial_norm, η_k)) + + @info " CG: $cg_iters iterations (residual: $(round(cg_residual, sigdigits=3)))" + + # Apply BC to solution + @cuda threads = threads blocks = blocks apply_dirichlet_kernel!(u, gpu_data.is_fixed) + + # Verify convergence + f_int = compute_residual_gpu!(gpu_data, u) + R_final = gpu_data.f_ext - f_int + @cuda threads = threads blocks = blocks apply_dirichlet_kernel!(R_final, gpu_data.is_fixed) + R_final_norm = sqrt(dot(R_final, R_final)) + + @info "Final residual norm: $R_final_norm" + + if R_final_norm < newton_tol + @info "✓ Converged in 1 Newton iteration (linear problem)" + return (Array(u), 1, total_cg_iters, R_final_norm, history) + else + @info "Residual still large - continuing with Newton iterations..." + end + end + + # Nonlinear iterations (if needed) + for newton_iter in 2:max_newton # Start from 2 since we did one solve above + # 1. Compute residual: R = f_ext - f_int(u) + # For linear elasticity: f_int(u) = K*u + f_int = compute_residual_gpu!(gpu_data, u) + R = gpu_data.f_ext - f_int + + # Apply Dirichlet BC to residual + threads = 256 + blocks = cld(n_dofs, threads) + @cuda threads = threads blocks = blocks apply_dirichlet_kernel!(R, gpu_data.is_fixed) + + R_norm = sqrt(dot(R, R)) + + @info "Newton iteration $newton_iter: ||R|| = $(R_norm)" + + # Check convergence + if R_norm < newton_tol + @info "✓ Newton converged in $newton_iter iterations!" + @info " Total CG iterations: $total_cg_iters" + if newton_iter > 1 + avg_cg = total_cg_iters / newton_iter + @info " Average CG per Newton: $(round(avg_cg, digits=1))" + end + + # For linear problems, just one CG solve + if newton_iter == 1 && length(history) == 0 + # Haven't done any CG yet - solve now + η_k = min(forcing_max, R_norm^forcing_power) + linear_tol = η_k * R_norm + + @info " LINEAR problem detected: solving K*u = R directly" + @info " CG tolerance: $linear_tol" + + Δu = CUDA.zeros(Float64, n_dofs) + cg_iters, cg_residual = cg_solve_matfree_gpu!( + Δu, R, gpu_data, u; + tol=linear_tol, + max_iter=max_cg_per_newton + ) + + u .= Δu # For linear: u = K^{-1} R + total_cg_iters += cg_iters + push!(history, (cg_iters, R_norm, η_k)) + + @info " CG: $cg_iters iterations (residual: $(round(cg_residual, sigdigits=3)))" + + # Apply BC to solution + @cuda threads = threads blocks = blocks apply_dirichlet_kernel!(u, gpu_data.is_fixed) + + return (Array(u), newton_iter, total_cg_iters, cg_residual, history) + end + + return (Array(u), newton_iter, total_cg_iters, R_norm, history) + end + + # 2. Adaptive forcing term (Eisenstat-Walker) + η_k = min(forcing_max, R_norm^forcing_power) + linear_tol = η_k * R_norm + + @info " Forcing term η = $(round(η_k, digits=3)) → linear tol = $(round(linear_tol, sigdigits=3))" + + # 3. Solve K(u)*Δu = -R inexactly (matrix-free CG) + Δu = CUDA.zeros(Float64, n_dofs) + + cg_iters, cg_residual = cg_solve_matfree_gpu!( + Δu, + -R, + gpu_data, + u; + tol=linear_tol, + max_iter=max_cg_per_newton + ) + + total_cg_iters += cg_iters + push!(history, (cg_iters, R_norm, η_k)) + + @info " CG: $cg_iters iterations (residual: $(round(cg_residual, sigdigits=3)))" + + # 4. Update solution + u .+= Δu + + # Apply Dirichlet BC to solution + @cuda threads = threads blocks = blocks apply_dirichlet_kernel!(u, gpu_data.is_fixed) + end # for newton_iter + + # Did not converge + @warn "Newton did not converge in $max_newton iterations" + @info " Final residual: $(history[end][2])" + @info " Total CG iterations: $total_cg_iters" + + return (Array(u), max_newton, total_cg_iters, history[end][2], history) +end # function solve_newton_krylov_gpu! + +# ============================================================================ +# Backend Interface Implementation +# ============================================================================ + +""" + JuliaFEM.initialize_backend(::GPU, physics::Physics{ElasticityPhysicsType}, time::Float64) + +Initialize GPU backend data from physics problem. + +Implements AbstractBackend interface. +""" +function JuliaFEM.initialize_backend(::GPU, physics::Physics{ElasticityPhysicsType}, time::Float64) + return initialize_gpu_data!(physics, time) +end + +""" + JuliaFEM.solve_backend!(data::ElasticityDataGPU, physics; tol, max_iter) + +Solve elasticity problem on GPU using inexact Newton-Krylov. + +Uses adaptive forcing terms (Eisenstat-Walker) to solve Newton and CG +iterations SIMULTANEOUSLY rather than in nested loops. + +For linear problems: Converges in 1 Newton iteration (becomes pure CG). +For nonlinear problems: Adaptively adjusts linear solve tolerance. + +Implements AbstractBackend interface. + +Returns: (u, newton_iterations, total_cg_iterations, residual, history) +""" +function JuliaFEM.solve_backend!(data::ElasticityDataGPU, physics::Physics{ElasticityPhysicsType}; + tol=1e-6, max_iter=1000, newton_tol=1e-6, max_newton=20, max_cg_per_newton=50) + # Compute external forces (Neumann BC) + @info "Computing external forces..." + fill!(data.f_ext, 0.0) + + n_surface = size(data.surface_nodes, 2) + if n_surface > 0 + threads = 256 + blocks = cld(n_surface, threads) + @cuda threads = threads blocks = blocks apply_surface_traction_kernel!( + data.f_ext, + data.surface_nodes, + data.surface_traction, + data.nodes + ) + end + + # Solve with inexact Newton-Krylov + @info "Solving with inexact Newton-Krylov..." + u, newton_iters, total_cg_iters, residual, history = solve_newton_krylov_gpu!( + data, physics; + newton_tol=newton_tol, + max_newton=max_newton, + max_cg_per_newton=max_cg_per_newton + ) + + # Return in the format expected by backend interface + # Include Newton-Krylov statistics + return (u, newton_iters, total_cg_iters, residual, history) +end + +# ============================================================================ +# Backward Compatibility +# ============================================================================ + +""" + solve_elasticity_gpu!(physics::Physics{ElasticityPhysicsType}; time, tol, max_iter) + +Backward compatibility wrapper for direct GPU solve calls. + +**Deprecated:** Use `solve!(physics; backend=GPU())` instead. +""" +function solve_elasticity_gpu!(physics::Physics{ElasticityPhysicsType}; time=0.0, tol=1e-6, max_iter=1000) + # Use the unified solve! interface + return JuliaFEM.solve!(physics; backend=GPU(), time=time, tol=tol, max_iter=max_iter) +end + +end # module JuliaFEMCUDAExt