diff --git a/S1/gsd123_#141/VICRegLoss_cuda.py b/S1/gsd123_#141/VICRegLoss_cuda.py new file mode 100644 index 00000000..7fa9eebe --- /dev/null +++ b/S1/gsd123_#141/VICRegLoss_cuda.py @@ -0,0 +1,227 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + +cuda_source = """ +#include +#include +#include + + +__device__ void atomicAddFloat(float* address, float val) { + atomicAdd(address, val); +} + + +__global__ void vicreg_stats_kernel( + const float* __restrict__ z1, + const float* __restrict__ z2, + float* z1_centered, + float* z2_centered, + float* loss_mse, + float* loss_std, + int batch_size, + int dim +) { + int d = blockIdx.x; // Current dimension + if (d >= dim) return; + + int tid = threadIdx.x; + int stride = blockDim.x; + + // 1. Calculate Means for column d + float sum1 = 0.0f; + float sum2 = 0.0f; + float diff_sum_sq = 0.0f; // For MSE: sum((z1-z2)^2) + + for (int n = tid; n < batch_size; n += stride) { + float v1 = z1[n * dim + d]; + float v2 = z2[n * dim + d]; + sum1 += v1; + sum2 += v2; + + float diff = v1 - v2; + diff_sum_sq += diff * diff; + } + + + __shared__ float s_sum1[256]; + __shared__ float s_sum2[256]; + __shared__ float s_diff[256]; + + s_sum1[tid] = sum1; + s_sum2[tid] = sum2; + s_diff[tid] = diff_sum_sq; + __syncthreads(); + + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) { + s_sum1[tid] += s_sum1[tid + s]; + s_sum2[tid] += s_sum2[tid + s]; + s_diff[tid] += s_diff[tid + s]; + } + __syncthreads(); + } + + float mean1 = s_sum1[0] / batch_size; + float mean2 = s_sum2[0] / batch_size; + + if (tid == 0) { + + atomicAddFloat(loss_mse, s_diff[0]); + } + + + + float var_sum1 = 0.0f; + float var_sum2 = 0.0f; + + for (int n = tid; n < batch_size; n += stride) { + float v1 = z1[n * dim + d]; + float v2 = z2[n * dim + d]; + + float c1 = v1 - mean1; + float c2 = v2 - mean2; + + z1_centered[n * dim + d] = c1; + z2_centered[n * dim + d] = c2; + + var_sum1 += c1 * c1; + var_sum2 += c2 * c2; + } + + + s_sum1[tid] = var_sum1; + s_sum2[tid] = var_sum2; + __syncthreads(); + + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) { + s_sum1[tid] += s_sum1[tid + s]; + s_sum2[tid] += s_sum2[tid + s]; + } + __syncthreads(); + } + + if (tid == 0) { + + float v1 = s_sum1[0] / (batch_size - 1); + float v2 = s_sum2[0] / (batch_size - 1); + + float std1 = sqrtf(v1 + 1e-4f); + float std2 = sqrtf(v2 + 1e-4f); + + float l1 = fmaxf(0.0f, 1.0f - std1); + float l2 = fmaxf(0.0f, 1.0f - std2); + + // Accumulate mean(relu(...)) -> sum / D + atomicAddFloat(loss_std, (l1 + l2) / dim); + } +} + + +__global__ void vicreg_cov_kernel( + const float* __restrict__ data, // z_centered [N, D] + float* loss_cov, + int batch_size, + int dim +) { + int row = blockIdx.y * blockDim.y + threadIdx.y; + int col = blockIdx.x * blockDim.x + threadIdx.x; + + if (row >= dim || col >= dim) return; + + // Compute Cov(row, col) + // C_ij = sum_n (z[n, row] * z[n, col]) / (N-1) + + float dot = 0.0f; + for (int n = 0; n < batch_size; ++n) { + dot += data[n * dim + row] * data[n * dim + col]; + } + dot /= (batch_size - 1); + + + + if (row != col) { + atomicAddFloat(loss_cov, (dot * dot) / dim); + } +} + +torch::Tensor vicreg_cuda_forward(torch::Tensor z1, torch::Tensor z2, + float lambda_param, float mu_param, float nu_param) { + int batch_size = z1.size(0); + int dim = z1.size(1); + + auto z1_c = z1.contiguous(); + auto z2_c = z2.contiguous(); + + auto z1_centered = torch::empty_like(z1_c); + auto z2_centered = torch::empty_like(z2_c); + + + auto loss_mse_t = torch::zeros({1}, z1.options()); + auto loss_std_t = torch::zeros({1}, z1.options()); + auto loss_cov_t = torch::zeros({1}, z1.options()); + + // 1. Stats Kernel + vicreg_stats_kernel<<>>( + z1_c.data_ptr(), + z2_c.data_ptr(), + z1_centered.data_ptr(), + z2_centered.data_ptr(), + loss_mse_t.data_ptr(), + loss_std_t.data_ptr(), + batch_size, + dim + ); + + // 2. Covariance Kernel + // Block size 16x16 = 256 threads + dim3 block(16, 16); + dim3 grid((dim + 15) / 16, (dim + 15) / 16); + + vicreg_cov_kernel<<>>( + z1_centered.data_ptr(), + loss_cov_t.data_ptr(), + batch_size, + dim + ); + + vicreg_cov_kernel<<>>( + z2_centered.data_ptr(), + loss_cov_t.data_ptr(), + batch_size, + dim + ); + + + + return lambda_param * (loss_mse_t / (batch_size * dim)) + + mu_param * loss_std_t + + nu_param * loss_cov_t; +} +""" + +cpp_source = """ +torch::Tensor vicreg_cuda_forward(torch::Tensor z1, torch::Tensor z2, + float lambda_param, float mu_param, float nu_param); +""" + +vicreg_module = load_inline( + name="vicreg_loss_opt", + cpp_sources=cpp_source, + cuda_sources=cuda_source, + functions=["vicreg_cuda_forward"], + verbose=False +) + + +class ModelNew(nn.Module): + def __init__(self, lambda_param, mu_param, nu_param): + super(ModelNew, self).__init__() + self.lambda_param = lambda_param + self.mu_param = mu_param + self.nu_param = nu_param + + def forward(self, z1, z2): + return vicreg_module.vicreg_cuda_forward(z1, z2, self.lambda_param, self.mu_param, self.nu_param) \ No newline at end of file diff --git a/S1/gsd123_#141/VICRegLoss_torch.py b/S1/gsd123_#141/VICRegLoss_torch.py new file mode 100644 index 00000000..e5cc953d --- /dev/null +++ b/S1/gsd123_#141/VICRegLoss_torch.py @@ -0,0 +1,47 @@ +import torch +import torch.nn as nn + + +class Model(nn.Module): + def __init__(self, lambda_param, mu_param, nu_param): + super(Model, self).__init__() + self.lambda_param = lambda_param + self.mu_param = mu_param + self.nu_param = nu_param + + def forward(self, z1: torch.Tensor, z2: torch.Tensor) -> torch.Tensor: + repr_loss = torch.nn.functional.mse_loss(z1, z2) + + z1_centered = z1 - z1.mean(dim=0) + z2_centered = z2 - z2.mean(dim=0) + + std_z1 = torch.sqrt(z1_centered.var(dim=0) + 1e-4) + std_z2 = torch.sqrt(z2_centered.var(dim=0) + 1e-4) + std_loss = torch.mean(torch.relu(1 - std_z1)) + torch.mean(torch.relu(1 - std_z2)) + + cov_z1 = torch.matmul(z1_centered.T, z1_centered) / (z1.shape[0] - 1) + cov_z2 = torch.matmul(z2_centered.T, z2_centered) / (z2.shape[0] - 1) + + cov_loss = (cov_z1.pow(2).sum() - cov_z1.diagonal().pow(2).sum()) / z1.shape[1] + cov_loss += (cov_z2.pow(2).sum() - cov_z2.diagonal().pow(2).sum()) / z2.shape[1] + + loss = self.lambda_param * repr_loss + self.mu_param * std_loss + self.nu_param * cov_loss + + return loss + + +batch_size = 16 +dim = 128 + + +def get_inputs(): + z1 = torch.randn(batch_size, dim) + z2 = torch.randn(batch_size, dim) + return [z1, z2] + + +def get_init_inputs(): + lambda_param = 25.0 + mu_param = 25.0 + nu_param = 1.0 + return [lambda_param, mu_param, nu_param] \ No newline at end of file diff --git a/S1/gsd123_#141/prompt.txt b/S1/gsd123_#141/prompt.txt new file mode 100644 index 00000000..1cbb094b --- /dev/null +++ b/S1/gsd123_#141/prompt.txt @@ -0,0 +1,76 @@ +You write custom CUDA kernels to replace the pytorch operators in the given GeGLU 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 chunk+gelu+elementwise_mul), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination. +Custom CUDA kernel extension via torch.utils.cpp_extension.load_inline + +VICReg (Variance-Invariance-Covariance Regularization) loss computation + +Multi-kernel design: statistics (mean/var) + covariance computation + +Dimension-wise parallelization for mean/variance calculation + +2D grid kernel for covariance matrix computation + +Three loss components: invariance (MSE), variance (std), covariance + +Shared memory reduction for column-wise statistics + +Atomic accumulation for loss components + +Centered representation storage for covariance reuse + +Numerical stability with epsilon in std computation + +Weighted loss combination with λ, μ, ν hyperparameters + + + + +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 + + +class Model(nn.Module): + def __init__(self, lambda_param, mu_param, nu_param): + super(Model, self).__init__() + self.lambda_param = lambda_param + self.mu_param = mu_param + self.nu_param = nu_param + + def forward(self, z1: torch.Tensor, z2: torch.Tensor) -> torch.Tensor: + repr_loss = torch.nn.functional.mse_loss(z1, z2) + + z1_centered = z1 - z1.mean(dim=0) + z2_centered = z2 - z2.mean(dim=0) + + std_z1 = torch.sqrt(z1_centered.var(dim=0) + 1e-4) + std_z2 = torch.sqrt(z2_centered.var(dim=0) + 1e-4) + std_loss = torch.mean(torch.relu(1 - std_z1)) + torch.mean(torch.relu(1 - std_z2)) + + cov_z1 = torch.matmul(z1_centered.T, z1_centered) / (z1.shape[0] - 1) + cov_z2 = torch.matmul(z2_centered.T, z2_centered) / (z2.shape[0] - 1) + + cov_loss = (cov_z1.pow(2).sum() - cov_z1.diagonal().pow(2).sum()) / z1.shape[1] + cov_loss += (cov_z2.pow(2).sum() - cov_z2.diagonal().pow(2).sum()) / z2.shape[1] + + loss = self.lambda_param * repr_loss + self.mu_param * std_loss + self.nu_param * cov_loss + + return loss + + +batch_size = 16 +dim = 128 + + +def get_inputs(): + z1 = torch.randn(batch_size, dim) + z2 = torch.randn(batch_size, dim) + return [z1, z2] + + +def get_init_inputs(): + lambda_param = 25.0 + mu_param = 25.0 + nu_param = 1.0 + return [lambda_param, mu_param, nu_param] \ No newline at end of file diff --git a/S1/gsd123_#141/run_code.py b/S1/gsd123_#141/run_code.py new file mode 100644 index 00000000..c13fc36a --- /dev/null +++ b/S1/gsd123_#141/run_code.py @@ -0,0 +1,77 @@ +########################################################### +# 性能和精度验证程序 +########################################################### +import torch +import torch.nn as nn +import time +from VICRegLoss_torch import Model, get_inputs, get_init_inputs +from VICRegLoss_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) + + 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 torch.relu 平均执行时间: {torch_time:.6f} 秒") + print(f"自定义 CUDA 内核 平均执行时间: {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() \ No newline at end of file