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

fixes mseloss #28

This commit is contained in:
ZZZJ 2025-11-12 09:29:33 +08:00
parent 79fe16bda1
commit beff3b7d1d
1 changed files with 99 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

143
S1/28/mseloss_cuda.py
View File

@ -1,7 +1,6 @@
# mseloss_cuda.py
import torch
from torch.utils.cpp_extension import load_inline
from mseloss_torch import BATCH_SIZE, DIM
class ModelNew(torch.nn.Module):
@ -12,59 +11,107 @@ class ModelNew(torch.nn.Module):
def _compile_cuda_kernel(self):
cpp_source = """
#include <torch/extension.h>
#include <cmath>
torch::Tensor mse_forward_cuda(torch::Tensor pred, torch::Tensor target);
"""
cuda_source = """
#include <cuda_runtime.h>
#include <cmath>
#include <device_launch_parameters.h>
#define BLOCK_SIZE 1024
#define VEC_SIZE 4
#define BLOCK_SIZE 256
#define VEC_SIZE 4
#define WARP_SIZE 32
__global__ void mse_fused_reduction_kernel(
// Warp-level reduction using shuffle instructions (faster than shared memory)
__device__ __forceinline__ float warp_reduce_sum(float val) {
#pragma unroll
for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xffffffff, val, offset);
}
return val;
}
// Optimized kernel: use float accumulation, warp primitives, and better memory access
__global__ void mse_optimized_kernel(
const float* __restrict__ pred,
const float* __restrict__ target,
double* __restrict__ output_sum,
float* __restrict__ output_sum,
int N_elements
) {
__shared__ double sh_sum[BLOCK_SIZE];
// Use float instead of double for accumulation (much faster on GPU)
float thread_sum = 0.0f;
double thread_sum = 0.0;
int N_vec = N_elements / VEC_SIZE;
int N_vec = N_elements / VEC_SIZE;
int grid_stride_vec = gridDim.x * blockDim.x;
const float4* __restrict__ pred4 = (const float4*)pred;
const float4* __restrict__ target4 = (const float4*)target;
const float4* __restrict__ pred4 = reinterpret_cast<const float4*>(pred);
const float4* __restrict__ target4 = reinterpret_cast<const float4*>(target);
// Process vectorized data with grid-stride loop
for (int idx_vec = blockIdx.x * blockDim.x + threadIdx.x;
idx_vec < N_vec;
idx_vec += grid_stride_vec)
{
float4 p4 = __ldg(&pred4[idx_vec]); // Use read-only cache
float4 t4 = __ldg(&target4[idx_vec]);
float4 p4 = pred4[idx_vec];
float4 t4 = target4[idx_vec];
double diff1 = (double)p4.x - (double)t4.x;
thread_sum += diff1 * diff1;
double diff2 = (double)p4.y - (double)t4.y;
thread_sum += diff2 * diff2;
// Compute differences and accumulate squared errors
float diff1 = p4.x - t4.x;
float diff2 = p4.y - t4.y;
float diff3 = p4.z - t4.z;
float diff4 = p4.w - t4.w;
double diff3 = (double)p4.z - (double)t4.z;
thread_sum += diff3 * diff3;
double diff4 = (double)p4.w - (double)t4.w;
thread_sum += diff4 * diff4;
// Use fused multiply-add for better performance
thread_sum = fmaf(diff1, diff1, thread_sum);
thread_sum = fmaf(diff2, diff2, thread_sum);
thread_sum = fmaf(diff3, diff3, thread_sum);
thread_sum = fmaf(diff4, diff4, thread_sum);
}
// Warp-level reduction (no shared memory needed for this step)
thread_sum = warp_reduce_sum(thread_sum);
// Use shared memory only for inter-warp reduction
__shared__ float warp_sums[BLOCK_SIZE / WARP_SIZE];
sh_sum[threadIdx.x] = thread_sum;
int lane = threadIdx.x % WARP_SIZE;
int warp_id = threadIdx.x / WARP_SIZE;
if (lane == 0) {
warp_sums[warp_id] = thread_sum;
}
__syncthreads();
// Final reduction by first warp only
if (warp_id == 0) {
thread_sum = (threadIdx.x < BLOCK_SIZE / WARP_SIZE) ? warp_sums[lane] : 0.0f;
thread_sum = warp_reduce_sum(thread_sum);
if (threadIdx.x == 0) {
output_sum[blockIdx.x] = thread_sum;
}
}
}
// Two-phase reduction kernel for final sum
__global__ void final_reduction_kernel(
const float* __restrict__ partial_sums,
float* __restrict__ output,
int n
) {
__shared__ float sh_sum[BLOCK_SIZE];
float sum = 0.0f;
for (int i = threadIdx.x; i < n; i += blockDim.x) {
sum += partial_sums[i];
}
sh_sum[threadIdx.x] = sum;
__syncthreads();
// Tree reduction in shared memory
#pragma unroll
for (int s = BLOCK_SIZE / 2; s > 0; s /= 2) {
if (threadIdx.x < s) {
sh_sum[threadIdx.x] += sh_sum[threadIdx.x + s];
@ -72,8 +119,8 @@ class ModelNew(torch.nn.Module):
__syncthreads();
}
if (threadIdx.x == 0) {
output_sum[blockIdx.x] = sh_sum[0];
if (threadIdx.x == 0) {
output[0] = sh_sum[0];
}
}
@ -84,38 +131,46 @@ class ModelNew(torch.nn.Module):
target = target.contiguous();
int N_elements = pred.numel();
TORCH_CHECK(N_elements % VEC_SIZE == 0,
"Total elements must be divisible by VEC_SIZE (4)");
if (N_elements % VEC_SIZE != 0) {
TORCH_CHECK(false, "Total elements must be divisible by VEC_SIZE (4) for optimization.");
}
// Adaptive grid size based on data size and GPU occupancy
const int block_size = BLOCK_SIZE;
const int grid_size = 256;
int num_sms;
cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, 0);
const int grid_size = min(num_sms * 4, (N_elements / VEC_SIZE + block_size - 1) / block_size);
auto partial_sum_output = torch::empty({grid_size}, pred.options().dtype(torch::kFloat64));
auto partial_sum = torch::empty({grid_size}, pred.options());
auto final_result = torch::empty({1}, pred.options());
mse_fused_reduction_kernel<<<grid_size, block_size>>>(
// Launch main kernel
mse_optimized_kernel<<<grid_size, block_size>>>(
pred.data_ptr<float>(),
target.data_ptr<float>(),
partial_sum_output.data_ptr<double>(),
partial_sum.data_ptr<float>(),
N_elements
);
// Launch final reduction kernel (all on GPU, no CPU sync)
final_reduction_kernel<<<1, block_size>>>(
partial_sum.data_ptr<float>(),
final_result.data_ptr<float>(),
grid_size
);
double total_sum = partial_sum_output.sum().item<double>();
// Divide by N to get mean (done on GPU)
final_result.div_(N_elements);
float mean_loss = (float)(total_sum / N_elements);
return torch::tensor(mean_loss, pred.options());
return final_result;
}
"""
self.mse_op = load_inline(
name="mse_fused_vectorized_op",
name="mse_optimized_op",
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=["mse_forward_cuda"],
extra_cuda_cflags=["-O3", "--use_fast_math"],
extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"],
verbose=True
)