06-CUDA 高级优化¶

// file: wmma_matmul.cu
// 使用WMMA API实现16x16矩阵乘法：D = A * B + C
// 编译: nvcc -arch=sm_70 -o wmma_matmul wmma_matmul.cu
// （sm_70=Volta, sm_80=Ampere, sm_90=Hopper）

#include <cuda_runtime.h>
#include <mma.h>        // WMMA头文件
#include <cuda_fp16.h>  // half精度
#include <cstdio>
#include <cstdlib>
#include <cmath>

using namespace nvcuda;

// ---- 矩阵维度：必须满足Tensor Core对齐要求 ----
constexpr int M = 16;
constexpr int N = 16;
constexpr int K = 16;

// ---- WMMA Kernel ----
__global__ void wmma_matmul_kernel(const half *A, const half *B,
                                   float *C, float *D,
                                   int m, int n, int k) {
    // 声明fragment
    // A: 行主序(row_major)，B: 列主序(col_major) —— WMMA要求
    wmma::fragment<wmma::matrix_a, M, N, K, half, wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b, M, N, K, half, wmma::col_major> b_frag;
    wmma::fragment<wmma::accumulator, M, N, K, float> acc_frag;
    wmma::fragment<wmma::accumulator, M, N, K, float> c_frag;

    // 加载A（m×k，行主序，leading dimension = k）
    wmma::load_matrix_sync(a_frag, A, K);

    // 加载B（k×n，列主序，leading dimension = k）
    wmma::load_matrix_sync(b_frag, B, K);

    // 加载偏置/累加器C（m×n，行主序，leading dimension = n）
    wmma::load_matrix_sync(c_frag, C, N, wmma::mem_row_major);

    // 执行MMA：acc = A * B + C
    wmma::mma_sync(acc_frag, a_frag, b_frag, c_frag);

    // 存储结果D（m×n，行主序）
    wmma::store_matrix_sync(D, acc_frag, N, wmma::mem_row_major);
}

// ---- CPU参考实现 ----
void cpu_matmul(const float *A, const float *B, const float *C,
                float *D, int m, int n, int k) {
    for (int i = 0; i < m; ++i)
        for (int j = 0; j < n; ++j) {
            float sum = C[i * n + j];
            for (int p = 0; p < k; ++p)
                sum += A[i * k + p] * B[p * n + j];
            D[i * n + j] = sum;
        }
}

int main() {
    // ---- 主机端数据准备 ----
    float h_A_fp32[M * K], h_B_fp32[K * N];
    float h_C[M * N], h_D_gpu[M * N], h_D_cpu[M * N];

    srand(42);
    for (int i = 0; i < M * K; ++i) h_A_fp32[i] = (float)(rand() % 5);
    for (int i = 0; i < K * N; ++i) h_B_fp32[i] = (float)(rand() % 5);
    for (int i = 0; i < M * N; ++i) h_C[i] = 0.0f;  // C初始化为0

    // 转FP16（用于GPU Tensor Core输入）
    half h_A_fp16[M * K], h_B_fp16[K * N];
    for (int i = 0; i < M * K; ++i) h_A_fp16[i] = __float2half(h_A_fp32[i]);
    // 注意：B需要转为列主序存储给WMMA
    for (int i = 0; i < K; ++i)
        for (int j = 0; j < N; ++j)
            h_B_fp16[j * K + i] = __float2half(h_B_fp32[i * N + j]); // 转置

    // ---- 设备端内存分配 ----
    half *d_A, *d_B;
    float *d_C, *d_D;
    cudaMalloc(&d_A, M * K * sizeof(half));
    cudaMalloc(&d_B, K * N * sizeof(half));
    cudaMalloc(&d_C, M * N * sizeof(float));
    cudaMalloc(&d_D, M * N * sizeof(float));

    cudaMemcpy(d_A, h_A_fp16, M * K * sizeof(half), cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B_fp16, K * N * sizeof(half), cudaMemcpyHostToDevice);
    cudaMemcpy(d_C, h_C, M * N * sizeof(float), cudaMemcpyHostToDevice);

    // ---- 启动Kernel（必须一个完整Warp = 32线程）----
    wmma_matmul_kernel<<<1, 32>>>(d_A, d_B, d_C, d_D, M, N, K);
    cudaDeviceSynchronize();

    cudaMemcpy(h_D_gpu, d_D, M * N * sizeof(float), cudaMemcpyDeviceToHost);

    // ---- CPU验证 ----
    cpu_matmul(h_A_fp32, h_B_fp32, h_C, h_D_cpu, M, N, K);

    float max_err = 0.0f;
    for (int i = 0; i < M * N; ++i) {
        float err = fabsf(h_D_gpu[i] - h_D_cpu[i]);
        if (err > max_err) max_err = err;
    }
    printf("Max error between GPU(Tensor Core) and CPU: %f\n", max_err);
    if (max_err < 1e-2f)
        printf("PASSED! Tensor Core result matches CPU.\n");
    else
        printf("FAILED! Max error too large.\n");

    // ---- 打印部分结果 ----
    printf("\nGPU result (first 4x4):\n");
    for (int i = 0; i < 4; ++i) {
        for (int j = 0; j < 4; ++j)
            printf("%8.1f", h_D_gpu[i * N + j]);
        printf("\n");
    }

    cudaFree(d_A); cudaFree(d_B); cudaFree(d_C); cudaFree(d_D);
    return 0;
}

