```
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

# Fused CUDA kernel that computes the full hinge loss:
#   mean(clamp(1 - predictions * targets, min=0))
# in a single pass (multiply + clamp + sum-reduction fused), followed by a
# scalar division. This avoids materializing any intermediate tensors.
hinge_loss_source = """
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <vector>
#include <algorithm>

__global__ void hinge_loss_sum_kernel(const float* __restrict__ preds,
                                      const float* __restrict__ tgts,
                                      float* __restrict__ out_sum,
                                      long long n) {
    __shared__ float shared[256];
    const int tid = threadIdx.x;
    const long long stride = (long long)blockDim.x * gridDim.x;
    float local = 0.0f;

    // Vectorized main loop (float4) for memory-bandwidth efficiency.
    const long long n4 = n >> 2;
    const float4* p4 = reinterpret_cast<const float4*>(preds);
    const float4* t4 = reinterpret_cast<const float4*>(tgts);
    for (long long i = (long long)blockIdx.x * blockDim.x + tid; i < n4; i += stride) {
        float4 p = p4[i];
        float4 t = t4[i];
        local += fmaxf(1.0f - p.x * t.x, 0.0f);
        local += fmaxf(1.0f - p.y * t.y, 0.0f);
        local += fmaxf(1.0f - p.z * t.z, 0.0f);
        local += fmaxf(1.0f - p.w * t.w, 0.0f);
    }
    // Scalar tail loop for the remaining (n % 4) elements.
    for (long long i = (n4 << 2) + (long long)blockIdx.x * blockDim.x + tid; i < n; i += stride) {
        local += fmaxf(1.0f - preds[i] * tgts[i], 0.0f);
    }

    // Block-level tree reduction in shared memory.
    shared[tid] = local;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) {
            shared[tid] += shared[tid + s];
        }
        __syncthreads();
    }

    // Accumulate per-block partial sums into a single global scalar.
    if (tid == 0) {
        atomicAdd(out_sum, shared[0]);
    }
}

torch::Tensor hinge_loss_cuda(torch::Tensor predictions, torch::Tensor targets) {
    TORCH_CHECK(predictions.is_cuda() && targets.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(predictions.scalar_type() == at::kFloat && targets.scalar_type() == at::kFloat,
                "Inputs must be float32");

    const c10::cuda::OptionalCUDAGuard device_guard(at::device_of(predictions));

    // Match PyTorch broadcasting semantics of predictions * targets.
    std::vector<torch::Tensor> broadcasted = torch::broadcast_tensors({predictions, targets});
    torch::Tensor preds = broadcasted[0].contiguous();
    torch::Tensor tgts = broadcasted[1].contiguous();

    const long long n = preds.numel();

    // 0-dim scalar accumulator on the GPU.
    torch::Tensor out = torch::scalar_tensor(0.0, preds.options());

    const int block_size = 256;
    long long needed = (n + block_size - 1) / block_size;
    int num_blocks = (int)std::min<long long>(std::max<long long>(needed, 1), 2048);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    hinge_loss_sum_kernel<<<num_blocks, block_size, 0, stream>>>(
        preds.data_ptr<float>(), tgts.data_ptr<float>(), out.data_ptr<float>(), n);

    // Divide by the number of elements to get the mean (0-dim tensor like torch.mean).
    return out / (double)n;
}
"""

hinge_loss_cpp_source = (
    "torch::Tensor hinge_loss_cuda(torch::Tensor predictions, torch::Tensor targets);"
)

# Compile the inline CUDA extension.
hinge_loss = load_inline(
    name="hinge_loss_fused",
    cpp_sources=hinge_loss_cpp_source,
    cuda_sources=hinge_loss_source,
    functions=["hinge_loss_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized Hinge Loss module using a fused custom CUDA kernel.

    The kernel fuses the multiply, clamp(min=0), and mean reduction into a
    single GPU pass, eliminating intermediate tensor allocations and extra
    kernel launches.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.hinge_loss = hinge_loss

    def forward(self, predictions, targets):
        if (predictions.is_cuda and targets.is_cuda
                and predictions.dtype == torch.float32
                and targets.dtype == torch.float32):
            return self.hinge_loss.hinge_loss_cuda(predictions, targets)
        # Fallback for CPU tensors or non-float32 dtypes.
        return torch.mean(torch.clamp(1 - predictions * targets, min=0))
```