Misam: Using ML in Dataflow Selection of Sparse-Sparse Matrix Multiplication

📄 arXiv: 2406.10166v2 📥 PDF

作者: Sanjali Yadav, Bahar Asgari

分类: cs.LG

发布日期: 2024-06-14 (更新: 2024-08-29)

备注: Accepted to ISCA 2024 MLArchSys workshop https://openreview.net/forum?id=A1V9FaZRbV

💡 一句话要点

Misam：利用机器学习进行稀疏矩阵乘法数据流选择优化

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 稀疏矩阵乘法 数据流选择 机器学习 决策树 深度强化学习 硬件加速器 性能优化

📋 核心要点

  1. 现有硬件加速器针对特定稀疏模式设计，数据流固定，无法有效处理多样化的稀疏矩阵乘法(SpGEMM)任务。
  2. 提出基于机器学习的自适应数据流选择方法，利用决策树和深度强化学习动态选择最优数据流方案。
  3. 实验结果表明，该方法在硬件加速器中动态选择数据流，性能提升可达28倍，优于传统启发式方法。

📝 摘要（中文）

稀疏矩阵乘法(SpGEMM)在科学计算、图分析和深度学习等领域至关重要。这些应用利用矩阵的稀疏性来减少存储和计算需求。然而，稀疏矩阵的不规则结构给性能优化带来了重大挑战。传统的硬件加速器针对具有固定数据流方案（内部、外部和行式）的特定稀疏模式量身定制，但当实际稀疏性偏离这些预定模式时，通常表现不佳。随着SpGEMM在具有不同稀疏特征的各个领域的应用扩展，对能够有效处理各种稀疏模式的硬件加速器的需求正在增加。本文提出了一种基于机器学习的方法，用于自适应地选择最适合具有不同稀疏模式的SpGEMM任务的数据流方案。通过采用决策树和深度强化学习，我们探索了这些技术在识别最佳数据流方案方面超越基于启发式方法的潜力。我们通过将模型的性能与启发式的性能进行比较来评估我们的模型，突出了每种方法的优缺点。我们的研究结果表明，在硬件加速器中使用机器学习进行动态数据流选择可以提供高达28倍的增益。

🔬 方法详解

问题定义：论文旨在解决稀疏矩阵乘法(SpGEMM)在不同稀疏模式下，传统硬件加速器因固定数据流方案而导致的性能瓶颈问题。现有方法依赖于针对特定稀疏模式优化的硬件加速器，当实际稀疏模式与预设模式不匹配时，性能会显著下降。因此，如何根据不同的稀疏模式动态选择最优的数据流方案是亟待解决的问题。

核心思路：论文的核心思路是利用机器学习技术，特别是决策树和深度强化学习，来预测并选择最适合当前SpGEMM任务的数据流方案。通过学习稀疏模式与数据流方案之间的关系，模型能够自适应地选择最优方案，从而提高SpGEMM的整体性能。这种方法避免了手动设计启发式规则的复杂性，并能够更好地适应各种稀疏模式。

技术框架：该方法包含以下主要阶段：1) 特征提取：从稀疏矩阵中提取能够表征其稀疏模式的特征，例如非零元素数量、分布等。2) 模型训练：使用提取的特征训练决策树或深度强化学习模型，学习稀疏模式与数据流方案之间的映射关系。3) 数据流选择：对于新的SpGEMM任务，提取其稀疏特征，并使用训练好的模型预测最优的数据流方案。4) 硬件加速：根据选择的数据流方案，配置硬件加速器执行SpGEMM计算。

关键创新：该论文的关键创新在于将机器学习技术应用于SpGEMM的数据流选择问题。与传统的启发式方法相比，机器学习方法能够自动学习稀疏模式与数据流方案之间的复杂关系，从而实现更优的性能。此外，该方法能够适应各种稀疏模式，无需针对每种模式手动设计优化策略。

关键设计：论文中使用了决策树和深度强化学习两种模型。对于决策树，需要选择合适的特征和分裂准则。对于深度强化学习，需要设计合适的奖励函数，以鼓励模型选择能够提高性能的数据流方案。具体的网络结构和参数设置在论文中可能有所描述，但摘要中未提供详细信息。此外，特征工程也是一个关键的设计环节，需要选择能够有效表征稀疏模式的特征。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，基于机器学习的动态数据流选择方法在SpGEMM任务中能够取得显著的性能提升，最高可达28倍。该方法优于传统的启发式方法，能够更好地适应各种稀疏模式，实现更优的性能。

🎯 应用场景

该研究成果可广泛应用于科学计算、图分析、深度学习等领域，尤其是在处理大规模稀疏矩阵乘法时。通过动态选择最优数据流方案，可以显著提高计算效率，降低能耗，加速相关应用的运行速度。未来，该技术有望集成到硬件加速器设计中，实现更加智能和高效的SpGEMM计算。

📄 摘要（原文）

Sparse matrix-matrix multiplication (SpGEMM) is a critical operation in numerous fields, including scientific computing, graph analytics, and deep learning. These applications exploit the sparsity of matrices to reduce storage and computational demands. However, the irregular structure of sparse matrices poses significant challenges for performance optimization. Traditional hardware accelerators are tailored for specific sparsity patterns with fixed dataflow schemes - inner, outer, and row-wise but often perform suboptimally when the actual sparsity deviates from these predetermined patterns. As the use of SpGEMM expands across various domains, each with distinct sparsity characteristics, the demand for hardware accelerators that can efficiently handle a range of sparsity patterns is increasing. This paper presents a machine learning based approach for adaptively selecting the most appropriate dataflow scheme for SpGEMM tasks with diverse sparsity patterns. By employing decision trees and deep reinforcement learning, we explore the potential of these techniques to surpass heuristic-based methods in identifying optimal dataflow schemes. We evaluate our models by comparing their performance with that of a heuristic, highlighting the strengths and weaknesses of each approach. Our findings suggest that using machine learning for dynamic dataflow selection in hardware accelerators can provide upto 28 times gains.