High-Order Spectral Element Methods for Wave Propagation on ARM Multicore CPU with SME: Optimizations and Implications

Abstract: Wave propagation based on the spectral element method (SEM) is a representative HPC workload, but existing SEM implementations are not well matched to emerging ARM multicore CPUs with Scalable Matrix Extension (SME). We present an SME-enabled optimization of \textsc{SPECFEM3D} on the emerging LX2 processor that combines an SME-aware batched small-matrix kernel for SEM tensor-product operators, a memory-aware hybrid MPI+OpenMP execution scheme for limited-HBM systems, and a dispersion-based iso-accuracy study of the $(h,p)$ tradeoff. At fixed polynomial order, the optimized implementation improves full-application performance by 4--6$\times$ over the original code and delivers clear gains over optimized non-SME CPU baselines. Beyond these implementation-level gains, our results suggest that SME shifts the performance-favorable operating point toward higher polynomial orders along the dispersion-based iso-accuracy frontier, further reducing time-to-solution and working-set size. These results indicate that SME affects not only kernel efficiency, but also the practical discretization tradeoff for SEM on modern ARM multicore platforms.

• The paper introduces an SME-tailored batched matrix multiplication kernel that boosts wave propagation performance in SEM.
• It demonstrates a shift in SEM discretization trade-offs, reducing cost scaling from O(p^4/h^3) to O(p^3/h^3) for high-order accuracy.
• Experimental results validate up to 7.7x kernel speedup and marked reductions in memory and time-to-solution for seismic applications.

High-Order Spectral Element Methods for Wave Propagation on ARM Multicore CPU with SME: Optimizations and Implications

Introduction and Motivation

This work addresses the intersection of high-order spectral element methods (SEM) for wave propagation and emerging ARM multicore architectures equipped with the Scalable Matrix Extension (SME). The primary application driver is large-scale seismic simulation (e.g., in earthquake hazard assessment and subsurface exploration), which is characterized by stringent accuracy requirements and computational demands. The spectral element method, due to its combination of geometric flexibility and high numerical fidelity utilizing tensor-product polynomial bases, is a standard in large-scale geophysical simulations. However, the alignment between SEM computational kernels—dominated by batched small matrix multiplications inherent in tensor-product operators—and the microarchitectural capabilities of new hardware extensions such as ARM SME has been historically suboptimal.

Architectural and Numerical Foundations

The implementation and experiments are centered on the LX2 processor, which is a dual-die ARMv9-A CPU featuring 304 cores and robust memory hierarchies, including both high-bandwidth memory (HBM) and large DDR capacity. Of critical significance is the ARM SME, which augments the core's vector capabilities with matrix-tile-based outer-product instructions for accelerating matrix multiplications. The ZA tile—a $64 \times 64$-byte matrix register—is the cornerstone of SME operations, enabling increased throughput for matrix-oriented workloads compared to conventional SIMD.

Figure 1: LX2 block diagram, showing NUMA domains, HBM, and position of SME vector/matrix engines.

High-order SEM utilizes tensor-product Gauss-Legendre-Lobatto interpolation over hexahedral meshes, with the stiffness operator applied via batched small matrix-matrix multiplications. For a typical SEM configuration, the sizes of these matrices are smaller than optimal for SME tile utilization, especially for moderate polynomial orders (e.g., $p=4 \ldots 8$) that are commonly used in practice. Understanding and overcoming this kernel-architecture mismatch is pivotal for exploiting the full computational potential of ARM multicore platforms.

Figure 2: Visualization of the SEM tensor-product basis derivative, highlighting axis-aligned differentiation as sequential batched matrix multiplications.

SME-Enabled Kernel Optimization

The authors introduce an SME-tailored batched matrix multiplication kernel that leverages software pipelining and SVE-based vector aggregation. Different tensor-product derivative directions are mapped onto SME tiles with explicit inter-slice pipelining, maximizing temporal overlap across tiles. For low-to-moderate $p$, aggregation of adjacent element slices within SVE vectors, together with register shuffling, enables packing multiple matrix operations into a single SME invocation, thus mitigating the inefficiency arising from the mismatch in operand shape versus tile size.

Figure 3: Schematic of the aggregated SME batched small matrix kernel, illustrating zone utilization and slice parallelism.

For the $\xi$-direction, where an in-memory transpose is required to enable contiguous tile loads, specialized in-register shuffle and transpose primitives minimize the penalty. Theoretical throughput bounds are derived as a function of $p$ and vector length, predicting near-linear scaling with $p$ up to full tile utilization.

SME Impact on Discretization and Application Performance

A central observation is that SME's improved efficiency at higher $p$ fundamentally alters the discretization trade-off. Historically, the computational cost of SEM stiffness operators has scaled as $O(p^4/h^3)$, favoring moderate $p$ and $h$ (element size) combinations for fixed accuracy. With SME, this cost scaling reduces to $p=4 \ldots 8$0 in the practical pre-tile-saturation regime, directly shifting the performance-favorable $p=4 \ldots 8$1 point towards higher polynomial orders.

