import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that computes Cross Entropy Loss for multi-class classification tasks.

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

    def forward(self, predictions, targets):
        return torch.nn.functional.cross_entropy(predictions, targets)

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

def get_inputs():
    return [torch.rand(batch_size, *input_shape), torch.randint(0, num_classes, (batch_size,))]

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>

at::Tensor cross_entropy_forward_cuda(at::Tensor predictions, at::Tensor targets);

at::Tensor cross_entropy_forward(at::Tensor predictions, at::Tensor targets) {
    return cross_entropy_forward_cuda(predictions, targets);
}
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>

static __device__ __forceinline__ float warp_reduce_sum(float v) {
    v += __shfl_down_sync(0xffffffff, v, 16);
    v += __shfl_down_sync(0xffffffff, v, 8);
    v += __shfl_down_sync(0xffffffff, v, 4);
    v += __shfl_down_sync(0xffffffff, v, 2);
    v += __shfl_down_sync(0xffffffff, v, 1);
    return v;
}

static __device__ __forceinline__ float warp_reduce_max(float v) {
    v = fmaxf(v, __shfl_down_sync(0xffffffff, v, 16));
    v = fmaxf(v, __shfl_down_sync(0xffffffff, v, 8));
    v = fmaxf(v, __shfl_down_sync(0xffffffff, v, 4));
    v = fmaxf(v, __shfl_down_sync(0xffffffff, v, 2));
    v = fmaxf(v, __shfl_down_sync(0xffffffff, v, 1));
    return v;
}

__global__ void init_scalar_kernel(float* out) {
    *out = 0.0f;
}

// One CTA computes one row.  Optimized for the benchmark shape (N,4096).
__global__ __launch_bounds__(256, 2)
void ce_mean_kernel(const float* __restrict__ x,
                    const long long* __restrict__ t,
                    float* __restrict__ out,
                    int n, int c) {
    const int row = blockIdx.x;
    const int tid = threadIdx.x;
    const float* rowp = x + ((size_t)row) * c;

    float m = -FLT_MAX;

    // Vectorized fast path for the fixed class count used by the benchmark.
    if (c == 4096) {
        const float4* row4 = reinterpret_cast<const float4*>(rowp);
        #pragma unroll
        for (int j = tid; j < 1024; j += 256) {
            float4 v = row4[j];
            m = fmaxf(m, v.x);
            m = fmaxf(m, v.y);
            m = fmaxf(m, v.z);
            m = fmaxf(m, v.w);
        }
    } else {
        for (int j = tid; j < c; j += blockDim.x) {
            m = fmaxf(m, rowp[j]);
        }
    }

    m = warp_reduce_max(m);
    __shared__ float smem[8];
    if ((tid & 31) == 0) smem[tid >> 5] = m;
    __syncthreads();
    m = (tid < 8) ? smem[tid] : -FLT_MAX;
    if (tid < 32) m = warp_reduce_max(m);
    __shared__ float row_max_s;
    if (tid == 0) row_max_s = m;
    __syncthreads();
    const float row_max = row_max_s;

    float s = 0.0f;
    if (c == 4096) {
        const float4* row4 = reinterpret_cast<const float4*>(rowp);
        #pragma unroll
        for (int j = tid; j < 1024; j += 256) {
            float4 v = row4[j];
            s += __expf(v.x - row_max);
            s += __expf(v.y - row_max);
            s += __expf(v.z - row_max);
            s += __expf(v.w - row_max);
        }
    } else {
        for (int j = tid; j < c; j += blockDim.x) {
            s += __expf(rowp[j] - row_max);
        }
    }

    s = warp_reduce_sum(s);
    if ((tid & 31) == 0) smem[tid >> 5] = s;
    __syncthreads();
    s = (tid < 8) ? smem[tid] : 0.0f;
    if (tid < 32) s = warp_reduce_sum(s);

    if (tid == 0) {
        const int cls = (int)t[row];
        const float correct = rowp[cls];
        const float loss = (logf(s) + row_max - correct) / (float)n;
        atomicAdd(out, loss);
    }
}

at::Tensor cross_entropy_forward_cuda(at::Tensor predictions, at::Tensor targets) {
    auto x = predictions.contiguous();
    auto y = targets.contiguous();
    const int n = (int)x.size(0);
    const int c = (int)x.size(1);
    auto out = at::empty({}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    init_scalar_kernel<<<1, 1, 0, stream>>>(out.data_ptr<float>());
    ce_mean_kernel<<<n, 256, 0, stream>>>(x.data_ptr<float>(), (const long long*)y.data_ptr<int64_t>(), out.data_ptr<float>(), n, c);
    return out;
}
'''

_ce_ext = load_inline(
    name='ce_fused_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['cross_entropy_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    verbose=False,
)

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

    def forward(self, predictions, targets):
        return _ce_ext.cross_entropy_forward(predictions, targets)