PaperlayerView exact revision on arXiv

Source-linked AI summary

Vectorizer: Vectorizing NumPy Programs with Shape-Guided Rewrite

Jingqian Liu, Xiaoyu Liu, Yuepeng Wang

arXiv:2609.08088v1cs.CLcs.SE

TL;DR

Efficient NumPy programming is difficult because vectorization requires careful reasoning about shapes, broadcasting, and indexing. The paper presents Vectorizer, which rewrites explicit loops into vectorized NumPy code using inside-out, shape- and dataflow-guided rules. Across 150 benchmarks, it vectorized 142 directly and 2 more after minor adaptations, with a 74.83× average speedup.

  • Problem

    Vectorizing loop-based NumPy programs is challenging because flexible APIs, broadcasting, advanced indexing, branches, and reductions require careful alignment and lack straightforward loop correspondences.

  • Method

    Vectorizer uses a domain-specific language and correct-by-construction inside-out rewrites guided by array shapes, maskedness, and dataflow analysis.

  • Results

    142 of 150 benchmarks were vectorized directly, and 2 more after minor adaptations; the resulting programs achieved a geometric mean speedup of 74.83×.

  • Takeaways & Limitations

    Vectorizer provides a consistently fast source-to-source approach for transforming many loop-based NumPy programs without search or symbolic reasoning.

  • Takeaways & Limitations

    Programs with strongly coupled loop-carried dependence are not generally vectorizable by the rewrite approach.

Abstract

from arXiv · show

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.

1 Introduction

Vectorizer targets the difficulty of manually aligning shapes, broadcasting, and indexing in loop-based NumPy code. It uses shape- and dataflow-guided rewrites to transform loops into vectorized statements while keeping the analysis tractable.

  • Motivation: Broadcasting and advanced indexing make efficient NumPy code difficult because programmers must correctly align axes and shapes.These challenges can lead programmers to use nested loops that incur runtime inefficiency.
  • Key insight: Many practical loops can be replaced by vectorized statements by expanding an innermost loop variable into an array and propagating its new axis.The approach treats outer-scope variables as fixed-shape free variables and rewrites expressions according to array shapes.
  • Approach: Vectorizer uses a domain-specific language and inside-out rewrite procedure guided by array shapes, maskedness, and dataflow analysis.Its rewrite rules are correct by construction, avoiding enumerative search or a separate equivalence-checking phase.
  • Workflow: The workflow repeatedly vectorizes inner loops and then applies optimizations such as indexing simplification and common subexpression elimination.The resulting code is produced as a runnable NumPy program after postprocessing.
  • Evaluation: 142 of 150 benchmarks were vectorized directly, with 2 more handled after minor adaptations.Vectorizer rewrote each NumPy function in 0.53 seconds on average, and the resulting programs were 74.83× faster on average than the originals.

2 Overview

The overview illustrates how explicit NumPy loops can be replaced by broadcasting-based array computations and shape-guided, inside-out rewrites. Branches require masked representations so only valid iterations contribute to vectorized updates.

  • Power computations: Broadcasting replaces explicit power-computation loops by combining expanded input arrays with an array of exponents.The vectorized function uses np.expand_dims(x,1) and exponents so NumPy performs the computations across aligned axes.
  • Inside-out vectorization: Vectorizing the innermost loop adds an axis to expressions containing its loop variable, allowing the new dimension to propagate through the program.Outer-scope variables remain fixed-shape free variables, while later axis arrangement enables the remaining loop to be vectorized.
  • Masked updates: Branches cannot replace loop variables directly with np.arange because some generated values correspond to iterations excluded by branch conditions.The transformation therefore uses masked arrays and masked updates to preserve dimensional structure while excluding invalid iterations.
  • Masked updates: Masked updates encode branch validity by masking restricted variables and applying updates only to unmasked indexed results.The DSL's make_masked operator constructs masked arrays and turns left-hand-side uses into masked updates.

3 Preliminaries

The preliminaries define the array-shape, broadcasting, advanced-indexing, and masked-array concepts used by the paper's DSL. These concepts determine how operands align, selections are formed, and invalid values are excluded from computations.

  • Shapes and broadcasting: An array shape records its size along each axis, with scalar values represented by the empty shape ().Axes are positions in the shape tuple, and their dimensionalities specify the corresponding sizes.
  • Shapes and broadcasting: Broadcasting implicitly expands arrays along size-1 and missing leading axes so compatible operands share a computation shape.Binary operators then perform element-wise computation on the broadcast operands.
  • Advanced indexing: Advanced indexing uses integer indexers that are broadcast to a common shape before selecting elements from the indexed base array.The selection result has the common broadcast shape of the indexers.
  • Masked arrays: MaskedArray pairs array data with a same-shaped boolean mask, omitting masked elements from computations and reductions.A reduction over entirely masked elements is masked, while np.ma.filled can replace masked entries with a specified value.