cudaEvent_t start, stop;
cudaEventCreate(&start);
cudaEventCreate(&stop);

cudaEventRecord(start, stream);
myKernel<<<grid, block, 0, stream>>>(...);
cudaEventRecord(stop, stream);
cudaEventSynchronize(stop);

float ms = 0;
cudaEventElapsedTime(&ms, start, stop);  // 精确GPU计时

// file: multi_stream_pipeline.cu
// 演示多Stream实现H2D + Compute + D2H三阶段Pipeline重叠
// 编译: nvcc -arch=sm_70 -o multi_stream multi_stream_pipeline.cu

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

// 简单Kernel：向量逐元素乘以常数
__global__ void scale_kernel(float *data, float factor, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        // 模拟较重计算
        float val = data[idx];
        for (int i = 0; i < 100; ++i)
            val = val * factor + 0.001f;
        data[idx] = val;
    }
}

int main() {
    const int NUM_STREAMS = 4;
    const int CHUNK_SIZE  = 1 << 20;  // 每个chunk 1M个float ≈ 4MB
    const int TOTAL_SIZE  = NUM_STREAMS * CHUNK_SIZE;
    const float FACTOR    = 1.0001f;

    // ---- 分配Pinned Host内存（必须！否则无法真正异步）----
    float *h_data;
    cudaMallocHost(&h_data, TOTAL_SIZE * sizeof(float));

    // 初始化数据
    for (int i = 0; i < TOTAL_SIZE; ++i)
        h_data[i] = (float)(i % 1000) * 0.001f;

    // ---- 分配Device内存 ----
    float *d_data;
    cudaMalloc(&d_data, TOTAL_SIZE * sizeof(float));

    // ---- 创建Stream和Event ----
    cudaStream_t streams[NUM_STREAMS];
    for (int i = 0; i < NUM_STREAMS; ++i)
        cudaStreamCreate(&streams[i]);

    cudaEvent_t start_event, stop_event;
    cudaEventCreate(&start_event);
    cudaEventCreate(&stop_event);

    // ======== 方式1: 单Stream（无重叠，作为基准）========
    cudaEventRecord(start_event, 0);
    cudaMemcpy(d_data, h_data, TOTAL_SIZE * sizeof(float),
               cudaMemcpyHostToDevice);

    int threads = 256;
    int blocks  = (TOTAL_SIZE + threads - 1) / threads;
    scale_kernel<<<blocks, threads>>>(d_data, FACTOR, TOTAL_SIZE);

    cudaMemcpy(h_data, d_data, TOTAL_SIZE * sizeof(float),
               cudaMemcpyDeviceToHost);
    cudaEventRecord(stop_event, 0);
    cudaEventSynchronize(stop_event);

    float single_ms = 0;
    cudaEventElapsedTime(&single_ms, start_event, stop_event);
    printf("Single Stream: %.3f ms\n", single_ms);

    // 重新初始化
    for (int i = 0; i < TOTAL_SIZE; ++i)
        h_data[i] = (float)(i % 1000) * 0.001f;

    // ======== 方式2: 多Stream Pipeline（H2D+Compute+D2H重叠）========
    cudaEventRecord(start_event, 0);

    int chunk_blocks = (CHUNK_SIZE + threads - 1) / threads;

    for (int i = 0; i < NUM_STREAMS; ++i) {
        int offset = i * CHUNK_SIZE;
        size_t bytes = CHUNK_SIZE * sizeof(float);

        // 阶段1: 异步H2D
        cudaMemcpyAsync(d_data + offset, h_data + offset, bytes,
                        cudaMemcpyHostToDevice, streams[i]);

        // 阶段2: 异步Kernel执行
        scale_kernel<<<chunk_blocks, threads, 0, streams[i]>>>(
            d_data + offset, FACTOR, CHUNK_SIZE);

        // 阶段3: 异步D2H
        cudaMemcpyAsync(h_data + offset, d_data + offset, bytes,
                        cudaMemcpyDeviceToHost, streams[i]);
    }

    cudaEventRecord(stop_event, 0);
    // 等待所有Stream完成
    for (int i = 0; i < NUM_STREAMS; ++i)
        cudaStreamSynchronize(streams[i]);
    cudaEventSynchronize(stop_event);

    float multi_ms = 0;
    cudaEventElapsedTime(&multi_ms, start_event, stop_event);
    printf("Multi Stream (%d streams): %.3f ms\n", NUM_STREAMS, multi_ms);
    printf("Speedup: %.2fx\n", single_ms / multi_ms);

    // ---- 清理 ----
    for (int i = 0; i < NUM_STREAMS; ++i)
        cudaStreamDestroy(streams[i]);
    cudaEventDestroy(start_event);
    cudaEventDestroy(stop_event);
    cudaFree(d_data);
    cudaFreeHost(h_data);

    return 0;
}

