From 6968605bfb30a9030c175b67ff88654a2450f7ec Mon Sep 17 00:00:00 2001 From: gsd <2396583337@qq.com> Date: Fri, 14 Nov 2025 10:57:42 +0800 Subject: [PATCH] fix CrossEntropyLoss #6 --- S1/gsd123_#6/CrossEntropyLoss_cuda.py | 156 +++++++++++++++++++++++++ S1/gsd123_#6/CrossEntropyLoss_torch.py | 26 +++++ S1/gsd123_#6/prompt.txt | 0 S1/gsd123_#6/run_code.py | 78 +++++++++++++ 4 files changed, 260 insertions(+) create mode 100644 S1/gsd123_#6/CrossEntropyLoss_cuda.py create mode 100644 S1/gsd123_#6/CrossEntropyLoss_torch.py create mode 100644 S1/gsd123_#6/prompt.txt create mode 100644 S1/gsd123_#6/run_code.py diff --git a/S1/gsd123_#6/CrossEntropyLoss_cuda.py b/S1/gsd123_#6/CrossEntropyLoss_cuda.py new file mode 100644 index 0000000..c91fa58 --- /dev/null +++ b/S1/gsd123_#6/CrossEntropyLoss_cuda.py @@ -0,0 +1,156 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + +BATCH_SIZE = 4096 +N_CLASSES = 1024 + + +class ModelNew(nn.Module): + + def __init__(self): + super().__init__() + self._compile_cuda_kernel() + + def _compile_cuda_kernel(self): + cpp_source = """ + #include + + // C++ 接口 + torch::Tensor cross_entropy_forward_cuda( + torch::Tensor logits, + torch::Tensor target + ); + """ + + cuda_source = """ + #include + #include + #include + #include + + #define BLOCK_SIZE 256 + + __global__ void cross_entropy_fused_kernel( + const float* __restrict__ logits_data, // (N, C) + const int64_t* __restrict__ target_data, // (N,) + float* __restrict__ loss_per_row_out, // (N,) + int N, + int C + ) { + // 当前处理的 Batch 索引 + int row_idx = blockIdx.x; + if (row_idx >= N) return; + + // 当前行的指针 + const float* row_logits = logits_data + row_idx * C; + int tid = threadIdx.x; + + // 共享内存：用于 Max 和 Sum 的归约 + __shared__ float s_data[BLOCK_SIZE]; + + float thread_max = -FLT_MAX; + + // Grid-Stride Loop 遍历类别 C + for (int c = tid; c < C; c += BLOCK_SIZE) { + float val = row_logits[c]; + if (val > thread_max) { + thread_max = val; + } + } + s_data[tid] = thread_max; + __syncthreads(); + + // 块内归约 (Max) + for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) { + if (tid < offset) { + if (s_data[tid + offset] > s_data[tid]) { + s_data[tid] = s_data[tid + offset]; + } + } + __syncthreads(); + } + + + float row_max_val = s_data[0]; + __syncthreads(); + + float thread_sum_exp = 0.0f; + + for (int c = tid; c < C; c += BLOCK_SIZE) { + float val = row_logits[c]; + thread_sum_exp += expf(val - row_max_val); + } + s_data[tid] = thread_sum_exp; + __syncthreads(); + + // 块内归约 (Sum) + for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) { + if (tid < offset) { + s_data[tid] += s_data[tid + offset]; + } + __syncthreads(); + } + + if (tid == 0) { + float row_sum_exp = s_data[0]; + float log_sum_exp = logf(row_sum_exp) + row_max_val; + + int64_t target_class = target_data[row_idx]; + float target_logit = row_logits[target_class]; + + // Cross Entropy Formula + loss_per_row_out[row_idx] = -target_logit + log_sum_exp; + } + } + + // C++ 封装函数 + torch::Tensor cross_entropy_forward_cuda( + torch::Tensor logits, + torch::Tensor target + ) { + TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor"); + TORCH_CHECK(target.is_cuda(), "target must be a CUDA tensor"); + TORCH_CHECK(logits.dim() == 2, "logits must be 2D"); + TORCH_CHECK(target.dim() == 1, "target must be 1D"); + + // 确保连续 + logits = logits.contiguous(); + target = target.contiguous(); + + int N = logits.size(0); // Batch Size + int C = logits.size(1); // Num Classes + + TORCH_CHECK(target.size(0) == N, "Target size mismatch"); + + auto losses = torch::empty({N}, logits.options()); + + // 启动配置: + // Grid: N (每个 Batch 一个 Block) + // Block: 256 + dim3 grid_dim(N); + dim3 block_dim(BLOCK_SIZE); + + cross_entropy_fused_kernel<<>>( + logits.data_ptr(), + target.data_ptr(), + losses.data_ptr(), + N, C + ); + + // 返回 Mean Reduction + return losses.mean(); + } + """ + + self.ce_op = load_inline( + name="cross_entropy_op_v1", + cpp_sources=cpp_source, + cuda_sources=cuda_source, + functions=["cross_entropy_forward_cuda"], + extra_cuda_cflags=["-O3"], + verbose=False + ) + + def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + return self.ce_op.cross_entropy_forward_cuda(logits, target) \ No newline at end of file diff --git a/S1/gsd123_#6/CrossEntropyLoss_torch.py b/S1/gsd123_#6/CrossEntropyLoss_torch.py new file mode 100644 index 0000000..fd8cef9 --- /dev/null +++ b/S1/gsd123_#6/CrossEntropyLoss_torch.py @@ -0,0 +1,26 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +BATCH_SIZE = 4096 +N_CLASSES = 1024 + + +class Model(nn.Module): + + def __init__(self): + super().__init__() + self.criterion = nn.CrossEntropyLoss(reduction='mean') + + def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + return self.criterion(logits, target) + + +def get_inputs(): + logits = torch.randn(BATCH_SIZE, N_CLASSES, dtype=torch.float32) + target = torch.randint(0, N_CLASSES, (BATCH_SIZE,), dtype=torch.long) + return [logits, target] + + +def get_init_inputs(): + return [] \ No newline at end of file diff --git a/S1/gsd123_#6/prompt.txt b/S1/gsd123_#6/prompt.txt new file mode 100644 index 0000000..e69de29 diff --git a/S1/gsd123_#6/run_code.py b/S1/gsd123_#6/run_code.py new file mode 100644 index 0000000..b9675b4 --- /dev/null +++ b/S1/gsd123_#6/run_code.py @@ -0,0 +1,78 @@ +import torch +import time +from CrossEntropyLoss_torch import Model, get_inputs, get_init_inputs +from CrossEntropyLoss_cuda import ModelNew + +def run_benchmark(): + if not torch.cuda.is_available(): + print("CUDA 不可用") + return + + device = torch.device("cuda") + + # 准备输入数据 + inputs = [x.cuda(device=device) for x in get_inputs()] + init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()] + + # 初始化模型 + torch_model = Model(*init_inputs).cuda() + cuda_model = ModelNew(*init_inputs).cuda() + + torch_model.eval() + cuda_model.eval() + + print("-------------------- 精度对齐验证 --------------------") + with torch.no_grad(): + # 预热GPU + _ = torch_model(*inputs) + _ = cuda_model(*inputs) + + # 正式测试 + output_torch = torch_model(*inputs) + output_cuda = cuda_model(*inputs) + + # 精度验证 + abs_diff = torch.abs(output_torch - output_cuda) + max_diff = torch.max(abs_diff).item() + mean_diff = torch.mean(abs_diff).item() + + if max_diff < 1e-4 and mean_diff < 1e-5: + print(f"✅ 精度对齐：最大误差 {max_diff:.6f}，平均误差 {mean_diff:.6f}") + precision_flag = True + else: + print(f"❌ 精度不一致：最大误差 {max_diff:.6f}，平均误差 {mean_diff:.6f}") + precision_flag = False + + print("\n-------------------- 性能加速比测试 --------------------") + num_iterations = 100 + + # 预热GPU + for _ in range(10): + _ = torch_model(*inputs) + _ = cuda_model(*inputs) + + # PyTorch模型计时 + torch.cuda.synchronize() + start_time = time.time() + for _ in range(num_iterations): + _ = torch_model(*inputs) + torch.cuda.synchronize() + torch_time = (time.time() - start_time) / num_iterations + + # 自定义CUDA内核计时 + torch.cuda.synchronize() + start_time = time.time() + for _ in range(num_iterations): + _ = cuda_model(*inputs) + torch.cuda.synchronize() + cuda_time = (time.time() - start_time) / num_iterations + + print(f"PyTorch内置Swish平均执行时间: {torch_time:.6f}秒") + print(f"自定义CUDA Swish平均执行时间: {cuda_time:.6f}秒") + speedup = torch_time / cuda_time if cuda_time > 0 else 0 + print(f"加速比 (Speedup): {speedup:.2f}x") + + return precision_flag, speedup + +if __name__ == "__main__": + precision_flag, speedup = run_benchmark() \ No newline at end of file