```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 keeps the proven linear grid-stride streaming layout (best
# DRAM behavior) but eliminates ALL broadcast-target loads and integer
# divisions from the hot loop:
#   * Key idea: choose the grid so that the float4 grid-stride is an exact
#     multiple of the row length Cv (= C/4). Then every grid-stride step
#     lands on the SAME column chunk, so each thread loads its broadcast
#     target float4 exactly once into registers (per-column targets) or
#     advances the per-row target pointer by a constant increment
#     (per-sample targets). The main loop issues nothing but streaming
#     prediction loads plus FMA/max.
#   * 8-deep unrolled float4 loads per thread for high memory-level
#     parallelism.
#   * Legacy vectorized/division kernels remain as fallbacks for huge row
#     lengths, 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>
#include <numeric>

#define BLOCK_SIZE 256
#define UNROLL 4
#define COLREG_UNROLL 8
#define MAX_DIMS 8
#define COLREG_MAX_UNIT 4096

// 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));
}

// =====================================================================
// Column-locked broadcast kernels.
// Host guarantees the grid-stride (gridDim.x * BLOCK_SIZE) is a multiple
// of Cv, so every stride step revisits the same column chunk.
//   * colreg: t is a per-column contiguous vector -> each thread loads ONE
//     float4 of t into registers before the loop; the loop issues only
//     streaming prediction loads (8 in flight) and FMA/max math.
//   * rowreg: t is a per-row scalar (stride sR) -> the row advances by a
//     compile-time-constant amount per stride step, so the t pointer is
//     bumped by dTR with no divisions in the loop.
// =====================================================================
template <typename IndexT>
__global__ void __launch_bounds__(BLOCK_SIZE)
hinge_bcast_colreg_kernel(const float4* __restrict__ p4,
                          const float4* __restrict__ t4,
                          float* __restrict__ out,
                          IndexT n4, IndexT Cv,
                          FastDiv fd, float inv_n) {
    const IndexT stride = (IndexT)gridDim.x * BLOCK_SIZE;
    IndexT i = (IndexT)blockIdx.x * BLOCK_SIZE + threadIdx.x;
    float local = 0.0f;
    if (i < n4) {
        IndexT r = fdiv_index(i, fd, Cv);
        float4 tv = __ldg(t4 + (i - r * Cv));
        for (; i + (IndexT)(COLREG_UNROLL - 1) * stride < n4;
               i += (IndexT)COLREG_UNROLL * stride) {
            float4 pv[COLREG_UNROLL];
            #pragma unroll
            for (int u = 0; u < COLREG_UNROLL; ++u)
                pv[u] = __ldcs(p4 + i + (IndexT)u * stride);
            #pragma unroll
            for (int u = 0; u < COLREG_UNROLL; ++u)
                local += hinge4(pv[u], tv);
        }
        for (; i < n4; i += stride)
            local += hinge4(__ldcs(p4 + i), tv);
    }
    local *= inv_n;
    local = blockReduceSum(local);
    if (threadIdx.x == 0) atomicAdd(out, local);
}

template <typename IndexT>
__global__ void __launch_bounds__(BLOCK_SIZE)
hinge_bcast_rowreg_kernel(const float4* __restrict__ p4,
                          const float* __restrict__ t,
                          float* __restrict__ out,
                          IndexT n4, IndexT Cv,
                          long long sR, long long dTR,
                          FastDiv fd, float inv_n) {
    const IndexT stride = (IndexT)gridDim.x * BLOCK_SIZE;
    IndexT i = (IndexT)blockIdx.x * BLOCK_SIZE + threadIdx.x;
    float local = 0.0f;
    if (i < n4) {
        IndexT r = fdiv_index(i, fd, Cv);
        const float* tp = t + (long long)r * sR;
        for (; i + (IndexT)3 * stride < n4;
               i += (IndexT)4 * stride, tp += 4 * dTR) {
            float4 a  = __ldcs(p4 + i);
            float4 b  = __ldcs(p4 + i + stride);
            float4 c2 = __ldcs(p4 + i + (IndexT)2 * stride);
            float4 d  = __ldcs(p4 + i + (IndexT)3 * stride);
            float ta = __ldg(tp);
            float tb = __ldg(tp + dTR);
            float tc = __ldg(tp + 2 * dTR);
            float td = __ldg(tp + 3 * dTR);
            local += (hinge4s(a, ta) + hinge4s(b, tb))
                   + (hinge4s(c2, tc) + hinge4s(d, td));
        }
        for (; i < n4; i += stride, tp += dTR)
            local += hinge4s(__ldcs(p4 + i), __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 (huge 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");

    // Non-float32 dtypes: keep exact PyTorch semantics via ATen.
    if (predictions.scalar_type() != at::kFloat ||
        targets.scalar_type() != at::kFloat) {
        auto bcnf = at::broadcast_tensors({predictions, targets});
        return at::mean(at::clamp(-at::mul(bcnf[0], bcnf[1]) + 1, 0));
    }

    // Match PyTorch broadcasting semantics exactly.
    auto bc = at::broadcast_tensors({predictions, targets});
    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;
                FastDiv fd = make_fastdiv(Cv);
                bool done = false;
                if (mode != 2) {
                    // Column-locked grid: stride = blocks*256 must be a
                    // multiple of Cv. Minimal block unit B0 = Cv/gcd(Cv,256).
                    long long g = std::gcd(Cv, (long long)BLOCK_SIZE);
                    long long B0 = Cv / g;
                    if (B0 <= COLREG_MAX_UNIT) {
                        int sm = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
                        long long cap = (long long)sm * 8;
                        long long need = (n4 + BLOCK_SIZE - 1) / BLOCK_SIZE;
                        long long j = cap / B0;
                        if (j < 1) j = 1;
                        long long uj = (need + B0 - 1) / B0;
                        if (uj < 1) uj = 1;
                        if (j > uj) j = uj;
                        long long blocks = B0 * j;
                        if (mode == 0) {
                            if (n4 <= I31)
                                hinge_bcast_colreg_kernel<unsigned int><<<(int)blocks, BLOCK_SIZE, 0, stream>>>(
                                    reinterpret_cast<const float4*>(pp),
                                    reinterpret_cast<const float4*>(tp),
                                    out.data_ptr<float>(),
                                    (unsigned int)n4, (unsigned int)Cv, fd, inv_n);
                            else
                                hinge_bcast_colreg_kernel<long long><<<(int)blocks, BLOCK_SIZE, 0, stream>>>(
                                    reinterpret_cast<const float4*>(pp),
                                    reinterpret_cast<const float4*>(tp),
                                    out.data_ptr<float>(),
                                    n4, Cv, fd, inv_n);
                        } else {
                            long long dTR = ((blocks * BLOCK_SIZE) / Cv) * sR;
                            if (n4 <= I31)
                                hinge_bcast_rowreg_kernel<unsigned int><<<(int)blocks, BLOCK_SIZE, 0, stream>>>(
                                    reinterpret_cast<const float4*>(pp), tp,
                                    out.data_ptr<float>(),
                                    (unsigned int)n4, (unsigned int)Cv,
                                    sR, dTR, fd, inv_n);
                            else
                                hinge_bcast_rowreg_kernel<long long><<<(int)blocks, BLOCK_SIZE, 0, stream>>>(
                                    reinterpret_cast<const float4*>(pp), tp,
                                    out.data_ptr<float>(),
                                    n4, Cv, sR, dTR, fd, inv_n);
                        }
                        done = true;
                    }
                }
                if (!done) {
                    // Linear vectorized fallback (huge Cv or general gather).
                    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_v5",
    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.
    The grid-stride is locked to a multiple of the row length so broadcast
    targets live in registers (per-column) or advance by a constant pointer
    bump (per-row), leaving only streaming prediction loads in the hot loop.
    """
    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)
```