forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish GaussianNLLLoss# 31' (#58) from uucoco/GPUCodeForces:GaussianNLLLoss into main
This commit is contained in:
commit
dfe343ddbd
|
|
@ -0,0 +1,169 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
import math
|
||||
|
||||
|
||||
N_BATCH = 128
|
||||
N_FEATURES = 512
|
||||
|
||||
|
||||
FULL = False
|
||||
EPS = 1e-6
|
||||
REDUCTION = 'mean'
|
||||
|
||||
|
||||
BLOCK_SIZE = 256
|
||||
|
||||
|
||||
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
|
||||
def __init__(self, full=False, eps=1e-6, reduction='mean'):
|
||||
super().__init__()
|
||||
self.full = full
|
||||
self.eps = eps
|
||||
self.reduction = reduction
|
||||
self.reduction_str = reduction
|
||||
|
||||
self.block_size = BLOCK_SIZE
|
||||
self._compile_cuda_kernel()
|
||||
|
||||
def _compile_cuda_kernel(self):
|
||||
|
||||
cpp_header = f"""
|
||||
#include <torch/extension.h>
|
||||
|
||||
// C++ 接口
|
||||
torch::Tensor gaussian_nll_loss_forward_cuda(
|
||||
torch::Tensor input,
|
||||
torch::Tensor target,
|
||||
torch::Tensor var,
|
||||
bool full_flag,
|
||||
float eps_val,
|
||||
std::string reduction
|
||||
);
|
||||
"""
|
||||
|
||||
cuda_source = f"""
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath> // for logf, fmaxf
|
||||
|
||||
#define BLOCK_SIZE {self.block_size}
|
||||
|
||||
#define CONST_TERM (0.5f * 1.8378770664f)
|
||||
|
||||
/*
|
||||
* GaussianNLLLoss 融合核函数 (Element-wise)
|
||||
*/
|
||||
__global__ void gaussian_nll_loss_fused_kernel(
|
||||
const float* __restrict__ input_data,
|
||||
const float* __restrict__ target_data,
|
||||
const float* __restrict__ var_data,
|
||||
float* __restrict__ output_data,
|
||||
int N_total,
|
||||
bool full_flag,
|
||||
float eps_val
|
||||
) {{
|
||||
for (int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
i < N_total;
|
||||
i += gridDim.x * blockDim.x)
|
||||
{{
|
||||
const float in = input_data[i];
|
||||
const float t = target_data[i];
|
||||
const float v = var_data[i];
|
||||
|
||||
const float v_clamped = fmaxf(v, eps_val);
|
||||
const float diff = in - t;
|
||||
float loss = 0.5f * (logf(v_clamped) + (diff * diff) / v_clamped);
|
||||
|
||||
if (full_flag) {{
|
||||
loss += CONST_TERM;
|
||||
}}
|
||||
|
||||
output_data[i] = loss;
|
||||
}}
|
||||
}}
|
||||
|
||||
// C++ 封装函数
|
||||
torch::Tensor gaussian_nll_loss_forward_cuda(
|
||||
torch::Tensor input,
|
||||
torch::Tensor target,
|
||||
torch::Tensor var,
|
||||
bool full_flag,
|
||||
float eps_val,
|
||||
std::string reduction
|
||||
) {{
|
||||
TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
|
||||
TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
|
||||
TORCH_CHECK(target.is_contiguous(), "target must be contiguous");
|
||||
TORCH_CHECK(var.is_contiguous(), "var must be contiguous");
|
||||
|
||||
TORCH_CHECK(input.sizes() == target.sizes(), "input and target shape mismatch");
|
||||
TORCH_CHECK(input.sizes() == var.sizes(), "input and var shape mismatch");
|
||||
|
||||
const int64_t N_total = input.numel();
|
||||
|
||||
auto output = torch::empty_like(input);
|
||||
|
||||
dim3 block_dim(BLOCK_SIZE);
|
||||
dim3 grid_dim((N_total + BLOCK_SIZE - 1) / BLOCK_SIZE);
|
||||
|
||||
gaussian_nll_loss_fused_kernel<<<grid_dim, block_dim>>>(
|
||||
input.data_ptr<float>(),
|
||||
target.data_ptr<float>(),
|
||||
var.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
N_total,
|
||||
full_flag,
|
||||
eps_val
|
||||
);
|
||||
|
||||
if (reduction == "mean") {{
|
||||
return output.mean();
|
||||
}} else if (reduction == "sum") {{
|
||||
return output.sum();
|
||||
}} else {{
|
||||
return output; // "none"
|
||||
}}
|
||||
}}
|
||||
"""
|
||||
|
||||
nvcc_flags = [
|
||||
'-O3',
|
||||
'--use_fast_math',
|
||||
'--expt-relaxed-constexpr'
|
||||
]
|
||||
|
||||
self.loss_op = load_inline(
|
||||
name="gaussian_nll_loss_op_v3_fixed_api",
|
||||
cpp_sources=cpp_header,
|
||||
cuda_sources=cuda_source,
|
||||
functions=["gaussian_nll_loss_forward_cuda"],
|
||||
extra_cuda_cflags=nvcc_flags,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor, var: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
|
||||
if target.size() != input.size():
|
||||
target = target.expand_as(input)
|
||||
if var.size() != input.size():
|
||||
var = var.expand_as(input)
|
||||
|
||||
|
||||
input_cont = input.contiguous()
|
||||
target_cont = target.contiguous()
|
||||
var_cont = var.contiguous()
|
||||
|
||||
return self.loss_op.gaussian_nll_loss_forward_cuda(
|
||||
input_cont,
|
||||
target_cont,
|
||||
var_cont,
|
||||
self.full,
|
||||
self.eps,
|
||||
self.reduction_str
|
||||
)
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import math
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# 常量定义
|
||||
# -------------------------------------------------------------
|
||||
N_BATCH = 128
|
||||
N_FEATURES = 512
|
||||
|
||||
# 损失函数参数
|
||||
FULL = False
|
||||
EPS = 1e-6
|
||||
REDUCTION = 'mean'
|
||||
|
||||
|
||||
# -------------------------------------------------------------
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
nn.GaussianNLLLoss 的纯 PyTorch 基准实现
|
||||
"""
|
||||
|
||||
def __init__(self, full=False, eps=1e-6, reduction='mean'):
|
||||
super().__init__()
|
||||
self.full = full
|
||||
self.eps = eps
|
||||
self.reduction = reduction
|
||||
|
||||
if self.full:
|
||||
self.const_term = 0.5 * math.log(2 * math.pi)
|
||||
else:
|
||||
self.const_term = 0.0
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor, var: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
# 1. 确保 var > eps。
|
||||
# torch.clamp(min=...) 等价于 max(var, eps)
|
||||
var_clamped = torch.clamp(var, min=self.eps)
|
||||
|
||||
# 2. 计算两个主要项
|
||||
term1_log = torch.log(var_clamped)
|
||||
term2_sq_err = (input - target).pow(2) / var_clamped
|
||||
|
||||
# 3. 组合
|
||||
# (N, *) 形状
|
||||
loss_unreduced = 0.5 * (term1_log + term2_sq_err) + self.const_term
|
||||
|
||||
# 4. 应用 Reduciton
|
||||
if self.reduction == 'mean':
|
||||
return loss_unreduced.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss_unreduced.sum()
|
||||
else: # 'none'
|
||||
return loss_unreduced
|
||||
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
生成 (N, D) 形状的输入
|
||||
"""
|
||||
input = torch.randn(N_BATCH, N_FEATURES, dtype=torch.float32)
|
||||
target = torch.randn(N_BATCH, N_FEATURES, dtype=torch.float32)
|
||||
|
||||
# Var 必须是正数
|
||||
var = torch.rand(N_BATCH, N_FEATURES, dtype=torch.float32) + EPS
|
||||
|
||||
return [input, target, var]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [FULL, EPS, REDUCTION]
|
||||
|
||||
|
|
@ -0,0 +1,102 @@
|
|||
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.
|
||||
|
||||
Technologies Used :
|
||||
|
||||
PyTorch: Deep learning framework.
|
||||
|
||||
CUDA: GPU acceleration for parallel computing.
|
||||
|
||||
C++/CUDA C++: High-performance kernel programming.
|
||||
|
||||
Inline C++/CUDA Extension (torch.utils.cpp_extension.load_inline): Just-In-Time (JIT) compilation of custom operators.
|
||||
|
||||
Fused Kernel: Combines multiple operations (logarithm, division, squaring, addition) into a single GPU kernel to reduce memory bandwidth and launch overhead.
|
||||
|
||||
Element-wise Parallelism: Each GPU thread handles an independent element of the input tensors.
|
||||
|
||||
Grid-Stride Loop: Efficiently processes data of arbitrary size using a fixed number of threads.
|
||||
|
||||
Math Operations: logf, fmaxf (fast math with --use_fast_math flag).
|
||||
|
||||
Tensor Contiguity Check: Ensures memory layout optimization.
|
||||
|
||||
Reduction Operations (Mean/Sum): Aggregates loss values in the kernel's C++ wrapper.
|
||||
|
||||
Memory Access Patterns: Uses __restrict__ keyword to hint at non-aliasing pointers for compiler optimization.
|
||||
|
||||
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
|
||||
import math
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# 常量定义
|
||||
# -------------------------------------------------------------
|
||||
N_BATCH = 128
|
||||
N_FEATURES = 512
|
||||
|
||||
# 损失函数参数
|
||||
FULL = False
|
||||
EPS = 1e-6
|
||||
REDUCTION = 'mean'
|
||||
|
||||
|
||||
# -------------------------------------------------------------
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
nn.GaussianNLLLoss 的纯 PyTorch 基准实现
|
||||
"""
|
||||
|
||||
def __init__(self, full=False, eps=1e-6, reduction='mean'):
|
||||
super().__init__()
|
||||
self.full = full
|
||||
self.eps = eps
|
||||
self.reduction = reduction
|
||||
|
||||
if self.full:
|
||||
self.const_term = 0.5 * math.log(2 * math.pi)
|
||||
else:
|
||||
self.const_term = 0.0
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor, var: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
# 1. 确保 var > eps。
|
||||
# torch.clamp(min=...) 等价于 max(var, eps)
|
||||
var_clamped = torch.clamp(var, min=self.eps)
|
||||
|
||||
# 2. 计算两个主要项
|
||||
term1_log = torch.log(var_clamped)
|
||||
term2_sq_err = (input - target).pow(2) / var_clamped
|
||||
|
||||
# 3. 组合
|
||||
# (N, *) 形状
|
||||
loss_unreduced = 0.5 * (term1_log + term2_sq_err) + self.const_term
|
||||
|
||||
# 4. 应用 Reduciton
|
||||
if self.reduction == 'mean':
|
||||
return loss_unreduced.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss_unreduced.sum()
|
||||
else: # 'none'
|
||||
return loss_unreduced
|
||||
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
生成 (N, D) 形状的输入
|
||||
"""
|
||||
input = torch.randn(N_BATCH, N_FEATURES, dtype=torch.float32)
|
||||
target = torch.randn(N_BATCH, N_FEATURES, dtype=torch.float32)
|
||||
|
||||
# Var 必须是正数
|
||||
var = torch.rand(N_BATCH, N_FEATURES, dtype=torch.float32) + EPS
|
||||
|
||||
return [input, target, var]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [FULL, EPS, REDUCTION]
|
||||
|
||||
|
|
@ -0,0 +1,77 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from GaussianNLLLoss_torch import Model, get_inputs, get_init_inputs
|
||||
from GaussianNLLLoss_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