4 Core Language

The core language is a high-level imperative DSL that captures core NumPy features, including bindings, updates, control flow, indexing, array construction, and maskedness. Its type analysis infers static shapes and maskedness, supporting formal soundness and a top-level vectorization procedure.

  • Syntax and Semantics: The DSL represents programs as functions built from statements such as bindings, updates, branches, loops, and sequential composition.Its syntax includes expressions, integral shape accesses, masked left-hand sides, and update statements.
  • Syntax and Semantics: Masked updates extend advanced-indexing updates with boolean mask indices and update selected elements only under the DSL’s masking conditions.Normal updates require an unmasked right-hand side, while masked updates use mask indices.
  • Syntax and Semantics: Its expressions cover NumPy-style indexing, shape access, array construction, operators, reductions, matrix multiplication, and dimension manipulation.The language also includes masked-array construction and replication operations.
  • Typing Shapes and Maskedness: The type environment maps variables to static shapes and maskedness, using ⊤ for masked arrays and ⊥ for unmasked arrays.The system supports fixed numbers of axes with dynamic axis lengths represented by symbols.
  • Vectorization Procedure: The top-level vectorization algorithm repeatedly analyzes shapes, selects an innermost loop, and rewrites the program while vectorizable loops remain.The type system’s soundness theorem states that execution of a well-typed program returns a value well-typed under the resulting environment.

5 Vectorizing with Type-Directed Rewrite

The approach rewrites programs into observationally equivalent, loop-free code by repeatedly vectorizing innermost loops using shape- and maskedness-aware rules. Its applicability is limited for programs with strongly coupled loop-carried dependence.

  • 5.1 Problem Statement: The goal is to transform a program P into an observationally equivalent loop-free program P′ with identical values and maskedness.The equivalence condition is stated for every evaluation environment under which P returns a value.
  • 5.2 Top-Level Algorithm: Vectorize repeatedly analyzes shapes, selects an innermost loop, rewrites it, and substitutes the result until no vectorizable loops remain.The procedure operates inside out, using inferred variable types to guide each rewrite.
  • 5.3 Type-Directed Rewrite: The core rewrite replaces a loop variable with an array of all iteration values and propagates the added axis through the loop body using broadcasting.The loop variable substitute is initialized with arange of the loop bound, while expressions and statements are rewritten compositionally.
  • 5.3 Type-Directed Rewrite: Expression rewrites preserve original broadcasting and indexing behavior by inserting axes according to inferred shapes and adding loop-variable indexers when needed.Binary operations, shape-sensitive operators, and advanced indexing each receive specialized axis adjustments.
  • 5.4 Rewriting Branches: Lifted branches are flattened and controlled with masked loop-variable substitutes so only iterations satisfying each condition contribute to the rewritten computation.Masked indexers are filled with zero and the corresponding result is masked to preserve the original branch semantics.
  • 5.6 Vectorizability: Programs with strongly coupled loop-carried dependence cannot generally be vectorized by the rewrite rules, so the approach is not complete for the full DSL.Cycles or backward dependence in dependence graphs generally prevent direct vectorization, aside from one special case.

6 Implementation

Vectorizer implements the rewrite approach using standard NumPy and Python's ast module, with a DSL that supports minor source adaptations. Postprocessing converts rewrite output into executable Python while removing common inefficiencies.

  • 6 Implementation: Vectorizer is implemented with standard NumPy and Python's ast module, making the tool lightweight and easy to run.Many practical NumPy constructs map directly to the DSL, so source programs often need only minor adaptations.
  • 6 Implementation: A postprocessing pipeline transforms rewritten programs into executable Python while eliminating masked-update overhead and other common inefficiencies.The passes simplify boolean masks and replace masked updates with executable NumPy constructs.

7 Evaluation

