[论文解读] Tensor processing primitives: a programming abstraction for efficiency and portability in deep learning workloads

Evangelos Georganas, Dhiraj Kalamkar+15 more|arXiv (Cornell University)|Nov 14, 2021

Parallel Computing and Optimization Techniques参考文献 19被引用 2

一句话总结

本文提出了张量处理原语（TPP），一种用于深度学习工作负载的可移植、高性能编程抽象，其使用一组精简的二维张量操作作为构建模块，实现复杂、高维的操作。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/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 of our approach using standalone kernels and end-to-end DL workloads expressed entirely via TPPs that outperform state-of-the-art implementations on multiple platforms.

研究动机与目标

• 解决尽管深度学习算法和硬件取得进展，但深度学习编程方法论停滞不前的问题。
• 克服现有深度学习库中平台特定、缺乏灵活性的内核的局限性。
• 提供一种高效率的编程抽象，实现复杂深度学习算子的可移植且高度优化的实现。
• 弥合深度学习系统开发中性能可移植性与开发人员生产力之间的差距。

提出的方法

• 定义一组最小化、平台无关的二维张量操作（虚拟张量指令集架构，virtual Tensor ISA）作为构建复杂深度学习工作负载的核心抽象。
• 利用这些原语作为构建模块，通过组合与分块（tiling）表达高维操作。
• 通过高度优化的、平台特定的内核实现TPP，以确保高性能。
• 将高层规格（TPP）与底层实现解耦，实现可移植性的同时保持性能。
• 端到端演示深度学习工作负载及独立内核，完全通过TPP表达。
• 在多个平台上评估性能，以验证效率与可移植性。

实验结果

研究问题

• RQ1 一组最小化、可移植的二维张量操作能否作为实现复杂深度学习工作负载的有效抽象？
• RQ2 TPP在多大程度上能够实现与手写优化、平台特定内核相当或更优的性能？
• RQ3 规格（TPP）与实现的分离如何同时实现可移植性与高性能？
• RQ4 是否能够仅使用TPP原语高效且可移植地表达端到端的深度学习工作负载？

主要发现

• 基于TPP的独立内核实现，在多个平台上均优于最先进的手写优化实现。
• 完全通过TPP表达的端到端深度学习工作负载，其性能与现有优化框架相比具有竞争力或更优。
• TPP的平台无关规格确保了在多样化硬件目标上的代码可移植性，且不牺牲性能。
• 该抽象通过允许开发人员使用一组小而一致的原语表达复杂操作，实现了高生产力。
• 规格与实现的分离使得能够针对每个平台对TPP内核进行针对性优化，同时保持接口兼容性。
• 该方法表明，在深度学习系统中，高性能与可移植性并非相互排斥。