Roadmap: 10⁹-element topology optimization on a single GPU (voxel backend, true matrix-free, multigrid)

[mohdibntarek]

Problem

TopOpt.jl currently caps out far below "gigantic" problems because every element's data is precomputed and held in RAM. At 10⁹ elements (1024³, 3D linear hex, Float64) the dominant consumers are:

Data	Where	Size @ 10⁹
Kes — one element stiffness SMatrix per element	src/TopOptProblems/elementinfo.jl	~4.6 TB
cell_dofs — dense Int matrix	src/TopOptProblems/metadata.jl	~192 GB
dof_cells — RaggedArray of Tuple{Int,Int}	src/TopOptProblems/metadata.jl	~384 GB
density-filter sparse Jacobian	src/CheqFilters/density_filter.jl	TBs (rmin-dependent)
u (displacements)	—	~24 GB
vars (densities)	—	~8 GB

Two further blockers compound this:

1. Ferrite itself is a setup-time RAM bottleneck. Grid/DofHandler/ConstraintHandler store all nodes + cell connectivity, so a 10⁹-element problem cannot even be constructed today.
2. The current matrix-free method is not element-free. CGMatrixFreeSolver avoids assembling the global K, but MatrixFreeOperator.mul! still reads the precomputed Kes vector (src/FEA/matrix_free_operator.jl), so all element matrices are still stored. We want element matrices recomputed on the fly (from a reference element for structured grids; from quadrature for general grids), never materialized.

The goal: a backend where the only resident data is u and vars plus a handful of stencils, so 1024³ fits in ~16 GB (Float32) on a single A100/H100 — the same approach as Aage et al. 2017 ("Giga-voxel computational morphogenesis").

Scope & target

• Hardware: single GPU (A100/H100-class).
• Objectives: compliance first, then thermal compliance, then stress-constrained (adjoint) — one abstraction for all three.
• Precision: Float32 on device by default, Float64 opt-in (the codebase already threads T = floattype(problem) everywhere).
• Non-goal for the solver: out-of-core/disk. Disk I/O is only for checkpointing (HDF5/JLD2) and VTK export; the Krylov iteration stays in aggregate memory.

Approach

1. VoxelProblem{dim,T} (new src/StructuredProblems/) — stores only (nels, sizes, E, ν, BC spec, rmin). No Ferrite. Element Ke, cell_dofs, node adjacency, and filter stencil are computed from a single reference element and recomputed on the fly — no Kes, no Metadata, no RaggedArray.

2. StencilOperator (<: AbstractMatrixOperator) — a 27-point stencil mul! (3D) with Dirichlet rows handled analytically. Recomputation policy: no element matrix is ever stored; for the structured path the stencil is a closed form, and for a general (Ferrite) path each mul! re-integrates the element matrices on the fly (memory-for-FLOPs tradeoff). Written in KernelAbstractions.jl; CUDA.jl is the first backend.

3. Geometric multigrid preconditioner — hand-rolled V-cycle on the structured grid; coarsest level solved directly (tiny system). GMG is load-bearing: SIMP's xmin=0.001 void regions make plain CG stall. Fallback: AMGCLWrap.jl (GPU algebraic multigrid).

4. Objective kernels (fused, hand-written rrules, no Zygote through the solver):

• compliance: dJ/dx_e = -dρ_e/dx_e · u_eᵀ K_e u_e (closed form, already the convention in src/Functions/compliance.jl).
• thermal: -T_eᵀ K_e T_e (scalar field, ~⅓ dofs).
• stress-constrained: adjoint solve K λ = ∂J/∂u reusing the same operator + GMG, then p-norm/KS aggregation kernels.

5. Iterative solver — swap IterativeSolvers.cg! (CPU-only) for Krylov.jl cg (GPU-capable, GMG pluggable). Mixed-precision: dot-product reductions in Float64 while the matvec stays Float32.

6. Filtering — replace the materialized sparse Jacobian with a stencil filter + transpose-stencil rrule.

Phased roadmap

• M0 — correctness spike, CPU, no Ferrite: VoxelProblem + StencilOperator + closed-form gradient on 64³, verified against FEASolver(DirectSolver, …). Proves the memory collapse + on-the-fly element matrices.
• M1 — GPU port: KernelAbstractions matvec + gradient + filter kernels; Krylov CG; Float32.
• M2 — GMG preconditioner; verify CG iteration count is mesh-independent (make-or-break).
• M3 — scale to 1024³ (occupancy, memory layout, fused kernels).
• M4 — thermal compliance (drop-in).
• M5 — stress-constrained (adjoint via reused solver + GMG; aggregation kernels).
• M6 — Nonconvex bridge (device-resident design, gradient transfer per iteration), checkpointing, streaming VTK, TopOptCUDAExt extension following the Makie/Flux weak-dep pattern.

New dependencies

KernelAbstractions, CUDA (weak dep), Krylov, HDF5/JLD2, optionally AMGCLWrap (weak dep) — each with a [compat] bound + Pkg.resolve().

Risks to spike early

1. Nonconvex at 10⁹ design vars — MMA subproblems are O(n) but hold several 8 GB CPU copies of x/gradient/bounds; verify Float32 or limited-copy operation.
2. CG reduction latency — global dot products are the throughput killer; single double-precision reduction per iteration.
3. Void-region conditioning — GMG must handle xmin contrast; AMGCLWrap is the fallback.
4. Hand-written rrules must stay AD-safe — no Zygote through Krylov; mirror the existing convention.

---

Prepared with assistance from deepseek-v4-pro via opencode.