用机器学习自动选线程数,让矩阵乘法快25%~40%。
A Machine Learning Approach Towards Runtime Optimisation of Matrix Multiplication
- 用机器学习模型实时预测最优线程数
- 在100MB内存下提速25%~40%
- 适合高性能计算中频繁调用GEMM的场景
通用矩阵乘法(GEMM)是科学计算的核心算法。单线程GEMM已通过分块和自调优等技术高度优化,但现代多核共享内存系统复杂,难以确定使多线程GEMM运行时间最短的线程数。本文提出一种概念验证方法,构建架构与数据结构感知的线性代数库(ADSALA),利用机器学习优化BLAS函数的运行时性能。具体而言,该方法基于训练数据,在运行时使用机器学习模型自动为给定的GEMM任务选择最优线程数。在两种不同高性能计算节点架构上测试:一个基于双路Intel Cascade Lake,另一个基于双路AMD Zen 3。结果显示,当GEMM内存占用在100MB以内时,相比传统BLAS实现,速度提升25%至40%。
原文摘要 · Abstract (English)
The GEneral Matrix Multiplication (GEMM) is one of the essential algorithms in scientific computing. Single-thread GEMM implementations are well-optimised with techniques like blocking and autotuning. However, due to the complexity of modern multi-core shared memory systems, it is challenging to determine the number of threads that minimises the multi-thread GEMM runtime. We present a proof-of-concept approach to building an Architecture and Data-Structure Aware Linear Algebra (ADSALA) software library that uses machine learning to optimise the runtime performance of BLAS routines. More specifically, our method uses a machine learning model on-the-fly to automatically select the optimal number of threads for a given GEMM task based on the collected training data. Test results on two different HPC node architectures, one based on a two-socket Intel Cascade Lake and the other on a two-socket AMD Zen 3, revealed a 25 to 40 per cent speedup compared to traditional GEMM implementations in BLAS when using GEMM of memory usage within 100 MB.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。