Matrix-Free 3D SIMP Topology Optimization with Fused Gather-GEMM-Scatter Kernels

Published 20 Apr 2026 in cs.CE and cs.DC | (2604.18020v1)

Abstract: The matrix-free gather-batched-GEMM-scatter pattern eliminates global stiffness assembly for three-dimensional SIMP topology optimization, but the conventional three-stage implementation forces avoidable DRAM traffic between stages. We present a single fused CUDA kernel, implemented through CuPy's runtime compilation interface, that performs gather, per-element stiffness multiplication, and scatter accumulation in one pass. On a single RTX 4090 (24 GB), the fused path reaches a problem-size-dependent 4.6-7.3x end-to-end SIMP wall-time speedup across 216k-4.9M cantilever elements and 4.4x on the 499,125-element torsion benchmark. Against the same-precision FP32 three-stage baseline, the fused path still yields 2.3-4.6x on cantilever and 2.8x on torsion. Isolated CUDA-event cantilever-operator measurements reach 8.9-13.8x per matvec call, while separate instrumented board-power traces at 216k and 1M show 3.2-4.9x lower energy than matched FP64 runs. A separate bridge stress test shows the same FP32-versus-FP64 three-stage trend under one distributed-load case; direct fused-kernel bridge benchmarks are not reported. We also evaluate a BF16 WMMA variant: a separate PyTorch BF16 GEMM proxy on matching tensor shapes yields 14.3x, but direct condition-number estimates of 6.1e5-2.3e6 across 64k-512k uniform-density test states imply BF16 conditioning products of 2.4e3-9.1e3, far above the 256 threshold, observed alongside BF16 iterative-refinement stagnation at the two tested inner tolerances.

Abstract PDF Upgrade to Chat

Authors (3)

Summary

The paper presents a fused gather–GEMM–scatter CUDA kernel that reduces intermediate DRAM traffic by nearly 1.8× and accelerates matrix–vector products by 6–7×.
It utilizes CuPy runtime compilation and BF16 tensor cores to optimize per-element stiffness computations in the 3D SIMP finite element method.
Numerical results demonstrate nearly identical solution quality with significant energy savings on large-scale cantilever problems compared to traditional methods.

Matrix-Free 3D SIMP Topology Optimization with Fused Gather–GEMM–Scatter Kernels: Technical Analysis

Introduction

The paper "Matrix-Free 3D SIMP Topology Optimization with Fused Gather-GEMM-Scatter Kernels" (2604.18020) presents a substantial advance in high-performance topology optimization (TO) for structural mechanics, focusing specifically on GPU-accelerated matrix-free finite element methods (FEM) and computational throughput. The authors develop a fused single-kernel implementation for the critical $K\mathbf{v}$ operation in 3D SIMP TO, deployed via CuPy runtime compilation and exploiting modern tensor core hardware. This essay provides a technical summary, numerical analysis, and discussion of the implications of the architecture and data observed in the paper.

Background: Challenges in Large-Scale 3D Topology Optimization

Standard 3D structural topology optimization, via the Solid Isotropic Material with Penalization (SIMP) method, is dominated by the sparse linear solves required for equilibrium at every iteration. The explicit assembly of global sparse stiffness matrices $K$ becomes unscalable due to memory requirements and bandwidth constraints on modern GPUs, especially as system sizes reach millions of elements [aage2017giga, aage2013topology, liu20143d]. Matrix-free operator application, which evaluates $K\mathbf{v}$ elementwise without explicit global assembly, is established as a remedy for these limitations [wang2025matrixfree, traff2023simple]. However, prior implementations typically decompose the operator into three GPU kernel launches—gather, batched GEMM, and scatter—that result in redundant global memory round-trips and DRAM bandwidth bottlenecks.

Fused Kernel Architecture and Implementation

The core contribution is the development of a fused single CUDA kernel, executing the gather, per-element stiffness multiply, and scatter in a one-pass approach, implemented through CuPy's runtime compilation (NVRTC) interface. This eliminates unnecessary DRAM traffic by keeping per-element intermediate vectors in registers and shared memory.

Figure 1: Comparison between the conventional three-stage pipeline (left) and the fused single-kernel pipeline (right), highlighting elimination of intermediate DRAM traffic.

This is achieved by leveraging the structure of the trilinear hexahedral ( $Q_1$ ) elements: the element stiffness matrix $K^{\mathrm{unit}}$ is identical across mesh elements (uniform mesh), allowing block-level shared memory broadcasting. The GPU kernel is launched with $n_{\mathrm{elem}}/128$ thread blocks (128 elements/block), and each thread handles the complete gather–GEMM–scatter sequence for a single element, with atomic accumulation to the global vector.

