写一段用cublas库进行批处理矩阵乘法的代码

以下是使用CUBLAS库进行批处理矩阵乘法的示例代码：

#include <stdio.h>
#include <stdlib.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>

#define N 3 // 矩阵大小
#define BATCH_SIZE 2 // 批处理大小

void printMatrix(float* matrix, int rows, int cols) {
    for (int i = 0; i < rows; i++) {
        for (int j = 0; j < cols; j++) {
            printf("%f ", matrix[i * cols + j]);
        }
        printf("\n");
    }
    printf("\n");
}

int main() {
    cublasHandle_t handle;
    cublasCreate(&handle);

    float* h_A[BATCH_SIZE]; // 批处理输入矩阵A
    float* h_B[BATCH_SIZE]; // 批处理输入矩阵B
    float* h_C[BATCH_SIZE]; // 批处理输出矩阵C

    float* d_A[BATCH_SIZE]; // GPU上的输入矩阵A
    float* d_B[BATCH_SIZE]; // GPU上的输入矩阵B
    float* d_C[BATCH_SIZE]; // GPU上的输出矩阵C

    // 为批处理矩阵分配内存
    for (int i = 0; i < BATCH_SIZE; i++) {
        h_A[i] = (float*)malloc(N * N * sizeof(float));
        h_B[i] = (float*)malloc(N * N * sizeof(float));
        h_C[i] = (float*)malloc(N * N * sizeof(float));
        cudaMalloc((void**)&d_A[i], N * N * sizeof(float));
        cudaMalloc((void**)&d_B[i], N * N * sizeof(float));
        cudaMalloc((void**)&d_C[i], N * N * sizeof(float));
    }

    // 初始化输入矩阵A和B
    for (int i = 0; i < BATCH_SIZE; i++) {
        for (int j = 0; j < N * N; j++) {
            h_A[i][j] = i + j;
            h_B[i][j] = i - j;
        }
    }

    // 将输入矩阵A和B从主机内存复制到GPU内存
    for (int i = 0; i < BATCH_SIZE; i++) {
        cudaMemcpy(d_A[i], h_A[i], N * N * sizeof(float), cudaMemcpyHostToDevice);
        cudaMemcpy(d_B[i], h_B[i], N * N * sizeof(float), cudaMemcpyHostToDevice);
    }

    const float alpha = 1.0f;
    const float beta = 0.0f;

    // 执行批处理矩阵乘法
    cublasSgemmBatched(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, N, N, &alpha, (const float**)d_A, N, (const float**)d_B, N, &beta, d_C, N, BATCH_SIZE);

    // 将输出矩阵C从GPU内存复制到主机内存
    for (int i = 0; i < BATCH_SIZE; i++) {
        cudaMemcpy(h_C[i], d_C[i], N * N * sizeof(float), cudaMemcpyDeviceToHost);
    }

    // 打印输入矩阵A和B
    printf("Matrix A:\n");
    for (int i = 0; i < BATCH_SIZE; i++) {
        printf("Batch %d:\n", i);
        printMatrix(h_A[i], N, N);
    }

    printf("Matrix B:\n");
    for (int i = 0; i < BATCH_SIZE; i++) {
        printf("Batch %d:\n", i);
        printMatrix(h_B[i], N, N);
    }

    // 打印输出矩阵C
    printf("Matrix C:\n");
    for (int i = 0; i < BATCH_SIZE; i++) {
        printf("Batch %d:\n", i);
        printMatrix(h_C[i], N, N);
    }

    // 释放内存
    for (int i = 0; i < BATCH_SIZE; i++) {
        free(h_A[i]);
        free(h_B[i]);
        free(h_C[i]);
        cudaFree(d_A[i]);
        cudaFree(d_B[i]);
        cudaFree(d_C[i]);
    }

    cublasDestroy(handle);

    return 0;
}

这个示例代码演示了如何使用CUBLAS库进行批处理矩阵乘法。它创建了两个批处理的输入矩阵A和B，并分配了相应的GPU内存。然后，它将输入矩阵从主机内存复制到GPU内存，并使用`cublasSgemmBatched`函数执行批处理矩阵乘法。最后，它将输出矩阵C从GPU内存复制到主机内存，并打印输入矩阵A和B以及输出矩阵C的内容。请确保您已正确安装CUDA和CUBLAS，并将代码与适当的编译器和库进行编译和链接。