Vectorizer was evaluated on 150 NumPy benchmarks against prior tools, measuring applicability, rewrite time, and performance improvement. It vectorized most benchmarks quickly and produced substantial speedups, while some failures and slowdowns remained.

  • Evaluation setup: 150 benchmarks from 12 datasets were used to evaluate Vectorizer, including translated programs and Stack Overflow examples.The evaluation examined explicit-loop NumPy programs and compared Vectorizer with Tenspiler and Tensorize.
  • Effectiveness: 142 of 150 benchmarks were vectorized directly, while 2 additional benchmarks succeeded after minor adaptations.The remaining 6 failures were primarily caused by complex loop-carried dependence.
  • Effectiveness: 70 of 150 benchmarks were solved by Tenspiler, whereas Tensorize generated output for 97 of 99 available MLIR benchmarks.Tensorize’s normalized outputs were judged semantically different from their sources for 15 benchmarks.
  • Efficiency: 0.53 seconds was Vectorizer’s average rewrite time per program, compared with 5.78 seconds for Tenspiler and 18.47 seconds for Tensorize.Vectorizer avoids search and symbolic reasoning, maintaining speed on semantically complex programs.
  • Performance improvement: Vectorizer slowed down 7 Stack Overflow benchmarks because materialized intermediates or predicated execution could outweigh loop-based memory and control-flow advantages.The authors identify workload-aware cost modeling as future work.
  • Performance improvement: 74.83× was Vectorizer’s geometric mean speedup over original loop-based implementations, with an additional 1.45× geometric mean speedup when paired with Numba.The performance comparison used generated inputs corresponding to 25,000,000 inner-loop iterations and median timings over 10 runs.

8 Related Work

Related work spans array DSLs, program synthesis, rewriting, and predicated execution. Vectorizer distinguishes itself through type-directed source-to-source rewrites rather than search or learning-based lifting.

  • Array DSLs: Array-programming systems range from NumPy and Halide to TACO and StableHLO, using different abstractions for efficient multidimensional computation.These systems motivate automated migration because differences in programming models complicate adoption of array DSLs.
  • Program synthesis: Program-synthesis approaches such as TF-Coder, C2TACO, Dexter, Tensorize, and Tenspiler use examples, symbolic reasoning, pattern matching, or search to lift programs.Their target systems and input languages vary across TensorFlow, TACO, Halide, and other array DSLs.
  • Program synthesis: Vectorizer uses type-directed rewrites to lift loop-based array programs without costly search, making the transformation more deterministic and cost-effective than learning-based approaches.Its rewrites target high-level NumPy programs directly.
  • Program rewrites: Prior rewriting systems optimize or migrate programs across representations, whereas Vectorizer applies source-to-source rewrites specifically to vectorize NumPy programs.The paper situates its approach among compiler optimization and high-performance-computing rewrite systems.
  • Predicated execution: Vectorizer represents conditional computation with masked arrays, paralleling predicated execution while targeting high-level NumPy source transformation.The related work connects this design to branch flattening and compiler if-conversion.

9 Conclusion

The paper concludes that Vectorizer combines correct-by-construction inside-out rewrites with shape and dataflow guidance to vectorize loop-based NumPy programs efficiently. Across 150 benchmarks, it achieved broad applicability, fast rewriting, and large average speedups.

  • Conclusion: Vectorizer uses correct-by-construction inside-out rewrites guided by array types and dataflow analysis.The approach avoids search and symbolic reasoning.
  • Conclusion: 142 of 150 benchmarks were vectorized directly, with 2 more succeeding after minor adaptations.This result summarizes the tool’s applicability across the evaluation set.
  • Conclusion: 0.53 seconds was the average time required to vectorize a benchmark.The reported timing summarizes the tool’s rewrite efficiency.
  • Conclusion: 74.83× was the geometric mean speedup of the resulting programs over the original implementations.This is the paper’s headline performance result.

A Formal Semantics of the Proposed DSL

The proposed DSL gives formal semantics for program execution, expression values, runtime shapes, and maskedness. Its evaluation environment tracks variable values and maskedness through scoped statement and expression rules.

  • Semantic framework: The DSL formalizes maskedness, statement execution, program evaluation, and expression evaluation through separate semantic rules.Figures 11–15 specify these components of the formal semantics.
  • Evaluation environment: A runtime variable state is represented as a value-maskedness pair, while Σ represents a stack of variable scopes.The semantics uses environments to evaluate statements, expressions, and program returns.
  • Auxiliary operators: Auxiliary operators define stack access, scope updates, top-level indexing, and other semantic operations used by the DSL rules.Auxiliary operators are distinguished from the DSL itself and support its formal definitions.
  • Expression semantics: Expression evaluation uses runtime shapes and indexing relations to specify how result elements correspond to operand elements.Most DSL operators use their established denotational meanings through JfK notation.
  • Maskedness semantics: Maskedness generally propagates through operators, with exceptions for masked indexers, matmul operands, and the second argument of filled.The DSL’s array indexing also requires unmasked integer indexers and a matching number of axes.

