SegFold: Accelerating Sparse GEMM with a Fine-Grained Dynamic Dataflow

2026-06-25

AI Summary

This work addresses the limitations of existing sparse general matrix-matrix multiplication (SpGEMM) accelerators, which rely on static dataflows that struggle to simultaneously achieve high data reuse and load balance. To overcome these inherent constraints, we propose SegFold, the first SpGEMM accelerator architecture incorporating fine-grained dynamic scheduling and dynamic remapping of partially completed work into its dataflow design. SegFold employs a customized memory controller to exploit data reuse within local windows and integrates dynamic task distribution with a reconfigurable merging network to ensure balanced workload allocation across processing units. Experimental results demonstrate that SegFold achieves an average speedup of 1.95× over state-of-the-art SpGEMM accelerators and up to 5.3× improvement compared to the best static configuration, across a wide range of sparsity patterns and matrix sizes.

Abstract

Generalized sparse matrix-matrix multiplication (SpGEMM) is critical in many domains. Existing CPUs and GPUs, as well as specialized accelerators, rely on static dataflows (e.g., inner product, outer product, Gustavson, etc.). Each static dataflow sacrifices some data reuse opportunities and imposes constraints on load balance. To address this inefficiency, we extend the typical SpGEMM dataflows by considering dynamism. Specifically, we add fine-grained dynamic scheduling to optimize reuse and reduce resource contention. We also develop dynamic remapping of partially completed work to improve load balance and parallelism. These ideas are formalized into a specific dataflow called Segment. To demonstrate Segment, we codesign a SpGEMM accelerator called SegFold. SegFold includes a memory controller that identifies fine-grained reuse opportunities in a local window of the stationary input array and exploits them through dynamic work assignment. It also incorporates a merge network that routes inputs to appropriate processing elements (PEs) for computation while dynamically remapping the work assigned to each PE to balance load. Across diverse densities and matrix sizes, SegFold achieves a geometric-mean $1.95\times$ speedup over state-of-the-art SpGEMM accelerators and $5.3\times$ over the best static dataflow configuration, demonstrating that adding dynamism to the dataflow design space unlocks reuse and load-balance gains that no static scheduling choice can achieve in isolation.

Problem

• SpGEMM
• static dataflow
• data reuse
• load balance
• sparse matrix multiplication

Innovation

• dynamic dataflow
• SpGEMM
• fine-grained scheduling
• load balancing
• memory reuse