[论文解读] 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 Techniques参考文献 47被引用 13

一句话总结

本文提出了张量处理原语（TPP），这是一种可移植的、高效的编程抽象，通过将一组最小化的二维张量操作作为“虚拟张量指令集架构”（Tensor ISA）暴露，以实现深度学习和高性能计算（HPC）工作负载的高性能、可移植实现。通过将算法规范与特定平台的优化解耦，TPP使开发人员和编译器能够以可移植方式表达复杂算子，同时在多种架构上实现卓越性能，在端到端深度学习工作负载和独立内核上均优于当前最先进实现。

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.

研究动机与目标

• 解决尽管深度学习工作负载和硬件快速进步，但深度学习编程方法论停滞不前的问题。
• 克服单体式、不可移植的深度学习库内核以及新型算子性能不佳的参考实现的局限性。
• 提供一种可移植的、高效率的抽象，使开发者无需掌握底层硬件知识，即可实现高效、平台优化的实现。
• 为开发人员和编译器实现算法规范与低级代码生成之间的清晰职责分离。
• 通过标准化的、最小化的算子集合，促进框架、编译器和硬件平台之间的复用与优化共享。

提出的方法

• 定义一组紧凑的、精度感知的二维张量原语（TPP），作为张量的虚拟指令集架构（ISA）。
• 设计TPP使其在规范层面与平台无关，但在实现层面高度优化，从而实现可移植代码并达到峰值性能。
• 使用TPP作为构建模块，表达深度学习和HPC工作负载中的复杂高维算子。
• 将TPP作为DL框架、厂商库（如oneDNN）和张量编译器（如MLIR方言）的下拉层集成到软件栈中。
• 针对目标架构，使用内建指令（如AVX512、SVE、NEON）实现TPP后端，以达到底层性能。
• 利用现有编译器基础设施（如MLIR）实现基于TPP的内核代码生成与优化。

实验结果

研究问题

• RQ1 一个最小化、可移植的二维张量操作集合，能否作为实现复杂、高性能深度学习和HPC工作负载的有效抽象？
• RQ2 与现有深度学习库内核或通用框架原语相比，通过TPP将算法规范与特定平台优化分离，是否能显著提升性能和可移植性？
• RQ3 在多种硬件平台上，TPP在真实世界深度学习和HPC工作负载中，与最先进实现相比，性能提升的幅度有多大？
• RQ4 TPP能否有效集成到MLIR、深度学习框架和厂商优化库等现有软件栈中？
• RQ5 使用虚拟张量ISA（TPP）如何实现优化共享，并减少编译器和系统之间的重复工作？

主要发现

• 完全基于TPP实现的端到端深度学习工作负载（如CNN、空洞卷积、DLRM、BERT和GNN）在多个平台上优于最先进实现。
• 使用TPP实现的独立内核性能优于优化的厂商库和参考实现。
• TPP抽象在保持底层性能的同时，实现了高生产力和可移植性，得益于平台特定的后端优化。
• TPP可原生集成到MLIR中作为下拉方言，实现在多个编译器栈上的高效代码生成与优化。
• 该方法减少了数据/机器学习科学家手动进行低级别优化的需求，加速了新型算子的开发与部署。
• 对多种ISA（x86、AArch64、ARMv8 SVE，以及未来GPU/POWER支持）的支持，证明了TPP抽象的可扩展性与可移植性。