TensorCore使用

最新推荐文章于 2024-11-09 07:00:00 发布

holyprince

最新推荐文章于 2024-11-09 07:00:00 发布

阅读量4.9k

点赞数 1

分类专栏： cuda编程高性能计算

本文链接：https://blog.csdn.net/wzh1026/article/details/103780784

版权

高性能计算同时被 2 个专栏收录

7 篇文章 3 订阅

订阅专栏

cuda编程

2 篇文章 0 订阅

订阅专栏

目前有两个库使用了TensorCore技术：cuBLAS与cuDNN。

cuBLAS主要加速了GEMM计算，cuDNN里主要用来加速卷积和RNN

而如何cuBLAS与cuDNN无法满足需求，也可以直接调用Tensor Cores进行编程。

Tensor Cores的调用API在nvcuda::wmma命名空间里。通过调用API进行数值初始化，进行矩阵乘加（MMA）操作然后再把数值存回内存。cuda官方给出了一个直接调用的例子，如果想要针对tensorcore进行高性能的优化，可以参考CUDA Toolkit 的cudaTensorCoreGemm。

TensorCore的使用方法：

头文件和命名空间

#include <mma.h>
using namespace nvcuda;

声明和初始化

tile矩阵乘法

// The only dimensions currently supported by WMMA
const int WMMA_M = 16;
const int WMMA_N = 16;
const int WMMA_K = 16;

__global__ void wmma_example(half *a, half *b, float *c, 
                             int M, int N, int K, 
                             float alpha, float beta) 
{

    // Leading dimensions. Packed with no transpositions.
    int lda = M;
    int ldb = K;
    int ldc = M;
    
    // Tile using a 2D grid
    int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;
    int warpN = (blockIdx.y * blockDim.y + threadIdx.y);

【1】使用TensorCore编程：https://devblogs.nvidia.com/programming-tensor-cores-cuda-9/

【2】TensorCore wiki ：https://github.com/maddyscientist/mixed-precision/wiki