forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish WelschLoss #20' (#162) from gsd123/GPUCodeForces:gsd20 into main
This commit is contained in:
commit
bba362d6d0
|
|
@ -0,0 +1,144 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
N, C, H, W = 32, 64, 56, 56
|
||||
|
||||
|
||||
class WelschLossCUDAOp(torch.autograd.Function):
|
||||
def forward(ctx, input, target, beta, reduction_id, op):
|
||||
if not input.is_cuda: input = input.cuda()
|
||||
if not target.is_cuda: target = target.cuda()
|
||||
input = input.contiguous()
|
||||
target = target.contiguous()
|
||||
output = op.welsch_loss_forward_cuda(input, target, beta, reduction_id)
|
||||
return output
|
||||
|
||||
def backward(ctx, grad_output):
|
||||
return grad_output, grad_output, None, None, None
|
||||
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, reduction='mean', beta=1.0):
|
||||
super().__init__()
|
||||
self.beta = float(beta)
|
||||
self.red_map = {'none': 0, 'mean': 1, 'sum': 2}
|
||||
if reduction not in self.red_map:
|
||||
raise ValueError("Invalid reduction")
|
||||
self.reduction_id = self.red_map[reduction]
|
||||
self._compile_cuda_kernel()
|
||||
|
||||
def _compile_cuda_kernel(self):
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
torch::Tensor welsch_loss_forward_cuda(torch::Tensor input, torch::Tensor target, float beta, int reduction);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <cuda_runtime.h>
|
||||
#include <device_functions.h>
|
||||
|
||||
__inline__ __device__ float warp_reduce(float v) {
|
||||
#pragma unroll
|
||||
for (int d = 16; d > 0; d >>= 1) v += __shfl_down_sync(0xffffffff, v, d);
|
||||
return v;
|
||||
}
|
||||
|
||||
__inline__ __device__ float block_reduce(float v) {
|
||||
__shared__ float s[32];
|
||||
int l = threadIdx.x & 31;
|
||||
int w = threadIdx.x >> 5;
|
||||
v = warp_reduce(v);
|
||||
if (l == 0) s[w] = v;
|
||||
__syncthreads();
|
||||
v = (threadIdx.x < (blockDim.x >> 5)) ? s[l] : 0.0f;
|
||||
if (w == 0) v = warp_reduce(v);
|
||||
return v;
|
||||
}
|
||||
|
||||
__global__ void welsch_kernel(const float* __restrict__ in, const float* __restrict__ tgt, float* __restrict__ out, int n, float k, int red) {
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
int n4 = n >> 2;
|
||||
float acc = 0.0f;
|
||||
|
||||
const float4* in4 = (const float4*)in;
|
||||
const float4* tgt4 = (const float4*)tgt;
|
||||
float4* out4 = (float4*)out;
|
||||
|
||||
for (int i = tid; i < n4; i += stride) {
|
||||
float4 iv = __ldg(&in4[i]);
|
||||
float4 tv = __ldg(&tgt4[i]);
|
||||
|
||||
float dx = iv.x - tv.x;
|
||||
float dy = iv.y - tv.y;
|
||||
float dz = iv.z - tv.z;
|
||||
float dw = iv.w - tv.w;
|
||||
|
||||
float lx = 1.0f - __expf(__fmul_rn(__fmul_rn(dx, dx), k));
|
||||
float ly = 1.0f - __expf(__fmul_rn(__fmul_rn(dy, dy), k));
|
||||
float lz = 1.0f - __expf(__fmul_rn(__fmul_rn(dz, dz), k));
|
||||
float lw = 1.0f - __expf(__fmul_rn(__fmul_rn(dw, dw), k));
|
||||
|
||||
if (red == 0) {
|
||||
out4[i] = make_float4(lx, ly, lz, lw);
|
||||
} else {
|
||||
acc += lx + ly + lz + lw;
|
||||
}
|
||||
}
|
||||
|
||||
int rem = n4 << 2;
|
||||
for (int i = rem + threadIdx.x; i < n; i += blockDim.x) {
|
||||
float d = in[i] - tgt[i];
|
||||
float l = 1.0f - __expf(__fmul_rn(__fmul_rn(d, d), k));
|
||||
if (red == 0) out[i] = l;
|
||||
else acc += l;
|
||||
}
|
||||
|
||||
if (red != 0) {
|
||||
acc = block_reduce(acc);
|
||||
if (threadIdx.x == 0) out[blockIdx.x] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor welsch_loss_forward_cuda(torch::Tensor input, torch::Tensor target, float beta, int reduction) {
|
||||
int64_t n = input.numel();
|
||||
auto opts = input.options();
|
||||
float k = -1.0f / (beta * beta);
|
||||
|
||||
const int bs = 256;
|
||||
const int gs = std::min((int)((n + (bs << 2) - 1) / (bs << 2)), 1024);
|
||||
const int fgs = std::max(gs, 1);
|
||||
|
||||
torch::Tensor output;
|
||||
if (reduction == 0) output = torch::empty_like(input);
|
||||
else output = torch::zeros({fgs}, opts);
|
||||
|
||||
welsch_kernel<<<fgs, bs>>>(input.data_ptr<float>(), target.data_ptr<float>(), output.data_ptr<float>(), n, k, reduction);
|
||||
|
||||
if (reduction != 0) {
|
||||
float s = output.sum().item<float>();
|
||||
if (reduction == 1) s /= n;
|
||||
return torch::tensor({s}, opts);
|
||||
}
|
||||
return output;
|
||||
}
|
||||
"""
|
||||
|
||||
self.op = load_inline(
|
||||
name='welsch_loss_cuda_v2',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['welsch_loss_forward_cuda'],
|
||||
extra_cuda_cflags=['-O3', '--use_fast_math'],
|
||||
verbose=False
|
||||
)
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
if isinstance(input, (list, tuple)): input = input[0]
|
||||
if isinstance(target, (list, tuple)): target = target[0]
|
||||
if not input.is_cuda: input = input.cuda()
|
||||
if not target.is_cuda: target = target.cuda()
|
||||
input = input.contiguous()
|
||||
target = target.contiguous()
|
||||
return WelschLossCUDAOp.apply(input, target, self.beta, self.reduction_id, self.op)
|
||||
|
|
@ -0,0 +1,53 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# 定义用于测试的假设常量
|
||||
N, C, H, W = 32, 64, 56, 56
|
||||
|
||||
|
||||
class WelschLoss(nn.Module):
|
||||
def __init__(self, reduction='mean', beta=1.0):
|
||||
super().__init__()
|
||||
self.reduction = reduction
|
||||
self.beta = float(beta)
|
||||
self.beta_sq = self.beta * self.beta
|
||||
if reduction not in ['none', 'mean', 'sum']:
|
||||
raise ValueError("Invalid reduction mode")
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
# Calculate squared difference: (x - y)^2
|
||||
diff_sq = (input - target) ** 2
|
||||
|
||||
# Calculate loss: 1 - exp(-(diff^2 / beta^2))
|
||||
loss = 1.0 - torch.exp(-diff_sq / self.beta_sq)
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
else:
|
||||
return loss
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, reduction='mean', beta=1.0):
|
||||
super().__init__()
|
||||
self.op = WelschLoss(reduction, beta)
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
# 适应 benchmark 输入格式 (如果输入是列表)
|
||||
if isinstance(input, (list, tuple)): input = input[0]
|
||||
if isinstance(target, (list, tuple)): target = target[0]
|
||||
|
||||
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', 1.0]
|
||||
|
|
@ -0,0 +1,108 @@
|
|||
You write custom CUDA kernels to replace the pytorch operators in the given GeGLU 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 chunk+gelu+elementwise_mul), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.
|
||||
|
||||
Core Optimization Techniques:
|
||||
|
||||
Performance Optimizations
|
||||
|
||||
Vectorized Processing - Uses float4 for 4-element vectorized loads/stores
|
||||
|
||||
Fast Math Compilation - --use_fast_math flag with -O3 for maximum speed
|
||||
|
||||
Efficient Grid Sizing - Optimized grid calculation using bit-shift operations
|
||||
|
||||
Read-Only Cache - Uses __ldg() intrinsic for cached memory access
|
||||
|
||||
Memory Optimizations
|
||||
|
||||
Vectorized Memory Access - Processes 4 elements per operation via float4
|
||||
|
||||
Memory Coalescing - Ensures contiguous memory access patterns
|
||||
|
||||
Minimal Memory Allocation - Only allocates necessary output tensors
|
||||
|
||||
Numerical Optimizations
|
||||
|
||||
Fast Math Operations - Uses __expf, __fmul_rn for optimized floating-point math
|
||||
|
||||
Pre-computed Constants - Calculates k = -1.0f / (beta * beta) once on host
|
||||
|
||||
Efficient Welsch Formula - Optimized computation: 1.0f - expf(d*d*k)
|
||||
|
||||
Kernel Design
|
||||
|
||||
Two-Phase Processing - Vectorized main loop + scalar remainder handling
|
||||
|
||||
Efficient Reduction - Warp-level and block-level reduction with shared memory
|
||||
|
||||
Flexible Output - Supports both element-wise and reduced outputs
|
||||
|
||||
Key Features
|
||||
|
||||
High Throughput - Vectorized processing maximizes memory bandwidth
|
||||
|
||||
Fast Exponential - Optimized Welsch loss computation using fast math
|
||||
|
||||
Efficient Reduction - Minimal synchronization in reduction steps
|
||||
|
||||
Remainder Handling - Properly processes non-multiple-of-4 elements
|
||||
|
||||
This implementation provides extremely efficient Welsch loss computation through extensive vectorization and fast math optimizations.
|
||||
|
||||
|
||||
|
||||
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 torch.nn.functional as F
|
||||
|
||||
# 定义用于测试的假设常量
|
||||
N, C, H, W = 32, 64, 56, 56
|
||||
|
||||
|
||||
class WelschLoss(nn.Module):
|
||||
def __init__(self, reduction='mean', beta=1.0):
|
||||
super().__init__()
|
||||
self.reduction = reduction
|
||||
self.beta = float(beta)
|
||||
self.beta_sq = self.beta * self.beta
|
||||
if reduction not in ['none', 'mean', 'sum']:
|
||||
raise ValueError("Invalid reduction mode")
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
# Calculate squared difference: (x - y)^2
|
||||
diff_sq = (input - target) ** 2
|
||||
|
||||
# Calculate loss: 1 - exp(-(diff^2 / beta^2))
|
||||
loss = 1.0 - torch.exp(-diff_sq / self.beta_sq)
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
else:
|
||||
return loss
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, reduction='mean', beta=1.0):
|
||||
super().__init__()
|
||||
self.op = WelschLoss(reduction, beta)
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
# 适应 benchmark 输入格式 (如果输入是列表)
|
||||
if isinstance(input, (list, tuple)): input = input[0]
|
||||
if isinstance(target, (list, tuple)): target = target[0]
|
||||
|
||||
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', 1.0]
|
||||
|
|
@ -0,0 +1,78 @@
|
|||
import torch
|
||||
import time
|
||||
from WelschLoss_torch import Model, get_inputs, get_init_inputs
|
||||
from WelschLoss_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