PaperlayerView exact revision on arXiv

Source-linked AI summary

Format Abstraction for Sparse Tensor Algebra Compilers

Stephen Chou, Fredrik Kjolstad, Saman Amarasinghe

arXiv:1804.10112v2cs.MScs.PL

TL;DR

Sparse tensor compilers struggle to efficiently support formats with fundamentally different data layouts without hard-coding every format combination. This paper introduces a capability-based level abstraction and modular code generator, achieving competitive performance across many formats and up to 3.6× faster direct COO matrix-vector multiplication than conversion to CSR.

  • Problem

    Supporting disparate tensor formats efficiently requires specialized code for format combinations, but hard-coding those combinations scales exponentially.

  • Method

    The paper defines a common level-format interface based on capabilities and properties, then composes six implementations into formats and generates kernels independently of specific formats.

  • Results

    The generated code is competitive with existing sparse libraries across sparse matrix and tensor computations, including COO matrix-vector multiplication up to 3.6× faster than conversion to CSR.

  • Takeaways & Limitations

    The modular design supports mixing formats for different computations and adding new formats without modifying the code generator, reducing data-translation costs.

  • Takeaways & Limitations

    The current technique does not yet support iteration-space tiling, shared-memory parallel code, accelerators, distributed systems, or additional formats such as DOK and LIL.

Abstract

from arXiv · show

This paper shows how to build a sparse tensor algebra compiler that is agnostic to tensor formats (data layouts). We develop an interface that describes formats in terms of their capabilities and properties, and show how to build a modular code generator where new formats can be added as plugins. We then describe six implementations of the interface that compose to form the dense, CSR/CSF, COO, DIA, ELL, and HASH tensor formats and countless variants thereof. With these implementations at hand, our code generator can generate code to compute any tensor algebra expression on any combination of the aforementioned formats. To demonstrate our technique, we have implemented it in the taco tensor algebra compiler. Our modular code generator design makes it simple to add support for new tensor formats, and the performance of the generated code is competitive with hand-optimized implementations. Furthermore, by extending taco to support a wider range of formats specialized for different application and data characteristics, we can improve end-user application performance. For example, if input data is provided in the COO format, our technique allows computing a single matrix-vector multiplication directly with the data in COO, which is up to 3.6$\times$ faster than by first converting the data to CSR.

1 INTRODUCTION

Sparse tensor algebra needs format-specific iteration, but supporting every disparate format through hard-coded strategies scales exponentially. The paper introduces a level-based abstraction and modular generator that supports varied formats while retaining competitive performance.

  • 1 INTRODUCTION: Supporting every combination of disparate tensor formats would require Θ(2^|F|) hard-coded code-generation strategies, preventing taco from scaling to formats such as COO and DIA.Different formats encode coordinates differently and therefore require distinct iteration code.
  • 1 INTRODUCTION: The compiler represents tensor dimensions through six composable per-dimension formats that express dense, CSR/CSF, COO, DIA, ELL, HASH, and many variants.Dense and compressed formats existed in taco; singleton, range, offset, and hashed formats are introduced here.
  • 1 INTRODUCTION: A common interface exposes level-format access methods and properties, allowing generated code to reason about capabilities rather than specific tensor formats.The interface hides how each level encodes a tensor dimension while retaining the information needed for access.
  • 1 INTRODUCTION: The modular generator emits efficient kernels for operands stored in any mix of formats and can be extended with new formats without modifying the generator.This lets users match formats to application and data characteristics.
  • 1 INTRODUCTION: Up to 3.6× faster COO matrix-vector multiplication is achieved by computing directly on COO data instead of first converting it to CSR.The technique is implemented as an extension to taco and generates code competitive with existing sparse libraries.

2 TENSOR STORAGE FORMATS

