import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs 3D Average Pooling.
    """
    def __init__(self, kernel_size: int, stride: int = None, padding: int = 0):
        """
        Initializes the Average Pooling layer.

        Args:
            kernel_size (int): Size of the kernel to apply pooling.
            stride (int, optional): Stride of the pooling operation. Defaults to None, which uses the kernel size.
            padding (int, optional): Padding to apply before pooling. Defaults to 0.
        """
        super(Model, self).__init__()
        self.avg_pool = nn.AvgPool3d(kernel_size=kernel_size, stride=stride, padding=padding)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Average Pooling to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, channels, depth, height, width).

        Returns:
            torch.Tensor: Output tensor with Average Pooling applied, shape depends on kernel_size, stride and padding.
        """
        return self.avg_pool(x)

batch_size = 16
channels = 32
depth = 128
height = 128
width = 256
kernel_size = 3
stride = 2
padding = 1

def get_inputs():
    x = torch.rand(batch_size, channels, depth, height, width)
    return [x]

def get_init_inputs():
    return [kernel_size, stride, padding]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = r'''
#include <torch/extension.h>

torch::Tensor avgpool3d_forward_cuda(torch::Tensor x, int64_t k, int64_t s, int64_t p);

torch::Tensor avgpool3d_forward(torch::Tensor x, int64_t k, int64_t s, int64_t p) {
    return avgpool3d_forward_cuda(x, k, s, p);
}
'''

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

__device__ __forceinline__ float ldg_f(const float* p) {
#if __CUDA_ARCH__ >= 350
    return __ldg(p);
#else
    return *p;
#endif
}

__global__ __launch_bounds__(128)
void avgpool3d_k3s2p1_tiled_kernel(const float* __restrict__ x,
                                    float* __restrict__ y,
                                    int N, int C, int D, int H, int W,
                                    int Do, int Ho, int Wo,
                                    int Dp, int Hp, int Wp,
                                    long long total_tiles) {
    long long tid = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= total_tiles) return;

    int wp = (int)(tid % Wp);
    long long q = tid / Wp;
    int hp = (int)(q % Hp);
    q /= Hp;
    int dp = (int)(q % Dp);
    long long nc = q / Dp;

    int od0 = dp << 1;
    int oh0 = hp << 1;
    int ow0 = wp << 1;
    int iz0 = (od0 << 1) - 1;
    int iy0 = (oh0 << 1) - 1;
    int ix0 = (ow0 << 1) - 1;

    const long long DHW = (long long)D * H * W;
    const long long HoWo = (long long)Ho * Wo;
    const long long base_in = nc * DHW;
    const long long base_out = nc * (long long)Do * HoWo;
    const float inv27 = 0.037037037037037035f;

    // Fast path for the overwhelmingly common, fully-valid 2x2x2 output tile.
    if (od0 + 1 < Do && oh0 + 1 < Ho && ow0 + 1 < Wo &&
        iz0 >= 0 && iy0 >= 0 && ix0 >= 0 &&
        iz0 + 4 < D && iy0 + 4 < H && ix0 + 4 < W) {

        float s000 = 0.f, s001 = 0.f, s010 = 0.f, s011 = 0.f;
        float s100 = 0.f, s101 = 0.f, s110 = 0.f, s111 = 0.f;

        #pragma unroll
        for (int zz = 0; zz < 5; ++zz) {
            float q00 = 0.f, q01 = 0.f, q10 = 0.f, q11 = 0.f;
            const long long zbase = base_in + (long long)(iz0 + zz) * H * W + (long long)iy0 * W + ix0;

            #pragma unroll
            for (int yy = 0; yy < 5; ++yy) {
                const float* row = x + zbase + (long long)yy * W;
                float v0 = ldg_f(row + 0);
                float v1 = ldg_f(row + 1);
                float v2 = ldg_f(row + 2);
                float v3 = ldg_f(row + 3);
                float v4 = ldg_f(row + 4);
                float a = v0 + v1 + v2;
                float b = v2 + v3 + v4;
                if (yy < 3) { q00 += a; q01 += b; }
                if (yy >= 2) { q10 += a; q11 += b; }
            }
            if (zz < 3) { s000 += q00; s001 += q01; s010 += q10; s011 += q11; }
            if (zz >= 2) { s100 += q00; s101 += q01; s110 += q10; s111 += q11; }
        }

        long long o_d0 = base_out + (long long)od0 * HoWo;
        long long o_d1 = o_d0 + HoWo;
        long long o_h0 = (long long)oh0 * Wo + ow0;
        long long o_h1 = o_h0 + Wo;
        y[o_d0 + o_h0]     = s000 * inv27;
        y[o_d0 + o_h0 + 1] = s001 * inv27;
        y[o_d0 + o_h1]     = s010 * inv27;
        y[o_d0 + o_h1 + 1] = s011 * inv27;
        y[o_d1 + o_h0]     = s100 * inv27;
        y[o_d1 + o_h0 + 1] = s101 * inv27;
        y[o_d1 + o_h1]     = s110 * inv27;
        y[o_d1 + o_h1 + 1] = s111 * inv27;
        return;
    }

    // Boundary/odd-size path, still count_include_pad=True (fixed divisor 27).
    #pragma unroll
    for (int dz = 0; dz < 2; ++dz) {
        int od = od0 + dz;
        if (od >= Do) continue;
        int izs = (od << 1) - 1;
        #pragma unroll
        for (int dy = 0; dy < 2; ++dy) {
            int oh = oh0 + dy;
            if (oh >= Ho) continue;
            int iys = (oh << 1) - 1;
            #pragma unroll
            for (int dx = 0; dx < 2; ++dx) {
                int ow = ow0 + dx;
                if (ow >= Wo) continue;
                int ixs = (ow << 1) - 1;
                float sum = 0.f;
                #pragma unroll
                for (int kz = 0; kz < 3; ++kz) {
                    int iz = izs + kz;
                    if ((unsigned)iz >= (unsigned)D) continue;
                    #pragma unroll
                    for (int ky = 0; ky < 3; ++ky) {
                        int iy = iys + ky;
                        if ((unsigned)iy >= (unsigned)H) continue;
                        const float* row = x + base_in + (long long)iz * H * W + (long long)iy * W;
                        #pragma unroll
                        for (int kx = 0; kx < 3; ++kx) {
                            int ix = ixs + kx;
                            if ((unsigned)ix < (unsigned)W) sum += ldg_f(row + ix);
                        }
                    }
                }
                y[base_out + (long long)od * HoWo + (long long)oh * Wo + ow] = sum * inv27;
            }
        }
    }
}

