import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that computes Kullback-Leibler Divergence for comparing two distributions.

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

    def forward(self, predictions, targets):
        return torch.nn.functional.kl_div(torch.log(predictions), targets, reduction='batchmean')

batch_size = 8192 * 2
input_shape = (8192 * 2,)
dim = 1

def get_inputs():
    scale = torch.rand(())
    return [(torch.rand(batch_size, *input_shape)*scale).softmax(dim=-1), torch.rand(batch_size, *input_shape).softmax(dim=-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 kl_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>

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

__inline__ __device__ float block_reduce_sum(float val) {
    static __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 + 31) / 32)) ? shared[lane] : 0.0f;
    if (wid == 0) val = warp_reduce_sum(val);
    return val;
}

__global__ void kl_reduce_kernel(const float* __restrict__ p,
                                 const float* __restrict__ t,
                                 float* __restrict__ partial,
                                 long long n4) {
    float sum = 0.0f;
    const float4* __restrict__ p4 = reinterpret_cast<const float4*>(p);
    const float4* __restrict__ t4 = reinterpret_cast<const float4*>(t);
    long long stride = (long long)blockDim.x * gridDim.x;
    for (long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
         idx < n4; idx += stride) {
        float4 pv = p4[idx];
        float4 tv = t4[idx];
        sum += tv.x * (logf(tv.x) - logf(pv.x));
        sum += tv.y * (logf(tv.y) - logf(pv.y));
        sum += tv.z * (logf(tv.z) - logf(pv.z));
        sum += tv.w * (logf(tv.w) - logf(pv.w));
    }
    sum = block_reduce_sum(sum);
    if (threadIdx.x == 0) partial[blockIdx.x] = sum;
}

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

torch::Tensor kl_forward(torch::Tensor predictions, torch::Tensor targets) {
    auto out = torch::empty({}, predictions.options());
    const long long n = predictions.numel();
    const int batch = predictions.size(0);
    const int threads = 512;
    const int blocks = 16384;
    auto partial = torch::empty({blocks}, predictions.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    kl_reduce_kernel<<<blocks, threads, 0, stream>>>(
        predictions.data_ptr<float>(), targets.data_ptr<float>(),
        partial.data_ptr<float>(), n / 4);
    final_reduce_kernel<<<1, 1024, 0, stream>>>(
        partial.data_ptr<float>(), out.data_ptr<float>(), blocks,
        1.0f / (float)batch);
    return out;
}
"""

_ext = load_inline(
    name='kl_fused_final_ext',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['kl_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.kl_forward(predictions, targets)