Tensor core 详解

// 包含 NVIDIA CUDA 的矩阵乘法和累加（WMMA）库，并使用 nvcuda 命名空间，以便访问 WMMA 接口。
#include <mma.h>
using namespace nvcuda;

// 核函数定义，用于执行矩阵乘法
__global__ void wmma_ker(half *a, half *b, float *c) {
    // 声明片段，用于存储矩阵数据
    wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::col_major> a_frag; // 矩阵A的片段，列主序
    wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::row_major> b_frag; // 矩阵B的片段，行主序
    wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag; // 用于累加的片段

    // 初始化输出片段c_frag为零
    wmma::fill_fragment(c_frag, 0.0f);

    // 加载输入矩阵数据到片段中
    wmma::load_matrix_sync(a_frag, a, 16);
    wmma::load_matrix_sync(b_frag, b, 16);

    // 执行矩阵乘法
    wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);

    // 将结果存储到输出矩阵中
    wmma::store_matrix_sync(c, c_frag, 16, wmma::mem_row_major);
}

// 详细注释说明以便理解

#include <mma.h>
#include <cuda_runtime.h>
#include <iostream>

using namespace nvcuda;

#define M 1024
#define N 1024
#define K 1024

__global__ void matrixMulKernel(const half* __restrict__ A, const half* __restrict__ B, float* __restrict__ C, int M, int N, int K) {
    // Declare shared memory to hold the tiles of A and B
    // 声明存储矩阵A和B片段的共享内存区域
    __shared__ half shared_A[16][16];
    __shared__ half shared_B[16][16];

    // Declare the fragments
    // 使用 WMMA 声明矩阵片段，a_frag 和 b_frag 分别用于存储矩阵 A 和 B 的片段，c_frag 用于累加结果。
    wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::col_major> a_frag;
    wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::row_major> b_frag;
    wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag;

    // 计算 warp 的位置
    
    int warpM = (blockIdx.y * blockDim.y + threadIdx.y) / 16;
    int warpN = (blockIdx.x * blockDim.x + threadIdx.x) / 16;

    wmma::fill_fragment(c_frag, 0.0f);

    for (int i = 0; i < K; i += 16) {
        int aRow = warpM * 16 + threadIdx.y % 16;
        int aCol = i + threadIdx.x % 16;
        int bRow = i + threadIdx.y % 16;
        int bCol = warpN * 16 + threadIdx.x % 16;

        // 计算共享内存中的行和列索引，并将矩阵 A 和 B 的数据加载到共享内存中
        if (aRow < M && aCol < K) {
            shared_A[threadIdx.y % 16][threadIdx.x % 16] = A[aRow * K + aCol];
        } else {
            shared_A[threadIdx.y % 16][threadIdx.x % 16] = 0.0;
        }

        if (bRow < K && bCol < N) {
            shared_B[threadIdx.y % 16][threadIdx.x % 16] = B[bRow * N + bCol];
        } else {
            shared_B[threadIdx.y % 16][threadIdx.x % 16] = 0.0;
        }

        __syncthreads();

        wmma::load_matrix_sync(a_frag, shared_A[0], 16);
        wmma::load_matrix_sync(b_frag, shared_B[0], 16);

        wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);

        __syncthreads();
    }
    // 计算输出矩阵 C 的行和列索引并将累加结果从 c_frag 存储到全局内存中的矩阵C
    int cRow = warpM * 16 + threadIdx.y % 16;
    int cCol = warpN * 16 + threadIdx.x % 16;

    if (cRow < M && cCol < N) {
        wmma::store_matrix_sync(C + cRow * N + cCol, c_frag, N,  wmma::mem_row_major);
    }
}

void matrixMultiply(const half* A, const half* B, float* C, int M, int N, int K) {
    dim3 threadsPerBlock(16, 16);
    dim3 blocksPerGrid((N + 15) / 16, (M + 15) / 16);

    matrixMulKernel<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, M, N, K);
    cudaDeviceSynchronize();
}

int main() {
    // Allocate and initialize host matrices
    half *h_A, *h_B;
    float *h_C;
    cudaMallocHost(&h_A, M * K * sizeof(half));
    cudaMallocHost(&h_B, K * N * sizeof(half));
    cudaMallocHost(&h_C, M * N * sizeof(float));

    for (int i = 0; i < M * K; i++) h_A[i] = __float2half(1.0f);
    for (int i = 0; i < K * N; i++) h_B[i] = __float2half(1.0f);
    for (int i = 0; i < M * N; i++) h_C[i] = 0.0f;

    // Allocate and initialize device matrices
    half *d_A, *d_B;
    float *d_C;
    cudaMalloc(&d_A, M * K * sizeof(half));
    cudaMalloc(&d_B, K * N * sizeof(half));
    cudaMalloc(&d_C, M * N * sizeof(float));

    cudaMemcpy(d_A, h_A, M * K * sizeof(half), cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, K * N * sizeof(half), cudaMemcpyHostToDevice);

    // Perform matrix multiplication on the device
    matrixMultiply(d_A, d_B, d_C, M, N, K);

    // Copy the result back to host
    cudaMemcpy(h_C, d_C, M * N * sizeof(float), cudaMemcpyDeviceToHost);

    // Print the result
    std::cout << "Result: " << std::endl;
    for (int i = 0; i < 10; i++) {
        for (int j = 0; j < 10; j++) {
            std::cout << h_C[i * N + j] << " ";
        }
        std::cout << std::endl;
    }

    // Cleanup
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);
    cudaFreeHost(h_A);
    cudaFreeHost(h_B);
    cudaFreeHost(h_C);

    return 0;
}