High-Performance Tensor Contraction without Transposition

TL;DR

This paper introduces TBLIS, a novel tensor contraction implementation that fuses transposition with internal operations, achieving near-matrix multiplication performance without extra workspace or explicit transpositions.

Contribution

The paper presents TBLIS, a tensor contraction implementation within the BLIS framework that eliminates the need for explicit transpositions and reduces overhead, improving performance and usability.

Findings

TBLIS achieves performance close to matrix multiplication.

In some cases, TBLIS outperforms traditional tensor contraction methods.

The approach simplifies tensor operations in scientific computing.

Abstract

Tensor computations--in particular tensor contraction (TC)--are important kernels in many scientific computing applications. Due to the fundamental similarity of TC to matrix multiplication (MM) and to the availability of optimized implementations such as the BLAS, tensor operations have traditionally been implemented in terms of BLAS operations, incurring both a performance and a storage overhead. Instead, we implement TC using the flexible BLIS framework, which allows for transposition (reshaping) of the tensor to be fused with internal partitioning and packing operations, requiring no explicit transposition operations or additional workspace. This implementation, TBLIS, achieves performance approaching that of MM, and in some cases considerably higher than that of traditional TC. Our implementation supports multithreading using an approach identical to that used for MM in BLIS, with similar performance characteristics. The complexity of managing tensor-to-matrix transformations is also handled automatically in our approach, greatly simplifying its use in scientific applications.