GPUCodeForces/S1/37/CosineEmbeddingLoss_cuda.py

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

N, C, H, W = 32, 64, 56, 56
EPS = 1e-8

assert (C * H * W) % 4 == 0, "Instance size (C*H*W) must be a multiple of 4"


class ModelNew(nn.Module):

    def __init__(self, margin=0.5):
        super().__init__()
        self.margin = margin
        self.eps = EPS
        self.block_size = 256
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>

        torch::Tensor cosine_loss_forward_cuda(
            torch::Tensor x1,
            torch::Tensor x2,
            torch::Tensor target,
            float margin,
            float eps,
            int N,
            int D);
        """

        cuda_source = f"""
        #include <cuda_runtime.h>
        #include <cmath>

        #define BLOCK_SIZE {self.block_size}
        #define WARP_SIZE 32
        #define ILP 4


        __inline__ __device__ 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;
        }}


        __global__ void cosine_embedding_kernel(
            const float* __restrict__ x1,
            const float* __restrict__ x2,
            const float* __restrict__ target,
            float* __restrict__ output,
            float margin,
            float eps,
            int D_vec // D / 4
        ) {{

            const int n_idx = blockIdx.x;


            const int offset = n_idx * D_vec * 4;

            const float4* x1_ptr = reinterpret_cast<const float4*>(x1 + offset);
            const float4* x2_ptr = reinterpret_cast<const float4*>(x2 + offset);

            float sum_dot = 0.0f;
            float sum_sq1 = 0.0f;
            float sum_sq2 = 0.0f;


            for (int i = threadIdx.x * ILP; i < D_vec; i += blockDim.x * ILP) {{
                float4 r1[ILP];
                float4 r2[ILP];


                #pragma unroll
                for (int k = 0; k < ILP; ++k) {{
                    if (i + k < D_vec) {{
                        r1[k] = __ldg(&x1_ptr[i + k]);
                        r2[k] = __ldg(&x2_ptr[i + k]);
                    }} else {{
                        r1[k] = make_float4(0.f, 0.f, 0.f, 0.f);
                        r2[k] = make_float4(0.f, 0.f, 0.f, 0.f);
                    }}
                }}

                // 2. 计算累加
                #pragma unroll
                for (int k = 0; k < ILP; ++k) {{
                    // Dot Product
                    sum_dot += r1[k].x * r2[k].x;
                    sum_dot += r1[k].y * r2[k].y;
                    sum_dot += r1[k].z * r2[k].z;
                    sum_dot += r1[k].w * r2[k].w;

                    // Norm Sq 1
                    sum_sq1 += r1[k].x * r1[k].x;
                    sum_sq1 += r1[k].y * r1[k].y;
                    sum_sq1 += r1[k].z * r1[k].z;
                    sum_sq1 += r1[k].w * r1[k].w;

                    // Norm Sq 2
                    sum_sq2 += r2[k].x * r2[k].x;
                    sum_sq2 += r2[k].y * r2[k].y;
                    sum_sq2 += r2[k].z * r2[k].z;
                    sum_sq2 += r2[k].w * r2[k].w;
                }}
            }}


            __shared__ float shared_data[32][3]; 

            int lane = threadIdx.x % WARP_SIZE;
            int wid = threadIdx.x / WARP_SIZE;


            sum_dot = warp_reduce_sum(sum_dot);
            sum_sq1 = warp_reduce_sum(sum_sq1);
            sum_sq2 = warp_reduce_sum(sum_sq2);


            if (lane == 0) {{
                shared_data[wid][0] = sum_dot;
                shared_data[wid][1] = sum_sq1;
                shared_data[wid][2] = sum_sq2;
            }}
            __syncthreads();


            if (wid == 0) {{
                // 读取
                sum_dot = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared_data[lane][0] : 0.0f;
                sum_sq1 = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared_data[lane][1] : 0.0f;
                sum_sq2 = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared_data[lane][2] : 0.0f;


                sum_dot = warp_reduce_sum(sum_dot);
                sum_sq1 = warp_reduce_sum(sum_sq1);
                sum_sq2 = warp_reduce_sum(sum_sq2);


                if (threadIdx.x == 0) {{
                    float norm1 = sqrtf(sum_sq1);
                    float norm2 = sqrtf(sum_sq2);
                    float cos_sim = sum_dot / (norm1 * norm2 + eps);

                    float t_val = target[n_idx]; 
                    float loss = 0.0f;

                    if (t_val == 1.0f) {{
                        loss = 1.0f - cos_sim;
                    }} else {{
                        loss = fmaxf(0.0f, cos_sim - margin);
                    }}

                    output[n_idx] = loss;
                }}
            }}
        }}

        torch::Tensor cosine_loss_forward_cuda(
            torch::Tensor x1,
            torch::Tensor x2,
            torch::Tensor target,
            float margin,
            float eps,
            int N,
            int D) 
        {{
            x1 = x1.contiguous();
            x2 = x2.contiguous();
            target = target.contiguous();

            auto output = torch::empty({{N}}, x1.options());

            int D_vec = D / 4;

            // Grid = Batch Size, Block = 256
            dim3 blocks(N);
            dim3 threads(BLOCK_SIZE);

            cosine_embedding_kernel<<<blocks, threads>>>(
                x1.data_ptr<float>(),
                x2.data_ptr<float>(),
                target.data_ptr<float>(),
                output.data_ptr<float>(),
                margin,
                eps,
                D_vec
            );

            return output;
        }}
        """

        self.op = load_inline(
            name='cosine_loss_cuda_v1',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['cosine_loss_forward_cuda'],
            extra_cuda_cflags=['-O3', '--use_fast_math'],
            verbose=False
        )

    def forward(self, x1: torch.Tensor, x2: torch.Tensor, target: torch.Tensor) -> torch.Tensor:

        x1_flat = x1.view(x1.size(0), -1)
        x2_flat = x2.view(x2.size(0), -1)

        if not x1_flat.is_cuda: x1_flat = x1_flat.cuda()
        if not x2_flat.is_cuda: x2_flat = x2_flat.cuda()
        if not target.is_cuda: target = target.cuda()

        N, D = x1_flat.shape

        out = self.op.cosine_loss_forward_cuda(
            x1_flat,
            x2_flat,
            target,
            self.margin,
            self.eps,
            N,
            D
        )

        return out.mean()