forked from ccf-ai-infra/GPUCodeForces
feat:add fused L1loss #45
This commit is contained in:
parent
f989885dde
commit
eab7de2cd5
|
|
@ -0,0 +1,221 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
l1_fused_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// 融合L1Loss内核 - 单步完成所有计算
|
||||
__global__ void l1_fused_kernel(
|
||||
const float* __restrict__ input,
|
||||
const float* __restrict__ target,
|
||||
float* __restrict__ loss_output,
|
||||
int size
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
// 局部累加器
|
||||
float local_l1 = 0.0f;
|
||||
|
||||
// 每个线程处理多个元素 - 融合diff+abs计算
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
for (int i = idx; i < size; i += stride) {
|
||||
float diff = input[i] - target[i]; // diff计算
|
||||
local_l1 += fabsf(diff); // abs计算 + 累加
|
||||
}
|
||||
|
||||
// 使用共享内存进行块内归约 - 融合sum计算
|
||||
extern __shared__ float shared_mem[];
|
||||
shared_mem[threadIdx.x] = local_l1;
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// 块内归约
|
||||
for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
|
||||
if (threadIdx.x < stride) {
|
||||
shared_mem[threadIdx.x] += shared_mem[threadIdx.x + stride];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// 每个block写入部分和到全局内存
|
||||
if (threadIdx.x == 0) {
|
||||
atomicAdd(loss_output, shared_mem[0]);
|
||||
}
|
||||
}
|
||||
|
||||
// 高级融合版本 - 多种优化策略
|
||||
__global__ void l1_fused_advanced_kernel(
|
||||
const float* __restrict__ input,
|
||||
const float* __restrict__ target,
|
||||
float* __restrict__ loss_output,
|
||||
int size,
|
||||
bool use_warp_reduce
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int warp_id = threadIdx.x / 32;
|
||||
int lane_id = threadIdx.x % 32;
|
||||
|
||||
// 局部累加器
|
||||
float local_l1 = 0.0f;
|
||||
|
||||
// 每个线程处理多个元素
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
for (int i = idx; i < size; i += stride) {
|
||||
float diff = input[i] - target[i];
|
||||
local_l1 += fabsf(diff);
|
||||
}
|
||||
|
||||
if (use_warp_reduce) {
|
||||
// 使用warp级归约 - 更高效的融合
|
||||
for (int offset = 16; offset > 0; offset /= 2) {
|
||||
local_l1 += __shfl_down_sync(0xffffffff, local_l1, offset);
|
||||
}
|
||||
|
||||
// 每个warp的第一个线程做原子操作
|
||||
if (lane_id == 0) {
|
||||
atomicAdd(loss_output, local_l1);
|
||||
}
|
||||
} else {
|
||||
// 使用共享内存归约
|
||||
extern __shared__ float shared_mem[];
|
||||
shared_mem[threadIdx.x] = local_l1;
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// 块内归约
|
||||
for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
|
||||
if (threadIdx.x < stride) {
|
||||
shared_mem[threadIdx.x] += shared_mem[threadIdx.x + stride];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
atomicAdd(loss_output, shared_mem[0]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 向量化融合版本 - 处理多个元素
|
||||
__global__ void l1_fused_vectorized_kernel(
|
||||
const float* __restrict__ input,
|
||||
const float* __restrict__ target,
|
||||
float* __restrict__ loss_output,
|
||||
int size
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
// 向量化处理 - 每次处理4个元素
|
||||
int vectorized_size = (size / 4) * 4;
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
|
||||
float local_l1 = 0.0f;
|
||||
|
||||
// 向量化主循环
|
||||
for (int i = idx * 4; i < vectorized_size; i += stride * 4) {
|
||||
// 加载4个元素
|
||||
float4 input_vec = *reinterpret_cast<const float4*>(&input[i]);
|
||||
float4 target_vec = *reinterpret_cast<const float4*>(&target[i]);
|
||||
|
||||
// 计算L1距离
|
||||
local_l1 += fabsf(input_vec.x - target_vec.x);
|
||||
local_l1 += fabsf(input_vec.y - target_vec.y);
|
||||
local_l1 += fabsf(input_vec.z - target_vec.z);
|
||||
local_l1 += fabsf(input_vec.w - target_vec.w);
|
||||
}
|
||||
|
||||
// 处理剩余元素
|
||||
for (int i = idx + vectorized_size; i < size; i += stride) {
|
||||
float diff = input[i] - target[i];
|
||||
local_l1 += fabsf(diff);
|
||||
}
|
||||
|
||||
// Warp级归约
|
||||
for (int offset = 16; offset > 0; offset /= 2) {
|
||||
local_l1 += __shfl_down_sync(0xffffffff, local_l1, offset);
|
||||
}
|
||||
|
||||
if (threadIdx.x % 32 == 0) {
|
||||
atomicAdd(loss_output, local_l1);
|
||||
}
|
||||
}
|
||||
|
||||
// 主函数：融合L1Loss
|
||||
torch::Tensor l1_fused_cuda(torch::Tensor input, torch::Tensor target, std::string mode = "fused") {
|
||||
TORCH_CHECK(input.scalar_type() == torch::kFloat32, "Input must be float32");
|
||||
TORCH_CHECK(target.scalar_type() == torch::kFloat32, "Target must be float32");
|
||||
TORCH_CHECK(input.sizes() == target.sizes(), "Input and target must have same shape");
|
||||
|
||||
auto input_contig = input.contiguous();
|
||||
auto target_contig = target.contiguous();
|
||||
int size = input_contig.numel();
|
||||
|
||||
// 创建输出张量
|
||||
auto l1_loss = torch::zeros(1, input.options());
|
||||
|
||||
const int block_size = 256;
|
||||
int num_blocks = min(1024, (size + block_size - 1) / block_size);
|
||||
|
||||
if (mode == "fused") {
|
||||
// 基础融合版本
|
||||
size_t shared_mem = block_size * sizeof(float);
|
||||
l1_fused_kernel<<<num_blocks, block_size, shared_mem>>>(
|
||||
input_contig.data_ptr<float>(),
|
||||
target_contig.data_ptr<float>(),
|
||||
l1_loss.data_ptr<float>(),
|
||||
size
|
||||
);
|
||||
} else if (mode == "advanced") {
|
||||
// 高级融合版本 - warp优化
|
||||
l1_fused_advanced_kernel<<<num_blocks, block_size>>>(
|
||||
input_contig.data_ptr<float>(),
|
||||
target_contig.data_ptr<float>(),
|
||||
l1_loss.data_ptr<float>(),
|
||||
size,
|
||||
true
|
||||
);
|
||||
} else if (mode == "vectorized") {
|
||||
// 向量化融合版本
|
||||
// 确保size是4的倍数
|
||||
int aligned_size = ((size + 3) / 4) * 4;
|
||||
int vectorized_num_blocks = min(1024, (aligned_size / 4 + block_size - 1) / block_size);
|
||||
|
||||
l1_fused_vectorized_kernel<<<vectorized_num_blocks, block_size>>>(
|
||||
input_contig.data_ptr<float>(),
|
||||
target_contig.data_ptr<float>(),
|
||||
l1_loss.data_ptr<float>(),
|
||||
size
|
||||
);
|
||||
}
|
||||
|
||||
return l1_loss;
|
||||
}
|
||||
"""
|
||||
|
||||
l1_fused_cpp_source = """
|
||||
torch::Tensor l1_fused_cuda(torch::Tensor input, torch::Tensor target, std::string mode);
|
||||
"""
|
||||
|
||||
# 编译融合CUDA代码
|
||||
l1_fused = load_inline(
|
||||
name="l1_fused",
|
||||
cpp_sources=l1_fused_cpp_source,
|
||||
cuda_sources=l1_fused_source,
|
||||
functions=["l1_fused_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, mode="fused"):
|
||||
super(ModelNew, self).__init__()
|
||||
self.mode = mode
|
||||
self.l1_fused = l1_fused
|
||||
|
||||
def forward(self, input, target):
|
||||
return self.l1_fused.l1_fused_cuda(input, target, self.mode)
|
||||
|
|
@ -0,0 +1,43 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
合理优化的PyTorch L1 Loss实现
|
||||
使用PyTorch内置函数，避免不必要的中间张量创建
|
||||
"""
|
||||
def __init__(self):
|
||||
super(Model, self).__init__()
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
使用PyTorch内置的l1_loss函数
|
||||
L1 Loss = |input - target|
|
||||
|
||||
Args:
|
||||
input (torch.Tensor): 预测值
|
||||
target (torch.Tensor): 真实值
|
||||
|
||||
Returns:
|
||||
torch.Tensor: L1 Loss标量值
|
||||
"""
|
||||
# 直接使用内置函数，让PyTorch处理优化
|
||||
return torch.nn.functional.l1_loss(
|
||||
input,
|
||||
target,
|
||||
reduction='sum'
|
||||
)
|
||||
|
||||
batch_size = 512
|
||||
num_features = 2000
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
生成合理的测试数据
|
||||
"""
|
||||
input_vals = torch.randn(batch_size, num_features)
|
||||
target_vals = torch.randn(batch_size, num_features)
|
||||
return [input_vals, target_vals]
|
||||
|
||||
def get_init_inputs():
|
||||
return [] # 没有特殊的初始化输入需求
|
||||
|
|
@ -0,0 +1,128 @@
|
|||
You write custom CUDA kernels to replace the pytorch operators in the 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 torch: The example given architecture is:
|
||||
|
||||
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():
|
||||
# randomly generate input tensors based on the model architecture
|
||||
a = torch.randn(1, 128).cuda()
|
||||
b = torch.randn(1, 128).cuda()
|
||||
return [a, b]
|
||||
|
||||
def get_init_inputs():
|
||||
# randomly generate tensors required for initialization based on the model architecture
|
||||
return []
|
||||
|
||||
|
||||
|
||||
The example new arch with custom CUDA kernels looks like this:
|
||||
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():
|
||||
# randomly generate input tensors based on the model architecture
|
||||
a = torch.randn(1, 128).cuda()
|
||||
b = torch.randn(1, 128).cuda()
|
||||
return [a, b]
|
||||
|
||||
def get_init_inputs():
|
||||
# randomly generate tensors required for initialization based on the model architecture
|
||||
return []
|
||||
|
||||
|
||||
|
||||
You are given the following architecture:
|
||||
|
||||
python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
合理优化的PyTorch L1 Loss实现
|
||||
使用PyTorch内置函数，避免不必要的中间张量创建
|
||||
"""
|
||||
def init(self):
|
||||
super(Model, self).init()
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
使用PyTorch内置的l1_loss函数
|
||||
L1 Loss = |input - target|
|
||||
|
||||
Args:
|
||||
input (torch.Tensor): 预测值
|
||||
target (torch.Tensor): 真实值
|
||||
|
||||
Returns:
|
||||
torch.Tensor: L1 Loss标量值
|
||||
"""
|
||||
# 直接使用内置函数，让PyTorch处理优化
|
||||
return torch.nn.functional.l1_loss(
|
||||
input,
|
||||
target,
|
||||
reduction='sum'
|
||||
)
|
||||
batch_size = 128
|
||||
num_features = 2000
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
生成合理的测试数据
|
||||
"""
|
||||
input_vals = torch.randn(batch_size, num_features)
|
||||
target_vals = torch.randn(batch_size, num_features)
|
||||
return [input_vals, target_vals]
|
||||
|
||||
def get_init_inputs():
|
||||
return [] # 没有特殊的初始化输入需求
|
||||
|
||||
|
||||
IMPORTANT FUSION REQUIREMENTS:
|
||||
|
||||
The current L1 Loss implementation involves multiple steps:
|
||||
1. Compute difference: diff = input - target
|
||||
2. Compute absolute value: abs_diff = |diff|
|
||||
3. Sum reduction: loss = sum(abs_diff)
|
||||
|
||||
YOUR TASK: Create a FUSED CUDA kernel that combines ALL these steps into a SINGLE kernel to maximize performance. The fusion should:
|
||||
|
||||
1. ELIMINATE intermediate tensor creation (no separate diff or abs_diff tensors)
|
||||
2. REDUCE memory bandwidth usage by computing diff and abs in-place
|
||||
3. MINIMIZE kernel launch overhead by using one kernel instead of multiple operations
|
||||
4. OPTIMIZE for GPU parallelism with efficient reduction strategies
|
||||
|
||||
Key fusion strategies to implement:
|
||||
- Compute (input - target) and |result| in the same thread loop
|
||||
- Use shared memory for efficient block-level reduction
|
||||
- Consider warp-level reduction for better performance
|
||||
- Implement vectorized memory access when possible
|
||||
|
||||
The fused kernel should achieve the same numerical result as the original PyTorch implementation while providing significant speedup through operator fusion.
|
||||
|
||||
Generate the complete CUDA implementation with:
|
||||
1. A fused kernel that combines diff + abs + sum operations
|
||||
2. Multiple optimization variants (basic fusion, warp-level, vectorized)
|
||||
3. Proper error checking and tensor validation
|
||||
4. Efficient memory access patterns
|
||||
5. Comprehensive performance optimizations
|
||||
|
|
@ -0,0 +1,75 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from l1_torchcode import Model, get_inputs, get_init_inputs
|
||||
from l1_fused_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 L1Loss 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA L1Loss 平均执行时间: {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