From 13eabeef3bc1e83552f7d0ff1745cfb8f81a6de2 Mon Sep 17 00:00:00 2001 From: Jukka Aho Date: Sat, 9 May 2026 16:30:27 +0300 Subject: [PATCH] refactor(ext): remove JuliaFEMCUDAExt.jl ext/JuliaFEMCUDAExt.jl | 934 ------------------------------------------------- 1 file changed, 934 deletions(-) --- ext/JuliaFEMCUDAExt.jl | 934 ----------------------------------------- 1 file changed, 934 deletions(-) delete mode 100644 ext/JuliaFEMCUDAExt.jl diff --git a/ext/JuliaFEMCUDAExt.jl b/ext/JuliaFEMCUDAExt.jl deleted file mode 100644 index cb2ff8d..0000000 --- a/ext/JuliaFEMCUDAExt.jl +++ /dev/null @@ -1,934 +0,0 @@ -# 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