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>

// Warp and block reductions for float sums.
__inline__ __device__ 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;
}

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

template<int BLOCK_SIZE>
__global__ __launch_bounds__(BLOCK_SIZE, 2)
void mse_partial_kernel(const float* __restrict__ pred,
                        const float* __restrict__ targ,
                        float* __restrict__ partial,
                        int64_t n) {
    float sum = 0.0f;
    int64_t tid = (int64_t)blockIdx.x * BLOCK_SIZE + threadIdx.x;
    int64_t stride = (int64_t)gridDim.x * BLOCK_SIZE;

    // Grid-stride loop with four independent accumulators improves ILP and reduces
    // rounding error a little versus one long dependency chain.
    float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;
    int64_t i = tid;
    int64_t step4 = stride * 4;
    for (; i + 3 * stride < n; i += step4) {
        float d0 = pred[i] - targ[i];
        float d1 = pred[i + stride] - targ[i + stride];
        float d2 = pred[i + 2 * stride] - targ[i + 2 * stride];
        float d3 = pred[i + 3 * stride] - targ[i + 3 * stride];
        s0 = fmaf(d0, d0, s0);
        s1 = fmaf(d1, d1, s1);
        s2 = fmaf(d2, d2, s2);
        s3 = fmaf(d3, d3, s3);
    }
    sum = (s0 + s1) + (s2 + s3);
    for (; i < n; i += stride) {
        float d = pred[i] - targ[i];
        sum = fmaf(d, d, sum);
    }

    sum = block_reduce_sum(sum);
    if (threadIdx.x == 0) partial[blockIdx.x] = sum;
}

template<int BLOCK_SIZE>
__global__ __launch_bounds__(BLOCK_SIZE, 1)
void mse_final_kernel(const float* __restrict__ partial,
                      float* __restrict__ out,
                      int nblocks,
                      float inv_n) {
    float sum = 0.0f;
    for (int i = threadIdx.x; i < nblocks; i += BLOCK_SIZE) {
        sum += partial[i];
    }
    sum = block_reduce_sum(sum);
    if (threadIdx.x == 0) out[0] = sum * inv_n;
}

torch::Tensor mse_forward(torch::Tensor predictions, torch::Tensor targets) {
    TORCH_CHECK(predictions.is_cuda(), "predictions must be a CUDA tensor");
    TORCH_CHECK(targets.is_cuda(), "targets must be a CUDA tensor");
    TORCH_CHECK(predictions.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(targets.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(predictions.numel() == targets.numel(), "input sizes must match");

    auto pred_c = predictions.contiguous();
    auto targ_c = targets.contiguous();
    const int64_t n = pred_c.numel();
    auto out = torch::empty({}, pred_c.options());
    if (n == 0) {
        return out;
    }

    constexpr int BLOCK = 256;
    int device = pred_c.get_device();
    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, device);
    // Enough blocks to cover the GPU well, while keeping final reduction tiny.
    int blocks = prop.multiProcessorCount * 8;
    int64_t min_blocks64 = (n + BLOCK - 1) / BLOCK;
    if (blocks > min_blocks64) blocks = (int)min_blocks64;
    if (blocks < 1) blocks = 1;
    // Cap partials to keep the second pass cheap. For the target problem this is
    // typically 800-1200 blocks depending on GPU.
    if (blocks > 4096) blocks = 4096;

    auto partial = torch::empty({blocks}, pred_c.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    mse_partial_kernel<BLOCK><<<blocks, BLOCK, 0, stream>>>(
        pred_c.data_ptr<float>(), targ_c.data_ptr<float>(), partial.data_ptr<float>(), n);
    mse_final_kernel<1024><<<1, 1024, 0, stream>>>(
        partial.data_ptr<float>(), out.data_ptr<float>(), blocks, 1.0f / (float)n);
    return out;
}
"""

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

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

    def forward(self, predictions, targets):
        return _ext.mse_forward(predictions, targets)