forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'optimized LogCoshLoss operator #2' (#200) from Lwh20070813/GPUCodeForces:lwh2 into main
This commit is contained in:
commit
df268306fd
|
|
@ -0,0 +1,171 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
N, C, H, W = 32, 64, 56, 56
|
||||
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, reduction='mean'):
|
||||
super().__init__()
|
||||
self.reduction = reduction
|
||||
self.red_map = {'none': 0, 'mean': 1, 'sum': 2}
|
||||
self.reduction_id = self.red_map[reduction]
|
||||
self.block_size = 256
|
||||
self._compile_cuda_kernel()
|
||||
|
||||
def _compile_cuda_kernel(self):
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
|
||||
torch::Tensor log_cosh_forward_cuda(
|
||||
torch::Tensor input,
|
||||
torch::Tensor target,
|
||||
int reduction);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
|
||||
#define BLOCK_SIZE 256
|
||||
#define WARP_SIZE 32
|
||||
|
||||
__inline__ __device__ float warp_reduce_sum(float val) {
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset /= 2) {
|
||||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
__inline__ __device__ float block_reduce_sum(float val) {
|
||||
__shared__ float shared[32];
|
||||
int lane = threadIdx.x % 32;
|
||||
int wid = threadIdx.x / 32;
|
||||
|
||||
val = warp_reduce_sum(val);
|
||||
if (lane == 0) shared[wid] = val;
|
||||
__syncthreads();
|
||||
|
||||
val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f;
|
||||
if (wid == 0) val = warp_reduce_sum(val);
|
||||
return val;
|
||||
}
|
||||
|
||||
__global__ void log_cosh_kernel(
|
||||
const float* __restrict__ input,
|
||||
const float* __restrict__ target,
|
||||
float* __restrict__ output,
|
||||
int n,
|
||||
int reduction
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
|
||||
float local_sum = 0.0f;
|
||||
float log_2 = 0.69314718056f;
|
||||
|
||||
float4* in_ptr = (float4*)input;
|
||||
float4* tgt_ptr = (float4*)target;
|
||||
float4* out_ptr = (float4*)output;
|
||||
|
||||
int vec_n = n / 4;
|
||||
|
||||
for (int i = idx; i < vec_n; i += stride) {
|
||||
float4 in_val = in_ptr[i];
|
||||
float4 tgt_val = tgt_ptr[i];
|
||||
|
||||
float diff[4];
|
||||
diff[0] = fabsf(in_val.x - tgt_val.x);
|
||||
diff[1] = fabsf(in_val.y - tgt_val.y);
|
||||
diff[2] = fabsf(in_val.z - tgt_val.z);
|
||||
diff[3] = fabsf(in_val.w - tgt_val.w);
|
||||
|
||||
float losses[4];
|
||||
#pragma unroll
|
||||
for(int k=0; k<4; ++k) {
|
||||
losses[k] = diff[k] + log1pf(expf(-2.0f * diff[k])) - log_2;
|
||||
}
|
||||
|
||||
if (reduction == 0) {
|
||||
float4 res;
|
||||
res.x = losses[0]; res.y = losses[1];
|
||||
res.z = losses[2]; res.w = losses[3];
|
||||
out_ptr[i] = res;
|
||||
} else {
|
||||
local_sum += losses[0] + losses[1] + losses[2] + losses[3];
|
||||
}
|
||||
}
|
||||
|
||||
int rem_start = vec_n * 4;
|
||||
for (int i = rem_start + idx; i < n; i += stride) {
|
||||
float diff = fabsf(input[i] - target[i]);
|
||||
float loss = diff + log1pf(expf(-2.0f * diff)) - log_2;
|
||||
|
||||
if (reduction == 0) {
|
||||
output[i] = loss;
|
||||
} else {
|
||||
local_sum += loss;
|
||||
}
|
||||
}
|
||||
|
||||
if (reduction != 0) {
|
||||
local_sum = block_reduce_sum(local_sum);
|
||||
if (threadIdx.x == 0) {
|
||||
atomicAdd(output, local_sum);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor log_cosh_forward_cuda(
|
||||
torch::Tensor input,
|
||||
torch::Tensor target,
|
||||
int reduction)
|
||||
{
|
||||
int64_t n = input.numel();
|
||||
auto options = input.options();
|
||||
|
||||
torch::Tensor output;
|
||||
if (reduction == 0) {
|
||||
output = torch::empty_like(input);
|
||||
} else {
|
||||
output = torch::zeros({1}, options);
|
||||
}
|
||||
|
||||
const int block_size = 256;
|
||||
const int grid_size = std::min((int)((n + block_size * 4 - 1) / (block_size * 4)), 1024);
|
||||
|
||||
log_cosh_kernel<<<grid_size, block_size>>>(
|
||||
input.data_ptr<float>(),
|
||||
target.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
n,
|
||||
reduction
|
||||
);
|
||||
|
||||
if (reduction == 1) {
|
||||
output.div_(n);
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
"""
|
||||
|
||||
self.op = load_inline(
|
||||
name='log_cosh_cuda_opt',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['log_cosh_forward_cuda'],
|
||||
extra_cuda_cflags=['-O3', '--use_fast_math'],
|
||||
verbose=False
|
||||
)
|
||||
|
||||
def forward(self, input, target):
|
||||
if not input.is_cuda: input = input.cuda()
|
||||
if not target.is_cuda: target = target.cuda()
|
||||
|
||||
input = input.contiguous()
|
||||
target = target.contiguous()
|
||||
|
||||
return self.op.log_cosh_forward_cuda(input, target, self.reduction_id)
|
||||
|
|
@ -0,0 +1,40 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import math
|
||||
|
||||
N, C, H, W = 32, 64, 56, 56
|
||||
|
||||
|
||||
class LogCoshLoss(nn.Module):
|
||||
def __init__(self, reduction='mean'):
|
||||
super().__init__()
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, input, target):
|
||||
diff = input - target
|
||||
loss = torch.abs(diff) + torch.nn.functional.softplus(-2. * torch.abs(diff)) - math.log(2.0)
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
return loss
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, reduction='mean'):
|
||||
super().__init__()
|
||||
self.op = LogCoshLoss(reduction)
|
||||
|
||||
def forward(self, input, target):
|
||||
return self.op(input, target)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
input = torch.randn(N, C, H, W, dtype=torch.float32)
|
||||
target = torch.randn(N, C, H, W, dtype=torch.float32)
|
||||
return [input, target]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return ['mean']
|
||||
|
|
@ -0,0 +1,82 @@
|
|||
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.
|
||||
|
||||
|
||||
Overview
|
||||
This implementation provides a highly optimized CUDA kernel for computing the Log-Cosh loss function, which is a smooth alternative to Mean Absolute Error (MAE) that is less sensitive to outliers than Mean Squared Error (MSE).
|
||||
|
||||
Mathematical Formulation
|
||||
The Log-Cosh loss is defined as:
|
||||
L(x, y) = log(cosh(x - y)) = |x-y| + log(1 + exp(-2|x-y|)) - log(2)
|
||||
|
||||
key Optimizations
|
||||
1. Vectorized Memory Access
|
||||
Uses float4 data type for coalesced memory operations
|
||||
|
||||
Processes 4 elements per thread simultaneously
|
||||
|
||||
Reduces memory transaction overhead by 75%
|
||||
|
||||
2. Numerical Stability
|
||||
Implements the stable formulation: |diff| + log1p(exp(-2*|diff|)) - log(2)
|
||||
|
||||
Avoids numerical overflow in cosh() calculation
|
||||
|
||||
Uses log1p() for accurate logarithm of (1 + x)
|
||||
|
||||
3. Parallel Reduction Strategy
|
||||
Warp-level reduction: 32-thread warp shuffle operations
|
||||
|
||||
Block-level reduction: Shared memory for intra-block reduction
|
||||
|
||||
Global reduction: Atomic operations for cross-block summation
|
||||
|
||||
4. Flexible Reduction Modes
|
||||
reduction=0: Element-wise output (no reduction)
|
||||
|
||||
reduction=1: Mean reduction (sum / n_elements)
|
||||
|
||||
reduction=2: Sum reduction
|
||||
|
||||
|
||||
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 math
|
||||
|
||||
N, C, H, W = 32, 64, 56, 56
|
||||
|
||||
|
||||
class LogCoshLoss(nn.Module):
|
||||
def __init__(self, reduction='mean'):
|
||||
super().__init__()
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, input, target):
|
||||
diff = input - target
|
||||
loss = torch.abs(diff) + torch.nn.functional.softplus(-2. * torch.abs(diff)) - math.log(2.0)
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
return loss
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, reduction='mean'):
|
||||
super().__init__()
|
||||
self.op = LogCoshLoss(reduction)
|
||||
|
||||
def forward(self, input, target):
|
||||
return self.op(input, target)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
input = torch.randn(N, C, H, W, dtype=torch.float32)
|
||||
target = torch.randn(N, C, H, W, dtype=torch.float32)
|
||||
return [input, target]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return ['mean']
|
||||
|
|
@ -0,0 +1,83 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import time
|
||||
from LogCoshLoss_torch import Model, get_inputs, get_init_inputs
|
||||
from LogCoshLoss_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()
|
||||
Loading…
Reference in New Issue