REVIEW 2 cited by
Operator Fusion in XLA: Analysis and Evaluation
Not yet reviewed by Pith; the record is open.
This paper has not been read by Pith yet. Machine review is queued; the pith claim, tier, and objections will appear here once it completes.
SPECIMEN: schema-true, not a live event
T0 review · schema-true
One-sentence machine reading of the paper's core claim.
pith:XXXXXXXX · record.json · timestamp
read the original abstract
Machine learning (ML) compilers are an active area of research because they offer the potential to automatically speedup tensor programs. Kernel fusion is often cited as an important optimization performed by ML compilers. However, there exists a knowledge gap about how XLA, the most common ML compiler, applies this nuanced optimization, what kind of speedup it can afford, and what low-level effects it has on hardware. Our paper aims to bridge this knowledge gap by studying key compiler passes of XLA's source code. Our evaluation on a reinforcement learning environment Cartpole shows how different fusion decisions in XLA are made in practice. Furthermore, we implement several XLA kernel fusion strategies that can achieve up to 10.56x speedup compared to our baseline implementation.
Forward citations
Cited by 2 Pith papers
-
Mambalaya: Einsum-Based Fusion Optimizations on State-Space Models
Mambalaya fuses the entire Mamba layer into one on-chip computation group, achieving simulated 4.9x prefill and 1.9x generation speedups over a MARCA-like baseline.
-
Nova: An End-to-End MLIR Compiler for Deep Learning
A search-free MLIR compiler that fuses whole training steps outperforms PyTorch by up to ~10% and cuts memory by up to 29%, fitting 144M-parameter MLPs on a 12 GB GPU.
Discussion (0). Continue with ORCID to comment.