Sparse tensor formats trade memory use, access patterns, assembly cost, and computational efficiency according to data structure and workload. Their differing coordinate encodings make general efficient code generation difficult, motivating scalable abstractions for format-aware iteration.

  • 2 TENSOR STORAGE FORMATS: No sparse tensor format is universally superior because the ideal choice depends on data structure, sparsity, computation, and hardware.Supporting many formats is desirable but requires specialized code for combinations of operand layouts.
  • 2.1 Survey of Tensor Formats: Dense arrays provide constant-time coordinate access but waste memory on zeros and may become impossible for tensors with many large dimensions.Dense storage explicitly stores every tensor component, including zeros.
  • 2.1 Survey of Tensor Formats: COO stores only nonzero coordinates and values in Θ(nnz) memory, while its append-oriented representation minimizes preprocessing cost.COO closely mirrors common tensor file formats but lacks efficient random access.
  • 2.1 Survey of Tensor Formats: HASH provides random access without explicitly storing zeros, but its poor ordered iteration restricts additive operations.COO has the opposite trade-off: efficient enumeration but inefficient random access.
  • 2.1 Survey of Tensor Formats: CSR removes redundant row coordinates with auxiliary position arrays, improving storage and performance for bandwidth-bound computations such as SpMV.Higher-order CSF compresses every dimension, while compressed formats are costly to assemble or modify.
  • 2.1 Survey of Tensor Formats: ELLPACK stores a bounded, equal number of components per row, exploiting regular sparsity patterns and exposing vectorization opportunities.It is suited to matrices such as those arising from well-formed meshes.
  • 2.2 Computing with Disparate Tensor Formats: Different coordinate encodings require different iteration strategies, so efficient computation needs specialized code for every combination of formats.Dense and CSR dimensions iterate differently, while COO requires simultaneous co-iteration and merging of coordinate arrays.
  • 2.2 Computing with Disparate Tensor Formats: Hard-coding those strategies causes Θ(2^|F|) growth in the number of cases, creating an exponential scalability barrier for format support.This blow-up effectively prevented earlier taco from supporting disparate formats such as COO and DIA.

3 TENSOR STORAGE ABSTRACTION

The abstraction represents tensor storage as coordinate hierarchies built from six per-dimension level formats. A capability-and-property interface lets code generation operate across formats without being hard-coded to particular data structures.

  • 3.1 Coordinate Hierarchies: A coordinate hierarchy has one level per tensor dimension, with root-to-leaf paths encoding tensor-component coordinates and nodes representing coordinates or unlabeled storage positions.Unlabeled nodes capture padding in physical storage, such as out-of-bounds diagonal segments in DIA.
  • 3.1 Coordinate Hierarchies: The hierarchy structure reflects each format’s memory encoding, so format-specific iteration patterns can be represented without the generator knowing the underlying tensor format.COO uses coordinate chains for complete nonzero coordinates, whereas CSR uses a tree structure with shared row ancestors.
  • 3.1 Coordinate Hierarchies: Dense, compressed, singleton, range, offset, and hashed levels encode coordinates through dimension ranges, segmented arrays, single coordinates, intervals, shifts, or hash maps.Compressed levels use coordinate arrays and segment bounds, while range and offset levels can implicitly encode structured coordinates.
  • 3.1 Coordinate Hierarchies: Six per-dimension level formats compose to represent common tensor formats and additional variants, including structured matrices cast as higher-order tensors.The framework also supports alternative compositions, such as a sparsely filled DIA variant using (dense, compressed, offset).
  • 3.1 Coordinate Hierarchies: Level capabilities and properties expose how to iterate, index, modify, and characterize levels, allowing the generator to manipulate storage through a common interface.Properties include full, ordered, unique, branchless, and compact; capabilities include iteration, locate, insert, and append.

Coordinate Value Iteration.

Coordinate value iteration enumerates possible coordinate values under given ancestors and uses coordinate access to find their positions. It generalizes dense-vector iteration through the level interface.

  • Coordinate Value Iteration: Coordinate value iteration returns bounds over coordinates with specified ancestors, then accesses each coordinate’s position while indicating whether that coordinate exists.The capability is exposed through coord_bounds and coord_access.
  • Coordinate Value Iteration: The capability generalizes dense-vector traversal by iterating coordinate values rather than directly iterating stored positions.Missing coordinates can be detected through the found result of coordinate access.
  • Coordinate Value Iteration: Code generation can remove existence checks when level properties guarantee that every iterated coordinate is present.The optimization follows from reasoning about level properties exposed by the abstraction.

Coordinate Position Iteration.

Coordinate position iteration traverses positions associated with a parent and retrieves the coordinate stored at each position. It generalizes sparse-vector traversal, including unlabeled positions.

  • Coordinate Position Iteration: Coordinate position iteration obtains a position range for a parent and accesses the coordinate encoded at each position through pos_bounds and pos_access.The returned found flag identifies whether a position is an actual child coordinate.
  • Coordinate Position Iteration: This capability generalizes sparse-vector iteration by traversing physical positions instead of enumerating coordinate values.It can represent positions that are not valid children or contain unlabeled storage nodes.

