import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that computes Hinge Loss for binary classification tasks.

    Parameters:
        None
    """
    def __init__(self):
        super(Model, self).__init__()

    def forward(self, predictions, targets):
        return torch.mean(torch.clamp(1 - predictions * targets, min=0))

batch_size = 32768
input_shape = (32768,)
dim = 1

def get_inputs():
    return [torch.rand(batch_size, *input_shape), torch.randint(0, 2, (batch_size,)).float() * 2 - 1]

def get_init_inputs():
    return []

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor hinge_forward(torch::Tensor predictions, torch::Tensor targets);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <stdint.h>

static __device__ __forceinline__ float hinge_val(float p, float t) {
    float v = fmaf(-p, t, 1.0f);
    return fmaxf(v, 0.0f);
}

__global__ void hinge_vec_pow2_kernel(const float* __restrict__ pred,
                                      const float* __restrict__ targ,
                                      float* __restrict__ partial,
                                      int64_t n4,
                                      int64_t cols) {
    constexpr int THREADS = 256;
    __shared__ float smem[THREADS];
    int tid = threadIdx.x;
    int64_t i0 = (int64_t)blockIdx.x * blockDim.x + tid;
    int64_t stride = (int64_t)gridDim.x * blockDim.x;
    int64_t mask = cols - 1;

    int64_t col = (i0 * 4) & mask;
    const float4 tv = *reinterpret_cast<const float4*>(targ + col);
    const float4* __restrict__ p4 = reinterpret_cast<const float4*>(pred);

    float sum = 0.0f;
    for (int64_t i = i0; i < n4; i += stride) {
        float4 pv = p4[i];
        float s0 = hinge_val(pv.x, tv.x);
        float s1 = hinge_val(pv.y, tv.y);
        float s2 = hinge_val(pv.z, tv.z);
        float s3 = hinge_val(pv.w, tv.w);
        sum += (s0 + s1) + (s2 + s3);
    }

    smem[tid] = sum;
    __syncthreads();
    for (int off = THREADS / 2; off > 0; off >>= 1) {
        if (tid < off) smem[tid] += smem[tid + off];
        __syncthreads();
    }
    if (tid == 0) partial[blockIdx.x] = smem[0];
}

__global__ void hinge_vec_generic_kernel(const float* __restrict__ pred,
                                         const float* __restrict__ targ,
                                         float* __restrict__ partial,
                                         int64_t n4,
                                         int64_t cols,
                                         int pow2) {
    constexpr int THREADS = 256;
    __shared__ float smem[THREADS];
    int tid = threadIdx.x;
    int64_t stride = (int64_t)gridDim.x * blockDim.x;
    int64_t mask = cols - 1;
    const float4* __restrict__ p4 = reinterpret_cast<const float4*>(pred);

    float sum = 0.0f;
    for (int64_t i = (int64_t)blockIdx.x * blockDim.x + tid; i < n4; i += stride) {
        int64_t base = i * 4;
        int64_t col = pow2 ? (base & mask) : (base % cols);
        float4 tv = *reinterpret_cast<const float4*>(targ + col);
        float4 pv = p4[i];
        float s0 = hinge_val(pv.x, tv.x);
        float s1 = hinge_val(pv.y, tv.y);
        float s2 = hinge_val(pv.z, tv.z);
        float s3 = hinge_val(pv.w, tv.w);
        sum += (s0 + s1) + (s2 + s3);
    }

    smem[tid] = sum;
    __syncthreads();
    for (int off = THREADS / 2; off > 0; off >>= 1) {
        if (tid < off) smem[tid] += smem[tid + off];
        __syncthreads();
    }
    if (tid == 0) partial[blockIdx.x] = smem[0];
}

__global__ void hinge_scalar_kernel(const float* __restrict__ pred,
                                    const float* __restrict__ targ,
                                    float* __restrict__ partial,
                                    int64_t n,
                                    int64_t cols,
                                    int pow2) {
    constexpr int THREADS = 256;
    __shared__ float smem[THREADS];
    int tid = threadIdx.x;
    int64_t stride = (int64_t)gridDim.x * blockDim.x;
    int64_t mask = cols - 1;

    float sum = 0.0f;
    for (int64_t i = (int64_t)blockIdx.x * blockDim.x + tid; i < n; i += stride) {
        int64_t col = pow2 ? (i & mask) : (i % cols);
        sum += hinge_val(pred[i], targ[col]);
    }

    smem[tid] = sum;
    __syncthreads();
    for (int off = THREADS / 2; off > 0; off >>= 1) {
        if (tid < off) smem[tid] += smem[tid + off];
        __syncthreads();
    }
    if (tid == 0) partial[blockIdx.x] = smem[0];
}

__global__ void final_reduce_kernel(const float* __restrict__ partial,
                                    float* __restrict__ out,
                                    int blocks,
                                    double denom) {
    constexpr int THREADS = 1024;
    __shared__ double smem[THREADS];
    int tid = threadIdx.x;
    double sum = 0.0;
    for (int i = tid; i < blocks; i += THREADS) {
        sum += (double)partial[i];
    }
    smem[tid] = sum;
    __syncthreads();
    for (int off = THREADS / 2; off > 0; off >>= 1) {
        if (tid < off) smem[tid] += smem[tid + off];
        __syncthreads();
    }
    if (tid == 0) out[0] = (float)(smem[0] / denom);
}

torch::Tensor hinge_forward(torch::Tensor predictions, torch::Tensor targets) {
    TORCH_CHECK(predictions.is_cuda(), "predictions must be CUDA");
    TORCH_CHECK(targets.is_cuda(), "targets must be CUDA");
    TORCH_CHECK(predictions.scalar_type() == torch::kFloat32, "float32 only");
    TORCH_CHECK(targets.scalar_type() == torch::kFloat32, "float32 only");

    const c10::cuda::CUDAGuard device_guard(predictions.device());
    auto out = torch::empty({}, predictions.options());

    int64_t n = predictions.numel();
    int64_t cols = targets.numel();
    TORCH_CHECK(cols > 0, "targets must be non-empty");

    const int threads = 256;
    bool pow2 = (cols & (cols - 1)) == 0;
    bool vec_ok = ((n & 3LL) == 0) && ((cols & 3LL) == 0) && (((uintptr_t)predictions.data_ptr<float>() & 15) == 0) && (((uintptr_t)targets.data_ptr<float>() & 15) == 0);
    int64_t work = vec_ok ? (n / 4) : n;
    int blocks = (int)std::min<int64_t>(4096, std::max<int64_t>(1, (work + threads - 1) / threads));

    auto partial = torch::empty({blocks}, predictions.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const float* pred = predictions.data_ptr<float>();
    const float* targ = targets.data_ptr<float>();
    float* part = partial.data_ptr<float>();

    if (vec_ok) {
        int stride_elems = blocks * threads * 4;
        if (pow2 && ((stride_elems & (cols - 1)) == 0)) {
            hinge_vec_pow2_kernel<<<blocks, threads, 0, stream>>>(pred, targ, part, n / 4, cols);
        } else {
            hinge_vec_generic_kernel<<<blocks, threads, 0, stream>>>(pred, targ, part, n / 4, cols, pow2 ? 1 : 0);
        }
    } else {
        hinge_scalar_kernel<<<blocks, threads, 0, stream>>>(pred, targ, part, n, cols, pow2 ? 1 : 0);
    }

    final_reduce_kernel<<<1, 1024, 0, stream>>>(part, out.data_ptr<float>(), blocks, (double)n);
    return out;
}
"""

_hinge_ext = load_inline(
    name="hinge_loss_inline_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["hinge_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()

    def forward(self, predictions, targets):
        return _hinge_ext.hinge_forward(predictions, targets)