xxq250
/
GPUCodeForces
forked from ccf-ai-infra/GPUCodeForces
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request 'finish cosineloss Operator #27' (#53) from ZZZJ/GPUCodeForces:cosineloss into main

This commit is contained in:
Kuohais 2025-11-12 09:42:45 +08:00
parent 7a9dc401f4 679036215e
commit c1cba4d39c
4 changed files with 467 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

317
S1/27/cosineloss_cuda.py Normal file
View File

@ -0,0 +1,317 @@
# cosineloss_cuda.py
import torch
from torch.utils.cpp_extension import load_inline
from cosineloss_torch import BATCH_SIZE, EMBEDDING_DIM, DIM, MARGIN
TOTAL_ELEMENTS = BATCH_SIZE * EMBEDDING_DIM
class ModelNew(torch.nn.Module):
def __init__(self):
super().__init__()
self._compile_cuda_kernel()
def _compile_cuda_kernel(self):
cpp_source = """
#include <torch/extension.h>
#include <cmath>
torch::Tensor cosine_forward_cuda(torch::Tensor x1, torch::Tensor x2, torch::Tensor y);
"""
cuda_source = """
#include <cuda_runtime.h>
#include <cmath>
#include <algorithm>
#define BLOCK_SIZE 1024
#define VEC_SIZE 4
#define MARGIN_VAL {margin_val}
struct CosineResult {{
double dot_sum;
double norm1_sq_sum;
double norm2_sq_sum;
}};
__global__ void cosine_pass1_kernel(
const float* __restrict__ x1,
const float* __restrict__ x2,
CosineResult* __restrict__ results,
int N_pairs,
int D_emb
) {{
int pair_idx = blockIdx.x;
if (pair_idx >= N_pairs) return;
__shared__ double sh_dot[BLOCK_SIZE];
__shared__ double sh_norm1[BLOCK_SIZE];
__shared__ double sh_norm2[BLOCK_SIZE];
double thread_dot_sum = 0.0;
double thread_norm1_sum = 0.0;
double thread_norm2_sum = 0.0;
int offset = pair_idx * D_emb;
int D_vec = D_emb / VEC_SIZE;
const float4* x1_4 = (const float4*)(x1 + offset);
const float4* x2_4 = (const float4*)(x2 + offset);
for (int d_vec = threadIdx.x; d_vec < D_vec; d_vec += blockDim.x) {{
float4 v1 = x1_4[d_vec];
float4 v2 = x2_4[d_vec];
// Dot Product
thread_dot_sum += (double)v1.x * (double)v2.x;
thread_dot_sum += (double)v1.y * (double)v2.y;
thread_dot_sum += (double)v1.z * (double)v2.z;
thread_dot_sum += (double)v1.w * (double)v2.w;
// Norm 1 Squared
thread_norm1_sum += (double)v1.x * (double)v1.x;
thread_norm1_sum += (double)v1.y * (double)v1.y;
thread_norm1_sum += (double)v1.z * (double)v1.z;
thread_norm1_sum += (double)v1.w * (double)v1.w;
// Norm 2 Squared
thread_norm2_sum += (double)v2.x * (double)v2.x;
thread_norm2_sum += (double)v2.y * (double)v2.y;
thread_norm2_sum += (double)v2.z * (double)v2.z;
thread_norm2_sum += (double)v2.w * (double)v2.w;
}}
sh_dot[threadIdx.x] = thread_dot_sum;
sh_norm1[threadIdx.x] = thread_norm1_sum;
sh_norm2[threadIdx.x] = thread_norm2_sum;
__syncthreads();
if (threadIdx.x < 512) {{
sh_dot[threadIdx.x] += sh_dot[threadIdx.x + 512];
sh_norm1[threadIdx.x] += sh_norm1[threadIdx.x + 512];
sh_norm2[threadIdx.x] += sh_norm2[threadIdx.x + 512];
}}
__syncthreads();
if (threadIdx.x < 256) {{
sh_dot[threadIdx.x] += sh_dot[threadIdx.x + 256];
sh_norm1[threadIdx.x] += sh_norm1[threadIdx.x + 256];
sh_norm2[threadIdx.x] += sh_norm2[threadIdx.x + 256];
}}
__syncthreads();
if (threadIdx.x < 128) {{
sh_dot[threadIdx.x] += sh_dot[threadIdx.x + 128];
sh_norm1[threadIdx.x] += sh_norm1[threadIdx.x + 128];
sh_norm2[threadIdx.x] += sh_norm2[threadIdx.x + 128];
}}
if (threadIdx.x < 64) {{
__syncthreads();
sh_dot[threadIdx.x] += sh_dot[threadIdx.x + 64];
sh_norm1[threadIdx.x] += sh_norm1[threadIdx.x + 64];
sh_norm2[threadIdx.x] += sh_norm2[threadIdx.x + 64];
}}
if (threadIdx.x < 32) {{
__syncthreads();
sh_dot[threadIdx.x] += sh_dot[threadIdx.x + 32];
sh_norm1[threadIdx.x] += sh_norm1[threadIdx.x + 32];
sh_norm2[threadIdx.x] += sh_norm2[threadIdx.x + 32];
}}
if (threadIdx.x < 16) {{
__syncthreads();
sh_dot[threadIdx.x] += sh_dot[threadIdx.x + 16];
sh_norm1[threadIdx.x] += sh_norm1[threadIdx.x + 16];
sh_norm2[threadIdx.x] += sh_norm2[threadIdx.x + 16];
}}
if (threadIdx.x < 8) {{
__syncthreads();
sh_dot[threadIdx.x] += sh_dot[threadIdx.x + 8];
sh_norm1[threadIdx.x] += sh_norm1[threadIdx.x + 8];
sh_norm2[threadIdx.x] += sh_norm2[threadIdx.x + 8];
}}
if (threadIdx.x < 4) {{
__syncthreads();
sh_dot[threadIdx.x] += sh_dot[threadIdx.x + 4];
sh_norm1[threadIdx.x] += sh_norm1[threadIdx.x + 4];
sh_norm2[threadIdx.x] += sh_norm2[threadIdx.x + 4];
}}
if (threadIdx.x < 2) {{
__syncthreads();
sh_dot[threadIdx.x] += sh_dot[threadIdx.x + 2];
sh_norm1[threadIdx.x] += sh_norm1[threadIdx.x + 2];
sh_norm2[threadIdx.x] += sh_norm2[threadIdx.x + 2];
}}
if (threadIdx.x == 0) {{
__syncthreads();
sh_dot[0] += sh_dot[1];
sh_norm1[0] += sh_norm1[1];
sh_norm2[0] += sh_norm2[1];
}}
if (threadIdx.x == 0) {{
results[pair_idx].dot_sum = sh_dot[0];
results[pair_idx].norm1_sq_sum = sh_norm1[0];
results[pair_idx].norm2_sq_sum = sh_norm2[0];
}}
}}
__global__ void cosine_final_kernel(
const CosineResult* __restrict__ pass1_results,
const float* __restrict__ y,
double* __restrict__ global_loss_sum,
int N_pairs
) {{
__shared__ double sh_loss_sum[BLOCK_SIZE];
double thread_loss_sum = 0.0;
for (int pair_idx = blockIdx.x * blockDim.x + threadIdx.x;
pair_idx < N_pairs;
pair_idx += gridDim.x * blockDim.x)
{{
double dot = pass1_results[pair_idx].dot_sum;
double norm1_sq = pass1_results[pair_idx].norm1_sq_sum;
double norm2_sq = pass1_results[pair_idx].norm2_sq_sum;
double label_y = (double)y[pair_idx];
double norm_prod = std::sqrt(norm1_sq * norm2_sq);
double cosine = (norm_prod > 1e-6) ? (dot / norm_prod) : 0.0;
if (label_y > 0) {{
thread_loss_sum += 1.0 - cosine;
}} else {{
thread_loss_sum += std::max(0.0, cosine - (double)MARGIN_VAL);
}}
}}
sh_loss_sum[threadIdx.x] = thread_loss_sum;
__syncthreads();
if (threadIdx.x < 512) {{
sh_loss_sum[threadIdx.x] += sh_loss_sum[threadIdx.x + 512];
}}
__syncthreads();
if (threadIdx.x < 256) {{
sh_loss_sum[threadIdx.x] += sh_loss_sum[threadIdx.x + 256];
}}
__syncthreads();
if (threadIdx.x < 128) {{
sh_loss_sum[threadIdx.x] += sh_loss_sum[threadIdx.x + 128];
}}
if (threadIdx.x < 64) {{
__syncthreads();
sh_loss_sum[threadIdx.x] += sh_loss_sum[threadIdx.x + 64];
}}
if (threadIdx.x < 32) {{
__syncthreads();
sh_loss_sum[threadIdx.x] += sh_loss_sum[threadIdx.x + 32];
}}
if (threadIdx.x < 16) {{
__syncthreads();
sh_loss_sum[threadIdx.x] += sh_loss_sum[threadIdx.x + 16];
}}
if (threadIdx.x < 8) {{
__syncthreads();
sh_loss_sum[threadIdx.x] += sh_loss_sum[threadIdx.x + 8];
}}
if (threadIdx.x < 4) {{
__syncthreads();
sh_loss_sum[threadIdx.x] += sh_loss_sum[threadIdx.x + 4];
}}
if (threadIdx.x < 2) {{
__syncthreads();
sh_loss_sum[threadIdx.x] += sh_loss_sum[threadIdx.x + 2];
}}
if (threadIdx.x == 0) {{
__syncthreads();
sh_loss_sum[0] += sh_loss_sum[1];
}}
if (threadIdx.x == 0) {{
global_loss_sum[blockIdx.x] = sh_loss_sum[0];
}}
}}
torch::Tensor cosine_forward_cuda(torch::Tensor x1, torch::Tensor x2, torch::Tensor y) {{
TORCH_CHECK(x1.is_cuda() && x2.is_cuda() && y.is_cuda(), "Inputs must be CUDA tensors");
x1 = x1.contiguous();
x2 = x2.contiguous();
y = y.contiguous();
int N_pairs = x1.size(0); // BATCH_SIZE
int D_emb = x1.size(1); // EMBEDDING_DIM
if (D_emb % VEC_SIZE != 0) {{
TORCH_CHECK(false, "Embedding dimension must be divisible by 4.");
}}
const int block_size_p1 = BLOCK_SIZE;
const int grid_size_p1 = N_pairs;
auto result_buffer = torch::empty({{N_pairs, 3}}, x1.options().dtype(torch::kFloat64));
cosine_pass1_kernel<<<grid_size_p1, block_size_p1>>>(
x1.data_ptr<float>(),
x2.data_ptr<float>(),
(CosineResult*)result_buffer.data_ptr<double>(),
N_pairs, D_emb
);
const int block_size_p2 = 256;
const int grid_size_p2 = 256;
auto final_loss_sum_buffer = torch::empty({{grid_size_p2}}, x1.options().dtype(torch::kFloat64));
cosine_final_kernel<<<grid_size_p2, block_size_p2>>>(
(CosineResult*)result_buffer.data_ptr<double>(),
y.data_ptr<float>(),
final_loss_sum_buffer.data_ptr<double>(),
N_pairs
);
double total_sum = final_loss_sum_buffer.sum().item<double>();
float mean_loss = (float)(total_sum / N_pairs);
return torch::tensor(mean_loss, x1.options());
}}
""".format(margin_val=MARGIN)
self.cos_op = load_inline(
name="cosine_fused_vectorized_op",
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=["cosine_forward_cuda"],
extra_cuda_cflags=["-O3", "--use_fast_math"],
verbose=True
)
def forward(self, x1: torch.Tensor, x2: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
return self.cos_op.cosine_forward_cuda(x1, x2, y)

28
S1/27/cosineloss_torch.py Normal file
View File

@ -0,0 +1,28 @@
# cosineloss_torch.py
import torch
import torch.nn as nn
import torch.nn.functional as F
BATCH_SIZE = 16
EMBEDDING_DIM = 256
DIM = BATCH_SIZE
MARGIN = 0.5
class Model(nn.Module):
def forward(self, x1: torch.Tensor, x2: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
return F.cosine_embedding_loss(x1, x2, y, margin=MARGIN, reduction='mean')
def get_inputs():
x1 = torch.randn(BATCH_SIZE, EMBEDDING_DIM, dtype=torch.float32)
x2 = torch.randn(BATCH_SIZE, EMBEDDING_DIM, dtype=torch.float32)
y = torch.randint(0, 2, size=(BATCH_SIZE,), dtype=torch.float32)
y[y == 0] = -1.0
return [x1, x2, y]
def get_init_inputs():
return []

34
S1/27/prompt.txt Normal file
View File

@ -0,0 +1,34 @@
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
BATCH_SIZE = 16
EMBEDDING_DIM = 256
DIM = BATCH_SIZE
MARGIN = 0.5
class Model(nn.Module):
def forward(self, x1: torch.Tensor, x2: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
return F.cosine_embedding_loss(x1, x2, y, margin=MARGIN, reduction='mean')
def get_inputs():
x1 = torch.randn(BATCH_SIZE, EMBEDDING_DIM, dtype=torch.float32)
x2 = torch.randn(BATCH_SIZE, EMBEDDING_DIM, dtype=torch.float32)
y = torch.randint(0, 2, size=(BATCH_SIZE,), dtype=torch.float32)
y[y == 0] = -1.0
return [x1, x2, y]
def get_init_inputs():
return []

88
S1/27/run_code.py Normal file
View File

@ -0,0 +1,88 @@
###########################################################
# 性能和精度验证程序
###########################################################
import torch
import torch.nn as nn
import time
from cosineloss_torch import Model, get_inputs, get_init_inputs
from cosineloss_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)
# 更严格的精度检查
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-05, atol=1e-05)
if precision_flag:
print("✅ 精度对齐：两个模型的输出结果非常接近。")
else:
print("❌ 精度不一致！")
print("\n-------------------- 性能加速比测试 --------------------")
num_iterations = 1000 # 增加迭代次数以获得更准确的时间测量
# Warm up
for _ in range(100):
_ = 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 (matmul + relu) 平均执行时间: {torch_time:.6f} 秒")
print(f"自定义 CUDA ReLU 平均执行时间: {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()