```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.
#
# The op is memory-bound at exactly 1x traffic (predictions read once), so
# this revision focuses on squeezing the remaining overheads while keeping
# the proven column-locked streaming kernel:
#   * Host-side broadcast resolution on raw sizes/strides (no temporary
#     expanded tensors are materialized), shaving host latency per call.
#   * Column-locked grid-stride: broadcast targets live in registers
#     (per-column) or advance by a constant pointer bump (per-row); the hot
#     loop issues nothing but streaming float4 prediction loads + FMA/max.
#   * Dual accumulators in the unrolled loops to shorten FP dependency
#     chains.
#   * Vectorized / scalar / generic-strided fallbacks keep every input
#     layout (broadcasts, non-contiguous, unaligned, >2^31, dim>8) 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
//     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 acc0 = 0.0f, acc1 = 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 += 2) {
                acc0 += hinge4(pv[u], tv);
                acc1 += hinge4(pv[u + 1], tv);
            }
        }
        for (; i < n4; i += stride)
            acc0 += hinge4(__ldcs(p4 + i), tv);
    }
    float local = (acc0 + acc1) * 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 acc0 = 0.0f, acc1 = 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);
            acc0 += hinge4s(a, ta) + hinge4s(c2, tc);
            acc1 += hinge4s(b, tb) + hinge4s(d, td);
        }
        for (; i < n4; i += stride, tp += dTR)
            acc0 += hinge4s(__ldcs(p4 + i), __ldg(tp));
    }
    float local = (acc0 + acc1) * inv_n;
    local = blockReduceSum(local);
    if (threadIdx.x == 0) atomicAdd(out, local);
}

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

template <typename IndexT>
__global__ void __launch_bounds__(BLOCK_SIZE)
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 __launch_bounds__(BLOCK_SIZE)
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 + (IndexT)(UNROLL - 1) * stride < n4; i += (IndexT)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 __launch_bounds__(BLOCK_SIZE)
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 __launch_bounds__(BLOCK_SIZE)
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(int sm, long long items, int per_thread) {
    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) {
    // Exact PyTorch semantics for anything outside the fast path.
    if (!predictions.is_cuda() || !targets.is_cuda() ||
        predictions.scalar_type() != at::kFloat ||
        targets.scalar_type() != at::kFloat) {
        return at::mean(at::clamp(-at::mul(predictions, targets) + 1, 0));
    }

    // ---- resolve broadcasting on raw sizes/strides (no copies) ----
    const int pd = (int)predictions.dim();
    const int td = (int)targets.dim();
    const int bnd = pd > td ? pd : td;
    if (bnd > MAX_DIMS) {
        return at::mean(at::clamp(-at::mul(predictions, targets) + 1, 0));
    }

    long long bshape[MAX_DIMS];
    long long sP[MAX_DIMS];
    long long sT[MAX_DIMS];
    for (int k = 0; k < bnd; ++k) {
        const int d = bnd - 1 - k;
        const long long ps = (k < pd) ? predictions.size(pd - 1 - k) : 1;
        const long long ts = (k < td) ? targets.size(td - 1 - k) : 1;
        long long pst = (k < pd) ? predictions.stride(pd - 1 - k) : 0;
        long long tst = (k < td) ? targets.stride(td - 1 - k) : 0;
        long long sz;
        if (ps == ts) {
            sz = ps;
        } else if (ps == 1) {
            sz = ts; pst = 0;
        } else if (ts == 1) {
            sz = ps; tst = 0;
        } else {
            // Not broadcastable: let ATen raise the canonical error.
            return at::mean(at::clamp(-at::mul(predictions, targets) + 1, 0));
        }
        bshape[d] = sz;
        sP[d] = pst;
        sT[d] = tst;
    }
    long long numel = 1;
    for (int d = 0; d < bnd; ++d) numel *= bshape[d];

    auto out = at::zeros({}, predictions.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;
    const int sm = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;

    const float* pp = predictions.data_ptr<float>();
    const float* tp = targets.data_ptr<float>();

    // Is the broadcast view contiguous starting at the base pointer?
    auto viewContig = [&](const long long* bs) -> bool {
        long long expect = 1;
        for (int d = bnd - 1; d >= 0; --d) {
            if (bshape[d] != 1) {
                if (bs[d] != expect) return false;
                expect *= bshape[d];
            }
        }
        return true;
    };

    bool pc = viewContig(sP);
    bool tc = viewContig(sT);
    // hinge is symmetric in p and t; make p the contiguous one if possible.
    if (!pc && tc) {
        std::swap(pp, tp);
        std::swap_ranges(sP, sP + MAX_DIMS, sT);
        pc = true;
        tc = false;
    }

    // ---- fast path 1: both contiguous ----
    if (pc && tc) {
        if ((numel & 3LL) == 0 && aligned16(pp) && aligned16(tp)) {
            long long n4 = numel >> 2;
            int blocks = grid_capped(sm, 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(sm, 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 && bnd >= 1) {
        const long long C = bshape[bnd - 1];
        const long long sC = sT[bnd - 1];
        long long sR = 0, P = 1;
        bool ok = true, sR_set = false;
        for (int d = bnd - 2; d >= 0; --d) {
            long long sz = bshape[d];
            long long st = sT[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) {
            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) {
                        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(sm, 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(sm, 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 strided fallback ----
    BroadcastInfo info;
    info.ndim = bnd;
    for (int d = 0; d < bnd; ++d) {
        info.shape[d]   = bshape[d];
        info.strideP[d] = sP[d];
        info.strideT[d] = sT[d];
    }
    int blocks = grid_capped(sm, numel, 1);
    hinge_strided_kernel<<<blocks, BLOCK_SIZE, 0, stream>>>(
        pp, tp, 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_v6",
    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.
    Broadcasting is resolved on raw strides (no temporary tensors), 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.hinge_loss_cuda

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