forked from ccf-ai-infra/GPUCodeForces
add batchnorm1d operator
This commit is contained in:
parent
2f77913144
commit
0ff6352146
|
|
@ -0,0 +1,168 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
batchnorm_source = r"""
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
|
||||
__global__ void batchnorm_forward_kernel(
|
||||
const float* __restrict__ x,
|
||||
const float* __restrict__ gamma,
|
||||
const float* __restrict__ beta,
|
||||
float* __restrict__ y,
|
||||
int batch,
|
||||
int features,
|
||||
float eps
|
||||
) {
|
||||
int feature = blockIdx.x;
|
||||
if (feature >= features) return;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
extern __shared__ float shared[];
|
||||
float* shm_sum = shared;
|
||||
float* shm_sq = shared + blockDim.x;
|
||||
|
||||
float sum = 0.0f;
|
||||
float sum_sq = 0.0f;
|
||||
|
||||
for (int row = tid; row < batch; row += blockDim.x) {
|
||||
float v = x[row * features + feature];
|
||||
sum += v;
|
||||
sum_sq += v * v;
|
||||
}
|
||||
|
||||
shm_sum[tid] = sum;
|
||||
shm_sq[tid] = sum_sq;
|
||||
__syncthreads();
|
||||
|
||||
for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
|
||||
if (tid < offset) {
|
||||
shm_sum[tid] += shm_sum[tid + offset];
|
||||
shm_sq[tid] += shm_sq[tid + offset];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
__shared__ float s_mean;
|
||||
__shared__ float s_inv_std;
|
||||
|
||||
if (tid == 0) {
|
||||
float mean = shm_sum[0] / batch;
|
||||
float var = shm_sq[0] / batch - mean * mean;
|
||||
var = var > 0.f ? var : 0.f;
|
||||
s_mean = mean;
|
||||
s_inv_std = rsqrtf(var + eps);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
float mean = s_mean;
|
||||
float inv_std = s_inv_std;
|
||||
float g = gamma[feature];
|
||||
float b = beta[feature];
|
||||
|
||||
for (int row = tid; row < batch; row += blockDim.x) {
|
||||
float v = x[row * features + feature];
|
||||
float norm = (v - mean) * inv_std;
|
||||
y[row * features + feature] = norm * g + b;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor batchnorm_cuda_forward(
|
||||
torch::Tensor x,
|
||||
torch::Tensor weight,
|
||||
torch::Tensor bias,
|
||||
double eps
|
||||
) {
|
||||
TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
|
||||
TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
|
||||
TORCH_CHECK(bias.is_cuda(), "bias must be a CUDA tensor");
|
||||
TORCH_CHECK(x.dtype() == torch::kFloat32, "only float32 tensors are supported");
|
||||
TORCH_CHECK(weight.dtype() == torch::kFloat32, "weight must be float32");
|
||||
TORCH_CHECK(bias.dtype() == torch::kFloat32, "bias must be float32");
|
||||
TORCH_CHECK(x.dim() == 2, "input must be 2D [batch, features]");
|
||||
TORCH_CHECK(weight.dim() == 1, "weight must be 1D");
|
||||
TORCH_CHECK(bias.dim() == 1, "bias must be 1D");
|
||||
TORCH_CHECK(x.size(1) == weight.size(0), "feature size mismatch between input and weight");
|
||||
TORCH_CHECK(weight.size(0) == bias.size(0), "weight and bias must have the same length");
|
||||
|
||||
auto x_contig = x.contiguous();
|
||||
auto weight_contig = weight.contiguous();
|
||||
auto bias_contig = bias.contiguous();
|
||||
|
||||
int batch = x_contig.size(0);
|
||||
int features = x_contig.size(1);
|
||||
|
||||
auto y = torch::empty_like(x_contig);
|
||||
|
||||
int threads = 256;
|
||||
if (batch < threads) {
|
||||
threads = 1;
|
||||
while (threads < batch) threads <<= 1;
|
||||
if (threads < 32) threads = 32;
|
||||
}
|
||||
size_t shared_mem = threads * 2 * sizeof(float);
|
||||
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
batchnorm_forward_kernel<<<features, threads, shared_mem, stream>>>(
|
||||
x_contig.data_ptr<float>(),
|
||||
weight_contig.data_ptr<float>(),
|
||||
bias_contig.data_ptr<float>(),
|
||||
y.data_ptr<float>(),
|
||||
batch,
|
||||
features,
|
||||
static_cast<float>(eps)
|
||||
);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
return y;
|
||||
}
|
||||
"""
|
||||
|
||||
batchnorm_cpp_source = r"""
|
||||
torch::Tensor batchnorm_cuda_forward(
|
||||
torch::Tensor x,
|
||||
torch::Tensor weight,
|
||||
torch::Tensor bias,
|
||||
double eps
|
||||
);
|
||||
"""
|
||||
|
||||
batchnorm_cuda = load_inline(
|
||||
name="batchnorm_cuda_ext",
|
||||
cpp_sources=batchnorm_cpp_source,
|
||||
cuda_sources=batchnorm_source,
|
||||
functions=["batchnorm_cuda_forward"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
"""
|
||||
Model performing matrix multiplication followed by custom CUDA BatchNorm and ReLU.
|
||||
"""
|
||||
def __init__(self, mat_weight: torch.Tensor, bn_weight: torch.Tensor, bn_bias: torch.Tensor, eps: float = 1e-5):
|
||||
super().__init__()
|
||||
if mat_weight.dim() != 2:
|
||||
raise ValueError("mat_weight must be a 2D tensor [input_dim, output_dim].")
|
||||
if bn_weight.dim() != 1 or bn_bias.dim() != 1:
|
||||
raise ValueError("BatchNorm weight and bias must be 1D.")
|
||||
if bn_weight.size(0) != mat_weight.size(1):
|
||||
raise ValueError("BatchNorm parameter size must match output_dim.")
|
||||
if bn_weight.size(0) != bn_bias.size(0):
|
||||
raise ValueError("BatchNorm weight and bias must share shape.")
|
||||
self.weight = nn.Parameter(mat_weight.clone())
|
||||
self.bn_weight = nn.Parameter(bn_weight.clone())
|
||||
self.bn_bias = nn.Parameter(bn_bias.clone())
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if not x.is_cuda:
|
||||
raise ValueError("Input must be a CUDA tensor.")
|
||||
if not self.weight.is_cuda:
|
||||
raise ValueError("Model weight must be on CUDA.")
|
||||
if not self.bn_weight.is_cuda or not self.bn_bias.is_cuda:
|
||||
raise ValueError("BatchNorm parameters must be on CUDA.")
|
||||
x = torch.matmul(x, self.weight)
|
||||
x = batchnorm_cuda.batchnorm_cuda_forward(x, self.bn_weight, self.bn_bias, self.eps)
|
||||
return torch.relu(x)
|
||||
|
|
@ -0,0 +1,41 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class BatchNormModel(nn.Module):
|
||||
"""
|
||||
Model that performs matrix multiplication followed by BatchNorm and ReLU activation.
|
||||
"""
|
||||
def __init__(self, weight, num_features=2048, eps=1e-5, momentum=0.1):
|
||||
super(BatchNormModel, self).__init__()
|
||||
self.weight = nn.Parameter(weight)
|
||||
# 设置 track_running_stats=False 使其始终使用当前批次统计量
|
||||
self.bn = nn.BatchNorm1d(num_features, eps=eps, momentum=momentum, track_running_stats=False)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Performs matrix multiplication, applies BatchNorm, then ReLU activation.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): Input tensor of shape [batch_size, input_dim]
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output tensor of shape [batch_size, output_dim]
|
||||
"""
|
||||
x = torch.matmul(x, self.weight)
|
||||
x = self.bn(x)
|
||||
return torch.relu(x)
|
||||
|
||||
# 添加别名以便在 run_code.py 中使用
|
||||
Model = BatchNormModel
|
||||
|
||||
batch_size = 16
|
||||
input_dim = 1024
|
||||
output_dim = 2048
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, input_dim)
|
||||
return [x]
|
||||
|
||||
def get_init_inputs():
|
||||
weight = torch.randn(input_dim, output_dim)
|
||||
return [weight]
|
||||
|
|
@ -0,0 +1,61 @@
|
|||
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, num_features) -> None:
|
||||
super().__init__()
|
||||
self.bn = nn.BatchNorm1d(num_features)
|
||||
|
||||
def forward(self, x):
|
||||
return self.bn(x)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(16, 2048).cuda()
|
||||
return [x]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [2048]
|
||||
```
|
||||
|
||||
You are given the following architecture to implement BatchNorm1d with custom CUDA kernel:
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""使用 PyTorch BatchNorm1d 的基准实现。"""
|
||||
|
||||
def __init__(self, num_features: int, eps: float = 1e-5, momentum: float = 0.1):
|
||||
super().__init__()
|
||||
self.bn = nn.BatchNorm1d(num_features, eps=eps, momentum=momentum)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""对输入做 BatchNorm，输出形状与输入一致。"""
|
||||
return self.bn(x)
|
||||
|
||||
|
||||
batch_size = 16
|
||||
feature_dim = 2048
|
||||
|
||||
|
||||
def get_inputs():
|
||||
x = torch.randn(batch_size, feature_dim)
|
||||
return [x]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [feature_dim]
|
||||
```
|
||||
|
|
@ -0,0 +1,99 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from batchnorm1d_torch import Model as TorchModel, get_inputs, get_init_inputs
|
||||
from batchnorm1d_cuda import ModelNew as CudaModel
|
||||
|
||||
def run_benchmark():
|
||||
# 检查 CUDA 是否可用
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
|
||||
return
|
||||
else:
|
||||
device = torch.device("cuda")
|
||||
|
||||
# 初始化模型
|
||||
# 获取 torch 版本的初始化参数
|
||||
torch_init_inputs = get_init_inputs()
|
||||
weight = torch_init_inputs[0].cuda(device=device)
|
||||
|
||||
# 为 BatchNorm 准备参数
|
||||
batch_size = 16
|
||||
input_dim = 1024
|
||||
output_dim = 2048
|
||||
bn_weight = torch.ones(output_dim, device=device, dtype=torch.float32)
|
||||
bn_bias = torch.zeros(output_dim, device=device, dtype=torch.float32)
|
||||
|
||||
# 初始化输入数据
|
||||
inputs = get_inputs()
|
||||
inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs]
|
||||
|
||||
# 初始化两个模型
|
||||
torch_model = TorchModel(weight.clone(), num_features=output_dim, eps=1e-5).cuda()
|
||||
cuda_model = CudaModel(weight.clone(), bn_weight.clone(), bn_bias.clone(), eps=1e-5).cuda()
|
||||
|
||||
torch_model.eval()
|
||||
cuda_model.eval()
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
output_torch = torch_model(*inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
|
||||
# 更严格的精度检查
|
||||
abs_diff = (output_torch - output_cuda).abs()
|
||||
max_diff = abs_diff.max().item()
|
||||
mean_diff = abs_diff.mean().item()
|
||||
|
||||
print(f"最大差异: {max_diff:.6f}")
|
||||
print(f"平均差异: {mean_diff:.6f}")
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03, atol=1e-03)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 1000 # 增加迭代次数以获得更准确的时间测量
|
||||
|
||||
# Warm up
|
||||
print("预热中...")
|
||||
for _ in range(100):
|
||||
with torch.no_grad():
|
||||
_ = torch_model(*inputs)
|
||||
_ = cuda_model(*inputs)
|
||||
|
||||
# PyTorch 模型计时
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
with torch.no_grad():
|
||||
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()
|
||||
with torch.no_grad():
|
||||
for _ in range(num_iterations):
|
||||
_ = cuda_model(*inputs)
|
||||
torch.cuda.synchronize()
|
||||
cuda_time = (time.time() - start_time) / num_iterations
|
||||
|
||||
print(f"PyTorch 平均执行时间: {torch_time*1000:.4f} 毫秒")
|
||||
print(f"自定义 CUDA BatchNorm 平均执行时间: {cuda_time*1000:.4f} 毫秒")
|
||||
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