// file: bank_conflict_demo.cu
// 演示Bank Conflict影响与padding优化
// 编译: nvcc -arch=sm_70 -o bank_conflict bank_conflict_demo.cu

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

constexpr int TILE_DIM = 32;
constexpr int BLOCK_ROWS = 8;

// ---- Naive转置（有Bank Conflict）----
__global__ void transpose_naive(float *odata, const float *idata,
                                 int width, int height) {
    // shared memory: 32×32, 列访问时32线程全部命中同一Bank
    __shared__ float tile[TILE_DIM][TILE_DIM];

    int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;
    int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;

    int index_in = xIndex + yIndex * width;

    // 加载：行访问，无Bank Conflict
    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS)
        tile[threadIdx.y + j][threadIdx.x] = idata[index_in + j * width];

    __syncthreads();

    // 写出转置位置
    xIndex = blockIdx.y * TILE_DIM + threadIdx.x;
    yIndex = blockIdx.x * TILE_DIM + threadIdx.y;

    // 列访问shared memory → 32-way Bank Conflict!
    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS)
        odata[xIndex + (yIndex + j) * height] =
            tile[threadIdx.x][threadIdx.y + j];
}

// ---- 优化转置（Padding消除Bank Conflict）----
__global__ void transpose_optimized(float *odata, const float *idata,
                                      int width, int height) {
    // +1 padding：错开Bank映射
    __shared__ float tile[TILE_DIM][TILE_DIM + 1];  // 关键优化！

    int xIndex = blockIdx.x * TILE_DIM + threadIdx.x;
    int yIndex = blockIdx.y * TILE_DIM + threadIdx.y;

    int index_in = xIndex + yIndex * width;

    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS)
        tile[threadIdx.y + j][threadIdx.x] = idata[index_in + j * width];

    __syncthreads();

    xIndex = blockIdx.y * TILE_DIM + threadIdx.x;
    yIndex = blockIdx.x * TILE_DIM + threadIdx.y;

    // 列访问时由于+1 padding，相邻线程访问不同Bank → 无冲突！
    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS)
        odata[xIndex + (yIndex + j) * height] =
            tile[threadIdx.x][threadIdx.y + j];
}

