diff --git a/S1/28/mseloss_cuda.py b/S1/28/mseloss_cuda.py index 9f24384..8a54816 100644 --- a/S1/28/mseloss_cuda.py +++ b/S1/28/mseloss_cuda.py @@ -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 - #include torch::Tensor mse_forward_cuda(torch::Tensor pred, torch::Tensor target); """ cuda_source = """ #include - #include + #include - #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(pred); + const float4* __restrict__ target4 = reinterpret_cast(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<<>>( + // Launch main kernel + mse_optimized_kernel<<>>( pred.data_ptr(), target.data_ptr(), - partial_sum_output.data_ptr(), + partial_sum.data_ptr(), N_elements ); + // Launch final reduction kernel (all on GPU, no CPU sync) + final_reduction_kernel<<<1, block_size>>>( + partial_sum.data_ptr(), + final_result.data_ptr(), + grid_size + ); - double total_sum = partial_sum_output.sum().item(); + // 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 )