Wafer-scale computing has emerged as an alternative solution to sustain performance scaling in the post-Moore era, driven by recent technology advancements such as chiplet integration. This enables considerable computing and storage capabilities on a single chip, making it capable of accommodating large machine learning models and datasets. Recent efforts have heralded the promise of wafer-scale architectures for deep learning inference and training. However, scaling the training of Graph Neural Networks in wafer-scale architecture remains a challenge and is relatively unexplored due to irregularities in gradient propagation as well as physical constraints from flat on-chip topologies. In this paper, we propose Aster, a topology-aware framework designed to efficiently support GNN training on arbitrary wafer-scale architectures. The proposed framework, as opposed to the current application or topology-specific heuristics, can be generalized to support any network topology and irregular GNN datasets. Specifically, we mathematically formulate commonly-seen network topologies in their geometric representation and prioritize communication efficiency during GNN workload partitioning and mapping. Based on the geometric representation, we propose a quadratic assignment problem solver to efficiently map irregular dataflows to a flat topology with reduced communication distance. The simulation results show that Aster can achieve performance speedup by$2.91 \times, 1.50 \times, 1.84 \times$, and$1.58 \times$in Mesh and speedup by$3.84 \times, 1.56 \times, 2.05 \times$, and$1.49 \times$in Torus on average compared to Mini-cut [1], ScalaGraph [2], ChunkV [3], and Chunk-E [4], respectively.

Sparse Matrix Dense Matrix Multiplication (SpMM) is a fundamental computation kernel across various domains, including scientific computing, machine learning, and graph processing.Despite extensive research, existing approaches optimize SpMM using loop transformations and linear algebra principles, which (1) poorly handle unstructured sparsity patterns, (2) rely on empirical methods to explore data reuse opportunities, and (3) enforce rigid coordinate alignment, compromising data locality.In this paper, we demonstrate that these limitations stem from the fundamental matrix representation and traditional dataflows of SpMM (e.g., inner-product, outer-product, and Gustavson).We propose Aquila, a graph transformation framework that reformulates SpMM computations as a graph optimization problem, leveraging graph theory to reinterpret tiling and dataflow.First, on the theoretical side, we introduce vertex decomposition and adaptive depth traversal (ADT) to enable non-contiguous tiling, where nonzero elements from discontinuous rows and columns are clustered by connectivity rather than following matrix dimensionality.This approach quantifies data reuse and improves data locality beyond traditional loop transformations while maintaining output equivalence.Second, on the algorithm side, we develop a pull-after-push (PaP) dataflow that simultaneously enhances the dense matrix data reuse while eliminating synchronization issues in output matrix accumulation.Third, building on our theoretical approach and dataflow, we present a versatile accelerator architecture that handles a variety of SpMM kernels with diverse data sizes and sparsity patterns in a unified architecture.Additionally, we introduce a bidirectional fiber tree (BFT) format to support the proposed graph-oriented dataflow in contrast to traditional column or row-major access.Evaluation across diverse sparse datasets shows Aquila achieves speedups of 4.3×, 3.4×, 3.7×, 2.9×, and 2.7× in execution time and up to 4.8× * Both authors contributed equally to this research.

Graph Neural Networks (GNNs) have been widely used to handle intricate graph-related problems, in which complex vertex and edge operations are performed in the form of message passing between vertices. Such complex GNN operations are highly dependent on the graph structure and can no longer be characterized as sparse-dense or general matrix multiplications. Consequently, current matrix-based data reuse and workload balancing optimizations have limited applicability to Message Passing-based GNN acceleration. In this paper, we leverage the mathematical insights from Gram Matrix to simultaneously exploit data reuse and workload balancing opportunities for message passing-based GNN accelerations. Upon this insight, we further propose a novel accelerator, named EGMA, that can efficiently facilitate a wide range of GNN models with improved data reuse and workload balance. Consequently, EGMA can achieve performance speedup by 1.57×, 1.72×, and 1.43× and energy reduction by 38.19%, 34.02%, and 24.54% on average compared to Betty, FlowGNN, and ReGNN, respectively.