pith. sign in

arxiv: 2606.00601 · v1 · pith:LGBAI4EZnew · submitted 2026-05-30 · 💻 cs.PL · cs.DC

ScanWeaver: Compiler-Driven Parallelization of Affine Recurrences via Associative Scan Lowering

Qiying Wu , Pavel Zolnikov This is my paper

Pith reviewed 2026-06-28 18:04 UTC · model grok-4.3

classification 💻 cs.PL cs.DC
keywords affine recurrencesassociative scanBlelloch scanMLIRGPU parallelizationselective state-space modelsMambacompiler lowering
0
0 comments X

The pith

ScanWeaver converts affine recurrences into associative scans that lower end-to-end through MLIR to parallel Blelloch execution on GPUs.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper introduces ScanWeaver as a compiler framework that takes recurrence computations, such as the input-dependent scans in selective state-space models, and rewrites them as associative scan operations. This rewrite allows systematic lowering via MLIR to compiler-generated Blelloch scans that execute in parallel on GPUs. The work treats the recurrence structure itself as a reusable compiler abstraction rather than a model-specific optimization target. Validation covers decomposition, lowering, code generation, and actual hardware execution, with results compared against sequential PyTorch, CUDA, and fused Mamba baselines. The central goal is to remove the sequential bottleneck while keeping the original linear-time semantics intact across a class of affine recurrences.

Core claim

ScanWeaver decomposes affine recurrences that arise in modern ML workloads into associative forms, then applies MLIR-based lowering to produce executable GPU programs that implement the Blelloch scan algorithm, preserving exact forward semantics for selective-scan workloads.

What carries the argument

Associative scan lowering that rewrites input-dependent affine recurrences into Blelloch-compatible parallel scans inside the MLIR infrastructure.

If this is right

  • Forward selective-scan workloads from models like Mamba can execute in parallel on GPUs using only compiler-generated code.
  • The same lowering path applies to other affine recurrences that share the same local recurrence structure.
  • MLIR becomes the single source for both the recurrence description and the final GPU binary.
  • Hand-written CUDA scan kernels can be replaced by verified compiler output for this class of computations.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The same decomposition technique could apply to recurrence patterns outside ML, such as certain dynamic programming or time-series algorithms.
  • Integration into standard ML frameworks might allow automatic parallelization of new recurrence-based layers without manual kernel development.
  • The approach opens a path to testing whether other scan variants or higher-order recurrences can be lowered through the same MLIR pipeline.

Load-bearing premise

The target affine recurrences keep their exact original semantics after decomposition into associative Blelloch scans, without any approximations or extra runtime checks that would change the computation.

What would settle it

Execute the generated GPU program on the same input sequences used by a sequential reference implementation and check whether every output element matches to machine precision; a single mismatch on any element falsifies semantic preservation.

Figures

Figures reproduced from arXiv: 2606.00601 by Pavel Zolnikov, Qiying Wu.

