Tool automatically converts loops into faster array operations in numpy
Vectorizer: Vectorizing NumPy Programs with Shape-Guided Rewrite
Computation and LanguageSoftware Engineering
Summary
Writing fast programs with NumPy, a popular tool for number crunching in Python, usually means avoiding explicit loops and using special array operations instead. These array operations are tricky for many programmers because they require careful thinking about how arrays fit together by shape and how data moves between them. The authors created a tool called Vectorizer that automatically rewrites slow loop-based code into fast array operations by analyzing shapes and data flow. This rewriting makes the programs run much faster without needing the programmer to do the complicated work themselves.
What this means in practice
- •For python software developers: Automatically improve the performance of Python programs using explicit loops on arrays by rewriting them into optimized NumPy vectorized operations.
- •For data engineering teams: Speed up numerical data processing pipelines by converting slow loop-based scripts to vectorized forms that run much faster in production environments.
Authors
Jingqian Liu, Xiaoyu Liu, Yuepeng Wang
Abstract
NumPy is a widely used Python library for numerical scientific computing, known for its declarative APIs and its optimized implementations. However, writing efficient NumPy programs, which often entails using vectorized array operations instead of explicit Python loops, may not be straightforward. This can be difficult for programmers who are accustomed to imperative array traversal, especially when vectorized API invocations require careful reasoning about shapes, broadcasting, and advanced indexing. This paper presents a rewrite-based approach for vectorizing Numpy programs with explicit loops over array data. Our approach vectorizes loops from the inside out, using array shapes and dataflow analysis to guide a source-to-source transformation that replaces loop bodies with vectorized statements. Following a set of rewrite rules that are correct by construction, our approach is consistently fast. We have implemented the approach as a tool called Vectorizer and evaluated it on 150 benchmarks collected from prior work and Stack Overflow. The evaluation shows that Vectorizer vectorizes 142 of the 150 benchmarks directly and 2 more after minor changes to the original benchmarks, with only 0.53 seconds on average to rewrite each one. The resulting programs are, on average, 74.83x faster than the original loop-based implementations.