int main() {
    const int WIDTH = 4096, HEIGHT = 4096;
    size_t bytes = WIDTH * HEIGHT * sizeof(float);

    float *h_input = (float *)malloc(bytes);
    float *h_output = (float *)malloc(bytes);
    for (int i = 0; i < WIDTH * HEIGHT; ++i) h_input[i] = (float)i;

    float *d_input, *d_output;
    cudaMalloc(&d_input, bytes);
    cudaMalloc(&d_output, bytes);
    cudaMemcpy(d_input, h_input, bytes, cudaMemcpyHostToDevice);

    dim3 grid(WIDTH / TILE_DIM, HEIGHT / TILE_DIM);
    dim3 block(TILE_DIM, BLOCK_ROWS);

    // 计时Naive版
    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);

    // Warmup
    transpose_naive<<<grid, block>>>(d_output, d_input, WIDTH, HEIGHT);
    cudaDeviceSynchronize();

    cudaEventRecord(start);
    for (int i = 0; i < 100; ++i)
        transpose_naive<<<grid, block>>>(d_output, d_input, WIDTH, HEIGHT);
    cudaEventRecord(stop);
    cudaEventSynchronize(stop);
    float naive_ms;
    cudaEventElapsedTime(&naive_ms, start, stop);

    // 计时优化版
    cudaEventRecord(start);
    for (int i = 0; i < 100; ++i)
        transpose_optimized<<<grid, block>>>(d_output, d_input, WIDTH, HEIGHT);
    cudaEventRecord(stop);
    cudaEventSynchronize(stop);
    float opt_ms;
    cudaEventElapsedTime(&opt_ms, start, stop);

    printf("Naive transpose:     %.3f ms (avg over 100 runs)\n", naive_ms / 100);
    printf("Optimized transpose: %.3f ms (avg over 100 runs)\n", opt_ms / 100);
    printf("Speedup: %.2fx\n", naive_ms / opt_ms);

    // 带宽计算（读+写各一次）
    float bw_naive = 2.0f * bytes / (naive_ms / 100.0f * 1e-3f) / 1e9f;
    float bw_opt   = 2.0f * bytes / (opt_ms / 100.0f * 1e-3f) / 1e9f;
    printf("Effective bandwidth: naive=%.1f GB/s, optimized=%.1f GB/s\n",
           bw_naive, bw_opt);

    cudaFree(d_input); cudaFree(d_output);
    free(h_input); free(h_output);
    cudaEventDestroy(start); cudaEventDestroy(stop);
    return 0;
}

// file: gemm_optimization.cu
// 5步GEMM优化全代码
// 编译: nvcc -arch=sm_70 -O3 -o gemm_opt gemm_optimization.cu -lcublas

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

// ================ Step 1: Naive GEMM ================
// 每个线程计算C的一个元素，直接从全局内存读A和B
// 性能瓶颈：全局内存访问被大量重复（A的每行被读N次，B的每列被读M次）
__global__ void gemm_naive(const float *A, const float *B, float *C,
                           int M, int N, int K) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < M && col < N) {
        float sum = 0.0f;
        for (int k = 0; k < K; ++k) {
            sum += A[row * K + k] * B[k * N + col];
        }
        C[row * N + col] = sum;
    }
}
// 性能：~300 GFLOPS (A100)，约峰值的1.5%

