forked from ccf-ai-infra/GPUCodeForces
finish HammingDistance #4
This commit is contained in:
parent
10eed82956
commit
7d7944fb64
|
|
@ -0,0 +1,154 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
import math
|
||||
|
||||
N_BATCH = 100
|
||||
D_VECTOR = 128
|
||||
|
||||
DIM = 1
|
||||
BLOCK_SIZE = 256
|
||||
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
|
||||
def __init__(self, dim=1):
|
||||
super().__init__()
|
||||
|
||||
if dim != 1:
|
||||
raise NotImplementedError("CUDA Kernel only supports dim=1 for (N,D) tensors")
|
||||
self.dim = dim
|
||||
self.block_size = BLOCK_SIZE
|
||||
self._compile_cuda_kernel()
|
||||
|
||||
def _compile_cuda_kernel(self):
|
||||
cpp_header = f"""
|
||||
#include <torch/extension.h>
|
||||
|
||||
// C++ 接口
|
||||
torch::Tensor hamming_distance_forward_cuda(
|
||||
torch::Tensor x1,
|
||||
torch::Tensor x2,
|
||||
int dim
|
||||
);
|
||||
"""
|
||||
|
||||
cuda_source = f"""
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
|
||||
#define BLOCK_SIZE {self.block_size}
|
||||
|
||||
/*
|
||||
* Hamming Distance 融合核函数
|
||||
* (输入: int64_t, 输出: float)
|
||||
*/
|
||||
__global__ void hamming_distance_fused_kernel(
|
||||
const int64_t* __restrict__ x1_data,
|
||||
const int64_t* __restrict__ x2_data,
|
||||
float* __restrict__ output_data,
|
||||
int N,
|
||||
int D
|
||||
) {{
|
||||
|
||||
__shared__ int s_data[BLOCK_SIZE];
|
||||
|
||||
const int n_idx = blockIdx.x; // [0, N-1]
|
||||
const int tid = threadIdx.x; // [0, BLOCK_SIZE-1]
|
||||
|
||||
const int64_t* p_in1 = x1_data + (int64_t)n_idx * D;
|
||||
const int64_t* p_in2 = x2_data + (int64_t)n_idx * D;
|
||||
|
||||
|
||||
int thread_sum = 0;
|
||||
|
||||
for (int d = tid; d < D; d += BLOCK_SIZE) {{
|
||||
if (p_in1[d] != p_in2[d]) {{
|
||||
thread_sum++;
|
||||
}}
|
||||
}}
|
||||
s_data[tid] = thread_sum;
|
||||
__syncthreads();
|
||||
|
||||
|
||||
for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {{
|
||||
if (tid < offset) {{
|
||||
s_data[tid] += s_data[tid + offset];
|
||||
}}
|
||||
__syncthreads();
|
||||
}}
|
||||
|
||||
|
||||
if (tid == 0) {{
|
||||
output_data[n_idx] = (float)s_data[0];
|
||||
}}
|
||||
}}
|
||||
|
||||
// C++ 封装函数
|
||||
torch::Tensor hamming_distance_forward_cuda(
|
||||
torch::Tensor x1,
|
||||
torch::Tensor x2,
|
||||
int dim
|
||||
) {{
|
||||
|
||||
TORCH_CHECK(x1.dim() == 2, "CUDA Kernel only supports 2D (N, D) input");
|
||||
TORCH_CHECK(dim == 1, "CUDA Kernel only supports dim=1");
|
||||
TORCH_CHECK(x1.is_cuda() && x2.is_cuda(), "Inputs must be CUDA");
|
||||
TORCH_CHECK(x1.is_contiguous() && x2.is_contiguous(), "Inputs must be contiguous");
|
||||
TORCH_CHECK(x1.sizes() == x2.sizes(), "Input shapes must match");
|
||||
|
||||
|
||||
TORCH_CHECK(x1.scalar_type() == torch::kLong, "Input 1 must be torch.long");
|
||||
TORCH_CHECK(x2.scalar_type() == torch::kLong, "Input 2 must be torch.long");
|
||||
|
||||
const int64_t N = x1.size(0);
|
||||
const int64_t D = x1.size(1);
|
||||
|
||||
|
||||
auto output_options = torch::TensorOptions()
|
||||
.dtype(torch::kFloat32)
|
||||
.device(x1.device());
|
||||
auto output = torch::empty({{N}}, output_options);
|
||||
|
||||
dim3 grid_dim(N);
|
||||
dim3 block_dim(BLOCK_SIZE);
|
||||
|
||||
hamming_distance_fused_kernel<<<grid_dim, block_dim>>>(
|
||||
x1.data_ptr<int64_t>(),
|
||||
x2.data_ptr<int64_t>(),
|
||||
output.data_ptr<float>(),
|
||||
N, D
|
||||
);
|
||||
|
||||
return output;
|
||||
}}
|
||||
"""
|
||||
|
||||
nvcc_flags = ['-O3']
|
||||
|
||||
# JIT (Just-In-Time) 编译
|
||||
self.hamming_op = load_inline(
|
||||
name="hamming_dist_op_v1",
|
||||
cpp_sources=cpp_header,
|
||||
cuda_sources=cuda_source,
|
||||
functions=["hamming_distance_forward_cuda"],
|
||||
extra_cuda_cflags=nvcc_flags,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
|
||||
# 确保连续性
|
||||
x1_cont = x1.contiguous()
|
||||
x2_cont = x2.contiguous()
|
||||
|
||||
# CUDA 核函数要求 Long 类型
|
||||
# (如果输入不是 Long, .to() 会创建一个副本)
|
||||
x1_long = x1_cont.to(torch.long)
|
||||
x2_long = x2_cont.to(torch.long)
|
||||
|
||||
return self.hamming_op.hamming_distance_forward_cuda(
|
||||
x1_long,
|
||||
x2_long,
|
||||
self.dim
|
||||
)
|
||||
|
|
@ -0,0 +1,54 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# 常量定义
|
||||
# -------------------------------------------------------------
|
||||
N_BATCH = 100
|
||||
D_VECTOR = 128
|
||||
|
||||
# 假设 dim=1
|
||||
DIM = 1
|
||||
|
||||
|
||||
# -------------------------------------------------------------
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
汉明距离 (Hamming Distance) 的纯 PyTorch 基准实现
|
||||
"""
|
||||
|
||||
def __init__(self, dim=1):
|
||||
super().__init__()
|
||||
# 假设 dim=1
|
||||
self.dim = dim
|
||||
|
||||
def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
|
||||
# x1, x2: (N, D), 假设为 long/int/bool 类型
|
||||
|
||||
# 1. 比较 (x1 != x2)
|
||||
# 如果 x1=[1, 0, 1], x2=[1, 1, 1]
|
||||
# diff=[False, True, False]
|
||||
diff = (x1 != x2)
|
||||
|
||||
# 2. 求和 (Sum)
|
||||
# False.sum() = 0, True.sum() = 1
|
||||
# sum([0, 1, 0]) = 1
|
||||
# 我们转换为 float 以匹配 CUDA 版本的输出类型
|
||||
return torch.sum(diff, dim=self.dim).to(torch.float32)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
生成两个 (N, D) 形状的 *整数* 输入
|
||||
(汉明距离的标准输入)
|
||||
"""
|
||||
# 随机生成 0 或 1
|
||||
input1 = torch.randint(0, 2, (N_BATCH, D_VECTOR), dtype=torch.long)
|
||||
input2 = torch.randint(0, 2, (N_BATCH, D_VECTOR), dtype=torch.long)
|
||||
return [input1, input2]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [DIM]
|
||||
|
|
@ -0,0 +1,91 @@
|
|||
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
|
||||
|
||||
Hamming Distance: Measure of dissimilarity between binary strings/vectors
|
||||
|
||||
Block-Level Parallel Reduction: Uses shared memory and tree reduction within thread blocks
|
||||
|
||||
One Block Per Sample: Each CUDA block processes one complete sample pair (N blocks for N samples)
|
||||
|
||||
Strided Memory Access: Threads process elements with stride equal to block size for coalesced memory access
|
||||
|
||||
Shared Memory Optimization: Uses __shared__ array for intermediate results and reduction
|
||||
|
||||
Tree Reduction Pattern: Binary tree reduction within block using __syncthreads()
|
||||
|
||||
Integer Arithmetic: Optimized for int64_t data type comparisons
|
||||
|
||||
Type Conversion: Converts final integer count to float32 output
|
||||
|
||||
Memory Coalescing: Ensures contiguous tensor layout for optimal memory access
|
||||
|
||||
Input Validation: Comprehensive type and shape checking in C++ wrapper
|
||||
|
||||
Grid-Stride Loops: Efficiently handles variable dimension sizes
|
||||
|
||||
Compiler Optimization: Uses -O3 flag for maximum performance
|
||||
|
||||
Tensor Contiguity Enforcement: Ensures optimal memory layout in PyTorch wrapper
|
||||
|
||||
Automatic Type Promotion: Converts input tensors to torch.long if needed
|
||||
|
||||
|
||||
|
||||
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_BATCH = 100
|
||||
D_VECTOR = 128
|
||||
|
||||
# 假设 dim=1
|
||||
DIM = 1
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
汉明距离 (Hamming Distance) 的纯 PyTorch 基准实现
|
||||
"""
|
||||
def __init__(self, dim=1):
|
||||
super().__init__()
|
||||
# 假设 dim=1
|
||||
self.dim = dim
|
||||
|
||||
def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
|
||||
# x1, x2: (N, D), 假设为 long/int/bool 类型
|
||||
|
||||
# 1. 比较 (x1 != x2)
|
||||
# 如果 x1=[1, 0, 1], x2=[1, 1, 1]
|
||||
# diff=[False, True, False]
|
||||
diff = (x1 != x2)
|
||||
|
||||
# 2. 求和 (Sum)
|
||||
# False.sum() = 0, True.sum() = 1
|
||||
# sum([0, 1, 0]) = 1
|
||||
# 我们转换为 float 以匹配 CUDA 版本的输出类型
|
||||
return torch.sum(diff, dim=self.dim).to(torch.float32)
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
生成两个 (N, D) 形状的 *整数* 输入
|
||||
(汉明距离的标准输入)
|
||||
"""
|
||||
# 随机生成 0 或 1
|
||||
input1 = torch.randint(0, 2, (N_BATCH, D_VECTOR), dtype=torch.long)
|
||||
input2 = torch.randint(0, 2, (N_BATCH, D_VECTOR), dtype=torch.long)
|
||||
return [input1, input2]
|
||||
|
||||
def get_init_inputs():
|
||||
return [DIM]
|
||||
|
|
@ -0,0 +1,78 @@
|
|||
import torch
|
||||
import time
|
||||
from HammingDistance_torch import Model, get_inputs, get_init_inputs
|
||||
from HammingDistance_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