import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that computes Smooth L1 (Huber) Loss for regression tasks.

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

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

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>

// CUDA declarations
torch::Tensor smooth_l1_forward_cuda(torch::Tensor predictions, torch::Tensor targets);

torch::Tensor smooth_l1_forward(torch::Tensor predictions, torch::Tensor targets) {
    return smooth_l1_forward_cuda(predictions, targets);
}
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <stdint.h>

#ifndef __CUDA_ARCH__
#define __CUDA_ARCH__ 0
#endif

static __device__ __forceinline__ float huber1(float d) {
    float a = fabsf(d);
    return (a < 1.0f) ? (0.5f * d * d) : (a - 0.5f);
}

__global__ void smooth_l1_reduce_kernel(const float* __restrict__ x,
                                        const float* __restrict__ y,
                                        float* __restrict__ partial,
                                        int64_t n) {
    extern __shared__ float smem[];
    const int tid = threadIdx.x;
    const int block_threads = blockDim.x;
    const int64_t stride = (int64_t)gridDim.x * block_threads;
    int64_t i = (int64_t)blockIdx.x * block_threads + tid;

    float sum = 0.0f;

    // grid-stride loop; unroll a little to reduce loop overhead
    for (; i + 3 * stride < n; i += 4 * stride) {
        float d0 = x[i] - y[i];
        float d1 = x[i + stride] - y[i + stride];
        float d2 = x[i + 2 * stride] - y[i + 2 * stride];
        float d3 = x[i + 3 * stride] - y[i + 3 * stride];
        sum += huber1(d0) + huber1(d1) + huber1(d2) + huber1(d3);
    }
    for (; i < n; i += stride) {
        sum += huber1(x[i] - y[i]);
    }

    smem[tid] = sum;
    __syncthreads();

    // block reduction in shared memory
    for (int s = block_threads >> 1; s >= 32; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }

    float v = smem[tid];
    if (tid < 32) {
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            v += __shfl_down_sync(0xffffffff, v, offset);
        }
        if (tid == 0) partial[blockIdx.x] = v;
    }
}

__global__ void final_reduce_kernel(const float* __restrict__ partial,
                                    float* __restrict__ out,
                                    int nblocks,
                                    float inv_n) {
    extern __shared__ float smem[];
    const int tid = threadIdx.x;
    float sum = 0.0f;
    for (int i = tid; i < nblocks; i += blockDim.x) {
        sum += partial[i];
    }
    smem[tid] = sum;
    __syncthreads();

    for (int s = blockDim.x >> 1; s >= 32; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }

    float v = smem[tid];
    if (tid < 32) {
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            v += __shfl_down_sync(0xffffffff, v, offset);
        }
        if (tid == 0) out[0] = v * inv_n;
    }
}

torch::Tensor smooth_l1_forward_cuda(torch::Tensor predictions, torch::Tensor targets) {
    TORCH_CHECK(predictions.is_cuda() && targets.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(predictions.scalar_type() == at::kFloat && targets.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(predictions.numel() == targets.numel(), "input sizes must match");

    auto x = predictions.contiguous();
    auto y = targets.contiguous();
    const int64_t n = x.numel();
    auto out = torch::empty({}, x.options());
    if (n == 0) {
        // PyTorch returns nan for mean over empty; set nan.
        out.fill_(NAN);
        return out;
    }

    const int threads = 256;
    // Enough blocks to expose parallelism, not so many that final reduction dominates.
    int blocks = (int)((n + threads - 1) / threads);
    if (blocks > 8192) blocks = 8192;
    if (blocks < 1) blocks = 1;

    auto partial = torch::empty({blocks}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    smooth_l1_reduce_kernel<<<blocks, threads, threads * sizeof(float), stream>>>(
        x.data_ptr<float>(), y.data_ptr<float>(), partial.data_ptr<float>(), n);
    final_reduce_kernel<<<1, 1024, 1024 * sizeof(float), stream>>>(
        partial.data_ptr<float>(), out.data_ptr<float>(), blocks, 1.0f / (float)n);
    return out;
}
"""

_ext = load_inline(
    name="smooth_l1_loss_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["smooth_l1_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.smooth_l1_forward(predictions, targets)