```
# 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 max(0, 1 - predictions * targets) element-wise
# and reduces to the mean in a single kernel launch (grid-stride loop +
# shared-memory block reduction + atomic accumulation of per-block means).
hinge_loss_cuda_source = """
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <limits>

#define BLOCK_SIZE 256
#define MAX_BLOCKS 4096

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

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

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

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

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

    // Accumulate this block's contribution to the mean
    if (threadIdx.x == 0) {
        atomicAdd(out, shared[0] / (float)n);
    }
}

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() == torch::kFloat32 &&
                targets.scalar_type() == torch::kFloat32, "Inputs must be float32");
    TORCH_CHECK(predictions.is_contiguous() && targets.is_contiguous(),
                "Inputs must be contiguous");
    TORCH_CHECK(predictions.numel() == targets.numel(),
                "Inputs must have the same number of elements");

    long n = predictions.numel();

    // Match torch.mean behavior on empty tensors (returns NaN)
    if (n == 0) {
        return torch::full({}, std::numeric_limits<float>::quiet_NaN(),
                           predictions.options());
    }

    auto out = torch::zeros({}, predictions.options());

    long num_blocks = (n + BLOCK_SIZE - 1) / BLOCK_SIZE;
    if (num_blocks > MAX_BLOCKS) num_blocks = MAX_BLOCKS;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    hinge_loss_kernel<<<num_blocks, BLOCK_SIZE, 0, stream>>>(
        predictions.data_ptr<float>(),
        targets.data_ptr<float>(),
        out.data_ptr<float>(),
        n);

    return out;
}
"""

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

# Compile the inline fused hinge-loss CUDA kernel
hinge_loss = load_inline(
    name="hinge_loss_fused",
    cpp_sources=hinge_loss_cpp_source,
    cuda_sources=hinge_loss_cuda_source,
    functions=["hinge_loss_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized hinge loss: fuses the multiply, subtract, clamp, and mean
    reduction into a single custom CUDA kernel, 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):
        # Fast path: contiguous float32 CUDA tensors -> fused kernel
        if (predictions.is_cuda and targets.is_cuda
                and predictions.dtype == torch.float32
                and targets.dtype == torch.float32):
            p, t = torch.broadcast_tensors(predictions, targets)
            p = p.contiguous().reshape(-1)
            t = t.contiguous().reshape(-1)
            return self.hinge_loss.hinge_loss_cuda(p, t)
        # Fallback for any other case
        return torch.mean(torch.clamp(1 - predictions * targets, min=0))
```