GPUCodeForces/S1/27/cosineloss_cuda.py

# 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)