forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish LogCoshLoss#42' (#70) from Lwh20070813/GPUCodeForces:LogCoshLoss into main
This commit is contained in:
commit
5c88be3d55
|
|
@ -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,77 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from LogCoshLoss_torch import Model, get_inputs, get_init_inputs
|
||||
from LogCoshLoss_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()
|
||||
Loading…
Reference in New Issue