Arithmetic and Memory Efficiency:

The fused kernel reduces the main DRAM memory traffic associated with intermediate per-element arrays ( $\mathbf{u}_{\text{elem}}$ , $\mathbf{f}_{\text{elem}}$ ) by 2× compared to the baseline.
Profiling shows effective DRAM traffic decreases from 768 bytes/element in the three-stage path to 384 bytes/element in the fused design, not accounting for index table and density field reads, which raise the practical reduction factor to ~1.8×.

Exploiting Tensor Cores: BF16 WMMA Variant

NVIDIA's Ada Lovelace tensor cores provide substantial raw throughput for matrix multiplications in BF16 precision, potentially doubling the FP32 throughput. The paper introduces a BF16 WMMA path for the kernel, padding the per-element $24\times24$ stiffness matrix to $32\times32$ , tiled as $K$ 0 blocks for tensor core execution.

Figure 2: WMMA-based BF16 tensor-core tiling of the per-element stiffness product—zero-padding to $K$ 1, processed as four $K$ 2 fragments per element.

While the BF16 path demonstrates up to 14.3× acceleration for the “GEMM stage” vs. FP64, the operator as a whole remains bandwidth-bound due to the gather/scatter bottlenecks, and overall wall-time speedup is capped at ~6–7×.

Numerical Results and Analysis

Operator Profiling:

Fused kernel achieves 6.0–6.6× per-matvec speedup over FP64 three-stage baseline for problems ranging $K$ 3– $K$ 4 elements (see Table 1 in paper).
Application to full SIMP-120 runs yields 4.6–7.3× wall-time speedup depending on problem size.

Wall Time and Energy:

Cantilever problems spanning $K$ 5 to $K$ 6 elements are solved in 17.5–997 seconds on an RTX 4090 for the fused kernel, vs. 80–6141s for the FP64 baseline.
Board-level energy measurements show 3.2–4.9× lower energy consumption for the fused path over FP64 runs.

Solution Quality:

Compliance and grayness (discreteness) metrics from the fused path are nearly identical (within $K$ 7 up to $K$ 8M elements, $K$ 9 at $K\mathbf{v}$ 0M elements) to FP64, confirming no numerical drift.
Repeatability and non-determinism induced by floating-point atomics are empirically negligible at engineering accuracy.

Precision Barrier in Mixed-Precision Solvers

The empirical studies and power-iteration-based conditioning estimates reveal that the SIMP stiffness matrices are extremely ill-conditioned:

$K\mathbf{v}$ 1 ranges from $K\mathbf{v}$ 2 (64k elements) to $K\mathbf{v}$ 3 (512k elements) under uniform density, with expected increases for late-phase SIMP penalties and higher contrast.
For BF16 arithmetic ( $K\mathbf{v}$ 4), the standard iterative refinement condition $K\mathbf{v}$ 5 is violated by more than an order of magnitude.
BF16 CG and even iterative refinement with BF16 inner solves stagnate with large compliance errors (44–54%).

The paper thus confirms, both theoretically and empirically, that direct use of BF16 in the CG loop is infeasible for practical 3D SIMP TO problems and that IR is insufficient to resolve this unless the operator's effective condition number is reduced.

Implications and Path Towards Multigrid and Smoothers

The observed throughput of BF16 tensor cores in the GEMM microbenchmarks motivates a future solution path: The use of the BF16 implementation as a fine-grid smoother inside a geometric multigrid V-cycle, where the effective condition number of the subproblem processed by the smoother can be controlled. Integration with geometric multigrid preconditioning is projected to resolve CG iteration bottlenecks, making it possible to exploit the BF16 path for interior smooths while maintaining FP32/FP64 elsewhere.

The software is implemented in pure Python with CuPy for GPU compute, supporting accessible deployment and extension by the engineering optimization community.

Conclusion

This work achieves a highly efficient matrix-free gather–GEMM–scatter kernel for 3D SIMP TO, validated through strong numerical and energy performance against the three-stage baseline. It meticulously quantifies the throughput and memory advantages enabled by runtime compilation and kernel fusion, along with the severe stability constraints imposed by SIMP-induced conditioning on low-precision (BF16) arithmetic. The results demonstrate that the memory bandwidth, not raw floating-point throughput, remains the principal limiter for matrix-free FEM on GPUs. The future integration of geometric multigrid algorithms—where BF16 smoothers can operate below their conditioning threshold—represents a compelling research direction, with the present implementation serving as a robust and extensible platform.

Markdown Report Issue