Figure 1
Figure 1. Figure 1: Overview of the ScanWeaver transformation. Sequential recurrences are rewritten into affine transition pairs, [PITH_FULL_IMAGE:figures/full_fig_p004_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: Affine scan reformulation converts sequential [PITH_FULL_IMAGE:figures/full_fig_p004_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: Concrete ScanWeaver lowering path from recur [PITH_FULL_IMAGE:figures/full_fig_p004_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: Latency scaling across sequence lengths for sequen [PITH_FULL_IMAGE:figures/full_fig_p005_4.png] view at source ↗
read the original abstract

Selective state-space models such as Mamba highlight the practical importance of input-dependent scan recurrences, which preserve linear-time sequence modeling while improving language modeling capabilities. However, these recurrences introduce stricter sequential dependencies than classical structured SSMs, limiting parallel execution on modern accelerators. We present \textbf{ScanWeaver}, a compiler framework that transforms recurrence-based computations into associative scan representations and lowers them end-to-end to executable GPU programs. We use Mamba-style selective scan as a motivating example of a broader class of affine recurrences that arise in modern ML workloads. Rather than targeting a single model family, ScanWeaver elevates this recurrence structure to a first-class compiler abstraction, enabling systematic MLIR-based lowering to compiler-generated Blelloch scan execution on GPUs. Across forward selective-scan workloads with matched local recurrence semantics, we validate affine recurrence decomposition, Blelloch lowering, MLIR GPU lowering, executable artifact generation, and actual GPU execution from generated MLIR. We benchmark the resulting ScanWeaver GPU execution against PyTorch and CUDA sequential baselines, and include the Mamba kernel as a fused production baseline for systems context.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

2 major / 0 minor

Summary. The paper presents ScanWeaver, a compiler framework that transforms recurrence-based computations (using Mamba-style selective scans as a motivating example of a broader class of affine recurrences) into associative scan representations. It performs end-to-end MLIR-based lowering to compiler-generated Blelloch scan execution on GPUs, with validation of decomposition, lowering, and actual GPU execution across forward selective-scan workloads that have matched local recurrence semantics, plus benchmarks against PyTorch, CUDA, and Mamba baselines.

Significance. If the central claim of semantics-preserving decomposition and lowering holds for the targeted class, the work would provide a systematic compiler-driven path to parallelize input-dependent linear recurrences in ML workloads on GPUs. The end-to-end lowering to executable MLIR artifacts and actual GPU execution is a concrete strength.

major comments (2)
  1. [Abstract] Abstract: the validation is described only as covering 'forward selective-scan workloads with matched local recurrence semantics,' with no quantitative results, error analysis, or details on equivalence checking. This leaves the central claim of general affine recurrence decomposition only partially supported.
  2. [Abstract (validation description)] The manuscript provides no formal argument, inductive proof, or counter-example analysis establishing that the decomposition rules yield bit-for-bit identical results for arbitrary affine recurrences (beyond cases where the local recurrence already matches the target associative form). This is load-bearing for the claim that ScanWeaver handles a 'broader class of affine recurrences.'

Simulated Author's Rebuttal

2 responses · 1 unresolved

We thank the referee for the feedback on the abstract and validation claims. We address the points regarding quantitative support and formal guarantees for the decomposition, clarifying the manuscript's empirical focus while acknowledging limitations in formal arguments.

read point-by-point responses

  1. Referee: [Abstract] Abstract: the validation is described only as covering 'forward selective-scan workloads with matched local recurrence semantics,' with no quantitative results, error analysis, or details on equivalence checking. This leaves the central claim of general affine recurrence decomposition only partially supported.

    Authors: The abstract intentionally scopes the validation to workloads with matched local recurrence semantics to emphasize that equivalence holds by construction of the decomposition rules. Quantitative performance results, including speedups over PyTorch and CUDA baselines and comparisons to the Mamba kernel, appear in the evaluation section. Equivalence is confirmed via end-to-end GPU execution producing matching outputs. We will revise the abstract to explicitly reference the evaluation section and note the empirical equivalence checks performed. revision: partial

  2. Referee: [Abstract (validation description)] The manuscript provides no formal argument, inductive proof, or counter-example analysis establishing that the decomposition rules yield bit-for-bit identical results for arbitrary affine recurrences (beyond cases where the local recurrence already matches the target associative form). This is load-bearing for the claim that ScanWeaver handles a 'broader class of affine recurrences.'

    Authors: The decomposition rules target the class of affine recurrences arising in ML workloads (exemplified by selective scans) that admit lowering to associative form; the 'matched local recurrence semantics' qualifier denotes precisely those cases where the rules apply and preserve semantics. The manuscript presents the rules with their rationale but does not include a formal inductive proof or exhaustive counter-example analysis, relying instead on the compiler implementation and empirical validation on representative workloads. revision: no

standing simulated objections not resolved
  • No formal inductive proof or counter-example analysis is provided for semantics preservation on arbitrary affine recurrences beyond the matched cases.

Circularity Check

0 steps flagged

No circularity: framework presented as independent lowering technique

full rationale

The provided abstract and manuscript description contain no equations, fitted parameters, or derivations that reduce claimed results to inputs by construction. ScanWeaver is described as a compiler framework that elevates recurrence structure to an MLIR abstraction and performs Blelloch lowering; validation is reported as empirical execution matching on workloads with 'matched local recurrence semantics.' No self-citation chains, ansatzes smuggled via prior work, or renaming of known results appear as load-bearing steps. The semantics-preservation claim is presented as an assumption of the approach rather than a derived quantity that collapses to the input data.

Axiom & Free-Parameter Ledger

0 free parameters · 2 axioms · 1 invented entities

The central claim rests on the domain assumption that affine recurrences admit exact associative decomposition and that MLIR can generate correct Blelloch-scan GPU code for them; no free parameters or invented physical entities are described.

axioms (2)
  • domain assumption Affine recurrences arising in ML workloads admit decomposition into associative scans that preserve exact local recurrence semantics.
    Invoked when the paper states that ScanWeaver transforms recurrence-based computations while matching local recurrence semantics.
  • domain assumption MLIR-based lowering to Blelloch scan produces executable GPU programs whose behavior matches the original recurrence.
    Required for the end-to-end claim of generating and running correct artifacts.
invented entities (1)
  • ScanWeaver compiler framework no independent evidence
    purpose: Elevates affine recurrence structure to a first-class abstraction for systematic lowering.
    New system introduced by the paper; no independent evidence outside the described implementation.

pith-pipeline@v0.9.1-grok · 5732 in / 1415 out tokens · 25179 ms · 2026-06-28T18:04:44.274381+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

13 extracted references · 2 canonical work pages · 2 internal anchors

  1. [1]

    Efficiently modeling long sequences with structured state spaces

    Albert Gu, Karan Goel, and Christopher R´e. Efficiently modeling long sequences with structured state spaces. International Conference on Learning Representations, 2021

    2021

  2. [2]

    Mamba: Linear-Time Sequence Modeling with Selective State Spaces

    Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces.arXiv preprint arXiv:2312.00752, 2023

    work page internal anchor Pith review Pith/arXiv arXiv 2023

  3. [3]

    Daniel Hillis and Guy L

    W. Daniel Hillis and Guy L. Steele Jr. Data parallel algorithms.Communications of the ACM, 29(12): 1170–1183, 1986

    1986

  4. [4]

    Blelloch

    Guy E. Blelloch. Prefix sums and their applications. In Synthesis of Parallel Algorithms. Morgan Kaufmann, 1990

    1990

  5. [5]

    Triton: An intermediate language and compiler for tiled neural network computations

    Philippe Tillet et al. Triton: An intermediate language and compiler for tiled neural network computations. Proceedings of Machine Learning and Systems, 2021

    2021

  6. [6]

    Cutlass: Fast linear algebra in cuda c++

    NVIDIA. Cutlass: Fast linear algebra in cuda c++. https://github.com/NVIDIA/ cutlass, 2024

    2024

  7. [7]

    Fu, Stefano Ermon, Atri Rudra, and Christopher R´e

    Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher R´e. Flashattention: Fast and memory- efficient exact attention with io-awareness.Advances in Neural Information Processing Systems, 2022

    2022

  8. [8]

    Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines.ACM SIGPLAN Notices, 48(6):519–530, 2013

    Jonathan Ragan-Kelley et al. Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines.ACM SIGPLAN Notices, 48(6):519–530, 2013

    2013

  9. [9]

    Tvm: An automated end-to-end op- timizing compiler for deep learning

    Tianqi Chen et al. Tvm: An automated end-to-end op- timizing compiler for deep learning. In13th USENIX Symposium on Operating Systems Design and Imple- mentation, pages 578–594, 2018

    2018

  10. [10]

    Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions

    Nicolas Vasilache et al. Tensor comprehensions: Framework-agnostic high-performance machine learn- ing abstractions. InarXiv preprint arXiv:1802.04730, 2018

    work page internal anchor Pith review Pith/arXiv arXiv 2018

  11. [11]

    Mlir: A compiler infrastructure for the end of moore’s law

    Chris Lattner et al. Mlir: A compiler infrastructure for the end of moore’s law. InLLVM Developers’ Meeting, 2020

    2020

  12. [12]

    IREE: Intermediate representation execution environment

    IREE Authors. IREE: Intermediate representation execution environment. https://github.com/ iree-org/iree, 2024

    2024

  13. [13]

    TensorIR: An abstraction for automatic tensorized program op- timization

    Siyuan Feng, Bohan Hou, Hongyi Jin, Wuwei Lin, Junru Shao, Ruihang Lai, Zihao Ye, Lianmin Zheng, Cody Hao Yu, Yong Yu, and Tianqi Chen. TensorIR: An abstraction for automatic tensorized program op- timization. InProceedings of the 28th ACM Interna- tional Conference on Architectural Support for Pro- gramming Languages and Operating Systems, 2023

    2023