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 <cuda.h>
#include <cuda_runtime.h>

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x " must be float32")

static __inline__ __device__ float warp_reduce_sum(float val) {
#pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__global__ void hinge_rows4_vec4_kernel(const float* __restrict__ pred,
                                       const float* __restrict__ target,
                                       float* __restrict__ partial,
                                       long long rows, int cols4) {
    long long base = ((long long)blockIdx.x) * 4;
    long long cols = ((long long)cols4) * 4;
    const float4* __restrict__ t4 = reinterpret_cast<const float4*>(target);
    const float4* __restrict__ p0 = reinterpret_cast<const float4*>(pred + (base + 0) * cols);
    const float4* __restrict__ p1 = reinterpret_cast<const float4*>(pred + (base + 1) * cols);
    const float4* __restrict__ p2 = reinterpret_cast<const float4*>(pred + (base + 2) * cols);
    const float4* __restrict__ p3 = reinterpret_cast<const float4*>(pred + (base + 3) * cols);
    bool b1 = base + 1 < rows, b2 = base + 2 < rows, b3 = base + 3 < rows;
    float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;
    for (int c = threadIdx.x; c < cols4; c += blockDim.x) {
        float4 tv = t4[c];
        float4 pv = p0[c];
        s0 += fmaxf(1.0f - pv.x * tv.x, 0.0f) + fmaxf(1.0f - pv.y * tv.y, 0.0f) +
              fmaxf(1.0f - pv.z * tv.z, 0.0f) + fmaxf(1.0f - pv.w * tv.w, 0.0f);
        if (b1) { pv = p1[c]; s1 += fmaxf(1.0f - pv.x * tv.x, 0.0f) + fmaxf(1.0f - pv.y * tv.y, 0.0f) + fmaxf(1.0f - pv.z * tv.z, 0.0f) + fmaxf(1.0f - pv.w * tv.w, 0.0f); }
        if (b2) { pv = p2[c]; s2 += fmaxf(1.0f - pv.x * tv.x, 0.0f) + fmaxf(1.0f - pv.y * tv.y, 0.0f) + fmaxf(1.0f - pv.z * tv.z, 0.0f) + fmaxf(1.0f - pv.w * tv.w, 0.0f); }
        if (b3) { pv = p3[c]; s3 += fmaxf(1.0f - pv.x * tv.x, 0.0f) + fmaxf(1.0f - pv.y * tv.y, 0.0f) + fmaxf(1.0f - pv.z * tv.z, 0.0f) + fmaxf(1.0f - pv.w * tv.w, 0.0f); }
    }

    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    s0 = warp_reduce_sum(s0); s1 = warp_reduce_sum(s1); s2 = warp_reduce_sum(s2); s3 = warp_reduce_sum(s3);
    __shared__ float sh[4][32];
    if (lane == 0) { sh[0][wid] = s0; sh[1][wid] = s1; sh[2][wid] = s2; sh[3][wid] = s3; }
    __syncthreads();
    if (wid == 0) {
        int nw = blockDim.x >> 5;
        s0 = (lane < nw) ? sh[0][lane] : 0.0f;
        s1 = (lane < nw) ? sh[1][lane] : 0.0f;
        s2 = (lane < nw) ? sh[2][lane] : 0.0f;
        s3 = (lane < nw) ? sh[3][lane] : 0.0f;
        s0 = warp_reduce_sum(s0); s1 = warp_reduce_sum(s1); s2 = warp_reduce_sum(s2); s3 = warp_reduce_sum(s3);
    }
    if (threadIdx.x == 0) {
        partial[base] = s0;
        if (b1) partial[base + 1] = s1;
        if (b2) partial[base + 2] = s2;
        if (b3) partial[base + 3] = s3;
    }
}

static __inline__ __device__ float block_reduce_sum(float val) {
    __shared__ float shared[32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    val = warp_reduce_sum(val);
    if (lane == 0) shared[wid] = val;
    __syncthreads();
    val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
    if (wid == 0) val = warp_reduce_sum(val);
    return val;
}

__global__ void hinge_rows_scalar_kernel(const float* __restrict__ pred,
                                         const float* __restrict__ target,
                                         float* __restrict__ partial,
                                         long long rows, long long cols) {
    long long r = (long long)blockIdx.x;
    if (r >= rows) return;
    const float* row = pred + r * cols;
    float sum = 0.0f;
    for (long long c = threadIdx.x; c < cols; c += blockDim.x) {
        sum += fmaxf(1.0f - row[c] * target[c], 0.0f);
    }
    sum = block_reduce_sum(sum);
    if (threadIdx.x == 0) partial[r] = sum;
}

__global__ void final_reduce_kernel(const float* __restrict__ partial,
                                    float* __restrict__ out,
                                    long long rows, float inv_total) {
    float sum = 0.0f;
    for (long long i = threadIdx.x; i < rows; i += blockDim.x) sum += partial[i];
    sum = block_reduce_sum(sum);
    if (threadIdx.x == 0) out[0] = sum * inv_total;
}

torch::Tensor hinge_forward(torch::Tensor predictions, torch::Tensor targets) {
    CHECK_CUDA(predictions); CHECK_CUDA(targets); CHECK_FLOAT(predictions); CHECK_FLOAT(targets);
    TORCH_CHECK(predictions.is_contiguous(), "predictions must be contiguous");
    TORCH_CHECK(targets.is_contiguous(), "targets must be contiguous");
    long long total = predictions.numel();
    long long cols = targets.numel();
    TORCH_CHECK(cols > 0 && total % cols == 0, "invalid shapes");
    long long rows = total / cols;
    auto partial = torch::empty({rows}, predictions.options());
    auto out = torch::empty({}, predictions.options());
    const int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if ((cols % 4 == 0) && (((uintptr_t)predictions.data_ptr<float>() & 15) == 0) && (((uintptr_t)targets.data_ptr<float>() & 15) == 0)) {
        hinge_rows4_vec4_kernel<<<(unsigned int)((rows + 3) / 4), threads, 0, stream>>>(predictions.data_ptr<float>(), targets.data_ptr<float>(), partial.data_ptr<float>(), rows, (int)(cols / 4));
    } else {
        hinge_rows_scalar_kernel<<<(unsigned int)rows, threads, 0, stream>>>(predictions.data_ptr<float>(), targets.data_ptr<float>(), partial.data_ptr<float>(), rows, cols);
    }
    final_reduce_kernel<<<1, threads, 0, stream>>>(partial.data_ptr<float>(), out.data_ptr<float>(), rows, 1.0f / (float)total);
    return out;
}
'''

_ext = load_inline(name='hinge_loss_ext_final', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['hinge_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
    def forward(self, predictions, targets):
        return _ext.hinge_forward(predictions, targets)