B Type Inference Rules

The type system infers array shapes and maskedness using operational semantics, broadcasting, and symbolic evaluation. These analyses provide the information needed to reason about program expressions and statements.

  • Inference framework: Static shape inference and maskedness inference are specified by complete rule sets for expressions.Figures 16 and 17 contain the complete rules for these analyses.
  • Inference framework: The type environment maps variables to fixed-rank shapes and maskedness, representing dimensions with symbols or integers.Masked arrays use ⊤, unmasked arrays use ⊥, and only the number of axes must be fixed.
  • Operational semantics: The semantics also define masked updates and program-level statement and expression behavior.The corresponding figures cover statement semantics, program semantics, expression-value evaluation, and maskedness evaluation.
  • Operational semantics: Runtime semantics define bindings, loops, indexing, and broadcasting over array values and shapes.The supplied rules include environment updates for bindings and loop execution, plus broadcasting for indexed operands.

C Proof of Soundness of Type Analysis

The soundness proofs show that static shape and maskedness analyses agree with runtime evaluation under the stated typing assumptions. The argument proceeds by structural induction over expression forms and culminates in soundness of type analysis.

  • Abstraction functions: The shape abstraction is the runtime Shape function, while maskedness maps masked values to ⊤ and unmasked values to ⊥.These abstractions connect runtime states to the static type environment.
  • Main soundness result: Theorem C.1 states that inferred shapes and maskedness equal the abstractions of runtime values and masks for well-annotated programs.The theorem assumes correctly abstracted input variables and concludes agreement for the returned program value.
  • Shape analysis: Shape soundness is established for literals, variables, indexing, routine calls, reductions, dimension expansion, matrix multiplication, arange, ones, and filled.Each inductive case aligns the corresponding runtime shape rule with its static inference rule.
  • Shape analysis: The shape rules preserve correctness through broadcasting, axis removal, axis insertion, symbolic dimensions, and matrix multiplication.The proof handles reductions, expand_dims or replicate, matmul, and symbolic constructors by matching static and runtime calculations.
  • Maskedness analysis: Lemma C.3 establishes soundness of maskedness analysis for expressions by structural induction.The base cases rely on the runtime maskedness of integral expressions and correctly typed variables.

D Rewrite Rules

The rewrite rules transform bindings, updates, reductions, indexing, variables, and branch-sensitive expressions while preserving shape and maskedness information. Their conditions control replication, masking, and axis adjustments during loop elimination.

  • Rule system: Figures 18 and 19 give the complete statement- and expression-level rewrite rules.The rules extend the informal rewrite procedure described earlier in the paper.
  • Statement rewrites: Bind rewrites replicate definitions when a bound expression depends on the current loop variable but is not lifted.Replication uses replicate or make_masked depending on branch context.
  • Statement rewrites: Reduction rewrites increment the reduced axis when the reduced expression gains a dimension during lifting.This preserves reduction over the corresponding logical axis after vectorization.
  • Expression rewrites: Variable rewrites add masking or replication when branch flattening exposes dependencies across defining branches or loop scopes.The rule handles both non-loop-local variables escaping branch restrictions and loop-local variables used in deeper branches.

E Proof of Correctness of Rewrite Rules

The correctness proof models vectorized execution as parallel execution of the innermost loop and shows that the rewrite rules preserve its results. It then lifts innermost-loop soundness to the complete Vectorize routine.

  • Loop rewrite soundness: Lemma E.1 proves soundness of loop rewrites by combining dependence-free parallelization with correctness of statement transformations.The proof separately handles loops without loop-carried dependence and loops containing one rewritable reduce statement.
  • Parallel semantics: The proof defines parallel evaluation and execution for innermost loops, including per-iteration arrays and merged updates.These definitions establish the semantic target that rewrites must reproduce.
  • Parallel semantics: Theorem E.4 shows that an innermost loop without loop-carried dependence has the same variable values under parallel execution as sequential execution.This theorem supplies the foundational equivalence used by the rewrite proof.
  • Program-level correctness: Theorem E.8 states that rewriting the innermost loop preserves both the returned value and maskedness for every evaluation environment.The rewritten program P′ satisfies v = v′ and μ = μ′.
  • Program-level correctness: Theorem E.9 states that Vectorize preserves program values and maskedness when input types are correctly annotated and the program is rewritable.This lifts innermost-loop correctness to the complete vectorization routine.
Loading 2609.08088v1…