[论文解读] High-Performance Tensor Contraction without Transposition

Devin A. Matthews|arXiv (Cornell University)|Jul 1, 2016

Parallel Computing and Optimization Techniques参考文献 33被引用 5

一句话总结

本文提出 TBLIS，一个高性能张量收缩库，通过将张量到矩阵的映射与基于 BLIS 的矩阵乘法内核融合，消除了显式转置操作。通过利用 BLIS 框架灵活的算法结构，作者实现了接近矩阵乘法级别的性能，在单核与多核工作负载中均优于传统方法（如 TTGT 和 BSMTC），尤其在复杂张量形状下表现更优。

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.

研究动机与目标

• 消除张量收缩中显式张量转置带来的性能与内存开销。
• 开发一个通用、高性能的张量收缩库，避免代码生成与运行时开销。
• 通过灵活可扩展的框架，实现与优化矩阵乘法相当的性能。
• 实现高效的多线程执行，同时最大限度减少用户在张量布局管理方面的干预。
• 提供一个生产就绪的开源解决方案（TBLIS），支持多种张量形状与收缩模式。

提出的方法

• 作者设计了两种新型张量到矩阵的映射方式——散列-矩阵张量收缩（SMTC）与块散列-矩阵张量收缩（BSMTC），使张量元素能够在矩阵内核中实现原地访问。
• 这些映射被集成到 BLIS 框架中，使得转置与数据打包可与内核计算融合，消除了单独的转置操作与额外的工作空间。
• 该方法利用了 BLIS 的模块化设计，将矩阵乘法分解为独立的算法组件，使自定义张量收缩算法可基于与 MM 相同的原子原语构建。
• 实现采用基于块的数据布局与优化的内存访问模式，以提升缓存重用与向量化能力，最大限度减少数据移动。
• 多线程通过复用 BLIS 矩阵乘法的线程模型实现，确保一致的性能扩展性。
• TBLIS 库自动处理张量到矩阵的索引转换，为用户提供低层次复杂性的抽象。

实验结果

研究问题

• RQ1 是否可以在不进行显式转置或额外工作空间的情况下，实现接近矩阵乘法性能的张量收缩？
• RQ2 基于 BLIS 的张量收缩性能与传统 TTGT 和 BSMTC 方法相比，在多种张量形状下表现如何？
• RQ3 所提出方法在计算密集型与通信密集型场景下，跨多核的可扩展性如何？
• RQ4 BLIS 框架的模块化设计是否能有效用于构建高性能、通用的张量收缩内核？
• RQ5 低效的数据访问模式对张量收缩性能有何影响，又该如何缓解？

主要发现

• BSMTC 算法在多种张量形状下性能接近峰值机器性能的 10% 以内，逼近矩阵乘法的效率。
• 在单核计算密集型场景中，BSMTC 与 BLIS 性能相近且接近峰值，而 TTT（TTGT）性能急剧下降，对较大张量的性能低于 10 GFLOPs。
• 在通信密集型场景中，BSMTC 维持在约 30 GFLOPs，而 TTT 降至 <10 GFLOPs，表明所提方法具有显著性能优势。
• 在多核执行中，BSMTC 相较 TTT 最高实现 21.1× 的加速比，加速比范围从 1.3× 到 21.1× 不等，具体取决于张量形状，展现出强大的可扩展性。
• 在若干情况下，BSMTC 表现优于 BLIS，可能归因于其内部使用了更优的矩阵乘法变体，凸显了算法定制的优势。
• TBLIS 库成功避免了 TTGT 的工作空间开销，并提供了一个生产就绪的、基于 BSD 许可证的开源实现。