diff --git a/demos/nodal_assembly_gpu.jl b/demos/nodal_assembly_gpu.jl new file mode 100644 index 0000000..f0e5b57 --- /dev/null +++ b/demos/nodal_assembly_gpu.jl @@ -0,0 +1,442 @@ +""" +Nodal Assembly GPU Implementation + +This is the GPU port of demos/nodal_assembly_cpu.jl using CUDA.jl. + +TWO-PHASE APPROACH: +1. compute_gp_data_kernel!() - Compute integration point stresses and material states +2. nodal_assembly_kernel!() - Assemble residual at nodes (matrix-free, no atomics!) + +Uses Tensors.jl throughout on GPU (CuArray{SymmetricTensor} works!) +""" + +using CUDA +using Tensors +using LinearAlgebra +using Printf + +# Material state for plasticity (GPU-compatible!) +struct PlasticState + ε_p::SymmetricTensor{2,3,Float64,6} # Plastic strain tensor + α::Float64 # Accumulated plastic strain +end + +# Material properties +struct Material + E::Float64 # Young's modulus + ν::Float64 # Poisson's ratio + σ_y::Float64 # Yield stress +end + +# Node-to-elements connectivity (CSR format) +struct NodeToElementsMap + ptr::CuArray{Int32,1} # Length: n_nodes + 1 + data::CuArray{Int32,1} # Length: total connections +end + +""" +Build CSR map: which elements touch each node? +""" +function build_node_to_elems_gpu(elements::Vector{NTuple{4,Int}}, n_nodes::Int) + # Count connections per node + counts = zeros(Int, n_nodes) + for elem in elements + for node in elem + counts[node] += 1 + end + end + + # Build CSR structure + ptr = cumsum([1; counts]) + data = Vector{Int32}(undef, sum(counts)) + + # Fill data array + offset = copy(ptr[1:end-1]) + for (elem_idx, elem) in enumerate(elements) + for node in elem + data[offset[node]] = elem_idx + offset[node] += 1 + end + end + + return NodeToElementsMap(CuArray(Int32.(ptr)), CuArray(data)) +end + +""" +Return mapping for von Mises perfect plasticity (using Tensors.jl on GPU!) + +This function works identically on CPU and GPU! +""" +@inline function return_mapping_tensor(ε_total::SymmetricTensor{2,3,T}, + state_old::PlasticState, + E, ν, σ_y) where T + # Elastic strain + ε_e = ε_total - state_old.ε_p + + # Elastic predictor + λ = E * ν / ((1 + ν) * (1 - 2ν)) + μ = E / (2(1 + ν)) + I = one(ε_e) + σ_trial = λ * tr(ε_e) * I + 2μ * ε_e + + # Deviatoric stress + σ_dev = dev(σ_trial) + σ_eq = sqrt(3 / 2 * (σ_dev ⊡ σ_dev)) + + # Yield function + f = σ_eq - σ_y + + if f <= T(0.0) + # Elastic + return (σ_trial, state_old) + else + # Plastic - radial return + Δγ = f / (3μ) + n = σ_dev / σ_eq + + σ = σ_trial - 2μ * Δγ * n + + # Update plastic state + Δε_p = Δγ * n + ε_p_new = state_old.ε_p + Δε_p + α_new = state_old.α + Δγ + + state_new = PlasticState(ε_p_new, α_new) + + return (σ, state_new) + end +end + +""" +PHASE 1 GPU KERNEL: Compute integration point data + +One thread per integration point! +""" +function compute_gp_data_kernel!( + σ_gp::CuDeviceArray{SymmetricTensor{2,3,Float64,6},1}, + states_new::CuDeviceArray{PlasticState,1}, + u::CuDeviceArray{Float64,1}, + nodes::CuDeviceArray{Float64,2}, # Shape: 3 × n_nodes + elements::CuDeviceArray{Int32,2}, # Shape: 4 × n_elems + states_old::CuDeviceArray{PlasticState,1}, + E, ν, σ_y +) + gp_idx = (blockIdx().x - 1) * blockDim().x + threadIdx().x + + if gp_idx <= length(σ_gp) + # Map GP to element and local GP + elem_idx = (gp_idx - 1) ÷ 4 + 1 # 4 GPs per Tet4 + # local_gp = (gp_idx - 1) % 4 + 1 # Not used yet (all GPs same for linear Tet4) + + # 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 (using Tensors.jl Vec!) + 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 (using Tensors.jl Vec!) + 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 (using tensor products!) + J = dN1_dxi ⊗ X1 + dN2_dxi ⊗ X2 + dN3_dxi ⊗ X3 + dN4_dxi ⊗ X4 + invJ = inv(J) + + # Physical derivatives (using tensor contractions!) + dN1_dx = invJ ⋅ dN1_dxi + dN2_dx = invJ ⋅ dN2_dxi + dN3_dx = invJ ⋅ dN3_dxi + dN4_dx = invJ ⋅ dN4_dxi + + # Strain (using tensor products and symmetric!) + ε = symmetric(dN1_dx ⊗ u1 + dN2_dx ⊗ u2 + dN3_dx ⊗ u3 + dN4_dx ⊗ u4) + + # Material state update (using Tensors.jl - works on GPU!) + state_old = states_old[gp_idx] + σ, state_new = return_mapping_tensor(ε, state_old, E, ν, σ_y) + + # Store results + σ_gp[gp_idx] = σ + states_new[gp_idx] = state_new + end + + return nothing +end + +""" +PHASE 2 GPU KERNEL: Nodal assembly (matrix-free, no atomics!) + +One thread per node! +""" +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} +) + node_idx = (blockIdx().x - 1) * blockDim().x + threadIdx().x + + if node_idx <= size(nodes, 2) + # Accumulate forces from all elements touching this node + f_node = zero(Vec{3,Float64}) + + # Gauss weights for Tet4 (standard 4-point quadrature) + gauss_weight = 1.0 / 24.0 + + # Shape derivatives in reference coordinates + 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)) + ) + + # Get element range for this node (CSR format) + elem_start = node_to_elems_ptr[node_idx] + elem_end = node_to_elems_ptr[node_idx+1] - 1 + + # Loop over touching elements + for elem_offset in elem_start:elem_end + elem_idx = node_to_elems_data[elem_offset] + + # Extract element nodes + n1 = elements[1, elem_idx] + n2 = elements[2, elem_idx] + n3 = elements[3, elem_idx] + n4 = elements[4, elem_idx] + + # Find local node index in element + 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 + + # Recompute geometry (matrix-free!) + 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) + + # Physical derivative for this node + dN_dx = invJ ⋅ dN_dxi[local_node] + + # Loop over Gauss points (4 per Tet4) + for local_gp in 1:4 + gp_idx = (elem_idx - 1) * 4 + local_gp + + # Get stress at this GP + σ = σ_gp[gp_idx] + + # Accumulate force (using tensor contraction!) + f_node += (dN_dx ⋅ σ) * (gauss_weight * detJ) + end + end + + # Write result (no atomics - this node is ours!) + 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 + +""" +Complete residual computation on GPU (two-phase approach) +""" +function compute_residual_gpu!( + r::CuArray{Float64,1}, + u::CuArray{Float64,1}, + nodes::CuArray{Float64,2}, + elements::CuArray{Int32,2}, + states_old::CuArray{PlasticState,1}, + mat::Material, + node_to_elems::NodeToElementsMap +) + n_gp = length(states_old) + n_nodes = size(nodes, 2) + + # Storage for integration point data (on GPU!) + σ_gp = CuArray{SymmetricTensor{2,3,Float64,6}}(undef, n_gp) + states_new = CuArray{PlasticState}(undef, n_gp) + + # Phase 1: Compute integration point data + threads = 256 + blocks = cld(n_gp, threads) + + @cuda threads = threads blocks = blocks compute_gp_data_kernel!( + σ_gp, states_new, + u, nodes, elements, + states_old, + mat.E, mat.ν, mat.σ_y + ) + + # Phase 2: Nodal assembly + fill!(r, 0.0) + + threads = 256 + blocks = cld(n_nodes, threads) + + @cuda threads = threads blocks = blocks nodal_assembly_kernel!( + r, σ_gp, + nodes, elements, + node_to_elems.ptr, node_to_elems.data + ) + + return r, states_new +end + +# ============================================================================ +# Test Setup +# ============================================================================ + +function main() + println("\n" * "="^70) + println("Nodal Assembly GPU Implementation - CUDA.jl + Tensors.jl") + println("="^70) + + # Check CUDA availability + if !CUDA.functional() + println("❌ CUDA not available! This demo requires a GPU.") + return + end + + println("\n✅ CUDA device: ", CUDA.name(CUDA.device())) + + # Single Tet4 element + nodes_cpu = Float64[ + 0.0 1.0 0.0 0.0; # X coordinates + 0.0 0.0 1.0 0.0; # Y coordinates + 0.0 0.0 0.0 1.0 # Z coordinates + ] + + elements_cpu = [(1, 2, 3, 4)] + elements_mat = Int32[e[i] for i in 1:4, e in elements_cpu] + + n_nodes = 4 + n_elems = 1 + n_gps = n_elems * 4 # 4 GPs per Tet4 + + # Material + mat = Material( + 210e3, # E = 210 GPa (steel) + 0.3, # ν = 0.3 + 250.0 # σ_y = 250 MPa + ) + + # Displacement (apply tension) + u_cpu = zeros(12) + u_cpu[4] = 0.01 # Move node 2 in X-direction + + # Initial states (all elastic) + states_old_cpu = [PlasticState(zero(SymmetricTensor{2,3,Float64}), 0.0) for _ in 1:n_gps] + + # Build node-to-elements map + println("\nBuilding node-to-elements map (CSR format)...") + node_to_elems = build_node_to_elems_gpu(elements_cpu, n_nodes) + + # Transfer to GPU + println("Transferring data to GPU...") + nodes_gpu = CuArray(nodes_cpu) + elements_gpu = CuArray(elements_mat) + u_gpu = CuArray(u_cpu) + states_old_gpu = CuArray(states_old_cpu) + r_gpu = CUDA.zeros(Float64, 12) + + # Compute residual on GPU + println("\n" * "-"^70) + println("Computing residual on GPU (two-phase nodal assembly)...") + println("-"^70) + + r_gpu, states_new_gpu = compute_residual_gpu!( + r_gpu, u_gpu, nodes_gpu, elements_gpu, + states_old_gpu, mat, node_to_elems + ) + + # Transfer results back to CPU + r_cpu = Array(r_gpu) + states_new_cpu = Array(states_new_gpu) + + println("\nResidual vector (internal forces):") + for i in 1:n_nodes + rx = r_cpu[3*i-2] + ry = r_cpu[3*i-1] + rz = r_cpu[3*i] + @printf("Node %d: [%12.6e, %12.6e, %12.6e]\n", i, rx, ry, rz) + end + + println("\nResidual norm: ", norm(r_cpu)) + + # Check material states + println("\n" * "-"^70) + println("Material States at Gauss Points:") + println("-"^70) + + for (gp_idx, state) in enumerate(states_new_cpu) + elem_idx = (gp_idx - 1) ÷ 4 + 1 + local_gp = (gp_idx - 1) % 4 + 1 + + status = state.α > 0.0 ? "Plastic" : "Elastic" + @printf("Elem %d, GP %d: %s (α = %.6e)\n", elem_idx, local_gp, status, state.α) + end + + # Test force balance + println("\n" * "-"^70) + println("Force Balance Check:") + println("-"^70) + + f_total = sum(reshape(r_cpu, 3, :), dims=2) + @printf("Sum of forces: [%.6e, %.6e, %.6e]\n", f_total[1], f_total[2], f_total[3]) + @printf("Should be ≈ zero for internal forces (tol: 1e-10)\n") + + if norm(f_total) < 1e-10 + println("✅ Force balance: PASSED") + else + println("❌ Force balance: FAILED") + end + + # Compare with CPU reference + println("\n" * "="^70) + println("Comparing with CPU reference (demos/nodal_assembly_cpu.jl)...") + println("="^70) + println("\nExpected residual norms should match!") + println(" CPU reference: 727.208...") + println(" GPU result: ", norm(r_cpu)) + + println("\n" * "="^70) + println("✅ GPU nodal assembly complete!") + println("="^70) + println("\nNext steps:") + println(" 1. Benchmark GPU vs CPU performance") + println(" 2. Scale to realistic mesh sizes (10K+ elements)") + println(" 3. Integrate with Newton-Krylov solver") + println(" 4. Add line search for convergence") + println(" 5. Add preconditioning (Chebyshev-Jacobi → GMG)") + println("="^70 * "\n") +end + +main()