forked from ccf-ai-infra/GPUCodeForces
FEAT:ADD dice_from_2d #108
This commit is contained in:
parent
f989885dde
commit
193f39bbee
|
|
@ -0,0 +1,112 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
dice_from_2d_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// Warp级归约版本 - 更高效的并行归约，并融合了Flatten操作
|
||||
__global__ void dice_from_2d_kernel_warp_reduction(
|
||||
const float* __restrict__ pred,
|
||||
const float* __restrict__ target,
|
||||
float* __restrict__ intersection,
|
||||
float* __restrict__ pred_sum,
|
||||
float* __restrict__ target_sum,
|
||||
int size
|
||||
) {
|
||||
int tid = threadIdx.x;
|
||||
int warp_id = (blockIdx.x * blockDim.x + tid) / 32;
|
||||
int lane_id = tid % 32;
|
||||
|
||||
// 每个warp处理一块数据
|
||||
int elements_per_warp = 1024; // 每个warp处理1K元素
|
||||
int warp_start = warp_id * elements_per_warp;
|
||||
int warp_end = min(warp_start + elements_per_warp, size);
|
||||
|
||||
// 寄存器累加
|
||||
float reg_intersection = 0.0f;
|
||||
float reg_pred_sum = 0.0f;
|
||||
float reg_target_sum = 0.0f;
|
||||
|
||||
// 处理数据 - 这里的size是总元素数，实现了隐式flatten
|
||||
for (int i = warp_start + lane_id; i < warp_end; i += 32) {
|
||||
float p = pred[i];
|
||||
float t = target[i];
|
||||
reg_intersection += p * t;
|
||||
reg_pred_sum += p;
|
||||
reg_target_sum += t;
|
||||
}
|
||||
|
||||
// Warp内归约 - 使用warp shuffle指令
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset /= 2) {
|
||||
reg_intersection += __shfl_down_sync(0xffffffff, reg_intersection, offset);
|
||||
reg_pred_sum += __shfl_down_sync(0xffffffff, reg_pred_sum, offset);
|
||||
reg_target_sum += __shfl_down_sync(0xffffffff, reg_target_sum, offset);
|
||||
}
|
||||
|
||||
// warp内第一个线程写入结果
|
||||
if (lane_id == 0) {
|
||||
atomicAdd(intersection, reg_intersection);
|
||||
atomicAdd(pred_sum, reg_pred_sum);
|
||||
atomicAdd(target_sum, reg_target_sum);
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor dice_from_2d_cuda(torch::Tensor pred, torch::Tensor target) {
|
||||
// --- 关键修改：直接计算总元素数，实现隐式flatten ---
|
||||
auto size = pred.numel();
|
||||
|
||||
// 创建中间结果张量
|
||||
auto intersection = torch::zeros(1, pred.options());
|
||||
auto pred_sum_tensor = torch::zeros(1, pred.options());
|
||||
auto target_sum_tensor = torch::zeros(1, pred.options());
|
||||
|
||||
// 使用Warp级归约版本
|
||||
const int block_size = 256;
|
||||
int num_warps_needed = (size + 1024 - 1) / 1024;
|
||||
int num_blocks = (num_warps_needed * 32 + block_size - 1) / block_size;
|
||||
|
||||
dice_from_2d_kernel_warp_reduction<<<num_blocks, block_size>>>(
|
||||
pred.data_ptr<float>(),
|
||||
target.data_ptr<float>(),
|
||||
intersection.data_ptr<float>(),
|
||||
pred_sum_tensor.data_ptr<float>(),
|
||||
target_sum_tensor.data_ptr<float>(),
|
||||
size
|
||||
);
|
||||
|
||||
// 计算最终DiceLoss
|
||||
float epsilon = 1e-6f;
|
||||
auto dice_score = (2.0f * intersection) / (pred_sum_tensor + target_sum_tensor + epsilon);
|
||||
auto dice_loss = 1.0f - dice_score;
|
||||
|
||||
return dice_loss;
|
||||
}
|
||||
"""
|
||||
|
||||
dice_from_2d_cpp_source = """
|
||||
torch::Tensor dice_from_2d_cuda(torch::Tensor pred, torch::Tensor target);
|
||||
"""
|
||||
|
||||
# 编译CUDA代码
|
||||
dice_from_2d = load_inline(
|
||||
name="dice_from_2d",
|
||||
cpp_sources=dice_from_2d_cpp_source,
|
||||
cuda_sources=dice_from_2d_source,
|
||||
functions=["dice_from_2d_cuda"],
|
||||
extra_cuda_cflags=[
|
||||
"-O3",
|
||||
"--use_fast_math",
|
||||
"-gencode=arch=compute_80,code=sm_80"
|
||||
],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super(ModelNew, self).__init__()
|
||||
self.dice_from_2d = dice_from_2d
|
||||
|
||||
def forward(self, pred, target):
|
||||
return self.dice_from_2d.dice_from_2d_cuda(pred, target)
|
||||
|
|
@ -0,0 +1,47 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
PyTorch基准实现：直接从2D输入计算Dice Loss
|
||||
"""
|
||||
def __init__(self):
|
||||
super(Model, self).__init__()
|
||||
|
||||
def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Applies Dice Loss to the 2D prediction and target tensors.
|
||||
|
||||
Args:
|
||||
pred (torch.Tensor): Prediction tensor of shape [N, C, H, W].
|
||||
target (torch.Tensor): Target tensor of same shape as pred.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Dice Loss value (scalar).
|
||||
"""
|
||||
# --- 第一步：在内部展平张量 ---
|
||||
pred_flat = pred.view(-1)
|
||||
target_flat = target.view(-1)
|
||||
|
||||
# --- 第二步：计算Dice Loss ---
|
||||
intersection = (pred_flat * target_flat).sum()
|
||||
pred_sum = pred_flat.sum()
|
||||
target_sum = target_flat.sum()
|
||||
|
||||
epsilon = 1e-6
|
||||
dice_score = (2.0 * intersection) / (pred_sum + target_sum + epsilon)
|
||||
dice_loss = 1.0 - dice_score
|
||||
|
||||
return dice_loss
|
||||
|
||||
batch_size = 32
|
||||
height, width = 256, 256
|
||||
channels = 1
|
||||
|
||||
def get_inputs():
|
||||
pred = torch.rand(batch_size, channels, height, width)
|
||||
target = torch.randint(0, 2, (batch_size, channels, height, width), dtype=torch.float32)
|
||||
return [pred, target]
|
||||
|
||||
def get_init_inputs():
|
||||
return [] # No special initialization inputs needed
|
||||
|
|
@ -0,0 +1,146 @@
|
|||
You write custom CUDA kernels to replace pytorch operators in given 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 matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.
|
||||
|
||||
Here's an example to show you the syntax of inline embedding custom CUDA operators in PyTorch. The example given architecture is a simple addition:
|
||||
|
||||
python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
def init(self) -> None:
|
||||
super().init()
|
||||
|
||||
def forward(self, a, b):
|
||||
return a + b
|
||||
def get_inputs():
|
||||
a = torch.randn(1, 128).cuda()
|
||||
b = torch.randn(1, 128).cuda()
|
||||
return [a, b]
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
||||
|
||||
|
||||
The example new architecture with a custom CUDA kernel looks like this:
|
||||
|
||||
python
|
||||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
add_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
global void add_kernel(const float* a, const float* b, float* out, int size) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < size) {
|
||||
out[idx] = a[idx] + b[idx];
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b) {
|
||||
auto out = torch::empty_like(a);
|
||||
int size = a.numel();
|
||||
const int block_size = 256;
|
||||
int num_blocks = (size + block_size - 1) / block_size;
|
||||
add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);
|
||||
return out;
|
||||
}
|
||||
"""
|
||||
|
||||
add_cpp_source = """
|
||||
torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b);
|
||||
"""
|
||||
|
||||
Compile the inline CUDA code
|
||||
add = load_inline(
|
||||
name="add",
|
||||
cpp_sources=add_cpp_source,
|
||||
cuda_sources=add_source,
|
||||
functions=["add_cuda"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def init(self):
|
||||
super(ModelNew, self).init()
|
||||
self.add = add
|
||||
|
||||
def forward(self, a, b):
|
||||
return self.add.add_cuda(a, b)
|
||||
|
||||
|
||||
---
|
||||
|
||||
Now, you are given the following PyTorch architecture to accelerate. The model computes the Dice Loss directly from 2D input tensors (e.g., from a convolutional layer) by first flattening them and then performing the loss calculation. This baseline implementation is efficient and uses PyTorch's highly optimized built-in functions for correctness and performance.
|
||||
|
||||
python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
PyTorch基准实现：直接从2D输入计算Dice Loss
|
||||
"""
|
||||
def __init__(self):
|
||||
super(Model, self).__init__()
|
||||
|
||||
def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Applies Dice Loss to the 2D prediction and target tensors.
|
||||
|
||||
Args:
|
||||
pred (torch.Tensor): Prediction tensor of shape [N, C, H, W].
|
||||
target (torch.Tensor): Target tensor of same shape as pred.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Dice Loss value (scalar).
|
||||
"""
|
||||
# --- 第一步：在内部展平张量 ---
|
||||
pred_flat = pred.view(-1)
|
||||
target_flat = target.view(-1)
|
||||
|
||||
# --- 第二步：计算Dice Loss ---
|
||||
intersection = (pred_flat * target_flat).sum()
|
||||
pred_sum = pred_flat.sum()
|
||||
target_sum = target_flat.sum()
|
||||
|
||||
epsilon = 1e-6
|
||||
dice_score = (2.0 * intersection) / (pred_sum + target_sum + epsilon)
|
||||
dice_loss = 1.0 - dice_score
|
||||
|
||||
return dice_loss
|
||||
|
||||
batch_size = 32
|
||||
height, width = 256, 256
|
||||
channels = 1
|
||||
|
||||
def get_inputs():
|
||||
pred = torch.rand(batch_size, channels, height, width)
|
||||
target = torch.randint(0, 2, (batch_size, channels, height, width), dtype=torch.float32)
|
||||
return [pred, target]
|
||||
|
||||
def get_init_inputs():
|
||||
return [] # No special initialization inputs needed
|
||||
|
||||
|
||||
|
||||
Your task is to generate the `ModelNew` architecture with a custom CUDA kernel that fuses the `flatten` operation with the Dice Loss calculation. The implementation must be highly optimized.
|
||||
|
||||
**CRITICAL REQUIREMENTS:**
|
||||
|
||||
1. **Performance Optimization & Fusion:**
|
||||
* **Operator Fusion:** The entire calculation (flattening the input tensors and then computing the Dice Loss) must be performed within a **single CUDA kernel**. The kernel should take the multi-dimensional tensors as input and directly output the scalar loss value.
|
||||
* The kernel should use a **Warp-level reduction** strategy for optimal performance. Warps should collaboratively iterate over the total number of elements in the tensors, effectively performing an "implicit flatten" within the kernel.
|
||||
|
||||
2. **Kernel Logic:**
|
||||
* The kernel should launch a grid of blocks where the total number of threads is sufficient to cover all elements in the input tensors (`pred.numel()`).
|
||||
* Each thread should process one element of the flattened tensors.
|
||||
* Use `__shfl_down_sync` for efficient warp-level reduction to compute the three required sums: `intersection`, `pred_sum`, and `target_sum`.
|
||||
* After the kernel completes, the host-side C++ wrapper function should perform the final scalar arithmetic to compute the Dice Loss from the three sums.
|
||||
|
||||
3. **Code Structure:** Follow the exact structure of the provided example, including `load_inline`, the CUDA source string, the C++ wrapper source string, and the `ModelNew` class. The `get_init_inputs` function must return `[]` to match the baseline.
|
||||
|
||||
4. **Compilation Flags:** Use `-O3` and `--use_fast_math` for maximum performance, as this is a common practice for high-throughput kernels like this. Avoid hardcoding compute capabilities to ensure portability.
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from dice_from_2d_torchcode import Model,get_inputs,get_init_inputs
|
||||
from dice_from_2d_cudacode 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 dice_from_2dLoss 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA dice_from_2dLoss 平均执行时间: {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()
|
||||
Loading…
Reference in New Issue