diff --git a/S1/gsd123_#9/circleloss_cuda.py b/S1/gsd123_#9/circleloss_cuda.py new file mode 100644 index 0000000..21323c1 --- /dev/null +++ b/S1/gsd123_#9/circleloss_cuda.py @@ -0,0 +1,298 @@ +# circleloss_cuda.py +import torch +import torch.nn.functional as F +from torch.utils.cpp_extension import load_inline +# 修复：从正确的文件导入 +from circleloss_torch import BATCH_SIZE, FEATURE_DIM, MARGIN, GAMMA + + +class ModelNew(torch.nn.Module): + + def __init__(self): + super().__init__() + self.margin = MARGIN + self.gamma = GAMMA + self._compile_cuda_kernel() + + def _compile_cuda_kernel(self): + cpp_source = """ + #include + + // C++ 接口 (保持不变) + torch::Tensor circleloss_forward_cuda( + torch::Tensor similarities, + torch::Tensor labels, + float margin_val, + float gamma_val + ); + """ + + cuda_source = """ + #include + #include + #include + #include // for int64_t + #include // For FLT_MAX + + #define BLOCK_SIZE 256 + + // ------------------------------------------------------------------ + // 阶段 1: 寻找 Logits 的最大值 + // ------------------------------------------------------------------ + __global__ void circleloss_find_max_kernel( + const float* __restrict__ similarities_data, + const int64_t* __restrict__ labels_data, + float* __restrict__ block_max_p_out, // (grid_size,) + float* __restrict__ block_max_n_out, // (grid_size,) + int n_elements, + int batch_size, + float margin_val, + float gamma_val + ) { + __shared__ float s_data_p[BLOCK_SIZE]; + __shared__ float s_data_n[BLOCK_SIZE]; + + float thread_max_p = -FLT_MAX; + float thread_max_n = -FLT_MAX; + + const float delta_p = 1.0f - margin_val; + const float delta_n = margin_val; + + int grid_stride = gridDim.x * blockDim.x; + + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; + idx < n_elements; + idx += grid_stride) + { + int i = idx / batch_size; + int j = idx % batch_size; + float s = similarities_data[idx]; + + if (labels_data[i] == labels_data[j]) { + // 正样本对 + float ap = fmaxf(0.0f, -s + 1.0f + margin_val); + float logit_p = -ap * (s - delta_p) * gamma_val; + thread_max_p = fmaxf(thread_max_p, logit_p); + } else { + // 负样本对 + float an = fmaxf(0.0f, s + margin_val); + float logit_n = an * (s - delta_n) * gamma_val; + thread_max_n = fmaxf(thread_max_n, logit_n); + } + } + + // --- 块内归约 (Max) - 正样本对 --- + s_data_p[threadIdx.x] = thread_max_p; + __syncthreads(); + for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) { + if (threadIdx.x < offset) { + s_data_p[threadIdx.x] = fmaxf(s_data_p[threadIdx.x], s_data_p[threadIdx.x + offset]); + } + __syncthreads(); + } + + // --- 块内归约 (Max) - 负样本对 --- + s_data_n[threadIdx.x] = thread_max_n; + __syncthreads(); + for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) { + if (threadIdx.x < offset) { + s_data_n[threadIdx.x] = fmaxf(s_data_n[threadIdx.x], s_data_n[threadIdx.x + offset]); + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + block_max_p_out[blockIdx.x] = s_data_p[0]; + block_max_n_out[blockIdx.x] = s_data_n[0]; + } + } + + + // ------------------------------------------------------------------ + // 阶段 2: 计算 Sum(Exp(Logit - Max)) + // ------------------------------------------------------------------ + __global__ void circleloss_sum_exp_diff_kernel( + const float* __restrict__ similarities_data, + const int64_t* __restrict__ labels_data, + float* __restrict__ block_sum_p_out, // (grid_size,) + float* __restrict__ block_sum_n_out, // (grid_size,) + float global_max_p, // 全局最大值 (标量) + float global_max_n, // 全局最大值 (标量) + int n_elements, + int batch_size, + float margin_val, + float gamma_val + ) { + __shared__ float s_data_p[BLOCK_SIZE]; + __shared__ float s_data_n[BLOCK_SIZE]; + + float thread_sum_p = 0.0f; + float thread_sum_n = 0.0f; + + const float delta_p = 1.0f - margin_val; + const float delta_n = margin_val; + + int grid_stride = gridDim.x * blockDim.x; + + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; + idx < n_elements; + idx += grid_stride) + { + int i = idx / batch_size; + int j = idx % batch_size; + float s = similarities_data[idx]; + + if (labels_data[i] == labels_data[j]) { + // 正样本对 + float ap = fmaxf(0.0f, -s + 1.0f + margin_val); + float logit_p = -ap * (s - delta_p) * gamma_val; + thread_sum_p += expf(logit_p - global_max_p); // 减去最大值 + } else { + // 负样本对 + float an = fmaxf(0.0f, s + margin_val); + float logit_n = an * (s - delta_n) * gamma_val; + thread_sum_n += expf(logit_n - global_max_n); // 减去最大值 + } + } + + // --- 块内归约 (Sum) - 正样本对 --- + s_data_p[threadIdx.x] = thread_sum_p; + __syncthreads(); + for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) { + if (threadIdx.x < offset) { + s_data_p[threadIdx.x] += s_data_p[threadIdx.x + offset]; + } + __syncthreads(); + } + + // --- 块内归约 (Sum) - 负样本对 --- + s_data_n[threadIdx.x] = thread_sum_n; + __syncthreads(); + for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) { + if (threadIdx.x < offset) { + s_data_n[threadIdx.x] += s_data_n[threadIdx.x + offset]; + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + block_sum_p_out[blockIdx.x] = s_data_p[0]; + block_sum_n_out[blockIdx.x] = s_data_n[0]; + } + } + + + // ------------------------------------------------------------------ + // C++ 封装函数 (现在执行两阶段逻辑) + // ------------------------------------------------------------------ + torch::Tensor circleloss_forward_cuda( + torch::Tensor similarities, + torch::Tensor labels, + float margin_val, + float gamma_val + ) { + // 检查 + TORCH_CHECK(similarities.is_cuda(), "Similarities tensor must be a CUDA tensor"); + TORCH_CHECK(labels.is_cuda(), "Labels tensor must be a CUDA tensor"); + + similarities = similarities.contiguous(); + labels = labels.contiguous(); + + TORCH_CHECK(labels.scalar_type() == torch::kInt64, "Labels tensor must be of type torch.long (int64_t)"); + + const int batch_size = labels.size(0); + const int n_elements = similarities.numel(); + + TORCH_CHECK(n_elements == batch_size * batch_size, "Similarities tensor has wrong size"); + + if (n_elements == 0) { + return torch::tensor(0.0f, similarities.options()); + } + + const int block_size = BLOCK_SIZE; + const int grid_size = std::max(1, (n_elements + block_size - 1) / block_size); + + // --- 阶段 1：运行 Find Max Kernel --- + auto block_max_p = torch::empty({grid_size}, similarities.options()); + auto block_max_n = torch::empty({grid_size}, similarities.options()); + + circleloss_find_max_kernel<<>>( + similarities.data_ptr(), + labels.data_ptr(), + block_max_p.data_ptr(), + block_max_n.data_ptr(), + n_elements, + batch_size, + margin_val, + gamma_val + ); + + // 在 C++ (GPU) 端找到全局最大值 + auto global_max_p_tensor = block_max_p.max(); + auto global_max_n_tensor = block_max_n.max(); + + // .item() 会导致 GPU -> CPU 同步，我们应尽量避免。 + // 但在这里我们需要这个值作为标量传递回下一个核函数。 + // 注意：一个更优的实现会使用 CUB 进行设备范围的归约， + // 但这对于 load_inline 来说太复杂了。 .max() 已经足够好了。 + const float global_max_p = global_max_p_tensor.item(); + const float global_max_n = global_max_n_tensor.item(); + + // --- 阶段 2：运行 Sum Exp Diff Kernel --- + auto block_sum_p = torch::empty({grid_size}, similarities.options()); + auto block_sum_n = torch::empty({grid_size}, similarities.options()); + + circleloss_sum_exp_diff_kernel<<>>( + similarities.data_ptr(), + labels.data_ptr(), + block_sum_p.data_ptr(), + block_sum_n.data_ptr(), + global_max_p, // 传递标量 + global_max_n, // 传递标量 + n_elements, + batch_size, + margin_val, + gamma_val + ); + + // --- 最终计算 (在 GPU 上) --- + + // 1. 对所有块的和进行求和 + auto global_sum_p = block_sum_p.sum(); + auto global_sum_n = block_sum_n.sum(); + + // 2. 稳定地计算 log(sum(exp(...))) + // logsumexp = max + log(sum(exp(x - max))) + auto log_sum_exp_p = global_max_p + torch::log(global_sum_p); + auto log_sum_exp_n = global_max_n + torch::log(global_sum_n); + + // 3. logsumexp_n + logsumexp_p + auto total_logit = log_sum_exp_p + log_sum_exp_n; + + // 4. 稳定的 F.softplus(x) = log(1 + exp(x)) + // 稳定的实现是: max(0, x) + log(1 + exp(-abs(x))) + auto zero_tensor = torch::tensor(0.0f, total_logit.options()); + auto max_val = torch::max(zero_tensor, total_logit); + auto loss = max_val + torch::log(1.0f + torch::exp(-torch::abs(total_logit))); + + return loss; + } + """ + + # JIT (Just-In-Time) 编译 + self.cl_op = load_inline( + name="circle_loss_op_v2_stable", + cpp_sources=cpp_source, + cuda_sources=cuda_source, + functions=["circleloss_forward_cuda"], + extra_cuda_cflags=["-O3"], + verbose=True # 设为 True 以便查看编译输出 + ) + + def forward(self, features: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + # 1. 执行优化的 matmul + # 假设输入的 features 已经是 L2 归一化的 + similarities = torch.matmul(features, features.t()) + + # 2. 调用我们编译好的、数值稳定的 CUDA C++ 函数 + return self.cl_op.circleloss_forward_cuda(similarities, labels, self.margin, self.gamma) \ No newline at end of file diff --git a/S1/gsd123_#9/circleloss_torch.py b/S1/gsd123_#9/circleloss_torch.py new file mode 100644 index 0000000..e97261f --- /dev/null +++ b/S1/gsd123_#9/circleloss_torch.py @@ -0,0 +1,60 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +# 定义常量 +BATCH_SIZE = 256 +FEATURE_DIM = 512 +MARGIN = 0.25 +GAMMA = 256 + + +class Model(nn.Module): + + def __init__(self): + super().__init__() + self.margin = MARGIN + self.gamma = GAMMA + + def forward(self, features: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + # Circle Loss的 PyTorch 实现 + # 假设 features 已经是 L2 归一化的 + # (B, D) @ (D, B) -> (B, B) + similarities = torch.matmul(features, features.t()) + + # 创建正负样本对的掩码 + mask_positive = labels.unsqueeze(1) == labels.unsqueeze(0) + mask_negative = labels.unsqueeze(1) != labels.unsqueeze(0) + + # 收集正样本对和负样本对的相似度 + # .masked_select() 会将张量展平 + sp = similarities[mask_positive] + sn = similarities[mask_negative] + + # 计算 Circle Loss 的 logits + # .detach() 用于停止梯度反向传播 + ap = torch.clamp_min(-sp.detach() + 1 + self.margin, min=0.) + an = torch.clamp_min(sn.detach() + self.margin, min=0.) + + delta_p = 1 - self.margin + delta_n = self.margin + + logit_p = -ap * (sp - delta_p) * self.gamma + logit_n = an * (sn - delta_n) * self.gamma + + # 使用 logsumexp 和 softplus 计算最终的 loss + # 这是 "unified" 版本的 loss + loss = F.softplus(torch.logsumexp(logit_n, dim=0) + torch.logsumexp(logit_p, dim=0)) + + return loss + + +def get_inputs(): + # 特征需要 L2 归一化 + features = F.normalize(torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32), p=2, dim=1) + labels = torch.randint(0, 10, (BATCH_SIZE,), dtype=torch.long) # 假设有 10 个类别 + return [features, labels] + + +def get_init_inputs(): + return [] \ No newline at end of file diff --git a/S1/gsd123_#9/prompt.txt b/S1/gsd123_#9/prompt.txt new file mode 100644 index 0000000..e12d5b2 --- /dev/null +++ b/S1/gsd123_#9/prompt.txt @@ -0,0 +1,158 @@ +You write custom CUDA kernels to replace the PyTorch operators in the given EvoNorm architecture to get speedups. +You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining normalization+affine_transform+nonlinear_gating), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination. + +Circle Loss CUDA Optimization with Two-Stage Reduction + +Core Algorithm +Loss Function: Implements Circle Loss for deep metric learning + +Key Formula: + +Positive pairs: logit_p = -α_p * (s_p - Δ_p) * γ + +Negative pairs: logit_n = α_n * (s_n - Δ_n) * γ + +Final loss: log(1 + exp(logsumexp_p + logsumexp_n)) + +CUDA Optimization Techniques +1. Two-Stage Parallel Reduction +Stage 1: Find global maximum of logits for numerical stability + +Kernel: circleloss_find_max_kernel + +Block-level max reduction with shared memory + +Grid-stride loops for load balancing + +Stage 2: Compute sum of exponentials with stability + +Kernel: circleloss_sum_exp_diff_kernel + +Subtract global max before exponentiation + +Block-level sum reduction + +2. Memory Access Optimization +Coalesced Memory Access: Sequential memory access patterns + +Shared Memory Utilization: + +s_data_p[BLOCK_SIZE] for positive pairs + +s_data_n[BLOCK_SIZE] for negative pairs + +Contiguous Tensors: Ensure input tensors are contiguous + +3. Numerical Stability Features +Log-Sum-Exp Trick: Subtract maximum before exponentiation + +Stable Softplus: max(0,x) + log(1 + exp(-abs(x))) + +Float Safety: Use FLT_MAX for initialization + +4. Parallelization Strategy +Grid-Stride Loops: Handle arbitrary input sizes + +Block Size: 256 threads per block for optimal occupancy + +Dynamic Grid Size: (n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE + +5. Kernel Design Patterns + +6. Performance Optimizations +Minimal CPU-GPU Synchronization: Avoid unnecessary .item() calls + +Batch Matrix Multiplication: Precompute similarity matrix + +Memory Pre-allocation: Pre-allocate output buffers + +7. Implementation Features +Template-Free Design: Optimized for float32 precision + +Error Checking: Comprehensive tensor validation + +Edge Case Handling: Empty tensor and size validation + +PyTorch Integration: Seamless tensor passing and gradient support + +Key CUDA Concepts Used +Shared Memory Reduction for parallel statistics computation + +Grid-Stride Loops for workload distribution + +Memory Coalescing for efficient global memory access + +Numerical Stability through careful floating-point handling + +Kernel Fusion combining multiple operations in single kernels + +Expected Performance Benefits +2-5x speedup over PyTorch implementation for large batch sizes + +Better numerical stability with log-sum-exp trick + +Scalable performance with increasing batch dimensions + +GPU utilization through optimized block and grid sizing + +This implementation provides a production-ready, numerically stable Circle Loss with significant performance improvements through careful CUDA optimization and parallel reduction patterns. +Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: +import torch +import torch.nn as nn +import torch.nn.functional as F + +# 定义常量 +BATCH_SIZE = 256 +FEATURE_DIM = 512 +MARGIN = 0.25 +GAMMA = 256 + +class Model(nn.Module): + + def __init__(self): + super().__init__() + self.margin = MARGIN + self.gamma = GAMMA + + def forward(self, features: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + + # Circle Loss的 PyTorch 实现 + # 假设 features 已经是 L2 归一化的 + # (B, D) @ (D, B) -> (B, B) + similarities = torch.matmul(features, features.t()) + + # 创建正负样本对的掩码 + mask_positive = labels.unsqueeze(1) == labels.unsqueeze(0) + mask_negative = labels.unsqueeze(1) != labels.unsqueeze(0) + + # 收集正样本对和负样本对的相似度 + # .masked_select() 会将张量展平 + sp = similarities[mask_positive] + sn = similarities[mask_negative] + + # 计算 Circle Loss 的 logits + # .detach() 用于停止梯度反向传播 + ap = torch.clamp_min(-sp.detach() + 1 + self.margin, min=0.) + an = torch.clamp_min(sn.detach() + self.margin, min=0.) + + delta_p = 1 - self.margin + delta_n = self.margin + + logit_p = -ap * (sp - delta_p) * self.gamma + logit_n = an * (sn - delta_n) * self.gamma + + # 使用 logsumexp 和 softplus 计算最终的 loss + # 这是 "unified" 版本的 loss + loss = F.softplus(torch.logsumexp(logit_n, dim=0) + torch.logsumexp(logit_p, dim=0)) + + return loss + +def get_inputs(): + + # 特征需要 L2 归一化 + features = F.normalize(torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32), p=2, dim=1) + labels = torch.randint(0, 10, (BATCH_SIZE,), dtype=torch.long) # 假设有 10 个类别 + return [features, labels] + +def get_init_inputs(): + return [] \ No newline at end of file diff --git a/S1/gsd123_#9/run_code.py b/S1/gsd123_#9/run_code.py new file mode 100644 index 0000000..6601ec4 --- /dev/null +++ b/S1/gsd123_#9/run_code.py @@ -0,0 +1,78 @@ +import torch +import time +from circleloss_torch import Model, get_inputs, get_init_inputs +from circleloss_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