forked from ccf-ai-infra/GPUCodeForces
fixed cosineloss #1
This commit is contained in:
parent
c6294b7773
commit
8207ecdfa3
|
|
@ -1,208 +0,0 @@
|
|||
# cosineloss_cuda_ultimate.py
|
||||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
from cosineloss_torch import BATCH_SIZE, EMBEDDING_DIM, DIM, MARGIN
|
||||
|
||||
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>
|
||||
|
||||
torch::Tensor cosine_forward_cuda(torch::Tensor x1, torch::Tensor x2, torch::Tensor y);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <cuda_runtime.h>
|
||||
#include <device_launch_parameters.h>
|
||||
|
||||
#define BLOCK_SIZE 256
|
||||
#define VEC_SIZE 4
|
||||
#define WARP_SIZE 32
|
||||
#define MARGIN_VAL {margin_val}f
|
||||
|
||||
// Fast warp reduction
|
||||
__device__ __forceinline__ float warp_reduce_sum(float val) {{
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {{
|
||||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||||
}}
|
||||
return val;
|
||||
}}
|
||||
|
||||
// Fast block reduction
|
||||
__device__ __forceinline__ float block_reduce_sum(float val) {{
|
||||
__shared__ float shared[WARP_SIZE];
|
||||
|
||||
int lane = threadIdx.x % WARP_SIZE;
|
||||
int wid = threadIdx.x / WARP_SIZE;
|
||||
|
||||
val = warp_reduce_sum(val);
|
||||
|
||||
if (lane == 0) shared[wid] = val;
|
||||
|
||||
__syncthreads();
|
||||
|
||||
val = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared[lane] : 0.0f;
|
||||
|
||||
if (wid == 0) val = warp_reduce_sum(val);
|
||||
|
||||
return val;
|
||||
}}
|
||||
|
||||
// Single-pass fused kernel - computes everything in one go
|
||||
__global__ void cosine_loss_fused_kernel(
|
||||
const float* __restrict__ x1,
|
||||
const float* __restrict__ x2,
|
||||
const float* __restrict__ y,
|
||||
float* __restrict__ output_loss,
|
||||
const int N_pairs,
|
||||
const int D_emb
|
||||
) {{
|
||||
const int pair_idx = blockIdx.x;
|
||||
|
||||
if (pair_idx >= N_pairs) return;
|
||||
|
||||
// Use float for better performance
|
||||
float thread_dot = 0.0f;
|
||||
float thread_norm1 = 0.0f;
|
||||
float thread_norm2 = 0.0f;
|
||||
|
||||
const int offset = pair_idx * D_emb;
|
||||
const int D_vec = D_emb / VEC_SIZE;
|
||||
|
||||
const float4* __restrict__ x1_4 = reinterpret_cast<const float4*>(x1 + offset);
|
||||
const float4* __restrict__ x2_4 = reinterpret_cast<const float4*>(x2 + offset);
|
||||
|
||||
// Vectorized computation
|
||||
for (int d_vec = threadIdx.x; d_vec < D_vec; d_vec += blockDim.x) {{
|
||||
float4 v1 = __ldg(&x1_4[d_vec]);
|
||||
float4 v2 = __ldg(&x2_4[d_vec]);
|
||||
|
||||
// Dot product with FMA
|
||||
thread_dot = fmaf(v1.x, v2.x, thread_dot);
|
||||
thread_dot = fmaf(v1.y, v2.y, thread_dot);
|
||||
thread_dot = fmaf(v1.z, v2.z, thread_dot);
|
||||
thread_dot = fmaf(v1.w, v2.w, thread_dot);
|
||||
|
||||
// Norm squared with FMA
|
||||
thread_norm1 = fmaf(v1.x, v1.x, thread_norm1);
|
||||
thread_norm1 = fmaf(v1.y, v1.y, thread_norm1);
|
||||
thread_norm1 = fmaf(v1.z, v1.z, thread_norm1);
|
||||
thread_norm1 = fmaf(v1.w, v1.w, thread_norm1);
|
||||
|
||||
thread_norm2 = fmaf(v2.x, v2.x, thread_norm2);
|
||||
thread_norm2 = fmaf(v2.y, v2.y, thread_norm2);
|
||||
thread_norm2 = fmaf(v2.z, v2.z, thread_norm2);
|
||||
thread_norm2 = fmaf(v2.w, v2.w, thread_norm2);
|
||||
}}
|
||||
|
||||
// Block-level reduction
|
||||
float dot = block_reduce_sum(thread_dot);
|
||||
float norm1_sq = block_reduce_sum(thread_norm1);
|
||||
float norm2_sq = block_reduce_sum(thread_norm2);
|
||||
|
||||
// Only thread 0 computes the final loss for this pair
|
||||
if (threadIdx.x == 0) {{
|
||||
float label = y[pair_idx];
|
||||
|
||||
// Compute cosine similarity with safe division
|
||||
float norm_prod = sqrtf(norm1_sq * norm2_sq);
|
||||
float cosine = (norm_prod > 1e-8f) ? (dot / norm_prod) : 0.0f;
|
||||
|
||||
// Compute loss based on label
|
||||
float loss;
|
||||
if (label > 0.0f) {{
|
||||
loss = 1.0f - cosine;
|
||||
}} else {{
|
||||
// Use fmaxf for max(0, cosine - margin)
|
||||
loss = fmaxf(0.0f, cosine - MARGIN_VAL);
|
||||
}}
|
||||
|
||||
output_loss[pair_idx] = loss;
|
||||
}}
|
||||
}}
|
||||
|
||||
// Final reduction kernel
|
||||
__global__ void final_sum_kernel(
|
||||
const float* __restrict__ losses,
|
||||
float* __restrict__ output,
|
||||
const int n
|
||||
) {{
|
||||
float sum = 0.0f;
|
||||
|
||||
for (int i = threadIdx.x; i < n; i += WARP_SIZE) {{
|
||||
sum += losses[i];
|
||||
}}
|
||||
|
||||
sum = warp_reduce_sum(sum);
|
||||
|
||||
if (threadIdx.x == 0) {{
|
||||
output[0] = sum;
|
||||
}}
|
||||
}}
|
||||
|
||||
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();
|
||||
|
||||
const int N_pairs = x1.size(0);
|
||||
const int D_emb = x1.size(1);
|
||||
|
||||
TORCH_CHECK(D_emb % VEC_SIZE == 0,
|
||||
"Embedding dimension must be divisible by 4");
|
||||
|
||||
// Allocate temp storage for per-pair losses
|
||||
auto per_pair_loss = torch::empty({{N_pairs}}, x1.options());
|
||||
auto final_result = torch::empty({{1}}, x1.options());
|
||||
|
||||
// Launch fused kernel - one block per pair
|
||||
const int block_size = BLOCK_SIZE;
|
||||
const int grid_size = N_pairs;
|
||||
|
||||
cosine_loss_fused_kernel<<<grid_size, block_size>>>(
|
||||
x1.data_ptr<float>(),
|
||||
x2.data_ptr<float>(),
|
||||
y.data_ptr<float>(),
|
||||
per_pair_loss.data_ptr<float>(),
|
||||
N_pairs,
|
||||
D_emb
|
||||
);
|
||||
|
||||
// Final reduction
|
||||
final_sum_kernel<<<1, WARP_SIZE>>>(
|
||||
per_pair_loss.data_ptr<float>(),
|
||||
final_result.data_ptr<float>(),
|
||||
N_pairs
|
||||
);
|
||||
|
||||
// Compute mean on GPU
|
||||
final_result.div_(N_pairs);
|
||||
|
||||
return final_result;
|
||||
}}
|
||||
""".format(margin_val=MARGIN)
|
||||
|
||||
self.cos_op = load_inline(
|
||||
name="cosine_ultimate_op",
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=["cosine_forward_cuda"],
|
||||
extra_cuda_cflags=[
|
||||
"-O3",
|
||||
"--use_fast_math",
|
||||
"-lineinfo",
|
||||
],
|
||||
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)
|
||||
|
|
@ -1,28 +0,0 @@
|
|||
# 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 []
|
||||
|
|
@ -1,34 +0,0 @@
|
|||
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 []
|
||||
|
|
@ -1,88 +0,0 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
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()
|
||||
Loading…
Reference in New Issue