DARPA Case Study: Up to 100x Faster Sparse Arrays | OpenTeams | OpenTeams

Case Study

Accelerating Scientific Computing for DARPA via PyData/Sparse and TACO Integration

Client: Defense Advanced Research Projects Agency (DARPA)

Overview

Most world data is stored in multi-dimensional arrays or tensors, much of which is inherently sparse. Traditional Python processing either treats sparse data as dense, incurring high computational overhead, or forces complex format transformations. Under DARPA contract HR001123C0139, OpenTeams and Quansight partnered with MIT to integrate PyData/Sparse with the Tensor Algebra Compiler (TACO) and Finch framework. The project successfully demonstrated a 100x speedup on sparse array benchmarks, while requiring <5% lines of code changed.

speedup on sparse array benchmarks

0 x

Lines of code changed

< 0 %

Project Details

The primary objective was to make sparse data processing in Python as accessible, performant, and device-agnostic as dense processing across standard APIs (NumPy, SciPy, scikit-learn). Key pillars included:

Intermediate Representation (IR)

Adopted TACO’s concrete index notation, iteration graphs, and merge lattices to optimize sparse loop kernels.

Target Architecture Code Generation

Enabled JIT compilation via Julia/LLVM and MLIR to target multiple architectures (x86-64, ARM, PowerPC, CUDA GPUs).

Python Array API Standard Adoption

Standardized interfaces across PyData/Sparse to enable drop-in interoperability for downstream libraries like SciPy and scikit-learn.

Key technical
accomplishments

Finch-TACO Integration

Integrated the Finch tensor compiler to support fine-grained level formats (CSR, CSC, COO, run-length encoding, block-sparse) on a per-dimension basis.

End-to-End Auto-Scheduling

Incorporated MIT’s cost-based Galley auto-scheduler to optimize whole-graph computations and fused operation execution.

GPU & Interoperability Protocols

Co-developed the zero-copy __binsparse__ protocol for fast GPU data transfer and SciPy/CuPy interchange.

Impact

Through automated JIT compilation and auto-scheduling, the framework achieved landmark speedups across key benchmarks, including a ~100x speedup over SciPy sparse for sampled matrix multiplication ((A @ B) * C) and a 40x speedup over Numba backends for matrix multiplication.

Speedup over SciPy sparse

~ 0 x

for sampled matrix multiplication ((A @ B) * C)

Speedup over Numba backends

0 x

for matrix multiplication

Make your Python
faster without rewriting it.

We build and maintain the open source tools your team already uses, and we can tune them for your workloads.