forked from ccf-ai-infra/GPUCodeForces
feat:add high performance euclidean #14
This commit is contained in:
parent
f876a28ada
commit
485ded3dff
|
|
@ -0,0 +1,210 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
euclidean_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// 基础版本 - 4元素展开
|
||||
__global__ void euclidean_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;
|
||||
float sum_squared = 0.0f;
|
||||
|
||||
// 4元素展开计算
|
||||
int dim = 0;
|
||||
for (; dim <= feature_dim - 4; dim += 4) {
|
||||
float x0 = x[base + dim];
|
||||
float x1 = x[base + dim + 1];
|
||||
float x2 = x[base + dim + 2];
|
||||
float x3 = x[base + dim + 3];
|
||||
|
||||
float y0 = y[base + dim];
|
||||
float y1 = y[base + dim + 1];
|
||||
float y2 = y[base + dim + 2];
|
||||
float y3 = y[base + dim + 3];
|
||||
|
||||
float diff0 = x0 - y0;
|
||||
float diff1 = x1 - y1;
|
||||
float diff2 = x2 - y2;
|
||||
float diff3 = x3 - y3;
|
||||
|
||||
sum_squared += diff0*diff0 + diff1*diff1 + diff2*diff2 + diff3*diff3;
|
||||
}
|
||||
|
||||
// 处理剩余元素
|
||||
for (; dim < feature_dim; dim++) {
|
||||
float diff = x[base + dim] - y[base + dim];
|
||||
sum_squared += diff * diff;
|
||||
}
|
||||
|
||||
distances[sample_idx] = sqrtf(sum_squared);
|
||||
}
|
||||
|
||||
// 向量化版本 - 使用float4
|
||||
__global__ void euclidean_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;
|
||||
|
||||
// 使用float4向量化
|
||||
const float4* x_vec = reinterpret_cast<const float4*>(x);
|
||||
const float4* y_vec = reinterpret_cast<const float4*>(y);
|
||||
|
||||
int feature_dim_vec = feature_dim / 4;
|
||||
int base_vec = sample_idx * feature_dim_vec;
|
||||
|
||||
float sum_squared = 0.0f;
|
||||
|
||||
// 向量化计算
|
||||
for (int dim_vec = 0; dim_vec < feature_dim_vec; dim_vec++) {
|
||||
float4 x_val = x_vec[base_vec + dim_vec];
|
||||
float4 y_val = y_vec[base_vec + dim_vec];
|
||||
|
||||
float diff_x = x_val.x - y_val.x;
|
||||
float diff_y = x_val.y - y_val.y;
|
||||
float diff_z = x_val.z - y_val.z;
|
||||
float diff_w = x_val.w - y_val.w;
|
||||
|
||||
sum_squared += diff_x*diff_x + diff_y*diff_y + diff_z*diff_z + diff_w*diff_w;
|
||||
}
|
||||
|
||||
// 处理剩余元素
|
||||
for (int dim = feature_dim_vec * 4; dim < feature_dim; dim++) {
|
||||
float diff = x[sample_idx * feature_dim + dim] - y[sample_idx * feature_dim + dim];
|
||||
sum_squared += diff * diff;
|
||||
}
|
||||
|
||||
distances[sample_idx] = sqrtf(sum_squared);
|
||||
}
|
||||
|
||||
// 融合版本 - 融合sqrt和快速数学
|
||||
__global__ void euclidean_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;
|
||||
|
||||
// 使用float4向量化
|
||||
const float4* x_vec = reinterpret_cast<const float4*>(x);
|
||||
const float4* y_vec = reinterpret_cast<const float4*>(y);
|
||||
|
||||
int feature_dim_vec = feature_dim / 4;
|
||||
int base_vec = sample_idx * feature_dim_vec;
|
||||
|
||||
float sum_squared = 0.0f;
|
||||
|
||||
// 融合计算 - 减少中间变量
|
||||
for (int dim_vec = 0; dim_vec < feature_dim_vec; dim_vec++) {
|
||||
float4 x_val = x_vec[base_vec + dim_vec];
|
||||
float4 y_val = y_vec[base_vec + dim_vec];
|
||||
|
||||
// 直接计算平方和，不存储中间结果
|
||||
sum_squared += (x_val.x - y_val.x) * (x_val.x - y_val.x) +
|
||||
(x_val.y - y_val.y) * (x_val.y - y_val.y) +
|
||||
(x_val.z - y_val.z) * (x_val.z - y_val.z) +
|
||||
(x_val.w - y_val.w) * (x_val.w - y_val.w);
|
||||
}
|
||||
|
||||
// 处理剩余元素
|
||||
for (int dim = feature_dim_vec * 4; dim < feature_dim; dim++) {
|
||||
float diff = x[sample_idx * feature_dim + dim] - y[sample_idx * feature_dim + dim];
|
||||
sum_squared += diff * diff;
|
||||
}
|
||||
|
||||
// 使用快速sqrt
|
||||
distances[sample_idx] = sqrtf(sum_squared);
|
||||
}
|
||||
|
||||
torch::Tensor euclidean_cuda(
|
||||
torch::Tensor x,
|
||||
torch::Tensor y,
|
||||
std::string mode = "vectorized"
|
||||
) {
|
||||
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") {
|
||||
// 融合版本
|
||||
euclidean_kernel_fused<<<batch_size, 1>>>(
|
||||
x_contig.data_ptr<float>(),
|
||||
y_contig.data_ptr<float>(),
|
||||
distances.data_ptr<float>(),
|
||||
batch_size,
|
||||
feature_dim
|
||||
);
|
||||
} else if (mode == "vectorized") {
|
||||
// 向量化版本（默认）
|
||||
euclidean_kernel_vectorized<<<batch_size, 1>>>(
|
||||
x_contig.data_ptr<float>(),
|
||||
y_contig.data_ptr<float>(),
|
||||
distances.data_ptr<float>(),
|
||||
batch_size,
|
||||
feature_dim
|
||||
);
|
||||
} else {
|
||||
// 基础版本
|
||||
euclidean_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;
|
||||
}
|
||||
"""
|
||||
|
||||
euclidean_cpp_source = """
|
||||
torch::Tensor euclidean_cuda(torch::Tensor x, torch::Tensor y, std::string mode);
|
||||
"""
|
||||
|
||||
# 编译CUDA代码
|
||||
euclidean = load_inline(
|
||||
name="euclidean",
|
||||
cpp_sources=euclidean_cpp_source,
|
||||
cuda_sources=euclidean_source,
|
||||
functions=["euclidean_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="vectorized"):
|
||||
super(ModelNew, self).__init__()
|
||||
self.mode = mode
|
||||
self.euclidean = euclidean
|
||||
|
||||
def forward(self, x, y):
|
||||
return self.euclidean.euclidean_cuda(x, y, self.mode)
|
||||
|
|
@ -0,0 +1,45 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Euclidean Distance implementation.
|
||||
Computes the Euclidean 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 Euclidean 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: Euclidean distances [batch_size]
|
||||
"""
|
||||
# Compute squared differences
|
||||
diff = x - y
|
||||
squared_diff = diff * diff
|
||||
|
||||
# Sum along feature dimension
|
||||
sum_squared = torch.sum(squared_diff, dim=1)
|
||||
|
||||
# Take square root
|
||||
distances = torch.sqrt(sum_squared)
|
||||
|
||||
return distances
|
||||
|
||||
batch_size = 256
|
||||
feature_dim = 512
|
||||
|
||||
def get_inputs():
|
||||
# Generate two sets of vectors
|
||||
x = torch.randn(batch_size, feature_dim)
|
||||
y = torch.randn(batch_size, feature_dim)
|
||||
return [x, y]
|
||||
|
||||
def get_init_inputs():
|
||||
return [] # No special initialization inputs needed
|
||||
|
|
@ -0,0 +1,288 @@
|
|||
# CUDA Euclidean Distance Optimization Expert Prompt
|
||||
|
||||
You are a CUDA optimization expert specializing in writing high-performance CUDA implementations for mathematical operations. Your task is to generate an optimized CUDA implementation for Euclidean Distance calculation that achieves significant speedup over PyTorch's native implementation.
|
||||
|
||||
## Core Optimization Principles
|
||||
|
||||
### 1. Performance Optimization Strategies
|
||||
- **Vectorization**: Prioritize using float4 for memory access vectorization to maximize memory bandwidth utilization
|
||||
- **Loop Unrolling**: Use 4-element unrolling to reduce loop overhead and improve instruction-level parallelism
|
||||
- **Memory Coalescing**: Ensure memory access patterns are coalesced and contiguous
|
||||
- **Single Block Strategy**: Use one block per sample to minimize synchronization overhead
|
||||
- **Fast Math**: Use --use_fast_math and fast mathematical functions (sqrtf)
|
||||
|
||||
### 2. Three-Tier Optimization Approach
|
||||
- **Basic Version**: 4-element loop unrolling for fundamental optimization
|
||||
- **Vectorized Version**: float4 vectorization for memory bandwidth optimization (recommended)
|
||||
- **Fused Version**: Fused computation with minimal intermediate variables
|
||||
|
||||
### 3. Memory Access Optimization
|
||||
- **Contiguous Memory**: Ensure input tensors are contiguous
|
||||
- **Vectorized Loading**: Use float4 pointer casting for 4-element simultaneous access
|
||||
- **Cache Efficiency**: Optimize for cache-friendly access patterns
|
||||
|
||||
## Algorithm Analysis
|
||||
|
||||
### Euclidean Distance Formula
|
||||
distance = √Σ(x_i - y_i)²
|
||||
|
||||
Computational Steps:
|
||||
|
||||
Calculate difference: diff = x_i - y_i
|
||||
Square the difference: diff²
|
||||
Sum all squared differences: Σ diff²
|
||||
Take square root: √sum
|
||||
|
||||
|
||||
### Optimization Opportunities
|
||||
- **High Compute Density**: Large number of arithmetic operations
|
||||
- **Memory Intensive**: Two large tensor read operations
|
||||
- **High Parallelism**: Each sample computation is independent
|
||||
- **Vectorization Friendly**: float4 vectorization shows significant benefits
|
||||
|
||||
## Code Template Structure
|
||||
|
||||
### CUDA Kernel Implementation
|
||||
cpp
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// Basic Version - 4-element unrolling
|
||||
global void euclidean_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;
|
||||
float sum_squared = 0.0f;
|
||||
|
||||
// 4-element unrolled computation
|
||||
int dim = 0;
|
||||
for (; dim <= feature_dim - 4; dim += 4) {
|
||||
float x0 = x[base + dim];
|
||||
float x1 = x[base + dim + 1];
|
||||
float x2 = x[base + dim + 2];
|
||||
float x3 = x[base + dim + 3];
|
||||
|
||||
float y0 = y[base + dim];
|
||||
float y1 = y[base + dim + 1];
|
||||
float y2 = y[base + dim + 2];
|
||||
float y3 = y[base + dim + 3];
|
||||
|
||||
float diff0 = x0 - y0;
|
||||
float diff1 = x1 - y1;
|
||||
float diff2 = x2 - y2;
|
||||
float diff3 = x3 - y3;
|
||||
|
||||
sum_squared += diff0*diff0 + diff1*diff1 + diff2*diff2 + diff3*diff3;
|
||||
}
|
||||
|
||||
// Handle remaining elements
|
||||
for (; dim < feature_dim; dim++) {
|
||||
float diff = x[base + dim] - y[base + dim];
|
||||
sum_squared += diff * diff;
|
||||
}
|
||||
|
||||
distances[sample_idx] = sqrtf(sum_squared);
|
||||
}
|
||||
|
||||
// Vectorized Version - float4 optimization (recommended)
|
||||
global void euclidean_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;
|
||||
|
||||
// Use float4 vectorization
|
||||
const float4* x_vec = reinterpret_cast<const float4*>(x);
|
||||
const float4* y_vec = reinterpret_cast<const float4*>(y);
|
||||
|
||||
int feature_dim_vec = feature_dim / 4;
|
||||
int base_vec = sample_idx * feature_dim_vec;
|
||||
|
||||
float sum_squared = 0.0f;
|
||||
|
||||
// Vectorized computation
|
||||
for (int dim_vec = 0; dim_vec < feature_dim_vec; dim_vec++) {
|
||||
float4 x_val = x_vec[base_vec + dim_vec];
|
||||
float4 y_val = y_vec[base_vec + dim_vec];
|
||||
|
||||
float diff_x = x_val.x - y_val.x;
|
||||
float diff_y = x_val.y - y_val.y;
|
||||
float diff_z = x_val.z - y_val.z;
|
||||
float diff_w = x_val.w - y_val.w;
|
||||
|
||||
sum_squared += diff_x*diff_x + diff_y*diff_y + diff_z*diff_z + diff_w*diff_w;
|
||||
}
|
||||
|
||||
// Handle remaining elements
|
||||
for (int dim = feature_dim_vec * 4; dim < feature_dim; dim++) {
|
||||
float diff = x[sample_idx * feature_dim + dim] - y[sample_idx * feature_dim + dim];
|
||||
sum_squared += diff * diff;
|
||||
}
|
||||
|
||||
distances[sample_idx] = sqrtf(sum_squared);
|
||||
}
|
||||
|
||||
// Fused Version - minimal intermediate variables
|
||||
global void euclidean_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;
|
||||
|
||||
// Use float4 vectorization
|
||||
const float4* x_vec = reinterpret_cast<const float4*>(x);
|
||||
const float4* y_vec = reinterpret_cast<const float4*>(y);
|
||||
|
||||
int feature_dim_vec = feature_dim / 4;
|
||||
int base_vec = sample_idx * feature_dim_vec;
|
||||
|
||||
float sum_squared = 0.0f;
|
||||
|
||||
// Fused computation - reduce intermediate variables
|
||||
for (int dim_vec = 0; dim_vec < feature_dim_vec; dim_vec++) {
|
||||
float4 x_val = x_vec[base_vec + dim_vec];
|
||||
float4 y_val = y_vec[base_vec + dim_vec];
|
||||
|
||||
// Direct computation without storing intermediate results
|
||||
sum_squared += (x_val.x - y_val.x) * (x_val.x - y_val.x) +
|
||||
(x_val.y - y_val.y) * (x_val.y - y_val.y) +
|
||||
(x_val.z - y_val.z) * (x_val.z - y_val.z) +
|
||||
(x_val.w - y_val.w) * (x_val.w - y_val.w);
|
||||
}
|
||||
|
||||
// Handle remaining elements
|
||||
for (int dim = feature_dim_vec * 4; dim < feature_dim; dim++) {
|
||||
float diff = x[sample_idx * feature_dim + dim] - y[sample_idx * feature_dim + dim];
|
||||
sum_squared += diff * diff;
|
||||
}
|
||||
|
||||
// Use fast sqrt
|
||||
distances[sample_idx] = sqrtf(sum_squared);
|
||||
}
|
||||
|
||||
torch::Tensor euclidean_cuda(
|
||||
torch::Tensor x,
|
||||
torch::Tensor y,
|
||||
std::string mode = “vectorized”
|
||||
) {
|
||||
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") {
|
||||
euclidean_kernel_fused<<<batch_size, 1>>>(
|
||||
x_contig.data_ptr<float>(),
|
||||
y_contig.data_ptr<float>(),
|
||||
distances.data_ptr<float>(),
|
||||
batch_size,
|
||||
feature_dim
|
||||
);
|
||||
} else if (mode == "vectorized") {
|
||||
euclidean_kernel_vectorized<<<batch_size, 1>>>(
|
||||
x_contig.data_ptr<float>(),
|
||||
y_contig.data_ptr<float>(),
|
||||
distances.data_ptr<float>(),
|
||||
batch_size,
|
||||
feature_dim
|
||||
);
|
||||
} else {
|
||||
euclidean_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;
|
||||
}
|
||||
|
||||
|
||||
|
||||
### Python Binding Template
|
||||
python
|
||||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
cpp_source = “”"
|
||||
torch::Tensor euclidean_cuda(torch::Tensor x, torch::Tensor y, std::string mode);
|
||||
“”"
|
||||
|
||||
euclidean = load_inline(
|
||||
name=“euclidean”,
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=[“euclidean_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=“vectorized”):
|
||||
super(ModelNew, self).init()
|
||||
self.mode = mode
|
||||
self.euclidean = euclidean
|
||||
|
||||
def forward(self, x, y):
|
||||
return self.euclidean.euclidean_cuda(x, y, self.mode)
|
||||
|
||||
|
||||
## Expected Performance Characteristics
|
||||
|
||||
### Performance Targets
|
||||
- **Basic Version**: 1.2-1.5x speedup
|
||||
- **Vectorized Version**: 1.4-1.6x speedup (recommended)
|
||||
- **Fused Version**: 1.3-1.5x speedup
|
||||
|
||||
### Optimal Usage
|
||||
- **Batch Size**: 64-256 for best performance
|
||||
- **Feature Dimension**: 128-1024 for optimal vectorization
|
||||
- **Data Type**: float32 required for vectorization
|
||||
- **Memory Layout**: Ensure contiguous tensors
|
||||
|
||||
## Key Success Factors
|
||||
|
||||
1. **Vectorization Priority**: float4 vectorization is the most critical optimization
|
||||
2. **Memory Coalescing**: Contiguous memory access patterns are essential
|
||||
3. **Single Block Strategy**: Minimizes synchronization overhead
|
||||
4. **Precision First**: Ensure numerical accuracy before optimization
|
||||
5. **Clean Implementation**: Avoid over-engineering for maintainability
|
||||
|
||||
## Task Requirements
|
||||
|
||||
Generate a complete CUDA implementation for Euclidean Distance calculation that:
|
||||
1. Achieves 1.4+ speedup over PyTorch native implementation
|
||||
2. Maintains numerical precision (rtol=1e-03, atol=1e-6)
|
||||
3. Provides three optimization modes (basic, vectorized, fused)
|
||||
4. Uses the recommended vectorized approach as default
|
||||
5. Includes proper error handling and memory safety checks
|
||||
6. Follows the exact structure and optimization strategies outlined above
|
||||
|
||||
The implementation should be production-ready with clear, maintainable code that demonstrates the optimization principles described.
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from euclidean_torchcode import Model, get_inputs, get_init_inputs
|
||||
from euclidean_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 euclidean 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA euclidean 平均执行时间: {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