[论文解读] Task-Based Algorithm for Matrix Multiplication: A Step Towards Block-Sparse Tensor Computing

Justus A. Calvin, Edward F. Valeev|arXiv (Cornell University)|Apr 20, 2015

Parallel Computing and Optimization Techniques参考文献 38被引用 7

一句话总结

本文提出了一种基于任务的可扩展通用矩阵乘法算法（SUMMA）重构方法，可在具有不规则、块稀疏数据结构的分布式内存系统上实现高效且可扩展的矩阵乘法。通过使用细粒度任务并行和多任务调度，该方法能够容忍非均匀矩阵分块带来的负载不平衡，并在最多262,144个核心上实现高达75%的效率，展示了在通用系统和高性能计算（HPC）系统上的优异性能。

ABSTRACT

Distributed-memory matrix multiplication (MM) is a key element of algorithms in many domains (machine learning, quantum physics). Conventional algorithms for dense MM rely on regular/uniform data decomposition to ensure load balance. These traits conflict with the irregular structure (block-sparse or rank-sparse within blocks) that is increasingly relevant for fast methods in quantum physics. To deal with such irregular data we present a new MM algorithm based on Scalable Universal Matrix Multiplication Algorithm (SUMMA). The novel features are: (1) multiple-issue scheduling of SUMMA iterations, and (2) fine-grained task-based formulation. The latter eliminates the need for explicit internodal synchronization; with multiple-iteration scheduling this allows load imbalance due to nonuniform matrix structure. For square MM with uniform and nonuniform block sizes (the latter simulates matrices with general irregular structure) we found excellent performance in weak and strong-scaling regimes, on commodity and high-end hardware.

研究动机与目标

• 解决在量子物理和张量计算中常见的不规则、块稀疏结构矩阵中，分布式内存矩阵乘法的负载不平衡问题。
• 开发一种通信避免、延迟容忍的矩阵乘法算法，可在不牺牲密集或均匀分块矩阵性能的前提下处理数据异质性。
• 通过基于任务的并行化重构标准密集MM算法，实现对不规则张量数据结构的高效计算。
• 在多种硬件上（包括通用集群和IBM BG/Q超算系统）展示强可扩展性和弱可扩展性性能。

提出的方法

• 采用细粒度任务并行编程模型重构标准SUMMA算法，实现计算与同步的解耦。
• 实现SUMMA迭代的多任务调度，以支持节点间计算与通信的重叠。
• 使用任务图表示矩阵乘法操作，支持动态负载均衡和对不规则数据分布的容错能力。
• 在基于任务的框架中复用现有的优化BLAS内核（如DGEMM）作为计算原语。
• 将算法集成到开源TiledArray库中，以支持通用张量计算，包括用户自定义类型和块结构。
• 应用数据过度分解技术，即使在块大小差异显著的情况下，也能减少每个进程的负载不平衡。

实验结果

研究问题

• RQ1 基于任务的SUMMA重构能否在保持密集矩阵高性能的同时，有效处理非均匀、块稀疏的矩阵结构？
• RQ2 在具有不规则数据的分布式内存矩阵乘法中，基于任务的并行化在负载均衡和通信隐藏方面有何改进？
• RQ3 SUMMA迭代的多任务调度在多大程度上能缓解数据异质性和网络延迟导致的性能下降？
• RQ4 在高端和通用HPC系统上，基于任务的SUMMA实现的强可扩展性和弱可扩展性表现如何？
• RQ5 该算法的效率在不同核心数和矩阵规模下与单节点峰值性能相比如何？

主要发现

• 在单节点上使用Intel MKL时，基于任务的SUMMA实现达到了理论峰值性能的90.98%，表现出极高的计算效率。
• 在IBM BG/Q系统上，无论矩阵大小和核心数量如何，该算法在非均匀分块条件下仍保持超过75%的效率。
• 在强可扩展极限下，当每个核心拥有四个或更少的矩阵块时，效率最低降至50%，这是由于任务粒度过细所致。
• 尽管在非均匀矩阵中最大块与最小块的大小比高达1:7.12，但通过数据过度分解，每个进程的负载不平衡被降低至1:1.35。
• 该算法在最多262,144个核心（2^18）上实现了有效扩展，无论采用均匀还是非均匀分块模式，性能均保持稳定。
• 该实现对硬件异质性和网络延迟具有鲁棒性，性能对拓扑结构和CPU时钟频率变化的敏感度较低。