Lecture 5: 并行模式三 —— 分块矩阵乘法：共享内存与数据复用 (对应 Lab 3: Tiled Matrix Multiply)

// (a) 静态共享内存：编译期确定大小，写在内核里，随内核一起编译进函数属性
__global__ void StaticKernel(const float* M, int Width)
{
    __shared__ float subTileM[32][32];        // 32*32*4 = 4 KB，每个 block 一份
    __shared__ float subTileN[32][32];
    subTileM[threadIdx.y][threadIdx.x] = M[0];
    __syncthreads();
}

// (b) 动态共享内存：大小在启动时以第三个 <<<>>> 参数给出，多个数组手工切分同一块
__global__ void DynamicKernel(const float* M, int Width)
{
    extern __shared__ float dynShared[];      // 只声明，不指定大小
    float (*subTileM)[33] = reinterpret_cast<float (*)[33]>(dynShared);
    float (*subTileN)[33] = reinterpret_cast<float (*)[33]>(dynShared + 32 * 33);
    subTileM[threadIdx.y][threadIdx.x] = M[0];
    __syncthreads();
}

// 启动动态版本：第三个 <<<>>> 参数是"每个 block 的动态共享内存字节数"
void launchDynamicVersion(dim3 grid, dim3 block, const float* dM, float* dP, int Width)
{
    int shmemBytes = 2 * 32 * 33 * (int)sizeof(float);        // 2 个 32x33 的 float 数组 = 8448 B
    // 若超过 48 KB（单个 block 静态共享内存上限），必须先 opt-in：
    // 例如 TILE=96 时 2*96*97*4 = 74496 B > 48 KB，就必须调用下面这一行。
    cudaFuncSetAttribute(DynamicKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmemBytes);
    DynamicKernel<<<grid, block, shmemBytes>>>(dM, Width);
    (void)dP;
}

#define BLOCK_SIZE 16        // 课程考题中 BLOCK_SIZE 是编译期常量，可取值 1 到 20

__global__ void BlockTranspose(float* A_elements, int A_width, int A_height)
{
    __shared__ float blockA[BLOCK_SIZE][BLOCK_SIZE];
    int baseIdx = blockIdx.x * BLOCK_SIZE + threadIdx.x;
    baseIdx += (blockIdx.y * BLOCK_SIZE + threadIdx.y) * A_width;
    blockA[threadIdx.y][threadIdx.x] = A_elements[baseIdx];
    A_elements[baseIdx] = blockA[threadIdx.x][threadIdx.y];   // 跨列读取 -> bank conflict
}

// 文件: tiled_matmul_basic.cu
// 编译: nvcc -O3 -arch=sm_80 tiled_matmul_basic.cu -o tiled_matmul_basic
// 运行: ./tiled_matmul_basic 2048      (矩阵边长，必须是 TILE_WIDTH 的整数倍)

#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <cuda_runtime.h>

#define CUDA_CHECK(call)                                                        \
    do {                                                                        \
        cudaError_t err__ = (call);                                             \
        if (err__ != cudaSuccess) {                                             \
            fprintf(stderr, "CUDA error %s:%d: %s\n", __FILE__, __LINE__,       \
                    cudaGetErrorString(err__));                                 \
            exit(EXIT_FAILURE);                                                 \
        }                                                                       \
    } while (0)

#define TILE_WIDTH 16

// A100 硬件档案（用于 Roofline 分析）
#define A100_PEAK_TFLOPS 19.5
#define A100_BW_GBPS     1555.0

// ---------- 朴素版本（讲义第 4-5 页）：每个输出元素 2K 次全局访存 ----------
__global__ void MatrixMulKernelNaive(const float* __restrict__ M,
                                     const float* __restrict__ N,
                                     float* __restrict__ P,
                                     int Width)
{
    int Row = blockIdx.y * blockDim.y + threadIdx.y;
    int Col = blockIdx.x * blockDim.x + threadIdx.x;
    if ((Row < Width) && (Col < Width)) {
        float Pvalue = 0.0f;
        for (int k = 0; k < Width; ++k)
            Pvalue += M[Row * Width + k] * N[k * Width + Col];
        P[Row * Width + Col] = Pvalue;
    }
}

// ---------- 分块版本（讲义第 34 页）：全局访存降低 TILE_WIDTH 倍 ----------
__global__ void MatrixMulKernelTiled(const float* __restrict__ M,
                                     const float* __restrict__ N,
                                     float* __restrict__ P,
                                     int Width)
{
    __shared__ float subTileM[TILE_WIDTH][TILE_WIDTH];
    __shared__ float subTileN[TILE_WIDTH][TILE_WIDTH];

    int bx = blockIdx.x;   int by = blockIdx.y;
    int tx = threadIdx.x;  int ty = threadIdx.y;

    // 该线程负责的 P 元素的行列（blockDim.x 与 blockDim.y 都等于 TILE_WIDTH）
    int Row = by * TILE_WIDTH + ty;
    int Col = bx * TILE_WIDTH + tx;

    float Pvalue = 0.0f;

    // 遍历计算该输出元素所需的全部 M/N tile（假定 Width 是 TILE_WIDTH 的整数倍）
    for (int m = 0; m < Width / TILE_WIDTH; ++m) {
        // 协作加载：每个线程搬 1 个 M 元素和 1 个 N 元素
        subTileM[ty][tx] = M[Row * Width + m * TILE_WIDTH + tx];
        subTileN[ty][tx] = N[(m * TILE_WIDTH + ty) * Width + Col];
        __syncthreads();                       // 屏障 1：确保 tile 已全部写满

        for (int k = 0; k < TILE_WIDTH; ++k)
            Pvalue += subTileM[ty][k] * subTileN[k][tx];

        __syncthreads();                       // 屏障 2：确保旧 tile 已被读完
    }

    P[Row * Width + Col] = Pvalue;
}

