[Paper Review] Tensor Processing Primitives: A Programming Abstraction for Efficiency and Portability in Deep Learning & HPC Workloads

Evangelos Georganas, Dhiraj Kalamkar+17 more|arXiv (Cornell University)|Apr 12, 2021

Parallel Computing and Optimization TechniquesComputer Science47 references13 citations

TL;DR

This paper introduces Tensor Processing Primitives (TPP), a portable, efficient programming abstraction that exposes a minimal set of 2D tensor operations as a 'virtual Tensor ISA' to enable high-performance, portable implementations of deep learning and HPC workloads. By decoupling algorithm specification from platform-specific optimization, TPPs allow developers and compilers to express complex operators portably while achieving superior performance across multiple architectures, outperforming state-of-the-art implementations in end-to-end DL workloads and standalone kernels.

ABSTRACT

During the past decade, novel Deep Learning (DL) algorithms, workloads and hardware have been developed to tackle a wide range of problems. Despite the advances in workload and hardware ecosystems, the programming methodology of DL systems is stagnant. DL workloads leverage either highly-optimized, yet platform-specific and inflexible kernels from DL libraries, or in the case of novel operators, reference implementations are built via DL framework primitives with underwhelming performance. This work introduces the Tensor Processing Primitives (TPP), a programming abstraction striving for efficient, portable implementation of DL workloads with high-productivity. TPPs define a compact, yet versatile set of 2D-tensor operators (or a virtual Tensor ISA), which subsequently can be utilized as building-blocks to construct complex operators on high-dimensional tensors. The TPP specification is platform-agnostic, thus code expressed via TPPs is portable, whereas the TPP implementation is highly-optimized and platform-specific. We demonstrate the efficacy and viability of our approach using standalone kernels and end-to-end DL & HPC workloads expressed entirely via TPPs that outperform state-of-the-art implementations on multiple platforms.

Motivation & Objective

• Address the stagnation in DL programming methodologies despite rapid advances in DL workloads and hardware.
• Overcome the limitations of monolithic, non-portable DL library kernels and underperforming reference implementations for novel operators.
• Provide a portable, high-productivity abstraction that enables efficient, platform-optimized implementations without requiring low-level hardware expertise.
• Enable a clean separation of concerns between algorithm specification and low-level code generation for both developers and compilers.
• Facilitate reuse and optimization sharing across frameworks, compilers, and hardware platforms through a standardized, minimal operator set.

Proposed method

• Define a compact, precision-aware set of 2D tensor primitives (TPPs) as a virtual instruction set architecture (ISA) for tensors.
• Design TPPs to be platform-agnostic in specification but highly optimized in implementation, enabling portable code with peak performance.
• Use TPPs as building blocks to express complex high-dimensional operators in DL and HPC workloads.
• Integrate TPPs into software stacks as a lowering layer for DL frameworks, vendor libraries (e.g., oneDNN), and Tensor Compilers (e.g., MLIR dialects).
• Implement TPP back-ends using intrinsic instructions (e.g., AVX512, SVE, NEON) for target architectures to achieve low-level performance.
• Leverage existing compiler infrastructure (e.g., MLIR) to enable code generation and optimization of TPP-based kernels.

Experimental results

Research questions

• RQ1 Can a minimal, portable set of 2D tensor operations serve as an effective abstraction for implementing complex, high-performance DL and HPC workloads?
• RQ2 How does the separation of algorithm specification (via TPPs) from platform-specific optimization improve performance and portability compared to existing DL library kernels or generic framework primitives?
• RQ3 To what extent can TPPs outperform state-of-the-art implementations in real-world DL and HPC workloads across diverse hardware platforms?
• RQ4 Can TPPs be effectively integrated into existing software stacks such as MLIR, DL frameworks, and vendor-optimized libraries?
• RQ5 How does the use of a virtual Tensor ISA (TPP) enable optimization sharing and reduce duplication across compilers and systems?

Key findings

• End-to-end DL workloads such as CNNs, dilated convolutions, DLRM, BERT, and GNNs implemented entirely via TPPs outperform state-of-the-art implementations on multiple platforms.
• Standalone kernel implementations using TPPs achieve superior performance compared to optimized vendor libraries and reference implementations.
• The TPP abstraction enables high productivity and portability while maintaining low-level performance through platform-specific back-end optimizations.
• TPPs can be natively integrated into MLIR as a lowering dialect, enabling efficient code generation and optimization across multiple compiler stacks.
• The approach reduces the need for manual low-level optimization by data/ML scientists, accelerating the development and deployment of novel operators.
• Support for diverse ISAs (x86, AArch64, ARMv8 SVE, and future GPU/POWER support) demonstrates the extensibility and portability of the TPP abstraction.