// ================ Step 2: Shared Memory Tiling ================
// 将矩阵分块（Tile），用shared memory缓存，减少全局内存访问
constexpr int TILE_SIZE = 32;

__global__ void gemm_shared_tiling(const float *A, const float *B, float *C,
                                    int M, int N, int K) {
    __shared__ float As[TILE_SIZE][TILE_SIZE];
    __shared__ float Bs[TILE_SIZE][TILE_SIZE];

    int row = blockIdx.y * TILE_SIZE + threadIdx.y;
    int col = blockIdx.x * TILE_SIZE + threadIdx.x;

    float sum = 0.0f;

    // 沿K维度分tile迭代
    for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; ++t) {
        // 协作加载A的tile到shared memory
        int a_col = t * TILE_SIZE + threadIdx.x;
        if (row < M && a_col < K)
            As[threadIdx.y][threadIdx.x] = A[row * K + a_col];
        else
            As[threadIdx.y][threadIdx.x] = 0.0f;

        // 协作加载B的tile到shared memory
        int b_row = t * TILE_SIZE + threadIdx.y;
        if (b_row < K && col < N)
            Bs[threadIdx.y][threadIdx.x] = B[b_row * N + col];
        else
            Bs[threadIdx.y][threadIdx.x] = 0.0f;

        __syncthreads();

        // 在shared memory中完成tile内的乘加
        for (int k = 0; k < TILE_SIZE; ++k)
            sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];

        __syncthreads();
    }

    if (row < M && col < N)
        C[row * N + col] = sum;
}
// 性能：~1500 GFLOPS（A100），约5x提速

// ================ Step 3: 共享内存分块加载 + 循环展开 ================
// 在 Step 2 基础上增加 #pragma unroll 循环展开，提高指令级并行度
constexpr int TILE3 = 32;
constexpr int FLOAT4_COUNT = TILE3 / 4;  // 每行8个float4

__global__ void gemm_vectorized(const float *A, const float *B, float *C,
                                 int M, int N, int K) {
    __shared__ float As[TILE3][TILE3];
    __shared__ float Bs[TILE3][TILE3];

    int row = blockIdx.y * TILE3 + threadIdx.y;
    int col = blockIdx.x * TILE3 + threadIdx.x;

    float sum = 0.0f;

    for (int t = 0; t < (K + TILE3 - 1) / TILE3; ++t) {
        int a_col = t * TILE3 + threadIdx.x;
        if (row < M && a_col < K)
            As[threadIdx.y][threadIdx.x] = A[row * K + a_col];
        else
            As[threadIdx.y][threadIdx.x] = 0.0f;

        // B的列方向是连续的 → 可以用float4加载
        int b_row = t * TILE3 + threadIdx.y;
        if (b_row < K && col < N)
            Bs[threadIdx.y][threadIdx.x] = B[b_row * N + col];
        else
            Bs[threadIdx.y][threadIdx.x] = 0.0f;

        __syncthreads();

        // 内层循环展开
        #pragma unroll
        for (int k = 0; k < TILE3; ++k)
            sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];

        __syncthreads();
    }

    if (row < M && col < N)
        C[row * N + col] = sum;
}
// 性能：~2000 GFLOPS，进一步提升内存带宽利用率

// ================ Step 4: Register Tiling ================
// 每个线程计算多个C元素（TM×TN的子矩阵），增加计算/访存比
constexpr int BM = 128;   // Block tile M
constexpr int BN = 128;   // Block tile N
constexpr int BK = 8;     // Block tile K
constexpr int TM = 8;     // Thread tile M（每线程计算8行）
constexpr int TN = 8;     // Thread tile N（每线程计算8列）