Figure 4: Contour shift on the iso-accuracy $p=4 \ldots 8$2 frontier: SME optimization reduces the cost curve, making higher $p=4 \ldots 8$3 regimes more favorable.

Higher $p=4 \ldots 8$4 allows a reduction in total global degrees of freedom—decreasing memory footprint and time-to-solution—while maintaining the same spatial dispersion accuracy. However, the minimum node spacing and resulting CFL constraints must be balanced against practical time-stepping. For high-accuracy scenarios such as those in reverse-time migration, the temporal-dispersion error bound is the primary constraint, further solidifying the advantage of high-order plus SME.

Parallelization and Execution Model

The memory pressure from duplicated per-process SEM data structures is particularly acute on high-core-count CPUs with relatively constrained HBM per NUMA node. The paper develops a hybrid MPI+OpenMP parallelization scheme, mapping MPI ranks to NUMA domains and exploiting thread concurrency within domains. Node update races, traditionally mitigated by atomics, are eliminated via graph coloring of the element adjacency graph, ensuring atomicity-free, conflict-free updates within color sets. Communication-computation overlap is achieved through a dedicated progress thread per MPI process, driving asynchronous exchanges and fully exposing hardware concurrency.

Figure 5: Workflow of hybrid process/thread parallelization, illustrating NUMA mapping, graph coloring, and communication threading.

Experimental Methodology and Results

The kernel-level analysis measures single-core FLOP rates for SME-enabled and baseline (SIMD/SVE) batched small matrix multiplications over a sweep of polynomial orders. The aggregated SME kernel reaches up to $p=4 \ldots 8$5 the performance of SVE at $p=4 \ldots 8$6, closely approaching the obtainable throughput bounds for the respective directions (especially $p=4 \ldots 8$7 and $p=4 \ldots 8$8 derivatives).

Figure 6: FLOPS comparison of SME-enabled vs. SIMD batched small matrix multiplications, as a function of $p=4 \ldots 8$9.

Wave-propagation accuracy is validated through an iso-accuracy analysis in homogeneous media. At stringent error targets ($p$0 at 8s propagation), larger $p$1 permits an upscaling of $p$2 while maintaining accuracy, with admissible element sizes scaling from 20m ($p$3) to 130m ($p$4).

Figure 7: Dispersion error as function of $p$5, highlighting the admissible domain for tight error thresholds.

End-to-end simulation metrics demonstrate that increasing $p$6 from 4 to 7 yields a $p$7 reduction in memory use and nearly $p$8 speedup; further increasing to $p$9 halves memory demand and reduces time-to-solution by an additional $\xi$0.

Figure 8: Application-level time-to-solution and memory footprint reduction achieved by SME-enabled higher-order SEM configurations.

Breakdown analyses isolate the performance contributions: vectorization and kernel fusion ($\xi$1-$\xi$2), graph coloring ($\xi$3-$\xi$4 depending on $\xi$5), and SME tile utilization (up to $\xi$6 on the stiffness kernel, $\xi$7 global).

Figure 9: Stacked contributions of each optimization stage to total speedup, with strongest effect from SME at high polynomial order.

Scaling experiments show that SME improves the competitive standing of ARM CPUs versus accelerators. For $\xi$8, performance matches A100 GPUs; at $\xi$9, A100 implementations exhaust shared memory, while SME enables continued scaling and $p$0 throughput improvement.

Figure 10: Intra- and inter-node scaling curves, showing weak-scaling near ideal up to available HBM and strong scaling dependent on local workload size.

Theoretical and Practical Implications

This research demonstrates that emerging matrix-oriented ISA extensions (e.g., ARM SME) can reshape not only kernel-level optimization strategies but also discretization regimes and parallel programming models in high-order PDE codes. Specifically, on memory-bandwidth-limited platforms, the optimal SEM configuration shifts decisively toward higher polynomial orders, further compressing memory requirements and wall time for large-scale wave simulation. The adoption of co-design strategies—spanning operand layout, pipelined SME-tile aggregation, and hybrid runtime models—enables the domain to transcend historical cost and scalability barriers on general-purpose architectures.

Practically, these optimizations position ARM multicore CPUs as credible, cost-effective alternatives to GPUs for high-order method-based workloads, with portability advantages for scientific codes. The results also suggest that continued increases in vector/tile widths in CPU ISAs will further extend the domain of affordable high-order discretization, amplifying the need for architecturally aware software infrastructure.

Conclusion

This work comprehensively optimizes high-order SEM for ARM multicore CPUs with SME, achieving large multiplicative gains at both the kernel and full-application levels for SPECFEM3D. Crucially, it establishes that SME alters not only hardware utilization but also the practical position of the $p$1 operating point, making higher-order discretizations superior in both cost and accuracy for modern memory-constrained multicore platforms. This has deep ramifications for future algorithm/hardware co-design in large-scale scientific simulation and motivates a reevaluation of performance models and discretization strategies in the exascale and post-exascale era.