```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 hinge loss: mean(clamp(1 - p * t, min=0)) in a single kernel pass.
#
# This revision targets the dominant broadcast workloads (a large contiguous
# prediction tensor against a small per-column or per-row target tensor):
#   * New "treg" kernels assign each thread one float4 column-chunk and a
#     contiguous slice of rows. The broadcast target value(s) are loaded ONCE
#     into registers and reused for every row, eliminating the per-element
#     target loads and all integer divisions from the inner loop.
#   * A flat (tile, column) thread decomposition keeps every thread busy and
#     every warp fully coalesced (>=128B contiguous per warp row-step) for any
#     column count, with no ragged-edge idle threads.
#   * Legacy vectorized/division kernels are kept as fallbacks for tiny column
#     counts, unaligned tails, general strided broadcasts, and the generic
#     N-D strided case, so all input layouts remain correct.
hinge_loss_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <algorithm>
#include <limits>
#include <cstdint>

#define BLOCK_SIZE 256
#define UNROLL 4
#define MAX_DIMS 8

// floor(v / d) for 0 <= v < 2^31, 1 <= d < 2^31 via precomputed magic constant.
struct FastDiv {
    unsigned long long M;
};

static inline FastDiv make_fastdiv(long long d) {
    FastDiv fd;
    fd.M = ((1ULL << 62) / (unsigned long long)d) + 1ULL;
    return fd;
}

__device__ __forceinline__ unsigned int fdiv_index(unsigned int v, const FastDiv& fd, unsigned int) {
    unsigned long long lo = (unsigned long long)v * fd.M;
    unsigned long long hi = __umul64hi((unsigned long long)v, fd.M);
    return (unsigned int)((hi << 2) | (lo >> 62));
}

__device__ __forceinline__ long long fdiv_index(long long v, const FastDiv&, long long d) {
    return v / d;
}

__device__ __forceinline__ float warpReduceSum(float v) {
    for (int offset = 16; offset > 0; offset >>= 1)
        v += __shfl_down_sync(0xffffffffu, v, offset);
    return v;
}

__device__ __forceinline__ float blockReduceSum(float v) {
    __shared__ float shared[BLOCK_SIZE / 32];
    int lane = threadIdx.x & 31;
    int wid  = threadIdx.x >> 5;
    v = warpReduceSum(v);
    if (lane == 0) shared[wid] = v;
    __syncthreads();
    v = (threadIdx.x < BLOCK_SIZE / 32) ? shared[threadIdx.x] : 0.0f;
    if (wid == 0) v = warpReduceSum(v);
    return v;
}

__device__ __forceinline__ float hinge_val(float pv, float tv) {
    float v = fmaf(-pv, tv, 1.0f);   // 1 - p * t
    // clamp(min=0) with NaN propagation to match torch.clamp semantics
    return v > 0.0f ? v : (v != v ? v : 0.0f);
}

__device__ __forceinline__ float hinge4(float4 p, float4 t) {
    return (hinge_val(p.x, t.x) + hinge_val(p.y, t.y))
         + (hinge_val(p.z, t.z) + hinge_val(p.w, t.w));
}

__device__ __forceinline__ float hinge4s(float4 p, float t) {
    return (hinge_val(p.x, t) + hinge_val(p.y, t))
         + (hinge_val(p.z, t) + hinge_val(p.w, t));
}

// =====================================================================
// Register-cached broadcast kernel.
// p is a contiguous (R, C) tensor viewed as (R, Cv) float4 rows.
// Thread flat id -> (row-tile y, column chunk c). Each thread owns one
// float4 column chunk for a contiguous run of rows [y*TR, min(R, y*TR+TR)).
// MODE 0: t is a per-column contiguous vector -> one float4, held in regs.
// MODE 1: t is a per-row scalar (stride sR)  -> one scalar load per row.
// Warps always span >=8 consecutive float4 column chunks (host guarantees
// Cv >= 8), so every row-step issues >=128B contiguous per warp.
// =====================================================================
template <typename IndexT, int MODE>
__global__ void hinge_bcast_treg_kernel(const float4* __restrict__ p4,
                                        const float* __restrict__ t,
                                        float* __restrict__ out,
                                        IndexT Cv, IndexT R, IndexT TR,
                                        IndexT tiles, long long sR,
                                        float inv_n) {
    IndexT id = (IndexT)blockIdx.x * BLOCK_SIZE + threadIdx.x;
    float local = 0.0f;
    IndexT y = id / Cv;
    if (y < tiles) {
        IndexT c = id - y * Cv;
        IndexT row0 = y * TR;
        IndexT len = R - row0;
        if (len > TR) len = TR;
        const float4* pp = p4 + row0 * Cv + c;
        if (MODE == 0) {
            float4 tv = __ldg(reinterpret_cast<const float4*>(t) + c);
            IndexT m = len;
            for (; m >= (IndexT)UNROLL; m -= (IndexT)UNROLL, pp += (IndexT)UNROLL * Cv) {
                float4 a = __ldcs(pp);
                float4 b = __ldcs(pp + Cv);
                float4 c2 = __ldcs(pp + (IndexT)2 * Cv);
                float4 d = __ldcs(pp + (IndexT)3 * Cv);
                local += (hinge4(a, tv) + hinge4(b, tv))
                       + (hinge4(c2, tv) + hinge4(d, tv));
            }
            for (; m > 0; --m, pp += Cv) local += hinge4(__ldcs(pp), tv);
        } else {
            const float* tp = t + (long long)row0 * sR;
            IndexT m = len;
            for (; m >= (IndexT)UNROLL; m -= (IndexT)UNROLL,
                   pp += (IndexT)UNROLL * Cv, tp += (long long)UNROLL * sR) {
                float4 a = __ldcs(pp);
                float4 b = __ldcs(pp + Cv);
                float4 c2 = __ldcs(pp + (IndexT)2 * Cv);
                float4 d = __ldcs(pp + (IndexT)3 * Cv);
                float ta = __ldg(tp);
                float tb = __ldg(tp + sR);
                float tc = __ldg(tp + (long long)2 * sR);
                float td = __ldg(tp + (long long)3 * sR);
                local += (hinge4s(a, ta) + hinge4s(b, tb))
                       + (hinge4s(c2, tc) + hinge4s(d, td));
            }
            for (; m > 0; --m, pp += Cv, tp += sR)
                local += hinge4s(__ldcs(pp), __ldg(tp));
        }
    }
    local *= inv_n;
    local = blockReduceSum(local);
    if (threadIdx.x == 0) atomicAdd(out, local);
}

// -------------------- both operands contiguous --------------------
template <typename IndexT>
__global__ void hinge_elem_vec4_kernel(const float4* __restrict__ p4,
                                       const float4* __restrict__ t4,
                                       float* __restrict__ out,
                                       IndexT n4, float inv_n) {
    IndexT stride = (IndexT)gridDim.x * BLOCK_SIZE;
    float local = 0.0f;
    IndexT i = (IndexT)blockIdx.x * BLOCK_SIZE + threadIdx.x;
    for (; i + (UNROLL - 1) * stride < n4; i += UNROLL * stride) {
        float4 pv[UNROLL], tv[UNROLL];
        #pragma unroll
        for (int u = 0; u < UNROLL; ++u) {
            pv[u] = __ldcs(p4 + i + (IndexT)u * stride);
            tv[u] = __ldcs(t4 + i + (IndexT)u * stride);
        }
        #pragma unroll
        for (int u = 0; u < UNROLL; ++u) local += hinge4(pv[u], tv[u]);
    }
    for (; i < n4; i += stride) {
        local += hinge4(__ldcs(p4 + i), __ldcs(t4 + i));
    }
    local *= inv_n;
    local = blockReduceSum(local);
    if (threadIdx.x == 0) atomicAdd(out, local);
}

template <typename IndexT>
__global__ void hinge_elem_scalar_kernel(const float* __restrict__ p,
                                         const float* __restrict__ t,
                                         float* __restrict__ out,
                                         IndexT n, float inv_n) {
    IndexT stride = (IndexT)gridDim.x * BLOCK_SIZE;
    float local = 0.0f;
    for (IndexT i = (IndexT)blockIdx.x * BLOCK_SIZE + threadIdx.x; i < n; i += stride)
        local += hinge_val(__ldcs(p + i), __ldcs(t + i));
    local *= inv_n;
    local = blockReduceSum(local);
    if (threadIdx.x == 0) atomicAdd(out, local);
}

// -------- linear broadcast fallback (small Cv / general gather) --------
// MODE 0: t per-column contiguous vector -> one float4 load.
// MODE 1: t per-row scalar (sC == 0)      -> one scalar load, splat.
// MODE 2: general 2D broadcast            -> four scalar gathers.
template <typename IndexT, int MODE>
__device__ __forceinline__ float4 load_t4(const float* __restrict__ t, IndexT v, IndexT Cv,
                                          long long sR, long long sC, const FastDiv& fd) {
    float4 r;
    if (MODE == 1) {
        IndexT row = fdiv_index(v, fd, Cv);
        float x = __ldg(t + (long long)row * sR);
        r.x = x; r.y = x; r.z = x; r.w = x;
    } else if (MODE == 0) {
        IndexT row = fdiv_index(v, fd, Cv);
        IndexT col = v - row * Cv;
        r = __ldg(reinterpret_cast<const float4*>(t + (long long)col * 4));
    } else {
        IndexT row = fdiv_index(v, fd, Cv);
        IndexT col = v - row * Cv;
        long long base = (long long)row * sR + (long long)col * 4 * sC;
        r.x = __ldg(t + base);
        r.y = __ldg(t + base + sC);
        r.z = __ldg(t + base + 2 * sC);
        r.w = __ldg(t + base + 3 * sC);
    }
    return r;
}

template <typename IndexT, int MODE>
__global__ void hinge_bcast_vec4_kernel(const float4* __restrict__ p4,
                                        const float* __restrict__ t,
                                        float* __restrict__ out,
                                        IndexT n4, IndexT Cv,
                                        long long sR, long long sC,
                                        FastDiv fd, float inv_n) {
    IndexT stride = (IndexT)gridDim.x * BLOCK_SIZE;
    float local = 0.0f;
    IndexT i = (IndexT)blockIdx.x * BLOCK_SIZE + threadIdx.x;
    for (; i + (UNROLL - 1) * stride < n4; i += UNROLL * stride) {
        float4 pv[UNROLL], tv[UNROLL];
        #pragma unroll
        for (int u = 0; u < UNROLL; ++u) pv[u] = __ldcs(p4 + i + (IndexT)u * stride);
        #pragma unroll
        for (int u = 0; u < UNROLL; ++u)
            tv[u] = load_t4<IndexT, MODE>(t, i + (IndexT)u * stride, Cv, sR, sC, fd);
        #pragma unroll
        for (int u = 0; u < UNROLL; ++u) local += hinge4(pv[u], tv[u]);
    }
    for (; i < n4; i += stride) {
        float4 pv = __ldcs(p4 + i);
        float4 tv = load_t4<IndexT, MODE>(t, i, Cv, sR, sC, fd);
        local += hinge4(pv, tv);
    }
    local *= inv_n;
    local = blockReduceSum(local);
    if (threadIdx.x == 0) atomicAdd(out, local);
}

template <typename IndexT>
__global__ void hinge_bcast_scalar_kernel(const float* __restrict__ p,
                                          const float* __restrict__ t,
                                          float* __restrict__ out,
                                          IndexT n, IndexT C,
                                          long long sR, long long sC,
                                          float inv_n) {
    IndexT stride = (IndexT)gridDim.x * BLOCK_SIZE;
    float local = 0.0f;
    for (IndexT i = (IndexT)blockIdx.x * BLOCK_SIZE + threadIdx.x; i < n; i += stride) {
        IndexT row = i / C;
        IndexT col = i - row * C;
        local += hinge_val(__ldcs(p + i), __ldg(t + (long long)row * sR + (long long)col * sC));
    }
    local *= inv_n;
    local = blockReduceSum(local);
    if (threadIdx.x == 0) atomicAdd(out, local);
}

// -------------------- generic strided fallback --------------------
struct BroadcastInfo {
    int ndim;
    long long shape[MAX_DIMS];
    long long strideP[MAX_DIMS];
    long long strideT[MAX_DIMS];
};

__global__ void hinge_strided_kernel(const float* __restrict__ p,
                                     const float* __restrict__ t,
                                     float* __restrict__ out,
                                     long long numel, float inv_n,
                                     BroadcastInfo info) {
    long long stride = (long long)gridDim.x * BLOCK_SIZE;
    float local = 0.0f;
    for (long long i = (long long)blockIdx.x * BLOCK_SIZE + threadIdx.x; i < numel; i += stride) {
        long long offP = 0, offT = 0, rem = i;
        #pragma unroll
        for (int d = MAX_DIMS - 1; d >= 0; --d) {
            if (d < info.ndim) {
                long long coord = rem % info.shape[d];
                rem /= info.shape[d];
                offP += coord * info.strideP[d];
                offT += coord * info.strideT[d];
            }
        }
        local += hinge_val(p[offP], t[offT]);
    }
    local *= inv_n;
    local = blockReduceSum(local);
    if (threadIdx.x == 0) atomicAdd(out, local);
}

static inline bool aligned16(const void* p) {
    return (reinterpret_cast<uintptr_t>(p) & 15) == 0;
}

// Grid sized to fill the device; grid-stride loops handle the rest.
static inline int grid_capped(long long items, int per_thread) {
    int sm = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    long long need = (items + (long long)BLOCK_SIZE * per_thread - 1) /
                     ((long long)BLOCK_SIZE * per_thread);
    long long cap = (long long)sm * 8;
    long long b = need < cap ? need : cap;
    if (b < 1) b = 1;
    return (int)b;
}

torch::Tensor hinge_loss_cuda(torch::Tensor predictions, torch::Tensor targets) {
    TORCH_CHECK(predictions.is_cuda() && targets.is_cuda(),
                "inputs must be CUDA tensors");

    auto p_in = predictions.scalar_type() == at::kFloat ? predictions
                                                        : predictions.to(at::kFloat);
    auto t_in = targets.scalar_type() == at::kFloat ? targets
                                                    : targets.to(at::kFloat);

    // Match PyTorch broadcasting semantics exactly.
    auto bc = at::broadcast_tensors({p_in, t_in});
    auto p = bc[0];
    auto t = bc[1];

    const long long numel = (long long)p.numel();
    auto out = at::zeros({}, p.options());
    if (numel == 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 / (float)numel;
    auto stream = at::cuda::getCurrentCUDAStream();
    const long long I31 = 2147483647LL;

    bool pc = p.is_contiguous();
    bool tc = t.is_contiguous();
    // hinge is symmetric in p and t; make p the contiguous one if possible.
    if (!pc && tc) { std::swap(p, t); pc = true; tc = false; }

    // ---- fast path 1: both contiguous ----
    if (pc && tc) {
        const float* pp = p.data_ptr<float>();
        const float* tp = t.data_ptr<float>();
        if ((numel & 3LL) == 0 && aligned16(pp) && aligned16(tp)) {
            long long n4 = numel >> 2;
            int blocks = grid_capped(n4, UNROLL);
            if (n4 <= I31)
                hinge_elem_vec4_kernel<unsigned int><<<blocks, BLOCK_SIZE, 0, stream>>>(
                    reinterpret_cast<const float4*>(pp),
                    reinterpret_cast<const float4*>(tp),
                    out.data_ptr<float>(), (unsigned int)n4, inv_n);
            else
                hinge_elem_vec4_kernel<long long><<<blocks, BLOCK_SIZE, 0, stream>>>(
                    reinterpret_cast<const float4*>(pp),
                    reinterpret_cast<const float4*>(tp),
                    out.data_ptr<float>(), n4, inv_n);
        } else {
            int blocks = grid_capped(numel, 1);
            if (numel <= I31)
                hinge_elem_scalar_kernel<unsigned int><<<blocks, BLOCK_SIZE, 0, stream>>>(
                    pp, tp, out.data_ptr<float>(), (unsigned int)numel, inv_n);
            else
                hinge_elem_scalar_kernel<long long><<<blocks, BLOCK_SIZE, 0, stream>>>(
                    pp, tp, out.data_ptr<float>(), numel, inv_n);
        }
        return out;
    }

    // ---- fast path 2: p contiguous, t collapses to offT = row*sR + col*sC ----
    if (pc && p.dim() >= 1 && p.dim() <= MAX_DIMS) {
        const int ndim = (int)p.dim();
        const long long C = p.size(ndim - 1);
        const long long sC = t.stride(ndim - 1);
        long long sR = 0, P = 1;
        bool ok = true, sR_set = false;
        for (int d = ndim - 2; d >= 0; --d) {
            long long sz = p.size(d);
            long long st = t.stride(d);
            if (sz != 1) {
                if (!sR_set) {
                    if (st % P != 0) { ok = false; break; }
                    sR = st / P;
                    sR_set = true;
                } else if (st != sR * P) { ok = false; break; }
            }
            P *= sz;
        }
        if (ok) {
            const float* pp = p.data_ptr<float>();
            const float* tp = t.data_ptr<float>();
            if ((C & 3LL) == 0 && aligned16(pp)) {
                int mode;
                if (sC == 0) mode = 1;                                  // per-row t
                else if (sC == 1 && sR == 0 && aligned16(tp)) mode = 0; // per-col t (vector)
                else mode = 2;                                          // general gather
                const long long n4 = numel >> 2;
                const long long Cv = C >> 2;
                const long long R = numel / C;
                if (mode != 2 && Cv >= 8) {
                    // Register-cached broadcast kernel: t held in registers,
                    // no divisions, fully coalesced, no idle threads.
                    int sm = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
                    long long target = (long long)sm * 2048;
                    long long tiles = 1;
                    if (Cv < target) {
                        tiles = (target + Cv - 1) / Cv;
                        long long maxTiles = (R + 15) / 16;
                        if (maxTiles < 1) maxTiles = 1;
                        if (tiles > maxTiles) tiles = maxTiles;
                        if (tiles > R) tiles = R;
                        if (tiles < 1) tiles = 1;
                    }
                    long long TR = (R + tiles - 1) / tiles;
                    tiles = (R + TR - 1) / TR;
                    long long total = Cv * tiles;
                    int blocks = (int)((total + BLOCK_SIZE - 1) / BLOCK_SIZE);
                    if (blocks < 1) blocks = 1;
                    if (n4 <= I31) {
                        const unsigned int Cvu = (unsigned int)Cv;
                        const unsigned int Ru = (unsigned int)R;
                        const unsigned int TRu = (unsigned int)TR;
                        const unsigned int tilesu = (unsigned int)tiles;
                        if (mode == 0)
                            hinge_bcast_treg_kernel<unsigned int, 0><<<blocks, BLOCK_SIZE, 0, stream>>>(
                                reinterpret_cast<const float4*>(pp), tp, out.data_ptr<float>(),
                                Cvu, Ru, TRu, tilesu, sR, inv_n);
                        else
                            hinge_bcast_treg_kernel<unsigned int, 1><<<blocks, BLOCK_SIZE, 0, stream>>>(
                                reinterpret_cast<const float4*>(pp), tp, out.data_ptr<float>(),
                                Cvu, Ru, TRu, tilesu, sR, inv_n);
                    } else {
                        if (mode == 0)
                            hinge_bcast_treg_kernel<long long, 0><<<blocks, BLOCK_SIZE, 0, stream>>>(
                                reinterpret_cast<const float4*>(pp), tp, out.data_ptr<float>(),
                                Cv, R, TR, tiles, sR, inv_n);
                        else
                            hinge_bcast_treg_kernel<long long, 1><<<blocks, BLOCK_SIZE, 0, stream>>>(
                                reinterpret_cast<const float4*>(pp), tp, out.data_ptr<float>(),
                                Cv, R, TR, tiles, sR, inv_n);
                    }
                } else {
                    // Linear vectorized fallback (tiny Cv or general gather).
                    FastDiv fd = make_fastdiv(Cv);
                    int blocks = grid_capped(n4, UNROLL);
                    if (n4 <= I31) {
                        const unsigned int n4u = (unsigned int)n4;
                        const unsigned int Cvu = (unsigned int)Cv;
                        if (mode == 0)
                            hinge_bcast_vec4_kernel<unsigned int, 0><<<blocks, BLOCK_SIZE, 0, stream>>>(
                                reinterpret_cast<const float4*>(pp), tp, out.data_ptr<float>(),
                                n4u, Cvu, sR, sC, fd, inv_n);
                        else if (mode == 1)
                            hinge_bcast_vec4_kernel<unsigned int, 1><<<blocks, BLOCK_SIZE, 0, stream>>>(
                                reinterpret_cast<const float4*>(pp), tp, out.data_ptr<float>(),
                                n4u, Cvu, sR, sC, fd, inv_n);
                        else
                            hinge_bcast_vec4_kernel<unsigned int, 2><<<blocks, BLOCK_SIZE, 0, stream>>>(
                                reinterpret_cast<const float4*>(pp), tp, out.data_ptr<float>(),
                                n4u, Cvu, sR, sC, fd, inv_n);
                    } else {
                        if (mode == 0)
                            hinge_bcast_vec4_kernel<long long, 0><<<blocks, BLOCK_SIZE, 0, stream>>>(
                                reinterpret_cast<const float4*>(pp), tp, out.data_ptr<float>(),
                                n4, Cv, sR, sC, fd, inv_n);
                        else if (mode == 1)
                            hinge_bcast_vec4_kernel<long long, 1><<<blocks, BLOCK_SIZE, 0, stream>>>(
                                reinterpret_cast<const float4*>(pp), tp, out.data_ptr<float>(),
                                n4, Cv, sR, sC, fd, inv_n);
                        else
                            hinge_bcast_vec4_kernel<long long, 2><<<blocks, BLOCK_SIZE, 0, stream>>>(
                                reinterpret_cast<const float4*>(pp), tp, out.data_ptr<float>(),
                                n4, Cv, sR, sC, fd, inv_n);
                    }
                }
            } else {
                int blocks = grid_capped(numel, 1);
                if (numel <= I31)
                    hinge_bcast_scalar_kernel<unsigned int><<<blocks, BLOCK_SIZE, 0, stream>>>(
                        pp, tp, out.data_ptr<float>(), (unsigned int)numel,
                        (unsigned int)C, sR, sC, inv_n);
                else
                    hinge_bcast_scalar_kernel<long long><<<blocks, BLOCK_SIZE, 0, stream>>>(
                        pp, tp, out.data_ptr<float>(), numel, C, sR, sC, inv_n);
            }
            return out;
        }
    }

    // ---- generic fallback ----
    if (p.dim() > MAX_DIMS) {
        return at::mean(at::clamp(-at::mul(p, t) + 1, 0));
    }
    BroadcastInfo info;
    info.ndim = (int)p.dim();
    for (int d = 0; d < info.ndim; ++d) {
        info.shape[d]   = p.size(d);
        info.strideP[d] = p.stride(d);
        info.strideT[d] = t.stride(d);
    }
    int blocks = grid_capped(numel, 1);
    hinge_strided_kernel<<<blocks, BLOCK_SIZE, 0, stream>>>(
        p.data_ptr<float>(), t.data_ptr<float>(), out.data_ptr<float>(),
        numel, inv_n, info);
    return out;
}
"""

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

# Compile the inline fused hinge-loss CUDA kernels.
hinge_loss_op = load_inline(
    name="hinge_loss_fused_v4",
    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):
    """
    Hinge Loss model backed by a single fused CUDA kernel that performs
    the elementwise hinge computation and the mean reduction in one pass.
    Broadcast targets are cached in registers (per-column vector or per-row
    scalar) so the main streaming loop issues only prediction loads.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.hinge_loss = hinge_loss_op

    def forward(self, predictions, targets):
        return self.hinge_loss.hinge_loss_cuda(predictions, targets)
```