__global__ void gemm_register_tiling(const float *A, const float *B, float *C,
                                      int M, int N, int K) {
    // Block内线程数 = (BM/TM) × (BN/TN) = 16 × 16 = 256
    const int thread_row = threadIdx.x / (BN / TN); // 0-15
    const int thread_col = threadIdx.x % (BN / TN); // 0-15

    // Shared memory分配
    __shared__ float As[BM][BK];   // 128×8
    __shared__ float Bs[BK][BN];   // 8×128

    // 指向当前Block负责的C子矩阵
    A += blockIdx.y * BM * K;
    B += blockIdx.x * BN;
    C += blockIdx.y * BM * N + blockIdx.x * BN;

    // 寄存器中存储TM×TN的局部结果
    float thread_results[TM][TN] = {0.0f};
    // 寄存器缓存A和B的列/行片段
    float reg_A[TM];
    float reg_B[TN];

    // 沿K维度分BK大小的tile迭代
    for (int bk = 0; bk < K; bk += BK) {
        // ---- 协作加载A的tile到shared memory ----
        // 每线程加载多个元素以填满 BM×BK 的shared memory
        for (int load_offset = 0; load_offset < BM * BK;
             load_offset += blockDim.x) {
            int idx = load_offset + threadIdx.x;
            if (idx < BM * BK) {
                int sm_row = idx / BK;
                int sm_col = idx % BK;
                int global_row = blockIdx.y * BM + sm_row;
                int global_col = bk + sm_col;
                if (global_row < M && global_col < K)
                    As[sm_row][sm_col] = A[sm_row * K + bk + sm_col];
                else
                    As[sm_row][sm_col] = 0.0f;
            }
        }

        // ---- 协作加载B的tile到shared memory ----
        for (int load_offset = 0; load_offset < BK * BN;
             load_offset += blockDim.x) {
            int idx = load_offset + threadIdx.x;
            if (idx < BK * BN) {
                int sm_row = idx / BN;
                int sm_col = idx % BN;
                int global_row = bk + sm_row;
                int global_col = blockIdx.x * BN + sm_col;
                if (global_row < K && global_col < N)
                    Bs[sm_row][sm_col] = B[(bk + sm_row) * N +
                                           blockIdx.x * BN + sm_col];
                else
                    Bs[sm_row][sm_col] = 0.0f;
            }
        }

        __syncthreads();

        // ---- 寄存器级计算 ----
        for (int dot_idx = 0; dot_idx < BK; ++dot_idx) {
            // 加载A的TM个元素到寄存器
            for (int i = 0; i < TM; ++i)
                reg_A[i] = As[thread_row * TM + i][dot_idx];

            // 加载B的TN个元素到寄存器
            for (int j = 0; j < TN; ++j)
                reg_B[j] = Bs[dot_idx][thread_col * TN + j];

            // 外积：TM × TN次FMA
            for (int i = 0; i < TM; ++i)
                for (int j = 0; j < TN; ++j)
                    thread_results[i][j] += reg_A[i] * reg_B[j];
        }

        __syncthreads();
    }

    // ---- 写回C ----
    for (int i = 0; i < TM; ++i)
        for (int j = 0; j < TN; ++j) {
            int global_row = blockIdx.y * BM + thread_row * TM + i;
            int global_col = blockIdx.x * BN + thread_col * TN + j;
            if (global_row < M && global_col < N)
                C[(thread_row * TM + i) * N +
                   thread_col * TN + j] = thread_results[i][j];
        }
}
// 性能：~6000-8000 GFLOPS（A100），计算/访存比从1提升到TM*TN=64

// ================ Step 5: Double Buffering ================
// 使用双缓冲：当前tile在计算时，下一个tile已经在加载
// 隐藏shared memory加载延迟

constexpr int DB_BM = 64;
constexpr int DB_BN = 64;
constexpr int DB_BK = 8;
constexpr int DB_TM = 4;
constexpr int DB_TN = 4;