// ------------------------------ 主机端辅助函数 ------------------------------
static void fillMatrix(float* A, int n, unsigned seed)
{
    for (int i = 0; i < n * n; ++i) {
        unsigned x = (unsigned)(i + seed * 7919u);
        x = x * 1103515245u + 12345u;
        A[i] = (float)((int)((x >> 16) & 0x7u) - 3) * 0.25f;   // 取值 -0.75 .. 1.00
    }
}

static void matmulCPU(const float* A, const float* B, float* C, int n)
{
    for (int i = 0; i < n; ++i)
        for (int j = 0; j < n; ++j) {
            float s = 0.0f;
            for (int k = 0; k < n; ++k) s += A[i * n + k] * B[k * n + j];
            C[i * n + j] = s;
        }
}

template <typename LaunchT>
static bool verifyKernel(const char* name, LaunchT launch, int n, dim3 block)
{
    size_t bytes = (size_t)n * (size_t)n * sizeof(float);
    float* hM   = (float*)malloc(bytes);
    float* hN   = (float*)malloc(bytes);
    float* hP   = (float*)malloc(bytes);
    float* hRef = (float*)malloc(bytes);
    fillMatrix(hM, n, 1u);
    fillMatrix(hN, n, 2u);
    matmulCPU(hM, hN, hRef, n);

    float *dM, *dN, *dP;
    CUDA_CHECK(cudaMalloc((void**)&dM, bytes));
    CUDA_CHECK(cudaMalloc((void**)&dN, bytes));
    CUDA_CHECK(cudaMalloc((void**)&dP, bytes));
    CUDA_CHECK(cudaMemcpy(dM, hM, bytes, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(dN, hN, bytes, cudaMemcpyHostToDevice));

    dim3 grid((n + block.x - 1) / block.x, (n + block.y - 1) / block.y);
    launch(dM, dN, dP, n, grid, block);
    CUDA_CHECK(cudaGetLastError());
    CUDA_CHECK(cudaDeviceSynchronize());
    CUDA_CHECK(cudaMemcpy(hP, dP, bytes, cudaMemcpyDeviceToHost));

    double maxErr = 0.0;
    for (int i = 0; i < n * n; ++i) {
        double d = fabs((double)hP[i] - (double)hRef[i]);
        if (d > maxErr) maxErr = d;
    }
    printf("[验证 %-6s] n=%d  最大绝对误差=%.3e  %s\n", name, n, maxErr,
           (maxErr < 1e-2) ? "PASS" : "FAIL");

    free(hM); free(hN); free(hP); free(hRef);
    CUDA_CHECK(cudaFree(dM)); CUDA_CHECK(cudaFree(dN)); CUDA_CHECK(cudaFree(dP));
    return maxErr < 1e-2;
}

template <typename LaunchT>
static double timeKernelMs(LaunchT launch, int iters)
{
    cudaEvent_t start, stop;
    CUDA_CHECK(cudaEventCreate(&start));
    CUDA_CHECK(cudaEventCreate(&stop));
    for (int i = 0; i < 3; ++i) launch();                    // 预热
    CUDA_CHECK(cudaEventRecord(start));
    for (int i = 0; i < iters; ++i) launch();
    CUDA_CHECK(cudaEventRecord(stop));
    CUDA_CHECK(cudaEventSynchronize(stop));
    float ms = 0.0f;
    CUDA_CHECK(cudaEventElapsedTime(&ms, start, stop));
    CUDA_CHECK(cudaEventDestroy(start));
    CUDA_CHECK(cudaEventDestroy(stop));
    return (double)ms / (double)iters;
}

int main(int argc, char** argv)
{
    int Width = (argc > 1) ? atoi(argv[1]) : 2048;
    if (Width % TILE_WIDTH != 0) {
        fprintf(stderr, "Width 必须是 TILE_WIDTH=%d 的整数倍\n", TILE_WIDTH);
        return EXIT_FAILURE;
    }

    // 1) 正确性验证（用小矩阵跑 CPU 参考实现）
    verifyKernel("naive", [](const float* M, const float* N, float* P, int W, dim3 g, dim3 b) {
        MatrixMulKernelNaive<<<g, b>>>(M, N, P, W);
    }, 512, dim3(16, 16));
    verifyKernel("tiled", [](const float* M, const float* N, float* P, int W, dim3 g, dim3 b) {
        MatrixMulKernelTiled<<<g, b>>>(M, N, P, W);
    }, 512, dim3(TILE_WIDTH, TILE_WIDTH));

    // 2) 分配大矩阵并计时
    size_t bytes = (size_t)Width * (size_t)Width * sizeof(float);
    float *hM = (float*)malloc(bytes), *hN = (float*)malloc(bytes);
    fillMatrix(hM, Width, 3u);
    fillMatrix(hN, Width, 4u);

    float *dM, *dN, *dP;
    CUDA_CHECK(cudaMalloc((void**)&dM, bytes));
    CUDA_CHECK(cudaMalloc((void**)&dN, bytes));
    CUDA_CHECK(cudaMalloc((void**)&dP, bytes));
    CUDA_CHECK(cudaMemcpy(dM, hM, bytes, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(dN, hN, bytes, cudaMemcpyHostToDevice));

    double flops = 2.0 * (double)Width * (double)Width * (double)Width;

    dim3 blockN(16, 16);
    dim3 gridN((Width + 15) / 16, (Width + 15) / 16);
    double msNaive = timeKernelMs([&] {
        MatrixMulKernelNaive<<<gridN, blockN>>>(dM, dN, dP, Width);
    }, 5);

    dim3 blockT(TILE_WIDTH, TILE_WIDTH);
    dim3 gridT(Width / TILE_WIDTH, Width / TILE_WIDTH);
    double msTiled = timeKernelMs([&] {
        MatrixMulKernelTiled<<<gridT, blockT>>>(dM, dN, dP, Width);
    }, 20);

    CUDA_CHECK(cudaGetLastError());

    // 3) 性能与 Roofline 打印
    printf("\n===== 性能 (Width = %d, 2*W^3 = %.3e FLOP) =====\n", Width, flops);
    printf("朴素  : %8.3f ms   %8.1f GFLOP/s   算术强度 0.25 FLOP/Byte (每 FLOP 4 Byte)\n",
           msNaive, flops / msNaive / 1e6);
    printf("分块  : %8.3f ms   %8.1f GFLOP/s   算术强度 T/4 = %.1f FLOP/Byte\n",
           msTiled, flops / msTiled / 1e6, TILE_WIDTH / 4.0);
    printf("加速比: %.2fx\n", msNaive / msTiled);

    double bwBoundNaive = A100_BW_GBPS * 0.25;                 // GB/s * FLOP/Byte = GFLOP/s
    double bwBoundTiled = A100_BW_GBPS * (TILE_WIDTH / 4.0);
    double balance      = A100_PEAK_TFLOPS * 1000.0 / A100_BW_GBPS;
    printf("\n===== Roofline (A100: %.1f TFLOPS, %.0f GB/s, 机器平衡点 %.2f FLOP/Byte) =====\n",
           A100_PEAK_TFLOPS, A100_BW_GBPS, balance);
    printf("朴素内存屋顶      : %8.1f GFLOP/s  (峰值占比 %.2f%%)\n",
           bwBoundNaive, 100.0 * bwBoundNaive / (A100_PEAK_TFLOPS * 1000.0));
    printf("分块内存屋顶      : %8.1f GFLOP/s  (峰值占比 %.2f%%)\n",
           bwBoundTiled, 100.0 * bwBoundTiled / (A100_PEAK_TFLOPS * 1000.0));
    printf("判定              : %s\n",
           (bwBoundTiled < A100_PEAK_TFLOPS * 1000.0) ? "仍然内存带宽受限（需要更大 tile / 线程粗化）"
                                                      : "计算受限");
    printf("理论全局流量      : 朴素 %.2f GB,  分块 %.2f GB (降低 %d 倍)\n",
           2.0 * (double)Width * Width * Width * 4.0 / 1e9,
           2.0 * (double)Width * Width * Width * 4.0 / TILE_WIDTH / 1e9, TILE_WIDTH);

    free(hM); free(hN);
    CUDA_CHECK(cudaFree(dM)); CUDA_CHECK(cudaFree(dN)); CUDA_CHECK(cudaFree(dP));
    return EXIT_SUCCESS;
}

// 文件: tiled_matmul_padded.cu
// 编译: nvcc -O3 -arch=sm_80 tiled_matmul_padded.cu -o tiled_matmul_padded
// 运行: ./tiled_matmul_padded 2048      (矩阵边长，必须是 TILE_WIDTH 的整数倍)

#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <cuda_runtime.h>

#define CUDA_CHECK(call)                                                        \
    do {                                                                        \
        cudaError_t err__ = (call);                                             \
        if (err__ != cudaSuccess) {                                             \
            fprintf(stderr, "CUDA error %s:%d: %s\n", __FILE__, __LINE__,       \
                    cudaGetErrorString(err__));                                 \
            exit(EXIT_FAILURE);                                                 \
        }                                                                       \
    } while (0)

#define TILE_WIDTH 32
#define PADDED_STRIDE (TILE_WIDTH + 1)

// ---- 内核 A：转置载入 + 无 padding。行步长 S = 32 => 跨列访问 32-way bank conflict ----
__global__ void MatrixMulKernelTransposeNoPad(const float* __restrict__ M,
                                              const float* __restrict__ N,
                                              float* __restrict__ P,
                                              int Width)
{
    __shared__ float subTileM[TILE_WIDTH][TILE_WIDTH];
    __shared__ float subTileN[TILE_WIDTH][TILE_WIDTH];

    int bx = blockIdx.x;  int by = blockIdx.y;
    int tx = threadIdx.x; int ty = threadIdx.y;
    int Row = by * TILE_WIDTH + ty;
    int Col = bx * TILE_WIDTH + tx;

    float Pvalue = 0.0f;

    for (int m = 0; m < Width / TILE_WIDTH; ++m) {
        // 转置写入：共享内存第 (tx,ty) 个元素保存全局第 (ty,tx) 个元素
        subTileM[tx][ty] = M[Row * Width + m * TILE_WIDTH + tx];
        subTileN[tx][ty] = N[(m * TILE_WIDTH + ty) * Width + Col];
        __syncthreads();

        for (int k = 0; k < TILE_WIDTH; ++k)
            Pvalue += subTileM[k][ty] * subTileN[tx][k];

        __syncthreads();
    }
    P[Row * Width + Col] = Pvalue;
}

// ---- 内核 B：转置载入 + 静态 padding。行步长 S = 33，gcd(33,32)=1 => 无冲突 ----
__global__ void MatrixMulKernelTransposePadStatic(const float* __restrict__ M,
                                                  const float* __restrict__ N,
                                                  float* __restrict__ P,
                                                  int Width)
{
    __shared__ float subTileM[TILE_WIDTH][PADDED_STRIDE];
    __shared__ float subTileN[TILE_WIDTH][PADDED_STRIDE];

    int bx = blockIdx.x;  int by = blockIdx.y;
    int tx = threadIdx.x; int ty = threadIdx.y;
    int Row = by * TILE_WIDTH + ty;
    int Col = bx * TILE_WIDTH + tx;

    float Pvalue = 0.0f;

    for (int m = 0; m < Width / TILE_WIDTH; ++m) {
        subTileM[tx][ty] = M[Row * Width + m * TILE_WIDTH + tx];
        subTileN[tx][ty] = N[(m * TILE_WIDTH + ty) * Width + Col];
        __syncthreads();

        for (int k = 0; k < TILE_WIDTH; ++k)
            Pvalue += subTileM[k][ty] * subTileN[tx][k];

        __syncthreads();
    }
    P[Row * Width + Col] = Pvalue;
}

// ---- 内核 C：动态共享内存 + padding（同一块共享内存手工切成两个带 padding 的二维数组）----
__global__ void MatrixMulKernelTransposePadDynamic(const float* __restrict__ M,
                                                   const float* __restrict__ N,
                                                   float* __restrict__ P,
                                                   int Width)
{
    extern __shared__ float dynShared[];
    float (*subTileM)[PADDED_STRIDE] = reinterpret_cast<float (*)[PADDED_STRIDE]>(dynShared);
    float (*subTileN)[PADDED_STRIDE] =
        reinterpret_cast<float (*)[PADDED_STRIDE]>(dynShared + TILE_WIDTH * PADDED_STRIDE);

    int bx = blockIdx.x;  int by = blockIdx.y;
    int tx = threadIdx.x; int ty = threadIdx.y;
    int Row = by * TILE_WIDTH + ty;
    int Col = bx * TILE_WIDTH + tx;

    float Pvalue = 0.0f;

    for (int m = 0; m < Width / TILE_WIDTH; ++m) {
        subTileM[tx][ty] = M[Row * Width + m * TILE_WIDTH + tx];
        subTileN[tx][ty] = N[(m * TILE_WIDTH + ty) * Width + Col];
        __syncthreads();

        for (int k = 0; k < TILE_WIDTH; ++k)
            Pvalue += subTileM[k][ty] * subTileN[tx][k];

        __syncthreads();
    }
    P[Row * Width + Col] = Pvalue;
}

// ------------------------------ 主机端辅助函数 ------------------------------
static void fillMatrix(float* A, int n, unsigned seed)
{
    for (int i = 0; i < n * n; ++i) {
        unsigned x = (unsigned)(i + seed * 7919u);
        x = x * 1103515245u + 12345u;
        A[i] = (float)((int)((x >> 16) & 0x7u) - 3) * 0.25f;
    }
}

static void matmulCPU(const float* A, const float* B, float* C, int n)
{
    for (int i = 0; i < n; ++i)
        for (int j = 0; j < n; ++j) {
            float s = 0.0f;
            for (int k = 0; k < n; ++k) s += A[i * n + k] * B[k * n + j];
            C[i * n + j] = s;
        }
}

template <typename LaunchT>
static bool verifyKernel(const char* name, LaunchT launch, int n, dim3 block)
{
    size_t bytes = (size_t)n * (size_t)n * sizeof(float);
    float* hM   = (float*)malloc(bytes);
    float* hN   = (float*)malloc(bytes);
    float* hP   = (float*)malloc(bytes);
    float* hRef = (float*)malloc(bytes);
    fillMatrix(hM, n, 1u);
    fillMatrix(hN, n, 2u);
    matmulCPU(hM, hN, hRef, n);

    float *dM, *dN, *dP;
    CUDA_CHECK(cudaMalloc((void**)&dM, bytes));
    CUDA_CHECK(cudaMalloc((void**)&dN, bytes));
    CUDA_CHECK(cudaMalloc((void**)&dP, bytes));
    CUDA_CHECK(cudaMemcpy(dM, hM, bytes, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(dN, hN, bytes, cudaMemcpyHostToDevice));

    dim3 grid((n + block.x - 1) / block.x, (n + block.y - 1) / block.y);
    launch(dM, dN, dP, n, grid, block);
    CUDA_CHECK(cudaGetLastError());
    CUDA_CHECK(cudaDeviceSynchronize());
    CUDA_CHECK(cudaMemcpy(hP, dP, bytes, cudaMemcpyDeviceToHost));

    double maxErr = 0.0;
    for (int i = 0; i < n * n; ++i) {
        double d = fabs((double)hP[i] - (double)hRef[i]);
        if (d > maxErr) maxErr = d;
    }
    printf("[验证 %-16s] n=%d  最大绝对误差=%.3e  %s\n", name, n, maxErr,
           (maxErr < 1e-2) ? "PASS" : "FAIL");

    free(hM); free(hN); free(hP); free(hRef);
    CUDA_CHECK(cudaFree(dM)); CUDA_CHECK(cudaFree(dN)); CUDA_CHECK(cudaFree(dP));
    return maxErr < 1e-2;
}

template <typename LaunchT>
static double timeKernelMs(LaunchT launch, int iters)
{
    cudaEvent_t start, stop;
    CUDA_CHECK(cudaEventCreate(&start));
    CUDA_CHECK(cudaEventCreate(&stop));
    for (int i = 0; i < 3; ++i) launch();
    CUDA_CHECK(cudaEventRecord(start));
    for (int i = 0; i < iters; ++i) launch();
    CUDA_CHECK(cudaEventRecord(stop));
    CUDA_CHECK(cudaEventSynchronize(stop));
    float ms = 0.0f;
    CUDA_CHECK(cudaEventElapsedTime(&ms, start, stop));
    CUDA_CHECK(cudaEventDestroy(start));
    CUDA_CHECK(cudaEventDestroy(stop));
    return (double)ms / (double)iters;
}

int main(int argc, char** argv)
{
    int Width = (argc > 1) ? atoi(argv[1]) : 2048;
    if (Width % TILE_WIDTH != 0) {
        fprintf(stderr, "Width 必须是 TILE_WIDTH=%d 的整数倍\n", TILE_WIDTH);
        return EXIT_FAILURE;
    }

    dim3 block(TILE_WIDTH, TILE_WIDTH);
    const int dynBytes = 2 * TILE_WIDTH * PADDED_STRIDE * (int)sizeof(float);   // 8448 B

    // 1) 正确性验证：三个内核在 n=512 上与 CPU 参考实现对比
    verifyKernel("no-pad", [](const float* M, const float* N, float* P, int W, dim3 g, dim3 b) {
        MatrixMulKernelTransposeNoPad<<<g, b>>>(M, N, P, W);
    }, 512, block);
    verifyKernel("pad-static", [](const float* M, const float* N, float* P, int W, dim3 g, dim3 b) {
        MatrixMulKernelTransposePadStatic<<<g, b>>>(M, N, P, W);
    }, 512, block);
    verifyKernel("pad-dynamic", [dynBytes](const float* M, const float* N, float* P, int W, dim3 g, dim3 b) {
        MatrixMulKernelTransposePadDynamic<<<g, b, dynBytes>>>(M, N, P, W);
    }, 512, block);

    // 2) 计时
    size_t bytes = (size_t)Width * (size_t)Width * sizeof(float);
    float *hM = (float*)malloc(bytes), *hN = (float*)malloc(bytes);
    fillMatrix(hM, Width, 3u);
    fillMatrix(hN, Width, 4u);

    float *dM, *dN, *dP;
    CUDA_CHECK(cudaMalloc((void**)&dM, bytes));
    CUDA_CHECK(cudaMalloc((void**)&dN, bytes));
    CUDA_CHECK(cudaMalloc((void**)&dP, bytes));
    CUDA_CHECK(cudaMemcpy(dM, hM, bytes, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(dN, hN, bytes, cudaMemcpyHostToDevice));

    double flops = 2.0 * (double)Width * (double)Width * (double)Width;
    dim3 grid(Width / TILE_WIDTH, Width / TILE_WIDTH);

    double msNoPad = timeKernelMs([&] {
        MatrixMulKernelTransposeNoPad<<<grid, block>>>(dM, dN, dP, Width);
    }, 20);
    double msPadS = timeKernelMs([&] {
        MatrixMulKernelTransposePadStatic<<<grid, block>>>(dM, dN, dP, Width);
    }, 20);
    double msPadD = timeKernelMs([&] {
        MatrixMulKernelTransposePadDynamic<<<grid, block, dynBytes>>>(dM, dN, dP, Width);
    }, 20);

    CUDA_CHECK(cudaGetLastError());

    printf("\n===== TILE_WIDTH=%d, Width=%d, 三个内核对比（每 block 共享内存：无 padding %d B，padding %d B）=====\n",
           TILE_WIDTH, Width,
           2 * TILE_WIDTH * TILE_WIDTH * (int)sizeof(float), dynBytes);
    printf("%-28s %10s %14s %12s\n", "内核", "耗时(ms)", "GFLOP/s", "相对最优");
    double best = msNoPad;
    if (msPadS < best) best = msPadS;
    if (msPadD < best) best = msPadD;
    printf("%-28s %10.3f %14.1f %11.2fx\n", "A 转置载入（无 padding）",
           msNoPad, flops / msNoPad / 1e6, best / msNoPad);
    printf("%-28s %10.3f %14.1f %11.2fx\n", "B 转置载入（静态 padding）",
           msPadS, flops / msPadS / 1e6, best / msPadS);
    printf("%-28s %10.3f %14.1f %11.2fx\n", "C 转置载入（动态 padding）",
           msPadD, flops / msPadD / 1e6, best / msPadD);
    printf("\nbank 分析：S=32 时 bank=(32*tx+ty)%%32=ty -> 32 路冲突；"
           "S=33 时 bank=(33*tx+ty)%%32=(tx+ty)%%32 -> 取遍 32 个 bank，无冲突\n");
    printf("算术强度 = TILE_WIDTH/4 = %.1f FLOP/Byte；"
           "A100 机器平衡点 = 12.54 FLOP/Byte\n", TILE_WIDTH / 4.0);

    free(hM); free(hN);
    CUDA_CHECK(cudaFree(dM)); CUDA_CHECK(cudaFree(dN)); CUDA_CHECK(cudaFree(dP));
    return EXIT_SUCCESS;
}

// 文件: tiled_matmul_coarsened.cu
// 编译: nvcc -O3 -arch=sm_80 tiled_matmul_coarsened.cu -o tiled_matmul_coarsened
// 运行: ./tiled_matmul_coarsened 2048        (任意边长；本程序也会自动测试 517 这种非整数倍尺寸)

#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <cuda_runtime.h>

#define CUDA_CHECK(call)                                                        \
    do {                                                                        \
        cudaError_t err__ = (call);                                             \
        if (err__ != cudaSuccess) {                                             \
            fprintf(stderr, "CUDA error %s:%d: %s\n", __FILE__, __LINE__,       \
                    cudaGetErrorString(err__));                                 \
            exit(EXIT_FAILURE);                                                 \
        }                                                                       \
    } while (0)

#define TILE_WIDTH   32      // 共享内存 tile 边长
#define THREADS_X    16      // 块内线程 x 维
#define THREADS_Y    16      // 块内线程 y 维（256 线程覆盖 32x32 输出块）
#define COARSE        2      // 每个线程算 COARSE x COARSE 个输出（寄存器 tile）

// ---------- 对照：示例 1 的分块内核（TILE=32、每线程 1 个输出，要求 Width % 32 == 0）----------
__global__ void MatrixMulKernelTiled32(const float* __restrict__ M,
                                       const float* __restrict__ N,
                                       float* __restrict__ P,
                                       int Width)
{
    __shared__ float subTileM[TILE_WIDTH][TILE_WIDTH];
    __shared__ float subTileN[TILE_WIDTH][TILE_WIDTH];
    int bx = blockIdx.x, by = blockIdx.y;
    int tx = threadIdx.x, ty = threadIdx.y;
    int Row = by * TILE_WIDTH + ty;
    int Col = bx * TILE_WIDTH + tx;
    float Pvalue = 0.0f;
    for (int m = 0; m < Width / TILE_WIDTH; ++m) {
        subTileM[ty][tx] = M[Row * Width + m * TILE_WIDTH + tx];
        subTileN[ty][tx] = N[(m * TILE_WIDTH + ty) * Width + Col];
        __syncthreads();
        for (int k = 0; k < TILE_WIDTH; ++k)
            Pvalue += subTileM[ty][k] * subTileN[k][tx];
        __syncthreads();
    }
    P[Row * Width + Col] = Pvalue;
}

// ---------- 粗化 + 边界处理：每个线程算 2x2 个输出，支持任意 Width ----------
__global__ void __launch_bounds__(THREADS_X * THREADS_Y)
MatrixMulKernelCoarsened(const float* __restrict__ M,
                         const float* __restrict__ N,
                         float* __restrict__ P,
                         int Width)
{
    __shared__ float subTileM[TILE_WIDTH][TILE_WIDTH + 1];   // +1 padding：行步长 33
    __shared__ float subTileN[TILE_WIDTH][TILE_WIDTH + 1];

    const int tx = threadIdx.x;          // 0 .. 15
    const int ty = threadIdx.y;          // 0 .. 15
    const int bx = blockIdx.x;
    const int by = blockIdx.y;

    // 该线程负责的 2x2 寄存器 tile 在 P 中的左上角
    const int rowBase = by * TILE_WIDTH + ty * COARSE;
    const int colBase = bx * TILE_WIDTH + tx * COARSE;

    float acc[COARSE][COARSE];
#pragma unroll
    for (int i = 0; i < COARSE; ++i)
#pragma unroll
        for (int j = 0; j < COARSE; ++j) acc[i][j] = 0.0f;

    const int numTiles = (Width + TILE_WIDTH - 1) / TILE_WIDTH;   // 向上取整，兼容任意尺寸

    for (int m = 0; m < numTiles; ++m) {
        // 协作加载：256 个线程 x 4 个元素 = 1024 = TILE_WIDTH^2，全部合并访存
#pragma unroll
        for (int t = 0; t < (TILE_WIDTH * TILE_WIDTH) / (THREADS_X * THREADS_Y); ++t) {
            const int off   = (ty * THREADS_X + tx) + t * (THREADS_X * THREADS_Y);  // 0..1023
            const int ldRow = off / TILE_WIDTH;      // tile 内的行 0..31
            const int ldCol = off % TILE_WIDTH;      // tile 内的列 0..31

            const int gRowM = by * TILE_WIDTH + ldRow;          // M 的全局行
            const int gColM = m * TILE_WIDTH + ldCol;           // M 的全局列
            const int gRowN = m * TILE_WIDTH + ldRow;           // N 的全局行
            const int gColN = bx * TILE_WIDTH + ldCol;          // N 的全局列

            subTileM[ldRow][ldCol] = (gRowM < Width && gColM < Width)
                                   ? M[gRowM * Width + gColM] : 0.0f;   // 越界填 0
            subTileN[ldRow][ldCol] = (gRowN < Width && gColN < Width)
                                   ? N[gRowN * Width + gColN] : 0.0f;   // 越界填 0
        }
        __syncthreads();                                    // 屏障 1：RAW

#pragma unroll 4
        for (int k = 0; k < TILE_WIDTH; ++k) {
            float mreg[COARSE], nreg[COARSE];
#pragma unroll
            for (int i = 0; i < COARSE; ++i) mreg[i] = subTileM[ty * COARSE + i][k];  // 广播访问
#pragma unroll
            for (int j = 0; j < COARSE; ++j) nreg[j] = subTileN[k][tx * COARSE + j];  // stride-1
#pragma unroll
            for (int i = 0; i < COARSE; ++i)
#pragma unroll
                for (int j = 0; j < COARSE; ++j)
                    acc[i][j] += mreg[i] * nreg[j];         // 4 次 FMA 只用 4 次共享内存读
        }
        __syncthreads();                                    // 屏障 2：WAR
    }

    // 写回：块外的线程只算不用（不写全局内存）
#pragma unroll
    for (int i = 0; i < COARSE; ++i)
#pragma unroll
        for (int j = 0; j < COARSE; ++j) {
            const int row = rowBase + i;
            const int col = colBase + j;
            if (row < Width && col < Width)
                P[row * Width + col] = acc[i][j];
        }
}

// ------------------------------ 主机端辅助函数 ------------------------------
static void fillMatrix(float* A, int n, unsigned seed)
{
    for (int i = 0; i < n * n; ++i) {
        unsigned x = (unsigned)(i + seed * 7919u);
        x = x * 1103515245u + 12345u;
        A[i] = (float)((int)((x >> 16) & 0x7u) - 3) * 0.25f;
    }
}

static void matmulCPU(const float* A, const float* B, float* C, int n)
{
    for (int i = 0; i < n; ++i)
        for (int j = 0; j < n; ++j) {
            float s = 0.0f;
            for (int k = 0; k < n; ++k) s += A[i * n + k] * B[k * n + j];
            C[i * n + j] = s;
        }
}

template <typename LaunchT>
static bool verifyKernel(const char* name, LaunchT launch, int n)
{
    size_t bytes = (size_t)n * (size_t)n * sizeof(float);
    float* hM   = (float*)malloc(bytes);
    float* hN   = (float*)malloc(bytes);
    float* hP   = (float*)malloc(bytes);
    float* hRef = (float*)malloc(bytes);
    fillMatrix(hM, n, 1u);
    fillMatrix(hN, n, 2u);
    matmulCPU(hM, hN, hRef, n);

    float *dM, *dN, *dP;
    CUDA_CHECK(cudaMalloc((void**)&dM, bytes));
    CUDA_CHECK(cudaMalloc((void**)&dN, bytes));
    CUDA_CHECK(cudaMalloc((void**)&dP, bytes));
    CUDA_CHECK(cudaMemcpy(dM, hM, bytes, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(dN, hN, bytes, cudaMemcpyHostToDevice));

    launch(dM, dN, dP, n);
    CUDA_CHECK(cudaGetLastError());
    CUDA_CHECK(cudaDeviceSynchronize());
    CUDA_CHECK(cudaMemcpy(hP, dP, bytes, cudaMemcpyDeviceToHost));

    double maxErr = 0.0;
    for (int i = 0; i < n * n; ++i) {
        double d = fabs((double)hP[i] - (double)hRef[i]);
        if (d > maxErr) maxErr = d;
    }
    printf("[验证 %-12s] n=%-5d 最大绝对误差=%.3e  %s\n", name, n, maxErr,
           (maxErr < 1e-2) ? "PASS" : "FAIL");

    free(hM); free(hN); free(hP); free(hRef);
    CUDA_CHECK(cudaFree(dM)); CUDA_CHECK(cudaFree(dN)); CUDA_CHECK(cudaFree(dP));
    return maxErr < 1e-2;
}

template <typename LaunchT>
static double timeKernelMs(LaunchT launch, int iters)
{
    cudaEvent_t start, stop;
    CUDA_CHECK(cudaEventCreate(&start));
    CUDA_CHECK(cudaEventCreate(&stop));
    for (int i = 0; i < 3; ++i) launch();
    CUDA_CHECK(cudaEventRecord(start));
    for (int i = 0; i < iters; ++i) launch();
    CUDA_CHECK(cudaEventRecord(stop));
    CUDA_CHECK(cudaEventSynchronize(stop));
    float ms = 0.0f;
    CUDA_CHECK(cudaEventElapsedTime(&ms, start, stop));
    CUDA_CHECK(cudaEventDestroy(start));
    CUDA_CHECK(cudaEventDestroy(stop));
    return (double)ms / (double)iters;
}

int main(int argc, char** argv)
{
    int Width = (argc > 1) ? atoi(argv[1]) : 2048;
    if (Width % TILE_WIDTH != 0) {
        printf("提示：Width=%d 不是 %d 的整数倍，仅测试粗化内核（对照内核需要整数倍）\n",
               Width, TILE_WIDTH);
    }

    dim3 block(THREADS_X, THREADS_Y);
    dim3 grid((Width + TILE_WIDTH - 1) / TILE_WIDTH,
              (Width + TILE_WIDTH - 1) / TILE_WIDTH);

    // 1) 正确性：整数倍尺寸 + 非整数倍尺寸（517 = 32*16 + 5）
    verifyKernel("tiled32", [&](const float* M, const float* N, float* P, int W) {
        dim3 g(W / TILE_WIDTH, W / TILE_WIDTH);
        MatrixMulKernelTiled32<<<g, dim3(TILE_WIDTH, TILE_WIDTH)>>>(M, N, P, W);
    }, 512);
    verifyKernel("coarsened", [&](const float* M, const float* N, float* P, int W) {
        dim3 g((W + TILE_WIDTH - 1) / TILE_WIDTH, (W + TILE_WIDTH - 1) / TILE_WIDTH);
        MatrixMulKernelCoarsened<<<g, block>>>(M, N, P, W);
    }, 512);
    verifyKernel("coarsened-517", [&](const float* M, const float* N, float* P, int W) {
        dim3 g((W + TILE_WIDTH - 1) / TILE_WIDTH, (W + TILE_WIDTH - 1) / TILE_WIDTH);
        MatrixMulKernelCoarsened<<<g, block>>>(M, N, P, W);
    }, 517);

    // 2) 计时
    size_t bytes = (size_t)Width * (size_t)Width * sizeof(float);
    float *hM = (float*)malloc(bytes), *hN = (float*)malloc(bytes);
    fillMatrix(hM, Width, 3u);
    fillMatrix(hN, Width, 4u);

    float *dM, *dN, *dP;
    CUDA_CHECK(cudaMalloc((void**)&dM, bytes));
    CUDA_CHECK(cudaMalloc((void**)&dN, bytes));
    CUDA_CHECK(cudaMalloc((void**)&dP, bytes));
    CUDA_CHECK(cudaMemcpy(dM, hM, bytes, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(dN, hN, bytes, cudaMemcpyHostToDevice));

    double flops = 2.0 * (double)Width * (double)Width * (double)Width;
    double msCoarse = 0.0, msTiled = 0.0;

    if (Width % TILE_WIDTH == 0) {
        dim3 gT(Width / TILE_WIDTH, Width / TILE_WIDTH);
        msTiled = timeKernelMs([&] {
            MatrixMulKernelTiled32<<<gT, dim3(TILE_WIDTH, TILE_WIDTH)>>>(dM, dN, dP, Width);
        }, 20);
    }
    msCoarse = timeKernelMs([&] {
        MatrixMulKernelCoarsened<<<grid, block>>>(dM, dN, dP, Width);
    }, 20);

    CUDA_CHECK(cudaGetLastError());

    printf("\n===== 性能对比 (Width=%d, 2*W^3=%.3e FLOP) =====\n", Width, flops);
    if (msTiled > 0.0)
        printf("%-34s %10.3f ms   %8.1f GFLOP/s  (按 19.5 TFLOPS 峰值计 %.1f%%)\n",
               "TILE=32 未粗化 (1024 线程/块)", msTiled, flops / msTiled / 1e6,
               100.0 * flops / msTiled / 1e6 / 19500.0);
    printf("%-34s %10.3f ms   %8.1f GFLOP/s  (按 19.5 TFLOPS 峰值计 %.1f%%)\n",
           "TILE=32 + 2x2 粗化 (256 线程/块)", msCoarse, flops / msCoarse / 1e6,
           100.0 * flops / msCoarse / 1e6 / 19500.0);
    if (msTiled > 0.0)
        printf("加速比: %.2fx\n", msTiled / msCoarse);
    printf("块数: %u x %u，每块共享内存 %d B\n",
           grid.x, grid.y, 2 * TILE_WIDTH * (TILE_WIDTH + 1) * (int)sizeof(float));

    free(hM); free(hN);
    CUDA_CHECK(cudaFree(dM)); CUDA_CHECK(cudaFree(dN)); CUDA_CHECK(cudaFree(dP));
    return EXIT_SUCCESS;
}