```python
# 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 implementation of:
#     mean(clamp(1 - predictions * targets, min=0))
# A single kernel computes the elementwise hinge loss and reduces it to the
# scalar mean, avoiding multiple kernel launches and intermediate tensors.
_hinge_loss_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/ExpandUtils.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <algorithm>
#include <cstdint>
#include <limits>

namespace {

constexpr int kBlockSize = 256;
constexpr int kMaxDims = 8;

struct BroadcastInfo {
    long long sizes[kMaxDims];
    long long stride_p[kMaxDims];
    long long stride_t[kMaxDims];
    int ndim;
};

// Block-wide sum reduction. Result is valid in thread 0 only.
__device__ __forceinline__ float block_reduce_sum(float sum) {
    // Warp-level reduction.
    for (int offset = 16; offset > 0; offset >>= 1) {
        sum += __shfl_down_sync(0xffffffffu, sum, offset);
    }
    __shared__ float sdata[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) {
        sdata[wid] = sum;
    }
    __syncthreads();
    if (wid == 0) {
        const int nwarps = (blockDim.x + 31) >> 5;
        sum = (lane < nwarps) ? sdata[lane] : 0.0f;
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            sum += __shfl_down_sync(0xffffffffu, sum, offset);
        }
    }
    return sum;
}

// Fast path: both operands contiguous with identical layout.
__global__ void hinge_contiguous_kernel(const float* __restrict__ pred,
                                        const float* __restrict__ targ,
                                        float* __restrict__ out,
                                        const long long n,
                                        const float inv_n,
                                        const int use_vec4) {
    const long long tid =
        static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x;
    const long long stride = static_cast<long long>(gridDim.x) * blockDim.x;

    float sum = 0.0f;
    if (use_vec4) {
        const long long n4 = n >> 2;
        const float4* pred4 = reinterpret_cast<const float4*>(pred);
        const float4* targ4 = reinterpret_cast<const float4*>(targ);
        for (long long k = tid; k < n4; k += stride) {
            const float4 p = pred4[k];
            const float4 t = targ4[k];
            sum += fmaxf(1.0f - p.x * t.x, 0.0f);
            sum += fmaxf(1.0f - p.y * t.y, 0.0f);
            sum += fmaxf(1.0f - p.z * t.z, 0.0f);
            sum += fmaxf(1.0f - p.w * t.w, 0.0f);
        }
        for (long long k = (n4 << 2) + tid; k < n; k += stride) {
            sum += fmaxf(1.0f - pred[k] * targ[k], 0.0f);
        }
    } else {
        for (long long k = tid; k < n; k += stride) {
            sum += fmaxf(1.0f - pred[k] * targ[k], 0.0f);
        }
    }

    sum = block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        atomicAdd(out, sum * inv_n);
    }
}

// General path: handles arbitrary broadcasting without materializing inputs.
__global__ void hinge_broadcast_kernel(const float* __restrict__ pred,
                                       const float* __restrict__ targ,
                                       float* __restrict__ out,
                                       const long long n,
                                       const float inv_n,
                                       const BroadcastInfo bc) {
    const long long tid =
        static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x;
    const long long stride = static_cast<long long>(gridDim.x) * blockDim.x;

    float sum = 0.0f;
    for (long long k = tid; k < n; k += stride) {
        long long rem = k;
        long long off_p = 0;
        long long off_t = 0;
        for (int d = bc.ndim - 1; d >= 0; --d) {
            const long long s = bc.sizes[d];
            const long long c = rem % s;
            rem /= s;
            off_p += c * bc.stride_p[d];
            off_t += c * bc.stride_t[d];
        }
        sum += fmaxf(1.0f - pred[off_p] * targ[off_t], 0.0f);
    }

    sum = block_reduce_sum(sum);
    if (threadIdx.x == 0) {
        atomicAdd(out, sum * inv_n);
    }
}

int pick_grid_size(long long work_items) {
    const long long blocks_needed = (work_items + kBlockSize - 1) / kBlockSize;
    const long long cap =
        static_cast<long long>(
            at::cuda::getCurrentDeviceProperties()->multiProcessorCount) * 8;
    long long grid = std::min(blocks_needed, cap);
    if (grid < 1) {
        grid = 1;
    }
    return static_cast<int>(grid);
}

}  // namespace

torch::Tensor hinge_loss_cuda(torch::Tensor pred, torch::Tensor targ) {
    TORCH_CHECK(pred.is_cuda() && targ.is_cuda(),
                "inputs must be CUDA tensors");
    TORCH_CHECK(pred.scalar_type() == at::kFloat &&
                targ.scalar_type() == at::kFloat,
                "inputs must be float32");

    const c10::cuda::CUDAGuard device_guard(pred.device());
    auto stream = at::cuda::getCurrentCUDAStream();

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

    // Match PyTorch broadcasting semantics of `predictions * targets`.
    const auto shape = at::infer_size(pred.sizes(), targ.sizes());
    auto pred_b = pred.expand(shape);
    auto targ_b = targ.expand(shape);

    const long long n = static_cast<long long>(pred_b.numel());
    if (n == 0) {
        // torch.mean of an empty tensor returns NaN.
        out.fill_(std::numeric_limits<float>::quiet_NaN());
        return out;
    }
    const float inv_n = 1.0f / static_cast<float>(n);

    bool contiguous = pred_b.is_contiguous() && targ_b.is_contiguous();
    if (!contiguous && static_cast<int>(shape.size()) > kMaxDims) {
        pred_b = pred_b.contiguous();
        targ_b = targ_b.contiguous();
        contiguous = true;
    }

    if (contiguous) {
        const float* p_ptr = pred_b.data_ptr<float>();
        const float* t_ptr = targ_b.data_ptr<float>();
        const bool aligned =
            ((reinterpret_cast<std::uintptr_t>(p_ptr) & 0xF) == 0) &&
            ((reinterpret_cast<std::uintptr_t>(t_ptr) & 0xF) == 0);
        const int use_vec4 = aligned ? 1 : 0;
        const long long work = use_vec4 ? ((n + 3) >> 2) : n;
        const int grid = pick_grid_size(work);
        hinge_contiguous_kernel<<<grid, kBlockSize, 0, stream>>>(
            p_ptr, t_ptr, out.data_ptr<float>(), n, inv_n, use_vec4);
        return out;
    }

    BroadcastInfo bc;
    bc.ndim = static_cast<int>(shape.size());
    for (int d = 0; d < bc.ndim; ++d) {
        bc.sizes[d]    = shape[d];
        bc.stride_p[d] = pred_b.stride(d);
        bc.stride_t[d] = targ_b.stride(d);
    }
    const int grid = pick_grid_size(n);
    hinge_broadcast_kernel<<<grid, kBlockSize, 0, stream>>>(
        pred_b.data_ptr<float>(), targ_b.data_ptr<float>(),
        out.data_ptr<float>(), n, inv_n, bc);
    return out;
}
"""

_hinge_loss_cpp_src = (
    "torch::Tensor hinge_loss_cuda(torch::Tensor pred, torch::Tensor targ);"
)

_hinge_loss_ext = load_inline(
    name="hinge_loss_fused_ext",
    cpp_sources=_hinge_loss_cpp_src,
    cuda_sources=_hinge_loss_cuda_src,
    functions=["hinge_loss_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Fused CUDA implementation of mean hinge loss for binary classification:
        mean(clamp(1 - predictions * targets, min=0))
    """

    def __init__(self):
        super(ModelNew, self).__init__()
        self._hinge_loss = _hinge_loss_ext.hinge_loss_cuda

    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(predictions, targets)
        # Reference fallback for any unexpected input configuration.
        return torch.mean(torch.clamp(1 - predictions * targets, min=0))
```