__global__ void gemm_double_buffer(const float *A, const float *B, float *C,
                                    int M, int N, int K) {
    // 双缓冲：两组shared memory交替使用
    __shared__ float As[2][DB_BM][DB_BK];
    __shared__ float Bs[2][DB_BK][DB_BN];

    const int thread_row = threadIdx.x / (DB_BN / DB_TN);
    const int thread_col = threadIdx.x % (DB_BN / DB_TN);

    float thread_results[DB_TM][DB_TN] = {0.0f};
    float reg_A[DB_TM], reg_B[DB_TN];

    int num_tiles = (K + DB_BK - 1) / DB_BK;

    // ---- 预加载第一个tile到buffer 0 ----
    for (int load_offset = 0; load_offset < DB_BM * DB_BK;
         load_offset += blockDim.x) {
        int idx = load_offset + threadIdx.x;
        if (idx < DB_BM * DB_BK) {
            int sr = idx / DB_BK, sc = idx % DB_BK;
            int gr = blockIdx.y * DB_BM + sr;
            int gc = sc;
            As[0][sr][sc] = (gr < M && gc < K) ? A[gr * K + gc] : 0.0f;
        }
    }
    for (int load_offset = 0; load_offset < DB_BK * DB_BN;
         load_offset += blockDim.x) {
        int idx = load_offset + threadIdx.x;
        if (idx < DB_BK * DB_BN) {
            int sr = idx / DB_BN, sc = idx % DB_BN;
            int gr = sr;
            int gc = blockIdx.x * DB_BN + sc;
            Bs[0][sr][sc] = (gr < K && gc < N) ? B[gr * N + gc] : 0.0f;
        }
    }
    __syncthreads();

    // ---- 主循环：计算当前buffer，预取下一个buffer ----
    for (int tile = 0; tile < num_tiles; ++tile) {
        int cur_buf = tile % 2;
        int nxt_buf = 1 - cur_buf;

        // 异步预加载下一个tile到nxt_buf
        if (tile + 1 < num_tiles) {
            int next_bk = (tile + 1) * DB_BK;
            for (int load_offset = 0; load_offset < DB_BM * DB_BK;
                 load_offset += blockDim.x) {
                int idx = load_offset + threadIdx.x;
                if (idx < DB_BM * DB_BK) {
                    int sr = idx / DB_BK, sc = idx % DB_BK;
                    int gr = blockIdx.y * DB_BM + sr;
                    int gc = next_bk + sc;
                    As[nxt_buf][sr][sc] = (gr < M && gc < K) ?
                                          A[gr * K + gc] : 0.0f;
                }
            }
            for (int load_offset = 0; load_offset < DB_BK * DB_BN;
                 load_offset += blockDim.x) {
                int idx = load_offset + threadIdx.x;
                if (idx < DB_BK * DB_BN) {
                    int sr = idx / DB_BN, sc = idx % DB_BN;
                    int gr = next_bk + sr;
                    int gc = blockIdx.x * DB_BN + sc;
                    Bs[nxt_buf][sr][sc] = (gr < K && gc < N) ?
                                          B[gr * N + gc] : 0.0f;
                }
            }
        }

        // 计算当前buffer
        for (int dot_idx = 0; dot_idx < DB_BK; ++dot_idx) {
            for (int i = 0; i < DB_TM; ++i)
                reg_A[i] = As[cur_buf][thread_row * DB_TM + i][dot_idx];
            for (int j = 0; j < DB_TN; ++j)
                reg_B[j] = Bs[cur_buf][dot_idx][thread_col * DB_TN + j];

            for (int i = 0; i < DB_TM; ++i)
                for (int j = 0; j < DB_TN; ++j)
                    thread_results[i][j] += reg_A[i] * reg_B[j];
        }

        __syncthreads();  // 确保预加载完成
    }

    // 写回C
    for (int i = 0; i < DB_TM; ++i)
        for (int j = 0; j < DB_TN; ++j) {
            int gr = blockIdx.y * DB_BM + thread_row * DB_TM + i;
            int gc = blockIdx.x * DB_BN + thread_col * DB_TN + j;
            if (gr < M && gc < N)
                C[gr * N + gc] = thread_results[i][j];
        }
}