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 avg_pool3d_forward(torch::Tensor x, int64_t kernel_size, int64_t stride, int64_t padding);
'''

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

#define BD 4
#define BH 8
#define BW 32
#define TD (2*BD+1)
#define TH (2*BH+1)
#define TW (2*BW+1)
#define TSIZE (TD*TH*TW)

__global__ __launch_bounds__(256, 4) void avgpool3d_k3s2p1_shared4_kernel(
        const float* __restrict__ x, float* __restrict__ y,
        int D, int H, int W, int Do, int Ho, int Wo, int C) {
    __shared__ float sh[TSIZE];

    int num_h_tiles = (Ho + BH - 1) / BH;
    int ht_dt = blockIdx.y;
    int ht = ht_dt % num_h_tiles;
    int dt = ht_dt / num_h_tiles;

    int ow_base = blockIdx.x * BW;
    int oh_base = ht * BH;
    int od_base = dt * BD;
    int nc = blockIdx.z;

    int iz0 = od_base * 2 - 1;
    int iy0 = oh_base * 2 - 1;
    int ix0 = ow_base * 2 - 1;
    long long nc_base = ((long long)nc) * D * H * W;

    bool full = (iz0 >= 0 && iy0 >= 0 && ix0 >= 0 &&
                 iz0 + TD <= D && iy0 + TH <= H && ix0 + TW <= W);

    for (int s = threadIdx.x; s < TSIZE; s += blockDim.x) {
        int lx = s % TW;
        int tmp = s / TW;
        int ly = tmp % TH;
        int lz = tmp / TH;
        int iz = iz0 + lz;
        int iy = iy0 + ly;
        int ix = ix0 + lx;
        float v = 0.0f;
        if (full || ((unsigned)iz < (unsigned)D && (unsigned)iy < (unsigned)H && (unsigned)ix < (unsigned)W)) {
            v = x[nc_base + ((long long)iz * H + iy) * W + ix];
        }
        sh[s] = v;
    }
    __syncthreads();

    int tid = threadIdx.x;
    int gw = tid & 7;          // 8 groups across W, four adjacent outputs per thread
    int tmp = tid >> 3;
    int lh = tmp & 7;
    int ld = tmp >> 3;

    int ow = ow_base + (gw << 2);
    int oh = oh_base + lh;
    int od = od_base + ld;

    if (od < Do && oh < Ho && ow < Wo) {
        int sx = gw << 3;
        int sy = lh << 1;
        int sz = ld << 1;
        float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;

        #pragma unroll
        for (int dz = 0; dz < 3; ++dz) {
            #pragma unroll
            for (int dy = 0; dy < 3; ++dy) {
                const float* p = sh + ((sz + dz) * TH + (sy + dy)) * TW + sx;
                float a0 = p[0], a1 = p[1], a2 = p[2], a3 = p[3], a4 = p[4];
                float a5 = p[5], a6 = p[6], a7 = p[7], a8 = p[8];
                s0 += a0 + a1 + a2;
                s1 += a2 + a3 + a4;
                s2 += a4 + a5 + a6;
                s3 += a6 + a7 + a8;
            }
        }

        long long out = (((long long)nc * Do + od) * Ho + oh) * Wo + ow;
        const float scale = 1.0f / 27.0f;
        y[out] = s0 * scale;
        if (ow + 1 < Wo) y[out + 1] = s1 * scale;
        if (ow + 2 < Wo) y[out + 2] = s2 * scale;
        if (ow + 3 < Wo) y[out + 3] = s3 * scale;
    }
}

__global__ void avgpool3d_generic_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 K, int S, int P, long long total) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride_grid = (long long)blockDim.x * gridDim.x;
    float inv = 1.0f / (float)(K * K * K);

    for (; idx < total; idx += stride_grid) {
        int ow = idx % Wo;
        long long t = idx / Wo;
        int oh = t % Ho; t /= Ho;
        int od = t % Do; t /= Do;
        int c = t % C;
        int n = t / C;

        int id0 = od * S - P;
        int ih0 = oh * S - P;
        int iw0 = ow * S - P;
        float sum = 0.0f;
        long long nc_base = (((long long)n * C + c) * D * H * W);

        for (int dz = 0; dz < K; ++dz) {
            int iz = id0 + dz;
            if ((unsigned)iz < (unsigned)D) {
                for (int dy = 0; dy < K; ++dy) {
                    int iy = ih0 + dy;
                    if ((unsigned)iy < (unsigned)H) {
                        for (int dx = 0; dx < K; ++dx) {
                            int ix = iw0 + dx;
                            if ((unsigned)ix < (unsigned)W) {
                                sum += x[nc_base + ((long long)iz * H + iy) * W + ix];
                            }
                        }
                    }
                }
            }
        }
        y[idx] = sum * inv;
    }
}

torch::Tensor avg_pool3d_forward(torch::Tensor x, int64_t kernel_size, int64_t stride, int64_t padding) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(x.dim() == 5, "x must be 5D");

    auto xc = x.contiguous();
    int N = (int)xc.size(0);
    int C = (int)xc.size(1);
    int D = (int)xc.size(2);
    int H = (int)xc.size(3);
    int W = (int)xc.size(4);
    int K = (int)kernel_size;
    int S = (int)stride;
    int P = (int)padding;

    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}, xc.options());

    long long total = (long long)N * C * Do * Ho * Wo;
    if (total == 0) return y;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (K == 3 && S == 2 && P == 1) {
        dim3 block(256);
        dim3 grid((Wo + BW - 1) / BW,
                  ((Ho + BH - 1) / BH) * ((Do + BD - 1) / BD),
                  N * C);
        avgpool3d_k3s2p1_shared4_kernel<<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), y.data_ptr<float>(), D, H, W, Do, Ho, Wo, C);
    } else {
        int threads = 256;
        int blocks = (int)((total + threads - 1) / threads);
        if (blocks > 65535) blocks = 65535;
        avgpool3d_generic_kernel<<<blocks, threads, 0, stream>>>(
            xc.data_ptr<float>(), y.data_ptr<float>(), N, C, D, H, W, Do, Ho, Wo, K, S, P, total);
    }
    return y;
}
'''

_ext = load_inline(
    name='avgpool3d_ext_final',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['avg_pool3d_forward'],
    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(stride) if stride is not None else int(kernel_size)
        self.padding = int(padding)

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