[Paper] 一种机器学习方法用于矩阵乘法的运行时优化

发布: 6天前 (2026年1月14日 GMT+8 11:28)

7 min read

原文: arXiv

Source: arXiv - 2601.09114v1

概述

本文提出了一种机器学习驱动的技术，用于在现代多核 CPU 上运行通用矩阵乘法（GEMM）例程时自动选择最佳线程数。通过即时训练轻量级模型，作者在两个当代 HPC 节点上实现了 25%‑40% 的加速，相较于流行 BLAS 库中使用的默认线程选择启发式方法。

ADSALA library prototype – 一个 “Architecture and Data‑Structure Aware Linear Algebra” 框架，将机器学习模型嵌入 BLAS 调用中。
On‑the‑fly model training – 在早期 GEMM 执行期间收集运行时特征（矩阵尺寸、缓存使用、核心拓扑），并持续优化回归模型。
Thread‑count prediction – 模型为每次 GEMM 调用输出最佳的 OpenMP 线程数，取代静态或朴素的启发式方法。
Cross‑architecture validation – 在 Intel Cascade Lake（2 × 18 核）和 AMD Zen 3（2 × 16 核）上的实验表明，对约 100 MiB 以下的矩阵工作负载可获得一致的性能提升。
Open‑source proof‑of‑concept – 作者公开代码和训练数据，支持可重复性研究并便于进一步扩展。

特征提取 – 对每个 GEMM 调用，库记录：
- 矩阵维度 (M, N, K)
- 估计工作集大小（字节）
- CPU 拓扑（插槽、核心数、超线程状态）
- 当前系统负载（可选）
模型选择 – 使用简单的决策树回归器（或梯度提升树）对这些特征进行训练，以映射到在短“探索阶段”中针对每个线程数测试得到的运行时间。
探索阶段 – 在新问题规模的前几次 GEMM 执行时，使用一系列线程数（例如 1、½ 核心、全部核心）运行。它们的运行时间用于填充训练集。
预测与部署 – 一旦模型达到预定义的置信阈值，它会预测后续 GEMM 调用的最佳线程数，然后使用该设置执行。
持续适应 – 如果出现新的矩阵尺寸或系统状态显著偏离训练分布，库会回退到一次简短的重新探索以更新模型。

整个流水线仅增加几毫秒的开销，与目标的多秒级 GEMM 核心相比可以忽略不计。

架构	矩阵占用空间 (≤ 100 MiB)	基准 BLAS（静态线程）	ADSALA（机器学习选定线程）	加速比
Intel Cascade Lake (2 × 18)	20 MiB – 100 MiB	1.0× (reference)	1.25× – 1.40×	25 % – 40 %
AMD Zen 3 (2 × 16)	20 MiB – 100 MiB	1.0×	1.28× – 1.38×	28 % – 38 %

性能关键库 – 科学 Python（NumPy、SciPy）、机器学习框架（TensorFlow、PyTorch）或自定义 HPC 内核的开发者可以嵌入 ADSALA 方法，实现自动、面向架构的线程调优，无需手动基准测试。
云端与容器环境 – 在虚拟化或容器化的场景中，进程可见的物理核心数量可能在运行时变化，自调优的 BLAS 能够即时适应，提供一致的吞吐量。
能源效率 – 当线程之间会争夺内存带宽时，使用更少的线程可以降低功耗，这对绿色计算计划是一个有吸引力的副作用。
部署简便 – 由于模型轻量（决策树）且在线训练，无需庞大的离线自动调优数据库；该库可以作为现有 BLAS 调用的即插即用替代品发布。

总体而言，本文展示了少量机器学习智能即可弥合手工调优 BLAS 库与不断变化的多核 CPU 架构之间的差距，为开发者提供了一条实现更快线性代数内核的实用路径。