```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 kernel: computes mean(clamp(1 - predictions * targets, min=0))
# in a single pass with no intermediate materialization. Supports full
# PyTorch broadcasting semantics via shape/stride arithmetic, and uses
# block-level tree reductions + double-precision atomic accumulation for
# numerical accuracy.
hinge_loss_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/AccumulateType.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <algorithm>
#include <cstdint>
#include <limits>
#include <type_traits>

#define HINGE_THREADS 256
#define HINGE_MAX_DIMS 8
#define HINGE_MAX_BLOCKS 2048

namespace hinge_cuda {

struct BroadcastParams {
    int64_t shape[HINGE_MAX_DIMS];
    int64_t ps[HINGE_MAX_DIMS];
    int64_t ts[HINGE_MAX_DIMS];
    int ndim;
};

void compute_broadcast_params(const at::Tensor& p, const at::Tensor& t,
                              BroadcastParams& bp, int64_t& total) {
    const auto p_sizes = p.sizes();
    const auto t_sizes = t.sizes();
    const auto p_strides = p.strides();
    const auto t_strides = t.strides();
    const int pd = (int)p_sizes.size();
    const int td = (int)t_sizes.size();
    int ndim = pd > td ? pd : td;
    if (ndim == 0) ndim = 1;  // scalar inputs
    TORCH_CHECK(ndim <= HINGE_MAX_DIMS, "Too many dimensions");
    total = 1;
    for (int i = 0; i < ndim; ++i) {
        const int pi = i - (ndim - pd);
        const int ti = i - (ndim - td);
        const int64_t psz = (pi >= 0) ? p_sizes[pi] : 1;
        const int64_t tsz = (ti >= 0) ? t_sizes[ti] : 1;
        TORCH_CHECK(psz == tsz || psz == 1 || tsz == 1,
                    "predictions and targets are not broadcastable");
        const int64_t sz = psz > tsz ? psz : tsz;
        bp.shape[i] = sz;
        bp.ps[i] = (psz == sz) ? ((pi >= 0) ? p_strides[pi] : 0) : 0;
        bp.ts[i] = (tsz == sz) ? ((ti >= 0) ? t_strides[ti] : 0) : 0;
        total *= sz;
    }
    for (int i = ndim; i < HINGE_MAX_DIMS; ++i) {
        bp.shape[i] = 1;
        bp.ps[i] = 0;
        bp.ts[i] = 0;
    }
    bp.ndim = ndim;
}

template <typename T>
__device__ __forceinline__ T warp_reduce_sum(T val) {
#pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_down_sync(0xffffffffu, val, offset);
    return val;
}

template <typename T>
__device__ __forceinline__ T block_reduce_sum(T val, int tid, int nthreads) {
    __shared__ T shared[HINGE_THREADS / 32];
    const int lane = tid & 31;
    const int wid = tid >> 5;
    val = warp_reduce_sum(val);
    if (lane == 0) shared[wid] = val;
    __syncthreads();
    const int nwarps = (nthreads + 31) >> 5;
    val = (tid < nwarps) ? shared[tid] : T(0);
    if (wid == 0) val = warp_reduce_sum(val);
    return val;
}

template <typename acc_t>
__device__ __forceinline__ acc_t hinge_val(acc_t a, acc_t b) {
    const acc_t v = acc_t(1) - a * b;
    return v > acc_t(0) ? v : acc_t(0);
}

// Fast path: both tensors flat & contiguous with identical numel.
template <typename scalar_t>
__global__ void hinge_flat_kernel(const scalar_t* __restrict__ p,
                                  const scalar_t* __restrict__ t,
                                  double* __restrict__ out,
                                  int64_t total) {
    using acc_t = typename at::acc_type<scalar_t, true>;
    const int64_t stride = (int64_t)gridDim.x * blockDim.x;
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    acc_t acc = acc_t(0);
    for (; idx < total; idx += stride) {
        acc += hinge_val<acc_t>(static_cast<acc_t>(p[idx]), static_cast<acc_t>(t[idx]));
    }
    acc = block_reduce_sum<acc_t>(acc, threadIdx.x, blockDim.x);
    if (threadIdx.x == 0) atomicAdd(out, (double)acc);
}

// Vectorized (float4) fast path for fp32.
__global__ void hinge_flat_vec4_kernel(const float* __restrict__ p,
                                       const float* __restrict__ t,
                                       double* __restrict__ out,
                                       int64_t nvec,
                                       int64_t total) {
    const float4* __restrict__ p4 = reinterpret_cast<const float4*>(p);
    const float4* __restrict__ t4 = reinterpret_cast<const float4*>(t);
    const int64_t stride = (int64_t)gridDim.x * blockDim.x;
    const int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    float acc = 0.f;
    for (int64_t i = idx; i < nvec; i += stride) {
        const float4 a = p4[i];
        const float4 b = t4[i];
        acc += hinge_val<float>(a.x, b.x);
        acc += hinge_val<float>(a.y, b.y);
        acc += hinge_val<float>(a.z, b.z);
        acc += hinge_val<float>(a.w, b.w);
    }
    for (int64_t i = (nvec << 2) + idx; i < total; i += stride) {
        acc += hinge_val<float>(p[i], t[i]);
    }
    acc = block_reduce_sum<float>(acc, threadIdx.x, blockDim.x);
    if (threadIdx.x == 0) atomicAdd(out, (double)acc);
}

// 1-D broadcast / strided path.
template <typename scalar_t>
__global__ void hinge_strided1d_kernel(const scalar_t* __restrict__ p,
                                       const scalar_t* __restrict__ t,
                                       double* __restrict__ out,
                                       int64_t total,
                                       int64_t ps0, int64_t ts0) {
    using acc_t = typename at::acc_type<scalar_t, true>;
    const int64_t stride = (int64_t)gridDim.x * blockDim.x;
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    acc_t acc = acc_t(0);
    for (; idx < total; idx += stride) {
        acc += hinge_val<acc_t>(static_cast<acc_t>(p[idx * ps0]),
                                static_cast<acc_t>(t[idx * ts0]));
    }
    acc = block_reduce_sum<acc_t>(acc, threadIdx.x, blockDim.x);
    if (threadIdx.x == 0) atomicAdd(out, (double)acc);
}

// 2-D broadcast path (division-free index math).
template <typename scalar_t>
__global__ void hinge_2d_kernel(const scalar_t* __restrict__ p,
                                const scalar_t* __restrict__ t,
                                double* __restrict__ out,
                                int64_t rows, int64_t cols,
                                int64_t ps0, int64_t ps1,
                                int64_t ts0, int64_t ts1) {
    using acc_t = typename at::acc_type<scalar_t, true>;
    const int64_t col0 = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    const int64_t row0 = (int64_t)blockIdx.y * blockDim.y + threadIdx.y;
    const int64_t cstride = (int64_t)gridDim.x * blockDim.x;
    const int64_t rstride = (int64_t)gridDim.y * blockDim.y;
    acc_t acc = acc_t(0);
    for (int64_t r = row0; r < rows; r += rstride) {
        const scalar_t* __restrict__ pr = p + r * ps0;
        const scalar_t* __restrict__ tr = t + r * ts0;
        for (int64_t c = col0; c < cols; c += cstride) {
            acc += hinge_val<acc_t>(static_cast<acc_t>(pr[c * ps1]),
                                    static_cast<acc_t>(tr[c * ts1]));
        }
    }
    const int tid = threadIdx.y * blockDim.x + threadIdx.x;
    acc = block_reduce_sum<acc_t>(acc, tid, blockDim.x * blockDim.y);
    if (tid == 0) atomicAdd(out, (double)acc);
}

// Generic N-D broadcast path.
template <typename scalar_t, typename index_t>
__global__ void hinge_generic_kernel(const scalar_t* __restrict__ p,
                                     const scalar_t* __restrict__ t,
                                     double* __restrict__ out,
                                     index_t total,
                                     BroadcastParams bp) {
    using acc_t = typename at::acc_type<scalar_t, true>;
    const index_t stride = (index_t)gridDim.x * blockDim.x;
    index_t idx = (index_t)blockIdx.x * blockDim.x + threadIdx.x;
    acc_t acc = acc_t(0);
    for (; idx < total; idx += stride) {
        index_t rem = idx;
        int64_t po = 0, to = 0;
#pragma unroll
        for (int d = HINGE_MAX_DIMS - 1; d >= 0; --d) {
            if (d < bp.ndim) {
                const index_t sh = (index_t)bp.shape[d];
                const index_t coord = rem % sh;
                rem = rem / sh;
                po += (int64_t)coord * bp.ps[d];
                to += (int64_t)coord * bp.ts[d];
            }
        }
        acc += hinge_val<acc_t>(static_cast<acc_t>(p[po]), static_cast<acc_t>(t[to]));
    }
    acc = block_reduce_sum<acc_t>(acc, threadIdx.x, blockDim.x);
    if (threadIdx.x == 0) atomicAdd(out, (double)acc);
}

}  // namespace hinge_cuda

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() == targets.scalar_type(), "Input dtypes must match");
    TORCH_CHECK(predictions.device() == targets.device(), "Inputs must be on the same device");
    const c10::cuda::OptionalCUDAGuard device_guard(predictions.device());

    hinge_cuda::BroadcastParams bp;
    int64_t total = 0;
    hinge_cuda::compute_broadcast_params(predictions, targets, bp, total);

    if (total == 0) {  // mean over an empty tensor is NaN, matching torch.mean
        return at::full({}, std::numeric_limits<double>::quiet_NaN(), predictions.options());
    }

    at::Tensor out = at::zeros({}, predictions.options().dtype(at::kDouble));
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const bool flat = predictions.numel() == total && targets.numel() == total &&
                      predictions.is_contiguous() && targets.is_contiguous();

    AT_DISPATCH_FLOATING_TYPES_AND_HALF(predictions.scalar_type(), "hinge_loss_cuda", [&] {
        const scalar_t* p_ptr = predictions.data_ptr<scalar_t>();
        const scalar_t* t_ptr = targets.data_ptr<scalar_t>();
        double* out_ptr = out.data_ptr<double>();

        if (flat) {
            if constexpr (std::is_same<scalar_t, float>::value) {
                const bool aligned =
                    ((reinterpret_cast<uintptr_t>(p_ptr) | reinterpret_cast<uintptr_t>(t_ptr)) & 0xF) == 0;
                if (aligned && total >= 4096) {
                    const int64_t nvec = total >> 2;
                    const int blocks = (int)std::min<int64_t>(
                        (nvec + HINGE_THREADS - 1) / HINGE_THREADS, HINGE_MAX_BLOCKS);
                    hinge_cuda::hinge_flat_vec4_kernel<<<blocks, HINGE_THREADS, 0, stream>>>(
                        reinterpret_cast<const float*>(p_ptr),
                        reinterpret_cast<const float*>(t_ptr), out_ptr, nvec, total);
                    return;
                }
            }
            const int blocks = (int)std::min<int64_t>(
                (total + HINGE_THREADS - 1) / HINGE_THREADS, HINGE_MAX_BLOCKS);
            hinge_cuda::hinge_flat_kernel<scalar_t><<<blocks, HINGE_THREADS, 0, stream>>>(
                p_ptr, t_ptr, out_ptr, total);
        } else if (bp.ndim == 1) {
            const int blocks = (int)std::min<int64_t>(
                (total + HINGE_THREADS - 1) / HINGE_THREADS, HINGE_MAX_BLOCKS);
            hinge_cuda::hinge_strided1d_kernel<scalar_t><<<blocks, HINGE_THREADS, 0, stream>>>(
                p_ptr, t_ptr, out_ptr, total, bp.ps[0], bp.ts[0]);
        } else if (bp.ndim == 2) {
            const dim3 block(32, 8);
            const int64_t rows = bp.shape[0];
            const int64_t cols = bp.shape[1];
            const unsigned gx = (unsigned)std::min<int64_t>(
                (cols + block.x - 1) / block.x, HINGE_MAX_BLOCKS);
            const unsigned gy = (unsigned)std::min<int64_t>(
                (rows + block.y - 1) / block.y, 4096);
            const dim3 grid(gx, gy);
            hinge_cuda::hinge_2d_kernel<scalar_t><<<grid, block, 0, stream>>>(
                p_ptr, t_ptr, out_ptr, rows, cols,
                bp.ps[0], bp.ps[1], bp.ts[0], bp.ts[1]);
        } else {
            const int blocks = (int)std::min<int64_t>(
                (total + HINGE_THREADS - 1) / HINGE_THREADS, HINGE_MAX_BLOCKS);
            if (total <= (int64_t)std::numeric_limits<int32_t>::max()) {
                hinge_cuda::hinge_generic_kernel<scalar_t, int32_t>
                    <<<blocks, HINGE_THREADS, 0, stream>>>(
                        p_ptr, t_ptr, out_ptr, (int32_t)total, bp);
            } else {
                hinge_cuda::hinge_generic_kernel<scalar_t, int64_t>
                    <<<blocks, HINGE_THREADS, 0, stream>>>(
                        p_ptr, t_ptr, out_ptr, total, bp);
            }
        }
    });
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out.div(total).to(predictions.scalar_type());
}
"""

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

hinge_loss_ext = load_inline(
    name="hinge_loss_fused_ext",
    cpp_sources=hinge_loss_cpp_source,
    cuda_sources=hinge_loss_cuda_source,
    functions=["hinge_loss_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    A model that computes Hinge Loss for binary classification tasks using a
    fused custom CUDA kernel (multiply + clamp + mean reduction in one pass,
    with full broadcasting support and no intermediate tensors).
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._ext = hinge_loss_ext

    def forward(self, predictions, targets):
        if (predictions.is_cuda and targets.is_cuda
                and predictions.device == targets.device
                and predictions.dtype == targets.dtype
                and predictions.dtype in (torch.float32, torch.float64, torch.float16)
                and not (torch.is_grad_enabled()
                         and (predictions.requires_grad or targets.requires_grad))):
            return self._ext.hinge_loss_cuda(predictions, targets)
        # Safe fallback preserving the original semantics in every other case.
        return torch.mean(torch.clamp(1 - predictions * targets, min=0))
```