Locate.

Locate provides random access to a coordinate within a hierarchy by searching among a parent’s children. Its properties describe level invariants that guide optimized generated code.

  • Locate: Locate searches for a requested child coordinate, returning its position and a found flag, and repeated calls can traverse a path to one tensor component.Efficient locate implementations benefit operands requiring random coordinate access.
  • Locate: Level properties such as full, ordered, unique, branchless, and compact describe storage invariants that the generator uses to emit optimized code.These properties distinguish, for example, sorted CSR columns from hash-map vectors and branchless COO levels from branching CSR levels.
  • Locate: Full, unique, ordered, branchless, and compact properties capture whether coordinates cover a dimension, avoid duplicates, follow an order, form chains, or occupy contiguous positions.The properties are format-dependent and may be configurable when they reflect application-level invariants rather than physical encoding.

Compact.

Assembly capabilities define how output tensor levels are initialized, modified, and finalized through fixed level functions.

  • Compact.: Assembly capabilities let the compiler construct result data structures without depending on a specific tensor format.The interface includes capabilities for iterating over and accessing coordinate hierarchy levels, plus inserting and appending coordinates.
  • Compact.: Insert capability places coordinates at arbitrary positions using locate, with initialization and finalization around the insertion process.The output level size is computed from its parent level when insertion is supported.

Append Capability.

The append capability builds output levels by adding coordinates in order and connecting them through the coordinate hierarchy.

  • Append Capability.: The append interface exposes append_coord for coordinates, append_edges for hierarchy links, and initialization and finalization functions.append_edges connects coordinates in the current level to a position in the previous level.
  • Append Capability.: Appending coordinates in order supports efficient assembly of hierarchical sparse outputs.The capability requires result coordinates to be appended in order.
  • Append Capability.: CSR output assembly repeatedly appends nonzero coordinates for the crd array and uses edge operations to build the pos array.Initialization occurs before computation, edges are added after each row, and finalization occurs at the end.

4 CODE GENERATION

The code-generation algorithm emits efficient tensor-algebra code for combinations of formats by reasoning about level capabilities, iteration spaces, and merge transformations.

  • 4 CODE GENERATION: The algorithm supports tensor formats expressible as compositions of level formats and extends prior code generation through capability- and property-based reasoning.This design targets many disparate formats rather than only dense and compressed levels.
  • 4 CODE GENERATION: Merging an indexed tensor expression requires iterating over the operands’ joint iteration space dimension by dimension.Matrix addition, for example, merges rows and then the corresponding dimensions of the operands.
  • 4 CODE GENERATION: Merge lattices encode the loops and sub-expressions needed to merge all input dimensions indexed by each variable.Each path through a lattice represents a possible runtime loop sequence for fully merging the inputs.
  • 4 CODE GENERATION: For CSR-plus-COO matrix addition, optimized merge lattices guide loops that incrementally merge operands until their dimensions are fully incorporated.The example assumes the COO operand has no empty row.
  • 4 CODE GENERATION: Merge optimization can co-iterate levels lacking locate while accessing other levels through locate, potentially reducing merge complexity when locate is constant-time.The benefit depends on operand ordering and whether levels are unordered.
  • 4 CODE GENERATION: Iterator conversion extends merging to unordered or non-unique levels by composing deduplication and reordering at runtime.The flowchart determines the conversions required for each operand in an intersection merge.

Deduplication.

Deduplication removes repeated coordinates before merging, while reordering creates ordered copies needed for co-iteration.

  • Deduplication.: Duplicate coordinates are aggregated into unique-coordinate iterators, with duplicate values summed at the bottom of coordinate hierarchies.At higher levels, child iterators are combined instead of directly aggregating values.
  • Deduplication.: Reordering stores an ordered copy of an unordered level so the compiler can merge it through co-iteration.The transformed iterator replaces the original unordered-level iterator during merging.

Reordering.

