FusedMM:面向图嵌入与图神经网络的统一 SDDMM-SpMM 核
机器学习
2021-10-28 v2 分布式、并行与集群计算
社会与信息网络
摘要
我们开发了一种融合矩阵乘法核,在称为 FusedMM 的单一操作下统一了采样稠密-稠密矩阵乘法(SDDMM)与稀疏-稠密矩阵乘法(SpMM)。借助用户自定义函数,FusedMM 可捕获流行图嵌入与 GNN 方法所需的几乎所有计算模式。FusedMM 比 Deep Graph Library 中等效核快一个数量级。FusedMM 的优越性能源于底层向量化核、合适的负载均衡方案以及对内存带宽的高效利用。FusedMM 可通过代码生成器调优性能,并在 Intel、AMD 和 ARM 处理器上表现同样良好。FusedMM 在不同处理器上将端到端图嵌入算法加速至多 28 倍。
引用
@article{arxiv.2011.06391,
title = {FusedMM: A Unified SDDMM-SpMM Kernel for Graph Embedding and Graph Neural Networks},
author = {Md. Khaledur Rahman and Majedul Haque Sujon and Ariful Azad},
journal= {arXiv preprint arXiv:2011.06391},
year = {2021}
}
备注
11 pages, published in IEEE IPDPS 2021