--- title: 'SegFold: Dynamic Dataflow for SpGEMM' url: https://www.emergentmind.com/topics/segfold type: topic --- # SegFold: Dynamic Dataflow for SpGEMM SegFold is a codesigned dataflow and accelerator for generalized sparse matrix–matrix multiplication (SpGEMM) that introduces fine-grained dynamism into both scheduling and mapping, departing from the three canonical static schedules—inner product, outer product, and Gustavson. In the formulation considered, SpGEMM computes $C_{m,n} = \sum_k A_{m,k} \times B_{k,n}$ for sparse $A$ and $B$. The central claim is that static dataflows commit in advance to a single reuse–imbalance trade-off, whereas SegFold adds sub-tile dynamic scheduling to optimize reuse and dynamic remapping of partially completed work to improve load balance and parallelism [2606.26701]. ## 1. Problem setting and motivation SegFold is motivated by the observation that the classical SpGEMM loop orders each privilege one form of locality while constraining another. Inner product, ordered as $M \to N \to K$, reuses $C$ partial sums because $K$ is innermost, but exhibits poor input reuse of $A$ and $B$ and is sensitive to row–column intersection balance. Outer product, ordered as $K \to M \to N$, reuses entire rows or columns of $A$ and $B$ because $K$ is outermost, but produces large, irregular partial-sum matrices $T$ and places $C$ reuse far apart in time. Gustavson, ordered as $M \to K \to N$, fully reuses rows of $A$ and reuses $C$ within each row, but does not reuse $B$ and still faces irregular $C$ reduction patterns [2606.26701]. The paper’s framing is that each static dataflow chooses a reuse pattern and a load-balance profile ahead of time. Real sparse patterns vary tile-to-tile and even within tiles, so a fixed schedule either loses reuse opportunities or suffers load imbalance. Recent adaptive schemes such as Spada and Flexagon are described as selecting among the three static schedules at tile granularity, while preserving a fixed execution order inside each tile. SegFold’s key insight is that adding dynamism at sub-tile granularity enables two simultaneous effects: reordering work within a window of $A$ to maximize row-wise reuse of $B$, and redistributing partial sums of $C$ across processing elements (PEs) on the fly to balance load [2606.26701]. A recurrent misconception in this area is to equate “adaptive” execution with any runtime choice among inner product, outer product, and Gustavson. SegFold’s contribution is narrower and more specific: it does not merely switch among preexisting loop nests, but extends the dataflow space by adding dynamic scheduling and dynamic mapping within the tile itself. This suggests that the proposal is best understood not as a new point in the static design space, but as a mechanism for relaxing the commitment to any single static schedule. ## 2. The Segment dataflow The dynamic dataflow formalized in the work is called **Segment**. Rather than fixing a single loop-nest order, Segment interleaves the $M\!-\!K$ and $K\!-\!N$ dimensions through two coordinated operations, called **SelectA** and **SegmentBC** [2606.26701]. SelectA performs fine-grained dynamic scheduling over an active window $W_k$ of up to $W_{\max}$ distinct $k$-values. At each cycle, it greedily selects up to $R_{\max}$ pairs $(m,k)$ from $A$ under two constraints: no two selected pairs may share the same $m$, to avoid $C$-row contention, and as many pairs as possible should share the same $k$, to maximize reuse of a row of $B$. The scheduler fills the active window with new $k$ values until $|W_k| = W_{\max}$, tracks completed $k$ values, and retires any $k$ whose entire row is done. The effect is to break both the static $M\!-\!K$ and $K\!-\!M$ nests and to permit partial-row streaming of $B$ by segmenting its $N$ dimension. SegmentBC performs on-the-fly element-wise redistribution after a batch of $(m,k)$ pairs has been selected. Each selected pair carries a partial row $\{B_{k,n}\}$, and each nonzero $B_{k,n}$ is injected into a virtual coordinate space $\mathcal V = X \times Y$, where $X$ indexes nonempty rows of $C$ and $Y$ indexes the current sorted nonzeros within that row. A time-varying injective mapping $$ f_t : (m,n) \mapsto (x,y) $$ is required to satisfy four conditions: injectivity, row saturation with no gaps in $Y$, column ordering so that $n_1 < n_2$ implies $y_1 < y_2$, and time-ascending motion in which entries move only to the right over time. When a new term $T_{m,n,k} = A_{m,k}\times B_{k,n}$ arrives, it either matches an existing $(m,n)$ entry and is reduced in place, or is inserted into sorted position with rightward shifts of existing entries. The work defines a segment displacement for each injected element and states that minimizing average displacement reduces network contention [2606.26701]. The significance of Segment lies in its attempt to combine multiple forms of reuse that are separated in conventional static schedules. The paper characterizes this as dynamic exploitation of reuse across $A$, $B$, and $C$, together with redistribution of reductions for better balance. A plausible implication is that Segment is not only a scheduling policy but also a structural constraint on how partial results may evolve over time, since sortedness, injectivity, and monotonic rightward movement are built into the mapping. ## 3. Microarchitectural realization in SegFold SegFold realizes Segment with a 16×16 PE array and a microarchitecture organized around a memory controller, an adaptive merge network, an index-to-PE mapper, and folding mechanisms [2606.26701]. The **memory controller and metadata path** store $A$ in column-major order and $B$ in row-major order, with DCSR plus a start-pointer for partial rows. The controller maintains the active window $W_k$, whose default size is up to 32, and scans the $A$-bitmask of size $O(W)$ each cycle using pure combinational logic to drive SelectA. It also coalesces $A$ and $B$ accesses to the on-chip cache and then to HBM2. Within each PE row, the **adaptive merge network** assigns one virtual row of $C$ to an ordered sequence of PEs whose stored column indices $c_i$ increase from left to right. For an incoming $B$ element with column index $b$, the network injects at a start position $s$ and then propagates right with simple per-PE comparator logic. If $b>c$, the element is forwarded rightward. If $b=c$, it accumulates and stops. If $b