The compiler chooses merge strategies from operand capabilities such as locate, ordering, and uniqueness, then generates specialized code for tensor-algebra expressions. Its recursive algorithm constructs merge lattices, emits access and assembly calls, and optimizes iteration through fused iterators.

  • Reordering: The most efficient merge strategy depends on whether operand levels support locate and whether they are ordered and unique.Without locate, the compiler co-iterates levels; with locate, it can iterate one operand and probe the other.
  • Reordering: Arbitrarily complex unions and intersections combine co-iteration over some operands with locate calls into the remaining operands.The operands to co-iterate are identified recursively from the target expression.
  • Reordering: The code-generation algorithm recursively processes index variables, builds merge lattices, and emits loops for the relevant coordinate hierarchy levels.It initializes iterators, computes result values, and assembles output indices through format-specific capability functions.
  • Reordering: The generated kernels iterate over input intersections and unions, compute values at joint coordinates, and specialize code for the operands’ formats.The algorithm calls access and assembly capabilities rather than relying on one fixed storage layout.
  • Reordering: Fusing iterators lets the compiler traverse multiple branchless coordinate levels with one loop, improving performance for formats such as COO.The optimization can iterate over two tensor dimensions simultaneously while preserving locate access to the output.

5 EVALUATION

The taco extension generates kernels across diverse sparse matrix and tensor formats, with performance competitive with existing libraries. Its broader format support also improves practical end-to-end performance when conversion costs or data characteristics make conventional layouts unsuitable.

  • 5 EVALUATION: The technique supports a wider range of formats and operations than the evaluated libraries, including both SoA and AoS COO variants.It supports formats such as DIA and ELL that some libraries omit, while TensorFlow supports only COO among the surveyed formats.
  • 5 EVALUATION: The generated code has performance competitive with existing libraries across sparse operations, including SpDM, sparse matrix addition, SpMV, and DIA SpMV.It is equal to or better than other libraries for SpDM and sparse matrix addition, while DIA SpMV is about 21% slower than MKL on average.
  • 5 EVALUATION: For COO SpMV, the generated kernel matches SciPy and MKL because all implement the same algorithm, while TensorFlow incurs overhead by treating the operation as SpDM.TensorFlow’s representation introduces a loop over a trivial column dimension for each input-vector access.
  • 5 EVALUATION: For sparse matrix addition, taco specializes code to operand order and can use narrower coordinate widths, reducing overhead relative to TensorFlow’s generic kernel.TensorFlow iterates over component coordinates and hard-codes 64-bit coordinates.
  • 5 EVALUATION: The generated code efficiently handles higher-order sparse tensor operations, outperforming TensorFlow and often the Tensor Toolbox while covering all five benchmark operations.Directly merging already sorted indices avoids re-sorting concatenated coordinates and reduces asymptotic work.
  • 5 EVALUATION: Direct COO SpMV can improve end-to-end performance for non-iterative applications by avoiding conversion to CSR when that preprocessing cost cannot be amortized.The benefit depends on computation and data characteristics, because COO computation itself can take longer than CSR due to higher memory traffic.

6 RELATED WORKS

Prior work spans sparse tensor formats, storage abstractions, and compilers, but existing approaches generally support fixed formats or limited structural compositions. This paper instead represents diverse tensor formats with six composable level formats sharing a common interface.

  • Six composable level formats represent a broad range of sparse matrix and higher-order tensor formats through a shared framework.
  • Earlier sparse-format abstractions described regular array partitions or dense/sparse dimensions, limiting support for unstructured matrices or formats beyond those data structures.
  • LL encoded sparse formats through nested lists and pairs, but changing the matrix format could require redefining the computation itself.
  • Other compiler approaches either transformed dense code into sparse code, generated runtime format conversions, or supported only fixed sets of standard sparse matrix formats.
  • Dense linear-algebra compilers and tensor-contraction systems provide related transformation techniques, but address dense computations rather than the paper’s broad sparse-format framework.

7 CONCLUSION AND FUTURE WORK

The paper presents a modular technique for generating tensor algebra kernels across disparate tensor formats. It identifies broader format support and parallel, accelerator, and distributed code generation as future directions.

  • The modular technique generates efficient tensor algebra kernels for disparate formats and allows new formats without modifying the code generator.
  • Future work includes supporting DOK, LIL, custom graph and accelerator-oriented formats, and formats exploiting structural or value symmetries.
  • Future work also includes shared-memory parallel code, iteration-space tiling, and targets such as GPUs and distributed-memory systems.
Loading 1804.10112v2…