forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'feat:add good performance Hamming #20' (#218) from wut0n/GPUCodeForces:Hamming into main
This commit is contained in:
commit
b8f9ea27f5
|
|
@ -0,0 +1,213 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
hamming_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// 基础版本 - 简单比较
|
||||
__global__ void hamming_kernel_basic(
|
||||
const float* __restrict__ x,
|
||||
const float* __restrict__ y,
|
||||
float* __restrict__ distances,
|
||||
int batch_size,
|
||||
int feature_dim
|
||||
) {
|
||||
int sample_idx = blockIdx.x;
|
||||
if (sample_idx >= batch_size) return;
|
||||
|
||||
int base = sample_idx * feature_dim;
|
||||
int distance = 0;
|
||||
|
||||
// Count different elements
|
||||
for (int i = 0; i < feature_dim; i++) {
|
||||
if (x[base + i] != y[base + i]) {
|
||||
distance++;
|
||||
}
|
||||
}
|
||||
|
||||
distances[sample_idx] = distance;
|
||||
}
|
||||
|
||||
// 向量化版本 - float4优化
|
||||
__global__ void hamming_kernel_vectorized(
|
||||
const float* __restrict__ x,
|
||||
const float* __restrict__ y,
|
||||
float* __restrict__ distances,
|
||||
int batch_size,
|
||||
int feature_dim
|
||||
) {
|
||||
int sample_idx = blockIdx.x;
|
||||
if (sample_idx >= batch_size) return;
|
||||
|
||||
int tid = threadIdx.x;
|
||||
int base = sample_idx * feature_dim;
|
||||
|
||||
// 使用共享内存进行归约
|
||||
extern __shared__ float shared_sum[];
|
||||
shared_sum[tid] = 0.0f;
|
||||
|
||||
// 每个线程处理4个元素（float4向量化）
|
||||
int stride = blockDim.x * 4;
|
||||
for (int dim = tid * 4; dim < feature_dim; dim += stride) {
|
||||
// 确保不越界
|
||||
if (dim + 3 < feature_dim) {
|
||||
float4 x_val = *reinterpret_cast<const float4*>(&x[base + dim]);
|
||||
float4 y_val = *reinterpret_cast<const float4*>(&y[base + dim]);
|
||||
|
||||
// 比较每个分量
|
||||
shared_sum[tid] += (x_val.x != y_val.x) + (x_val.y != y_val.y) +
|
||||
(x_val.z != y_val.z) + (x_val.w != y_val.w);
|
||||
} else {
|
||||
// 处理剩余元素
|
||||
for (int i = dim; i < feature_dim; i++) {
|
||||
if (x[base + i] != y[base + i]) {
|
||||
shared_sum[tid] += 1.0f;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// 块内归约求和
|
||||
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
shared_sum[tid] += shared_sum[tid + stride];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// 第一个线程写入结果
|
||||
if (tid == 0) {
|
||||
distances[sample_idx] = shared_sum[0];
|
||||
}
|
||||
}
|
||||
|
||||
// 融合版本 - 单kernel完成所有计算
|
||||
__global__ void hamming_kernel_fused(
|
||||
const float* __restrict__ x,
|
||||
const float* __restrict__ y,
|
||||
float* __restrict__ distances,
|
||||
int batch_size,
|
||||
int feature_dim
|
||||
) {
|
||||
int sample_idx = blockIdx.x;
|
||||
if (sample_idx >= batch_size) return;
|
||||
|
||||
int tid = threadIdx.x;
|
||||
int base = sample_idx * feature_dim;
|
||||
|
||||
// 使用共享内存存储部分结果
|
||||
extern __shared__ float shared_data[];
|
||||
shared_data[tid] = 0.0f;
|
||||
|
||||
// 每个线程处理多个元素
|
||||
int elements_per_thread = (feature_dim + blockDim.x - 1) / blockDim.x;
|
||||
int start_idx = tid * elements_per_thread;
|
||||
int end_idx = min(start_idx + elements_per_thread, feature_dim);
|
||||
|
||||
// 计算分配给这个线程的元素
|
||||
for (int i = start_idx; i < end_idx; i++) {
|
||||
if (x[base + i] != y[base + i]) {
|
||||
shared_data[tid] += 1.0f;
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// 归约求最终结果
|
||||
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
shared_data[tid] += shared_data[tid + stride];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// 第一个线程写入结果
|
||||
if (tid == 0) {
|
||||
distances[sample_idx] = shared_data[0];
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor hamming_cuda(
|
||||
torch::Tensor x,
|
||||
torch::Tensor y,
|
||||
std::string mode = "fused"
|
||||
) {
|
||||
TORCH_CHECK(x.scalar_type() == torch::kFloat32, "X must be float32");
|
||||
TORCH_CHECK(y.scalar_type() == torch::kFloat32, "Y must be float32");
|
||||
|
||||
auto x_contig = x.contiguous();
|
||||
auto y_contig = y.contiguous();
|
||||
|
||||
int batch_size = x_contig.size(0);
|
||||
int feature_dim = x_contig.size(1);
|
||||
|
||||
auto distances = torch::zeros({batch_size}, x.options());
|
||||
|
||||
if (mode == "fused") {
|
||||
// 融合版本 - 推荐
|
||||
const int block_size = 256;
|
||||
size_t shared_mem = block_size * sizeof(float);
|
||||
|
||||
hamming_kernel_fused<<<batch_size, block_size, shared_mem>>>(
|
||||
x_contig.data_ptr<float>(),
|
||||
y_contig.data_ptr<float>(),
|
||||
distances.data_ptr<float>(),
|
||||
batch_size,
|
||||
feature_dim
|
||||
);
|
||||
} else if (mode == "vectorized") {
|
||||
// 向量化版本
|
||||
const int block_size = 64; // 每个线程处理4个元素
|
||||
size_t shared_mem = block_size * sizeof(float);
|
||||
|
||||
hamming_kernel_vectorized<<<batch_size, block_size, shared_mem>>>(
|
||||
x_contig.data_ptr<float>(),
|
||||
y_contig.data_ptr<float>(),
|
||||
distances.data_ptr<float>(),
|
||||
batch_size,
|
||||
feature_dim
|
||||
);
|
||||
} else {
|
||||
// 基础版本
|
||||
hamming_kernel_basic<<<batch_size, 1>>>(
|
||||
x_contig.data_ptr<float>(),
|
||||
y_contig.data_ptr<float>(),
|
||||
distances.data_ptr<float>(),
|
||||
batch_size,
|
||||
feature_dim
|
||||
);
|
||||
}
|
||||
|
||||
return distances;
|
||||
}
|
||||
"""
|
||||
|
||||
hamming_cpp_source = """
|
||||
torch::Tensor hamming_cuda(torch::Tensor x, torch::Tensor y, std::string mode);
|
||||
"""
|
||||
|
||||
# Compile the inline CUDA code
|
||||
hamming = load_inline(
|
||||
name="hamming",
|
||||
cpp_sources=hamming_cpp_source,
|
||||
cuda_sources=hamming_source,
|
||||
functions=["hamming_cuda"],
|
||||
extra_cuda_cflags=[
|
||||
"-O3",
|
||||
"--use_fast_math",
|
||||
"-gencode=arch=compute_80,code=sm_80"
|
||||
],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, mode="fused"):
|
||||
super(ModelNew, self).__init__()
|
||||
self.mode = mode
|
||||
self.hamming = hamming
|
||||
|
||||
def forward(self, x, y):
|
||||
return self.hamming.hamming_cuda(x, y, self.mode)
|
||||
|
|
@ -0,0 +1,49 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Hamming Distance implementation.
|
||||
Computes the Hamming distance between two sets of vectors.
|
||||
"""
|
||||
def __init__(self):
|
||||
super(Model, self).__init__()
|
||||
|
||||
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute Hamming distance between x and y.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): First set of vectors [batch_size, feature_dim]
|
||||
y (torch.Tensor): Second set of vectors [batch_size, feature_dim]
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Hamming distances [batch_size]
|
||||
"""
|
||||
# Input validation
|
||||
if x.shape != y.shape:
|
||||
raise ValueError(f"Input tensors must have the same shape, got {x.shape} and {y.shape}")
|
||||
|
||||
if x.dim() != 2:
|
||||
raise ValueError(f"Input tensors must be 2D, got {x.dim()}D")
|
||||
|
||||
# Compute Hamming distance: count of different elements
|
||||
# Step 1: Compare elements (x != y gives boolean tensor)
|
||||
diff = (x != y)
|
||||
|
||||
# Step 2: Convert to float and sum along feature dimension
|
||||
distance = torch.sum(diff.float(), dim=1)
|
||||
|
||||
return distance
|
||||
|
||||
batch_size = 256
|
||||
feature_dim = 512
|
||||
|
||||
def get_inputs():
|
||||
# Generate two sets of integer vectors (0 or 1 for simplicity)
|
||||
x = torch.randint(0, 2, (batch_size, feature_dim)).float()
|
||||
y = torch.randint(0, 2, (batch_size, feature_dim)).float()
|
||||
return [x, y]
|
||||
|
||||
def get_init_inputs():
|
||||
return [] # No special initialization inputs needed
|
||||
|
|
@ -0,0 +1,138 @@
|
|||
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) -> None:
|
||||
super().init()
|
||||
|
||||
def forward(self, a, b):
|
||||
return a + b
|
||||
def get_inputs():
|
||||
# randomly generate input tensors based on the model architecture
|
||||
a = torch.randn(1, 128).cuda()
|
||||
b = torch.randn(1, 128).cuda()
|
||||
return [a, b]
|
||||
|
||||
def get_init_inputs():
|
||||
# randomly generate tensors required for initialization based on the model architecture
|
||||
return []
|
||||
|
||||
|
||||
|
||||
The example new arch with custom CUDA kernels looks like this:
|
||||
|
||||
python
|
||||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
relu_source = “”"
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
global void relu_kernel(const float* x, float* y, int size) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < size) {
|
||||
y[idx] = fmaxf(x[idx], 0.f);
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor relu_cuda(torch::Tensor x) {
|
||||
auto size = x.numel();
|
||||
auto y = torch::empty_like(x);
|
||||
const int block_size = 256;
|
||||
int num_blocks = (size + block_size - 1) / block_size;
|
||||
relu_kernel<<<num_blocks, block_size>>>(x.data_ptr<float>(), y.data_ptr<float>(), size);
|
||||
return y;
|
||||
}
|
||||
“”"
|
||||
|
||||
relu_cpp_source = “”"
|
||||
torch::Tensor relu_cuda(torch::Tensor x);
|
||||
“”"
|
||||
|
||||
Compile the inline CUDA code
|
||||
relu = load_inline(
|
||||
name=“relu”,
|
||||
cpp_sources=relu_cpp_source,
|
||||
cuda_sources=relu_source,
|
||||
functions=[“relu_cuda”],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def init(self):
|
||||
super(ModelNew, self).init()
|
||||
self.relu = relu # The module containing the kernel
|
||||
|
||||
def forward(self, x):
|
||||
return self.relu.relu_cuda(x)
|
||||
def get_inputs():
|
||||
# randomly generate input tensors based on the model architecture
|
||||
a = torch.randn(1, 128).cuda()
|
||||
b = torch.randn(1, 128).cuda()
|
||||
return [a, b]
|
||||
|
||||
def get_init_inputs():
|
||||
# randomly generate tensors required for initialization based on the model architecture
|
||||
return []
|
||||
|
||||
|
||||
|
||||
You are given the following architecture:
|
||||
|
||||
python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
“”"
|
||||
Hamming Distance implementation.
|
||||
Computes the Hamming distance between two sets of vectors.
|
||||
“”"
|
||||
def init(self):
|
||||
super(Model, self).init()
|
||||
|
||||
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute Hamming distance between x and y.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): First set of vectors [batch_size, feature_dim]
|
||||
y (torch.Tensor): Second set of vectors [batch_size, feature_dim]
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Hamming distances [batch_size]
|
||||
"""
|
||||
# Input validation
|
||||
if x.shape != y.shape:
|
||||
raise ValueError(f"Input tensors must have the same shape, got {x.shape} and {y.shape}")
|
||||
|
||||
if x.dim() != 2:
|
||||
raise ValueError(f"Input tensors must be 2D, got {x.dim()}D")
|
||||
|
||||
# Compute Hamming distance: count of different elements
|
||||
# Step 1: Compare elements (x != y gives boolean tensor)
|
||||
diff = (x != y)
|
||||
|
||||
# Step 2: Convert to float and sum along feature dimension
|
||||
distance = torch.sum(diff.float(), dim=1)
|
||||
|
||||
return distance
|
||||
batch_size = 256
|
||||
feature_dim = 512
|
||||
|
||||
def get_inputs():
|
||||
# Generate two sets of integer vectors (0 or 1 for simplicity)
|
||||
x = torch.randint(0, 2, (batch_size, feature_dim)).float()
|
||||
y = torch.randint(0, 2, (batch_size, feature_dim)).float()
|
||||
return [x, y]
|
||||
|
||||
def get_init_inputs():
|
||||
return [] # No special initialization inputs needed
|
||||
|
||||
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from hamming_torchcode import Model, get_inputs, get_init_inputs
|
||||
from hamming_cudacode 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 mahalanobis 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA mahalanobis 平均执行时间: {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