import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that computes the Mean Squared Error loss for regression tasks.

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

    def forward(self, predictions, targets):
        return torch.mean((predictions - targets) ** 2)

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

def get_inputs():
    scale = torch.rand(())
    return [torch.rand(batch_size, *input_shape)*scale, torch.rand(batch_size, *input_shape)]

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 mse_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>
#include <cstdint>

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")

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

template<int BLOCK>
__device__ __forceinline__ void store_block_sum(float sum, float* __restrict__ partial) {
    int tid = threadIdx.x;
    sum = warp_reduce_sum(sum);
    __shared__ float warp_sums[BLOCK / 32];
    if ((tid & 31) == 0) warp_sums[tid >> 5] = sum;
    __syncthreads();
    float block_sum = (tid < (BLOCK / 32)) ? warp_sums[tid] : 0.0f;
    if (tid < 32) {
        block_sum = warp_reduce_sum(block_sum);
        if (tid == 0) partial[blockIdx.x] = block_sum;
    }
}

static __device__ __forceinline__ float4 load4(const float4* __restrict__ p, int64_t i) { return p[i]; }

template<int BLOCK>
__global__ void mse_stage1_vec4_kernel(const float4* __restrict__ p4,
                                       const float4* __restrict__ t4,
                                       float* __restrict__ partial,
                                       int64_t n4) {
    int tid = threadIdx.x;
    int64_t stride = (int64_t)gridDim.x * BLOCK;
    int64_t idx = (int64_t)blockIdx.x * BLOCK + tid;
    float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;
    int64_t stride4 = stride << 2;
    for (; idx + 3 * stride < n4; idx += stride4) {
        float4 a0 = load4(p4, idx);            float4 b0 = load4(t4, idx);
        float4 a1 = load4(p4, idx + stride);   float4 b1 = load4(t4, idx + stride);
        float4 a2 = load4(p4, idx + 2*stride); float4 b2 = load4(t4, idx + 2*stride);
        float4 a3 = load4(p4, idx + 3*stride); float4 b3 = load4(t4, idx + 3*stride);
        float d;
        d = a0.x - b0.x; s0 = fmaf(d,d,s0); d = a0.y - b0.y; s0 = fmaf(d,d,s0); d = a0.z - b0.z; s0 = fmaf(d,d,s0); d = a0.w - b0.w; s0 = fmaf(d,d,s0);
        d = a1.x - b1.x; s1 = fmaf(d,d,s1); d = a1.y - b1.y; s1 = fmaf(d,d,s1); d = a1.z - b1.z; s1 = fmaf(d,d,s1); d = a1.w - b1.w; s1 = fmaf(d,d,s1);
        d = a2.x - b2.x; s2 = fmaf(d,d,s2); d = a2.y - b2.y; s2 = fmaf(d,d,s2); d = a2.z - b2.z; s2 = fmaf(d,d,s2); d = a2.w - b2.w; s2 = fmaf(d,d,s2);
        d = a3.x - b3.x; s3 = fmaf(d,d,s3); d = a3.y - b3.y; s3 = fmaf(d,d,s3); d = a3.z - b3.z; s3 = fmaf(d,d,s3); d = a3.w - b3.w; s3 = fmaf(d,d,s3);
    }
    float sum = (s0 + s1) + (s2 + s3);
    for (; idx < n4; idx += stride) {
        float4 a = p4[idx]; float4 b = t4[idx];
        float d = a.x - b.x; sum = fmaf(d,d,sum);
        d = a.y - b.y; sum = fmaf(d,d,sum);
        d = a.z - b.z; sum = fmaf(d,d,sum);
        d = a.w - b.w; sum = fmaf(d,d,sum);
    }
    store_block_sum<BLOCK>(sum, partial);
}

template<int BLOCK>
__global__ void mse_stage1_scalar_kernel(const float* __restrict__ p,
                                  const float* __restrict__ t,
                                  float* __restrict__ partial,
                                  int64_t n) {
    int tid = threadIdx.x;
    int64_t idx = (int64_t)blockIdx.x * BLOCK + tid;
    int64_t stride = (int64_t)gridDim.x * BLOCK;
    float sum = 0.0f;
    for (; idx < n; idx += stride) {
        float d = p[idx] - t[idx];
        sum = fmaf(d, d, sum);
    }
    store_block_sum<BLOCK>(sum, partial);
}

template<int BLOCK>
__global__ void mse_stage2_kernel(const float* __restrict__ partial,
                                  float* __restrict__ out,
                                  int blocks,
                                  int64_t n) {
    int tid = threadIdx.x;
    float sum = 0.0f;
    for (int i = tid; i < blocks; i += BLOCK) sum += partial[i];
    sum = warp_reduce_sum(sum);
    __shared__ float warp_sums[BLOCK / 32];
    if ((tid & 31) == 0) warp_sums[tid >> 5] = sum;
    __syncthreads();
    float total = (tid < (BLOCK / 32)) ? warp_sums[tid] : 0.0f;
    if (tid < 32) {
        total = warp_reduce_sum(total);
        if (tid == 0) out[0] = total / (float)n;
    }
}

torch::Tensor mse_forward(torch::Tensor predictions, torch::Tensor targets) {
    CHECK_CUDA(predictions); CHECK_CUDA(targets);
    TORCH_CHECK(predictions.scalar_type() == torch::kFloat32, "predictions must be float32");
    TORCH_CHECK(targets.scalar_type() == torch::kFloat32, "targets must be float32");
    TORCH_CHECK(predictions.numel() == targets.numel(), "size mismatch");
    TORCH_CHECK(predictions.is_contiguous(), "predictions must be contiguous");
    TORCH_CHECK(targets.is_contiguous(), "targets must be contiguous");
    const int64_t n = predictions.numel();
    auto out = torch::empty({}, predictions.options());
    if (n == 0) return out;
    constexpr int BLOCK = 256;
    const float* p = predictions.data_ptr<float>();
    const float* t = targets.data_ptr<float>();
    uintptr_t addr = reinterpret_cast<uintptr_t>(p) | reinterpret_cast<uintptr_t>(t);
    bool use_vec4 = ((addr & 15) == 0) && ((n & 3LL) == 0);
    int64_t items = use_vec4 ? (n >> 2) : n;
    int blocks = (int)((items + BLOCK - 1) / BLOCK);
    if (blocks > 2048) blocks = 2048;
    auto partial = torch::empty({blocks}, predictions.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (use_vec4) {
        mse_stage1_vec4_kernel<BLOCK><<<blocks, BLOCK, 0, stream>>>(reinterpret_cast<const float4*>(p), reinterpret_cast<const float4*>(t), partial.data_ptr<float>(), n >> 2);
    } else {
        mse_stage1_scalar_kernel<BLOCK><<<blocks, BLOCK, 0, stream>>>(p, t, partial.data_ptr<float>(), n);
    }
    mse_stage2_kernel<BLOCK><<<1, BLOCK, 0, stream>>>(partial.data_ptr<float>(), out.data_ptr<float>(), blocks, n);
    return out;
}
"""

_ext = load_inline(name="mse_mean_ext_final", cpp_sources=_cpp_src, cuda_sources=_cuda_src,
                   functions=["mse_forward"], verbose=False,
                   extra_cuda_cflags=["-O3", "--use_fast_math"],
                   extra_cflags=["-O3"])

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