细粒度任务并行GPU通用矩阵乘

摘要

稠密线性代数运算对模式识别和生物信息等许多实际应用至关重要,而通用矩阵乘(GEMM)在稠密线性代数运算中占有基础地位.在cuBLAS与MAGMA中,GEMM被实现为若干kernel函数,对大规模GEMM计算能够取得达到很高的性能.然而,现有实现对批量的小规模GEMM计算性能则较为有限.而且,现有实现也不能在多个具有不同性能的GPU之间自动扩展并达到负载均衡.本文提出任务并行式GEMM(TPGEMM),用细粒度任务并行的方式实现批量矩阵乘和多GPU矩阵乘.一个或多个GEMM的计算能够被拆分为多个任务,动态地调度到一个或多个GPU上.TPGEMM避免了为批量矩阵乘启动多个kernel函数的开销,对批量矩阵乘能够取得显著高于cuBLAS与MAGMA的性能.在低开销细粒度任务调度的基础上,TPGEMM支持单个GEMM计算在多个GPU间的自动并行,在一台具有4个不同性能GPU的工作站上取得了接近100％的扩展效率.