softmax的cuda编程代码

Softmax函数是一种常用的归一化函数，在深度学习中常用于神经网络的最后一层，将每个输入值转换为概率分布。CUDA编程是针对NVIDIA GPU的并行计算环境，下面是使用CUDA编写的简单softmax函数的示例：

#include <cuda_runtime.h>
__global__ void softmax(float* input, float* output, int n_classes, int size)
{
    // 获取当前线程ID
    int thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (thread_idx < size) {
        float sum_exp = 0.0f;

        // 计算前向传播的exp和累加和
        for (int i = 0; i < n_classes; ++i) {
            float val = input[thread_idx * n_classes + i];
            output[thread_idx * n_classes + i] = exp(val);
            sum_exp += output[thread_idx * n_classes + i];
        }

        // 归一化，将结果设置为概率值
        output[thread_idx * n_classes] = 1.0f / sum_exp;
        
        // 非原子操作更新其他元素的概率
        for (int i = 1; i < n_classes; ++i) {
            output[thread_idx * n_classes + i] *= sum_exp - output[thread_idx * n_classes];
        }
    }
}

// 主函数中调用softmax函数
void softmaxKernel(float* h_input, float* d_output, int n_classes, int size, cudaStream_t stream)
{
    dim3 threadsPerBlock(256, 1); // CUDA块大小
    dim3 numBlocks(size / threadsPerBlock.x + (size % threadsPerBlock.x != 0 ? 1 : 0), 1);

    // 将host数据复制到device
    cudaMemcpy(d_input, h_input, size * n_classes * sizeof(float), cudaMemcpyHostToDevice, stream);

    // 调用kernel
    softmax<<<numBlocks, threadsPerBlock>>>(d_input, d_output, n_classes, size, stream);

    // 将device数据复制回host
    cudaMemcpy(h_output, d_output, size * n_classes * sizeof(float), cudaMemcpyDeviceToHost, stream);
}