```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-2 improvements over turn 1:
#  * The generic N-D broadcast kernel (64-bit div/mod per element per dim) was
#    the bottleneck (~10.5 ms). It is replaced by specialized 2-D broadcast
#    kernels that need NO per-element div/mod: a 2-D grid walks rows/columns
#    directly with pointer arithmetic.
#  * The 2-D kernels are templated on whether each operand's column stride is
#    0 (broadcast along the row) or 1 (row-contiguous), and use float4 loads
#    when alignment permits. This covers all common patterns:
#      contiguous x contiguous, matrix x row-vector, matrix x col-vector,
#      col-vector x row-vector (outer), anything x scalar, 1-D broadcasts.
#  * A fully general strided 2-D kernel and the old N-D kernel remain as
#    fallbacks for arbitrary inputs.
_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;
}

// Fast path: both operands contiguous with identical layout.
__global__ void __launch_bounds__(kBlockSize)
hinge_contiguous_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 sum = 0.0f;
    if (use_vec4) {
        const long long n4 = n >> 2;
        const float4* pred4 = reinterpret_cast<const float4*>(pred);
        const float4* targ4 = reinterpret_cast<const float4*>(targ);
        for (long long k = tid; k < n4; k += stride) {
            const float4 p = pred4[k];
            const float4 t = targ4[k];
            sum += fmaxf(1.0f - p.x * t.x, 0.0f);
            sum += fmaxf(1.0f - p.y * t.y, 0.0f);
            sum += fmaxf(1.0f - p.z * t.z, 0.0f);
            sum += fmaxf(1.0f - p.w * t.w, 0.0f);
        }
        for (long long k = (n4 << 2) + tid; k < n; k += stride) {
            sum += fmaxf(1.0f - pred[k] * targ[k], 0.0f);
        }
    } else {
        for (long long k = tid; k < n; k += stride) {
            sum += fmaxf(1.0f - pred[k] * targ[k], 0.0f);
        }
    }

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

// 2-D broadcast kernel, no div/mod anywhere.
// SP1/ST1: column stride of pred/targ, restricted to 0 or 1 (compile time).
// VEC4: use float4 loads for the row-contiguous operands (compile time).
// Element (r, c): pred[r * sp0 + c * SP1] * targ[r * st0 + c * ST1].
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 += fmaxf(1.0f - pv.x * tv.x, 0.0f);
                sum += fmaxf(1.0f - pv.y * tv.y, 0.0f);
                sum += fmaxf(1.0f - pv.z * tv.z, 0.0f);
                sum += fmaxf(1.0f - pv.w * tv.w, 0.0f);
            }
            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;
}

void launch_contiguous(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;
    long long grid = (work + kBlockSize - 1) / kBlockSize;
    const long long cap = static_cast<long long>(sm_count()) * 16;
    if (grid > cap) grid = cap;
    if (grid < 1) grid = 1;
    hinge_contiguous_kernel<<<static_cast<int>(grid), kBlockSize, 0, stream>>>(
        p, t, o, n, inv_n, use_vec4);
}

dim3 pick_grid_2d(long long chunks_per_row, long long R) {
    const long long target = static_cast<long long>(sm_count()) * 16;
    long long gx = (chunks_per_row + kBlockSize - 1) / kBlockSize;
    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 dim3 grid = pick_grid_2d(chunks, R);
    if (vec4) {
        hinge_2d_kernel<SP1, ST1, true><<<grid, kBlockSize, 0, stream>>>(
            p, t, o, R, C, inv_n, sp0, st0);
    } else {
        hinge_2d_kernel<SP1, ST1, false><<<grid, kBlockSize, 0, stream>>>(
            p, t, o, R, C, inv_n, sp0, st0);
    }
}

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

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

    auto out = torch::zeros({}, pred.options());

    // Match PyTorch broadcasting semantics of `predictions * targets`.
    const auto shape = at::infer_size(pred.sizes(), targ.sizes());
    auto pred_b = pred.expand(shape);
    auto targ_b = targ.expand(shape);

    const long long n = static_cast<long long>(pred_b.numel());
    if (n == 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 / static_cast<float>(n);
    const int ndim = static_cast<int>(shape.size());

    const float* p_ptr = pred_b.data_ptr<float>();
    const float* t_ptr = targ_b.data_ptr<float>();
    float* o_ptr = out.data_ptr<float>();

    // Fast path: identical contiguous layouts.
    if (pred_b.is_contiguous() && targ_b.is_contiguous()) {
        launch_contiguous(p_ptr, t_ptr, o_ptr, n, inv_n, stream);
        return out;
    }

    // Specialized 2-D broadcast paths (covers the common cases and any
    // 1-D broadcast by treating it as 1 x C).
    if (ndim <= 2) {
        long long R, C, sp0, sp1, st0, st1;
        if (ndim == 2) {
            R = shape[0]; C = shape[1];
            sp0 = pred_b.stride(0); sp1 = pred_b.stride(1);
            st0 = targ_b.stride(0); st1 = targ_b.stride(1);
        } else {
            R = 1; C = shape[0];
            sp0 = 0; sp1 = pred_b.stride(0);
            st0 = 0; st1 = targ_b.stride(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 dim3 grid = pick_grid_2d(C, R);
            hinge_2d_general_kernel<<<grid, kBlockSize, 0, stream>>>(
                p_ptr, t_ptr, o_ptr, R, C, inv_n, sp0, sp1, st0, st1);
        }
        return out;
    }

    // Very high-rank broadcasts: materialize contiguously and use fast path.
    if (ndim > kMaxDims) {
        auto pc = pred_b.contiguous();
        auto tc = targ_b.contiguous();
        launch_contiguous(pc.data_ptr<float>(), tc.data_ptr<float>(),
                          o_ptr, n, inv_n, stream);
        return out;
    }

    // Generic N-D fallback.
    BroadcastInfo bc;
    bc.ndim = ndim;
    for (int d = 0; d < ndim; ++d) {
        bc.sizes[d]    = shape[d];
        bc.stride_p[d] = pred_b.stride(d);
        bc.stride_t[d] = targ_b.stride(d);
    }
    long long grid = (n + kBlockSize - 1) / kBlockSize;
    const long long cap = static_cast<long long>(sm_count()) * 16;
    if (grid > cap) grid = cap;
    if (grid < 1) grid = 1;
    hinge_broadcast_kernel<<<static_cast<int>(grid), kBlockSize, 0, stream>>>(
        p_ptr, t_ptr, o_ptr, n, inv_n, bc);
    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_v2",
    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"],
)


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__()
        self._hinge_loss = _hinge_loss_ext.hinge_loss_cuda

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