diff --git a/S1/hli28146_#58/focaltverskyloss_cuda.py b/S1/hli28146_#58/focaltverskyloss_cuda.py new file mode 100644 index 00000000..74f873ae --- /dev/null +++ b/S1/hli28146_#58/focaltverskyloss_cuda.py @@ -0,0 +1,208 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + +cpp_source = """ +#include + +torch::Tensor focal_tversky_loss_cuda_forward( + const torch::Tensor& y_pred, + const torch::Tensor& y_true, + float alpha, + float beta, + float gamma, + float epsilon); +""" + +cuda_source = """ +#include +#include +#include + +#define BLOCK_SIZE 256 +#define WARP_SIZE 32 + +struct __align__(16) Float4 { + float x, y, z, w; +}; + +template +__device__ __forceinline__ T warp_reduce_sum(T val) { + #pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + return val; +} + +// Reduce 3 values simultaneously +__device__ __forceinline__ void block_reduce_sum_3( + float& val1, float& val2, float& val3) +{ + static __shared__ float shared1[32]; + static __shared__ float shared2[32]; + static __shared__ float shared3[32]; + + int lane = threadIdx.x % WARP_SIZE; + int wid = threadIdx.x / WARP_SIZE; + + val1 = warp_reduce_sum(val1); + val2 = warp_reduce_sum(val2); + val3 = warp_reduce_sum(val3); + + if (lane == 0) { + shared1[wid] = val1; + shared2[wid] = val2; + shared3[wid] = val3; + } + __syncthreads(); + + // Last warp reduction (assuming block size 256) + val1 = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared1[lane] : 0.0f; + val2 = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared2[lane] : 0.0f; + val3 = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared3[lane] : 0.0f; + + if (wid == 0) { + val1 = warp_reduce_sum(val1); + val2 = warp_reduce_sum(val2); + val3 = warp_reduce_sum(val3); + } +} + +// Sigmoid +__device__ __forceinline__ float fast_sigmoid(float x) { + return 1.0f / (1.0f + __expf(-x)); +} + +// Fused Kernel +__global__ void focal_tversky_kernel( + float* __restrict__ output, + const float* __restrict__ logits, + const float* __restrict__ targets, + int spatial_size, + int num_blocks, + float alpha, + float beta, + float gamma, + float epsilon) +{ + int slice_idx = blockIdx.x; + if (slice_idx >= num_blocks) return; + + int offset = slice_idx * spatial_size; + const float* slice_logits = logits + offset; + const float* slice_targets = targets + offset; + + float tp_sum = 0.0f; + float fn_sum = 0.0f; + float fp_sum = 0.0f; + + int tid = threadIdx.x; + int stride = blockDim.x; + + int i = tid * 4; + while (i < spatial_size) { + if (i + 4 <= spatial_size) { + Float4 l_vec = reinterpret_cast(&slice_logits[i])[0]; + Float4 t_vec = reinterpret_cast(&slice_targets[i])[0]; + + // Unroll manually + float p, g; + + p = fast_sigmoid(l_vec.x); g = t_vec.x; + tp_sum += p * g; fn_sum += g * (1.0f - p); fp_sum += p * (1.0f - g); + + p = fast_sigmoid(l_vec.y); g = t_vec.y; + tp_sum += p * g; fn_sum += g * (1.0f - p); fp_sum += p * (1.0f - g); + + p = fast_sigmoid(l_vec.z); g = t_vec.z; + tp_sum += p * g; fn_sum += g * (1.0f - p); fp_sum += p * (1.0f - g); + + p = fast_sigmoid(l_vec.w); g = t_vec.w; + tp_sum += p * g; fn_sum += g * (1.0f - p); fp_sum += p * (1.0f - g); + + } else { + for (int k = 0; k < 4 && i + k < spatial_size; ++k) { + float p = fast_sigmoid(slice_logits[i+k]); + float g = slice_targets[i+k]; + tp_sum += p * g; + fn_sum += g * (1.0f - p); + fp_sum += (1.0f - g) * p; + } + } + i += stride * 4; + } + + block_reduce_sum_3(tp_sum, fn_sum, fp_sum); + + if (tid == 0) { + float denominator = tp_sum + alpha * fn_sum + beta * fp_sum + epsilon; + float ti = (tp_sum + epsilon) / denominator; + + float loss = __powf(1.0f - ti, gamma); + output[slice_idx] = loss; + } +} + +torch::Tensor focal_tversky_loss_cuda_forward( + const torch::Tensor& logits, + const torch::Tensor& targets, + float alpha, + float beta, + float gamma, + float epsilon) +{ + TORCH_CHECK(logits.is_cuda() && targets.is_cuda(), "Inputs must be CUDA"); + TORCH_CHECK(logits.is_contiguous() && targets.is_contiguous(), "Inputs must be contiguous"); + + int batch_size = logits.size(0); + int channels = logits.size(1); + + int64_t total_elements = logits.numel(); + int64_t num_slices = batch_size * channels; + int spatial_size = total_elements / num_slices; + + auto output = torch::empty({num_slices}, logits.options()); + + focal_tversky_kernel<<>>( + output.data_ptr(), + logits.data_ptr(), + targets.data_ptr(), + spatial_size, + num_slices, + alpha, + beta, + gamma, + epsilon + ); + + return output.mean(); +} +""" + +class ModelNew(nn.Module): + def __init__(self, alpha=0.7, beta=0.3, gamma=0.75, epsilon=1e-7): + super(ModelNew, self).__init__() + self.alpha = alpha + self.beta = beta + self.gamma = gamma + self.epsilon = epsilon + + self.op = load_inline( + name='focal_tversky_op', + cpp_sources=cpp_source, + cuda_sources=cuda_source, + functions=['focal_tversky_loss_cuda_forward'], + verbose=False, + extra_cuda_cflags=['-O3'] + ) + + def forward(self, y_pred, y_true): + return self.op.focal_tversky_loss_cuda_forward( + y_pred.contiguous(), + y_true.contiguous(), + self.alpha, + self.beta, + self.gamma, + self.epsilon + ) \ No newline at end of file diff --git a/S1/hli28146_#58/focaltverskyloss_torch.py b/S1/hli28146_#58/focaltverskyloss_torch.py new file mode 100644 index 00000000..92f483c4 --- /dev/null +++ b/S1/hli28146_#58/focaltverskyloss_torch.py @@ -0,0 +1,62 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +BATCH_SIZE = 16 +CHANNELS = 4 +DEPTH = 32 +HEIGHT = 128 +WIDTH = 128 +SHAPE = (BATCH_SIZE, CHANNELS, DEPTH, HEIGHT, WIDTH) + +ALPHA = 0.7 +BETA = 0.3 +GAMMA = 0.75 +EPSILON = 1e-7 + +class FocalTverskyLoss(nn.Module): + ''' + FOCAL TVERSKY LOSS(https://arxiv.org/pdf/1810.07842) + The input tensors are expected to have a shape of (B, N, H, W, L), where: + B is the batch size + N is the number of channels + H, W, L represent the depth, height, and width of the volumes, respectively + ''' + def __init__(self, alpha=0.7, beta=0.3, gamma=0.75, epsilon=1e-7): + super(FocalTverskyLoss, self).__init__() + self.alpha = alpha + self.beta = beta + self.gamma = gamma + self.epsilon = epsilon + + def forward(self, y_pred, y_true): + # y_pred: logits + # y_true: binary targets (0 or 1) + y_pred = torch.sigmoid(y_pred) + + # Reduction over spatial dims (2, 3, 4) + tp = (y_true * y_pred).sum(dim=(2, 3, 4)) + fn = (y_true * (1 - y_pred)).sum(dim=(2, 3, 4)) + fp = ((1 - y_true) * y_pred).sum(dim=(2, 3, 4)) + + tversky_index = (tp + self.epsilon) / (tp + self.alpha * fn + self.beta * fp + self.epsilon) + + loss = (1 - tversky_index).pow(self.gamma) + + return loss.mean() + +class Model(nn.Module): + def __init__(self, alpha=0.7, beta=0.3, gamma=0.75, epsilon=1e-7): + super(Model, self).__init__() + self.loss_fn = FocalTverskyLoss(alpha, beta, gamma, epsilon) + + def forward(self, y_pred, y_true): + return self.loss_fn(y_pred, y_true) + +def get_inputs(): + y_pred = torch.randn(SHAPE, dtype=torch.float32) + y_true = torch.randint(0, 2, SHAPE, dtype=torch.float32) + return [y_pred.contiguous(), y_true.contiguous()] + +def get_init_inputs(): + return [ALPHA, BETA, GAMMA, EPSILON] \ No newline at end of file diff --git a/S1/hli28146_#58/prompt.txt b/S1/hli28146_#58/prompt.txt new file mode 100644 index 00000000..e957ab23 --- /dev/null +++ b/S1/hli28146_#58/prompt.txt @@ -0,0 +1,92 @@ +Write a custom CUDA kernel to optimize `Focal Tversky Loss` for 3D segmentation. + +Formula: +TI_c = (TP_c + epsilon) / (TP_c + alpha * FN_c + beta * FP_c + epsilon) +Loss = Mean( (1 - TI_c)^gamma ) + +Parameters: +- alpha, beta, gamma: Loss weighting parameters. +- epsilon: Smoothing factor for numerical stability. + +Optimization Strategy: Fused Block-per-Channel Reduction + +1. Flattened View: Treat the input (B, C, D, H, W) as `B * C` independent slices. + +2. Block-per-Slice Parallelism: Launch `B * C` CUDA blocks. Each block reduces one spatial slice to compute TP, FN, and FP simultaneously. + +3. Fused Accumulation: + - Load logit and target using vectorized `float4`. + - Compute `p = sigmoid(logit)` in register. + - Accumulate TP, FN, FP in registers. + +4. Shared Memory Reduction: Perform parallel reduction for the three accumulators. + +5. Final Calculation: Thread 0 uses the reduced TP, FN, FP and the passed `epsilon` to calculate the loss for the slice. + +6. Global Reduction: The C++ wrapper performs the final mean reduction. + +Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: + +```python +import torch +import torch.nn as nn +import torch.nn.functional as F + +BATCH_SIZE = 16 +CHANNELS = 4 +DEPTH = 32 +HEIGHT = 128 +WIDTH = 128 +SHAPE = (BATCH_SIZE, CHANNELS, DEPTH, HEIGHT, WIDTH) + +ALPHA = 0.7 +BETA = 0.3 +GAMMA = 0.75 +EPSILON = 1e-7 + +class FocalTverskyLoss(nn.Module): + ''' + FOCAL TVERSKY LOSS(https://arxiv.org/pdf/1810.07842) + The input tensors are expected to have a shape of (B, N, H, W, L), where: + B is the batch size + N is the number of channels + H, W, L represent the depth, height, and width of the volumes, respectively + ''' + def __init__(self, alpha=0.7, beta=0.3, gamma=0.75, epsilon=1e-7): + super(FocalTverskyLoss, self).__init__() + self.alpha = alpha + self.beta = beta + self.gamma = gamma + self.epsilon = epsilon + + def forward(self, y_pred, y_true): + # y_pred: logits + # y_true: binary targets (0 or 1) + y_pred = torch.sigmoid(y_pred) + + # Reduction over spatial dims (2, 3, 4) + tp = (y_true * y_pred).sum(dim=(2, 3, 4)) + fn = (y_true * (1 - y_pred)).sum(dim=(2, 3, 4)) + fp = ((1 - y_true) * y_pred).sum(dim=(2, 3, 4)) + + tversky_index = (tp + self.epsilon) / (tp + self.alpha * fn + self.beta * fp + self.epsilon) + + loss = (1 - tversky_index).pow(self.gamma) + + return loss.mean() + +class Model(nn.Module): + def __init__(self, alpha=0.7, beta=0.3, gamma=0.75, epsilon=1e-7): + super(Model, self).__init__() + self.loss_fn = FocalTverskyLoss(alpha, beta, gamma, epsilon) + + def forward(self, y_pred, y_true): + return self.loss_fn(y_pred, y_true) + +def get_inputs(): + y_pred = torch.randn(SHAPE, dtype=torch.float32) + y_true = torch.randint(0, 2, SHAPE, dtype=torch.float32) + return [y_pred.contiguous(), y_true.contiguous()] + +def get_init_inputs(): + return [ALPHA, BETA, GAMMA, EPSILON] \ No newline at end of file diff --git a/S1/hli28146_#58/run_code.py b/S1/hli28146_#58/run_code.py new file mode 100644 index 00000000..3c414060 --- /dev/null +++ b/S1/hli28146_#58/run_code.py @@ -0,0 +1,74 @@ +########################################################### +# 性能和精度验证程序 +########################################################### +import torch +import torch.nn as nn +import time +from focaltverskyloss_torch import Model,get_inputs,get_init_inputs +from focaltverskyloss_cuda import ModelNew + +def run_benchmark(): + # 检查 CUDA 是否可用 + if not torch.cuda.is_available(): + print("CUDA 不可用,请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。") + return + else: + device = torch.device("cuda") + + # 初始化模型 + init_inputs = get_init_inputs() + init_inputs = [ + x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs + ] + inputs = get_inputs() + inputs = [ + x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs + ] + + torch_model = Model(*init_inputs).cuda() + cuda_model = ModelNew(*init_inputs).cuda() + + torch_model.eval() + cuda_model.eval() + + print("-------------------- 精度对齐验证 --------------------") + with torch.no_grad(): + output_torch = torch_model( *inputs) + output_cuda = cuda_model(*inputs) + + precision_flag = torch.allclose(output_torch, output_cuda,rtol=1e-03) + if precision_flag: + print("✅ 精度对齐:两个模型的输出结果非常接近。") + else: + print("❌ 精度不一致!") + + print("\n-------------------- 性能加速比测试 --------------------") + num_iterations = 100 + + # 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 torch.relu 平均执行时间: {torch_time:.6f} 秒") + print(f"自定义 CUDA 内核 平均执行时间: {cuda_time:.6f} 秒") + speedup = 0 + if cuda_time > 0: + speedup = torch_time / cuda_time + print(f"加速比 (Speedup): {speedup:.2f}x") + else: + print("CUDA 内核执行时间为0,无法计算加速比。") + return precision_flag,speedup +if __name__ == "__main__": + precision_flag,speedup = run_benchmark() \ No newline at end of file