```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))
#
# Turn-4 improvements over turn 3:
#  * Lower host-side latency on the hot path: broadcast shape/stride metadata
#    is computed directly on stack arrays instead of building two expanded
#    tensor views (at::infer_size + expand), and the output accumulator is
#    created with empty + a 4-byte cudaMemsetAsync instead of the much
#    heavier torch::zeros dispatch + fill-kernel launch. The main reduction
#    kernel is bandwidth-bound (~95% of DRAM peak), so shaving exposed CPU
#    launch overhead is the remaining win.
#  * Same set of specialized kernels (streaming, flat-modulo broadcast row,
#    2-D no-div kernels, generic N-D) that proved correct and fast in turn 3,
#    plus an added same-device guard for robustness.
_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) {
    #pragma unroll
    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;
}

__device__ __forceinline__ float hinge4(const float4 p, const float4 t) {
    return fmaxf(1.0f - p.x * t.x, 0.0f)
         + fmaxf(1.0f - p.y * t.y, 0.0f)
         + fmaxf(1.0f - p.z * t.z, 0.0f)
         + fmaxf(1.0f - p.w * t.w, 0.0f);
}

// Both operands contiguous with identical layout (read-once streaming).
__global__ void __launch_bounds__(kBlockSize)
hinge_stream_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 sum0 = 0.0f, sum1 = 0.0f;
    if (use_vec4) {
        const long long n4 = n >> 2;
        const float4* p4 = reinterpret_cast<const float4*>(pred);
        const float4* t4 = reinterpret_cast<const float4*>(targ);
        long long k = tid;
        for (; k + stride < n4; k += 2 * stride) {
            const float4 p0 = __ldcs(p4 + k);
            const float4 t0 = __ldcs(t4 + k);
            const float4 p1 = __ldcs(p4 + k + stride);
            const float4 t1 = __ldcs(t4 + k + stride);
            sum0 += hinge4(p0, t0);
            sum1 += hinge4(p1, t1);
        }
        if (k < n4) {
            sum0 += hinge4(__ldcs(p4 + k), __ldcs(t4 + k));
        }
        for (long long e = (n4 << 2) + tid; e < n; e += stride) {
            sum0 += fmaxf(1.0f - __ldcs(pred + e) * __ldcs(targ + e), 0.0f);
        }
    } else {
        for (long long k = tid; k < n; k += stride) {
            sum0 += fmaxf(1.0f - __ldcs(pred + k) * __ldcs(targ + k), 0.0f);
        }
    }

    const float sum = block_reduce_sum(sum0 + sum1);
    if (threadIdx.x == 0) {
        atomicAdd(out, sum * inv_n);
    }
}

// Flat kernel for: one operand contiguous ("str", read once), the other a
// broadcast row ("bcast") of length C, where element k maps to bcast[k % C].
// The modulo is maintained incrementally (no per-element division).
__global__ void __launch_bounds__(kBlockSize)
hinge_flat_mod_kernel(const float* __restrict__ str,
                      const float* __restrict__ bcast,
                      float* __restrict__ out,
                      const long long n,
                      const long long C,
                      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 sum0 = 0.0f, sum1 = 0.0f;
    if (use_vec4) {
        const long long n4 = n >> 2;
        const long long c4 = C >> 2;
        const float4* s4 = reinterpret_cast<const float4*>(str);
        const float4* b4 = reinterpret_cast<const float4*>(bcast);
        const long long js  = stride % c4;
        const long long js2 = (2 * stride) % c4;
        long long j  = tid % c4;
        long long j2 = j + js;
        if (j2 >= c4) { j2 -= c4; }
        long long k = tid;
        for (; k + stride < n4; k += 2 * stride) {
            const float4 p0 = __ldcs(s4 + k);
            const float4 t0 = __ldg(b4 + j);
            const float4 p1 = __ldcs(s4 + k + stride);
            const float4 t1 = __ldg(b4 + j2);
            sum0 += hinge4(p0, t0);
            sum1 += hinge4(p1, t1);
            j  += js2; if (j  >= c4) { j  -= c4; }
            j2 += js2; if (j2 >= c4) { j2 -= c4; }
        }
        if (k < n4) {
            sum0 += hinge4(__ldcs(s4 + k), __ldg(b4 + j));
        }
        for (long long e = (n4 << 2) + tid; e < n; e += stride) {
            sum0 += fmaxf(1.0f - __ldcs(str + e) * __ldg(bcast + (e % C)),
                          0.0f);
        }
    } else {
        const long long js = stride % C;
        long long j = tid % C;
        for (long long k = tid; k < n; k += stride) {
            sum0 += fmaxf(1.0f - __ldcs(str + k) * __ldg(bcast + j), 0.0f);
            j += js;
            if (j >= C) { j -= C; }
        }
    }

    const float sum = block_reduce_sum(sum0 + sum1);
    if (threadIdx.x == 0) {
        atomicAdd(out, sum * inv_n);
    }
}

// 2-D broadcast kernel, no div/mod anywhere (generic fallback for the
// patterns the flat kernels do not cover).
template<int SP1, int ST1, bool VEC4>
__global__ void __launch_bounds__(kBlockSize)
hinge_2d_kernel(const float* __restrict__ pred,
                const float* __restrict__ targ,
                float* __restrict__ out,
                const long long R, const long long C,
                const float inv_n,
                const long long sp0, const long long st0) {
    const long long cid =
        static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x;
    const long long cstride =
        static_cast<long long>(gridDim.x) * blockDim.x;

    float sum = 0.0f;
    for (long long r = blockIdx.y; r < R; r += gridDim.y) {
        const float* pr = pred + r * sp0;
        const float* tr = targ + r * st0;
        float ps = 0.0f, ts = 0.0f;
        if (SP1 == 0) { ps = pr[0]; }
        if (ST1 == 0) { ts = tr[0]; }
        if (VEC4) {
            const long long n4 = C >> 2;
            const float4* pr4 = reinterpret_cast<const float4*>(pr);
            const float4* tr4 = reinterpret_cast<const float4*>(tr);
            for (long long c4 = cid; c4 < n4; c4 += cstride) {
                float4 pv, tv;
                if (SP1 == 1) { pv = pr4[c4]; }
                else          { pv = make_float4(ps, ps, ps, ps); }
                if (ST1 == 1) { tv = tr4[c4]; }
                else          { tv = make_float4(ts, ts, ts, ts); }
                sum += hinge4(pv, tv);
            }
            for (long long c = (n4 << 2) + cid; c < C; c += cstride) {
                const float pv = (SP1 == 1) ? pr[c] : ps;
                const float tv = (ST1 == 1) ? tr[c] : ts;
                sum += fmaxf(1.0f - pv * tv, 0.0f);
            }
        } else {
            for (long long c = cid; c < C; c += cstride) {
                const float pv = (SP1 == 1) ? pr[c] : ps;
                const float tv = (ST1 == 1) ? tr[c] : ts;
                sum += fmaxf(1.0f - pv * tv, 0.0f);
            }
        }
    }

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

// General 2-D kernel for arbitrary (e.g. sliced) strides, still no div/mod.
__global__ void __launch_bounds__(kBlockSize)
hinge_2d_general_kernel(const float* __restrict__ pred,
                        const float* __restrict__ targ,
                        float* __restrict__ out,
                        const long long R, const long long C,
                        const float inv_n,
                        const long long sp0, const long long sp1,
                        const long long st0, const long long st1) {
    const long long cid =
        static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x;
    const long long cstride =
        static_cast<long long>(gridDim.x) * blockDim.x;

    float sum = 0.0f;
    for (long long r = blockIdx.y; r < R; r += gridDim.y) {
        const float* pr = pred + r * sp0;
        const float* tr = targ + r * st0;
        for (long long c = cid; c < C; c += cstride) {
            sum += fmaxf(1.0f - pr[c * sp1] * tr[c * st1], 0.0f);
        }
    }

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

// General N-D fallback: handles arbitrary broadcasting without materializing.
__global__ void __launch_bounds__(kBlockSize)
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 sm_count() {
    return at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
}

long long grid_cap() {
    return static_cast<long long>(sm_count()) * 16;
}

long long flat_grid(long long work) {
    long long grid = (work + kBlockSize - 1) / kBlockSize;
    const long long cap = grid_cap();
    if (grid > cap) grid = cap;
    if (grid < 1) grid = 1;
    return grid;
}

void launch_stream(const float* p, const float* t, float* o,
                   long long n, float inv_n, cudaStream_t stream) {
    const bool aligned =
        ((reinterpret_cast<std::uintptr_t>(p) & 0xF) == 0) &&
        ((reinterpret_cast<std::uintptr_t>(t) & 0xF) == 0);
    const int use_vec4 = aligned ? 1 : 0;
    const long long work = use_vec4 ? ((n + 3) >> 2) : n;
    const long long grid = flat_grid(work);
    hinge_stream_kernel<<<static_cast<int>(grid), kBlockSize, 0, stream>>>(
        p, t, o, n, inv_n, use_vec4);
}

void launch_flat_mod(const float* str, const float* bcast, float* o,
                     long long n, long long C, float inv_n,
                     cudaStream_t stream) {
    const bool aligned =
        ((reinterpret_cast<std::uintptr_t>(str) & 0xF) == 0) &&
        ((reinterpret_cast<std::uintptr_t>(bcast) & 0xF) == 0);
    const bool vec4 = aligned && ((C & 3) == 0);
    const long long work = vec4 ? ((n + 3) >> 2) : n;
    const long long grid = flat_grid(work);
    hinge_flat_mod_kernel<<<static_cast<int>(grid), kBlockSize, 0, stream>>>(
        str, bcast, o, n, C, inv_n, vec4 ? 1 : 0);
}

int pick_block_2d(long long chunks_per_row) {
    if (chunks_per_row >= 256) return 256;
    if (chunks_per_row >= 128) return 128;
    if (chunks_per_row >= 64)  return 64;
    return 32;
}

dim3 pick_grid_2d(long long chunks_per_row, long long R, int block) {
    const long long target = grid_cap();
    long long gx = (chunks_per_row + block - 1) / block;
    if (gx > target) gx = target;
    if (gx < 1) gx = 1;
    long long gy = target / gx;
    if (gy > R) gy = R;
    if (gy < 1) gy = 1;
    if (gy > 65535) gy = 65535;
    return dim3(static_cast<unsigned>(gx), static_cast<unsigned>(gy));
}

template<int SP1, int ST1>
void launch_2d_simple(const float* p, const float* t, float* o,
                      long long R, long long C, float inv_n,
                      long long sp0, long long st0, bool vec4,
                      cudaStream_t stream) {
    const long long chunks = vec4 ? ((C + 3) >> 2) : C;
    const int block = pick_block_2d(chunks);
    const dim3 grid = pick_grid_2d(chunks, R, block);
    if (vec4) {
        hinge_2d_kernel<SP1, ST1, true><<<grid, block, 0, stream>>>(
            p, t, o, R, C, inv_n, sp0, st0);
    } else {
        hinge_2d_kernel<SP1, ST1, false><<<grid, block, 0, stream>>>(
            p, t, o, R, C, inv_n, sp0, st0);
    }
}

// If the operand described by (strides, sizes) always reads element
// data[(k % row_len)] at flat output index k, return row_len; else -1.
long long bcast_row_len(const long long* strides, const long long* sizes,
                        int nd) {
    long long row_len = 1;
    for (int d = 0; d < nd; ++d) {
        const long long sz = sizes[d];
        const long long st = strides[d];
        if (d == nd - 1) {
            if (sz == 1) {
                // contributes nothing
            } else if (st == 1) {
                row_len = sz;
            } else {
                return -1;
            }
        } else {
            if (sz == 1 || st == 0) {
                // contributes nothing
            } else {
                return -1;
            }
        }
    }
    return row_len;
}

}  // 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");
    TORCH_CHECK(pred.get_device() == targ.get_device(),
                "inputs must be on the same CUDA device");

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

    // empty + tiny memset: avoids the heavier zeros() dispatch + fill kernel.
    auto out = torch::empty({}, pred.options());
    float* o_ptr = out.data_ptr<float>();

    const int nd_p = static_cast<int>(pred.dim());
    const int nd_t = static_cast<int>(targ.dim());
    const int nd = nd_p > nd_t ? nd_p : nd_t;

    const float* p_ptr = pred.data_ptr<float>();
    const float* t_ptr = targ.data_ptr<float>();

    // Fast path: identical contiguous layouts (covers 0-dim as well).
    if (nd_p == nd_t && pred.is_contiguous() && targ.is_contiguous() &&
        pred.sizes() == targ.sizes()) {
        const long long n = static_cast<long long>(pred.numel());
        if (n == 0) {
            out.fill_(std::numeric_limits<float>::quiet_NaN());
            return out;
        }
        cudaMemsetAsync(o_ptr, 0, sizeof(float), stream);
        launch_stream(p_ptr, t_ptr, o_ptr, n,
                      1.0f / static_cast<float>(n), stream);
        return out;
    }

    if (nd <= kMaxDims) {
        // Manually compute broadcasted sizes/strides (right-aligned), which
        // is cheaper than materializing two expanded tensor views.
        long long sz[kMaxDims], sp[kMaxDims], st[kMaxDims];
        long long n = 1;
        for (int i = 0; i < nd; ++i) {
            const int dp = nd_p - 1 - i;
            const int dt = nd_t - 1 - i;
            const long long psz = (dp >= 0) ? pred.size(dp) : 1;
            const long long tsz = (dt >= 0) ? targ.size(dt) : 1;
            long long pst = (dp >= 0) ? pred.stride(dp) : 0;
            long long tst = (dt >= 0) ? targ.stride(dt) : 0;
            long long osz;
            if (psz == tsz) {
                osz = psz;
            } else if (psz == 1) {
                osz = tsz;
                pst = 0;
            } else if (tsz == 1) {
                osz = psz;
                tst = 0;
            } else {
                TORCH_CHECK(false,
                            "predictions and targets are not broadcastable");
                osz = 1;
            }
            sz[nd - 1 - i] = osz;
            sp[nd - 1 - i] = pst;
            st[nd - 1 - i] = tst;
            n *= osz;
        }
        if (n == 0) {
            // torch.mean of an empty tensor returns NaN.
            out.fill_(std::numeric_limits<float>::quiet_NaN());
            return out;
        }
        cudaMemsetAsync(o_ptr, 0, sizeof(float), stream);
        const float inv_n = 1.0f / static_cast<float>(n);

        // "Full coverage contiguous": operand exactly matches the broadcast
        // shape and is contiguous, so output element k reads operand[k].
        auto covers = [&](const at::Tensor& x, int nd_x) -> bool {
            if (nd_x != nd) return false;
            for (int i = 0; i < nd; ++i) {
                if (x.size(i) != sz[i]) return false;
            }
            return x.is_contiguous();
        };
        const bool p_full = covers(pred, nd_p);
        const bool t_full = covers(targ, nd_t);

        // Flat paths: one operand contiguous, the other a broadcast row
        // (covers matrix x row-vector, outer-product style, x scalar, and any
        //  higher-rank variant of these, e.g. (A,B,C) x (C,)).
        if (nd >= 1 && p_full) {
            const long long rl = bcast_row_len(st, sz, nd);
            if (rl >= 0) {
                launch_flat_mod(p_ptr, t_ptr, o_ptr, n, rl, inv_n, stream);
                return out;
            }
        }
        if (nd >= 1 && t_full) {
            const long long rl = bcast_row_len(sp, sz, nd);
            if (rl >= 0) {
                // hinge(p, t) is symmetric in its two operands.
                launch_flat_mod(t_ptr, p_ptr, o_ptr, n, rl, inv_n, stream);
                return out;
            }
        }

        // Specialized 2-D broadcast paths (column vectors, sliced layouts).
        if (nd <= 2) {
            long long R, C, sp0, sp1, st0, st1;
            if (nd == 2) {
                R = sz[0]; C = sz[1];
                sp0 = sp[0]; sp1 = sp[1];
                st0 = st[0]; st1 = st[1];
            } else if (nd == 1) {
                R = 1; C = sz[0];
                sp0 = 0; sp1 = sp[0];
                st0 = 0; st1 = st[0];
            } else {
                // Two scalars: n == 1.
                R = 1; C = 1;
                sp0 = sp1 = st0 = st1 = 0;
            }
            const bool simple =
                (sp1 == 0 || sp1 == 1) && (st1 == 0 || st1 == 1);
            if (simple) {
                // float4 needs 16B-aligned row starts for row-contiguous
                // operands.
                const bool p_ok = (sp1 == 0) ||
                    (((reinterpret_cast<std::uintptr_t>(p_ptr) & 0xF) == 0) &&
                     ((sp0 & 3) == 0));
                const bool t_ok = (st1 == 0) ||
                    (((reinterpret_cast<std::uintptr_t>(t_ptr) & 0xF) == 0) &&
                     ((st0 & 3) == 0));
                const bool vec4 = p_ok && t_ok;
                const int key =
                    static_cast<int>(sp1) * 2 + static_cast<int>(st1);
                switch (key) {
                    case 0:
                        launch_2d_simple<0, 0>(p_ptr, t_ptr, o_ptr, R, C,
                                               inv_n, sp0, st0, vec4, stream);
                        break;
                    case 1:
                        launch_2d_simple<0, 1>(p_ptr, t_ptr, o_ptr, R, C,
                                               inv_n, sp0, st0, vec4, stream);
                        break;
                    case 2:
                        launch_2d_simple<1, 0>(p_ptr, t_ptr, o_ptr, R, C,
                                               inv_n, sp0, st0, vec4, stream);
                        break;
                    default:
                        launch_2d_simple<1, 1>(p_ptr, t_ptr, o_ptr, R, C,
                                               inv_n, sp0, st0, vec4, stream);
                        break;
                }
            } else {
                const int block = pick_block_2d(C);
                const dim3 grid = pick_grid_2d(C, R, block);
                hinge_2d_general_kernel<<<grid, block, 0, stream>>>(
                    p_ptr, t_ptr, o_ptr, R, C, inv_n, sp0, sp1, st0, st1);
            }
            return out;
        }

        // Generic N-D fallback (no materialization).
        BroadcastInfo bc;
        bc.ndim = nd;
        for (int d = 0; d < nd; ++d) {
            bc.sizes[d]    = sz[d];
            bc.stride_p[d] = sp[d];
            bc.stride_t[d] = st[d];
        }
        const long long grid = flat_grid(n);
        hinge_broadcast_kernel<<<static_cast<int>(grid), kBlockSize, 0,
                                 stream>>>(p_ptr, t_ptr, o_ptr, n, inv_n, bc);
        return out;
    }

    // Very high-rank broadcasts: materialize contiguously and use fast path.
    {
        const auto shape = at::infer_size(pred.sizes(), targ.sizes());
        auto pc = pred.expand(shape).contiguous();
        auto tc = targ.expand(shape).contiguous();
        const long long n = static_cast<long long>(pc.numel());
        if (n == 0) {
            out.fill_(std::numeric_limits<float>::quiet_NaN());
            return out;
        }
        cudaMemsetAsync(o_ptr, 0, sizeof(float), stream);
        launch_stream(pc.data_ptr<float>(), tc.data_ptr<float>(), o_ptr, n,
                      1.0f / static_cast<float>(n), stream);
        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_v4",
    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"],
)

_HINGE_LOSS = _hinge_loss_ext.hinge_loss_cuda


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__()

    def forward(self, predictions, targets):
        if (predictions.is_cuda and targets.is_cuda
                and predictions.dtype == torch.float32
                and targets.dtype == torch.float32):
            return _HINGE_LOSS(predictions, targets)
        # Reference fallback for any unexpected input configuration.
        return torch.mean(torch.clamp(1 - predictions * targets, min=0))
```