```
# 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: computes mean(clamp(1 - predictions * targets, min=0))
# in a single pass with a grid-stride loop, block-level shared-memory
# reduction, and atomic accumulation into a scalar.
hinge_loss_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void hinge_loss_kernel(const float* __restrict__ preds,
                                  const float* __restrict__ targets,
                                  double* __restrict__ out,
                                  long n) {
    __shared__ double shared[256];

    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long stride = (long)gridDim.x * blockDim.x;

    double sum = 0.0;
    for (long i = idx; i < n; i += stride) {
        float v = 1.0f - preds[i] * targets[i];
        sum += (double)fmaxf(v, 0.0f);
    }

    shared[threadIdx.x] = sum;
    __syncthreads();

    // Block reduction
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            shared[threadIdx.x] += shared[threadIdx.x + s];
        }
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        atomicAdd(out, shared[0]);
    }
}

torch::Tensor hinge_loss_cuda(torch::Tensor preds, torch::Tensor targets) {
    TORCH_CHECK(preds.is_cuda() && targets.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(preds.scalar_type() == torch::kFloat32, "preds must be float32");
    TORCH_CHECK(targets.scalar_type() == torch::kFloat32, "targets must be float32");
    TORCH_CHECK(preds.numel() == targets.numel(), "Shapes must match after broadcast");

    const long n = preds.numel();
    auto out = torch::zeros({}, preds.options().dtype(torch::kFloat64));

    if (n == 0) {
        return out.to(torch::kFloat32);
    }

    const int block_size = 256;
    long num_blocks = (n + block_size - 1) / block_size;
    if (num_blocks > 4096) num_blocks = 4096;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    hinge_loss_kernel<<<(int)num_blocks, block_size, 0, stream>>>(
        preds.data_ptr<float>(),
        targets.data_ptr<float>(),
        out.data_ptr<double>(),
        n);

    auto mean = out / (double)n;
    return mean.to(torch::kFloat32);
}
"""

hinge_loss_cpp_source = (
    "#include <ATen/cuda/CUDAContext.h>\n"
    "torch::Tensor hinge_loss_cuda(torch::Tensor preds, torch::Tensor targets);"
)

hinge_loss = load_inline(
    name="hinge_loss_fused",
    cpp_sources=hinge_loss_cpp_source,
    cuda_sources="#include <ATen/cuda/CUDAContext.h>\n" + hinge_loss_source,
    functions=["hinge_loss_cuda"],
    verbose=True,
)


class ModelNew(nn.Module):
    """
    Optimized hinge loss model using a single fused CUDA kernel that
    performs the multiply, subtract, clamp, and mean reduction in one pass.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.hinge_loss = hinge_loss

    def forward(self, predictions, targets):
        # Handle broadcasting exactly like the eager implementation, then
        # run the fully fused kernel on contiguous buffers.
        p, t = torch.broadcast_tensors(predictions, targets)
        p = p.contiguous()
        t = t.contiguous()
        return self.hinge_loss.hinge_loss_cuda(p, t)
```