__global__ __launch_bounds__(256)
void avgpool3d_generic_kernel(const float* __restrict__ x,
                              float* __restrict__ y,
                              int D, int H, int W, int Do, int Ho, int Wo,
                              int k, int s, int p, long long total) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = (int)(idx % Wo);
    long long q = idx / Wo;
    int oh = (int)(q % Ho); q /= Ho;
    int od = (int)(q % Do); q /= Do;
    long long nc = q;
    long long base = nc * (long long)D * H * W;
    int izs = od * s - p;
    int iys = oh * s - p;
    int ixs = ow * s - p;
    float sum = 0.f;
    for (int kz = 0; kz < k; ++kz) {
        int iz = izs + kz;
        if ((unsigned)iz >= (unsigned)D) continue;
        for (int ky = 0; ky < k; ++ky) {
            int iy = iys + ky;
            if ((unsigned)iy >= (unsigned)H) continue;
            const float* row = x + base + (long long)iz * H * W + (long long)iy * W;
            for (int kx = 0; kx < k; ++kx) {
                int ix = ixs + kx;
                if ((unsigned)ix < (unsigned)W) sum += ldg_f(row + ix);
            }
        }
    }
    y[idx] = sum / (float)(k * k * k);
}

torch::Tensor avgpool3d_forward_cuda(torch::Tensor x, int64_t k64, int64_t s64, int64_t p64) {
    TORCH_CHECK(x.is_cuda(), "input must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(x.dim() == 5, "expected 5D NCDHW input");

    int N = (int)x.size(0), C = (int)x.size(1), D = (int)x.size(2), H = (int)x.size(3), W = (int)x.size(4);
    int k = (int)k64, s = (int)s64, p = (int)p64;
    int Do = (D + 2 * p - k) / s + 1;
    int Ho = (H + 2 * p - k) / s + 1;
    int Wo = (W + 2 * p - k) / s + 1;
    auto y = torch::empty({N, C, Do, Ho, Wo}, x.options());

    const float* xp = x.data_ptr<float>();
    float* yp = y.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (k == 3 && s == 2 && p == 1) {
        int Dp = (Do + 1) >> 1;
        int Hp = (Ho + 1) >> 1;
        int Wp = (Wo + 1) >> 1;
        long long total_tiles = (long long)N * C * Dp * Hp * Wp;
        dim3 block(128);
        dim3 grid((unsigned int)((total_tiles + block.x - 1) / block.x));
        avgpool3d_k3s2p1_tiled_kernel<<<grid, block, 0, stream>>>(xp, yp, N, C, D, H, W, Do, Ho, Wo, Dp, Hp, Wp, total_tiles);
    } else {
        long long total = (long long)N * C * Do * Ho * Wo;
        dim3 block(256);
        dim3 grid((unsigned int)((total + block.x - 1) / block.x));
        avgpool3d_generic_kernel<<<grid, block, 0, stream>>>(xp, yp, D, H, W, Do, Ho, Wo, k, s, p, total);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

_ext = load_inline(
    name='avgpool3d_custom_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['avgpool3d_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, kernel_size: int, stride: int = None, padding: int = 0):
        super(ModelNew, self).__init__()
        self.kernel_size = int(kernel_size)
        self.stride = int(kernel_size if stride is None else stride)
        self.padding = int(padding)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.avgpool3d_forward(x, self.kernel_size, self.stride, self.padding)