import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D convolution, max pooling, log sum exp, and ReLU activation.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super(Model, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.max_pool = nn.MaxPool3d(kernel_size=2, stride=2)

    def forward(self, x):
        """
        Args:
            x: Input tensor of shape (batch_size, in_channels, depth, height, width)
        Returns:
            Output tensor of shape (batch_size, out_channels, depth', height', width')
        """
        x = self.conv(x)
        x = self.max_pool(x)
        x = torch.logsumexp(x, dim=1, keepdim=True)
        x = torch.relu(x)
        return x

batch_size = 4
in_channels = 32
out_channels = 64
depth, height, width = 32, 128, 128
kernel_size = 3
stride = 1
padding = 1

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

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

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

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor fused_pool_lse_forward(torch::Tensor x);
'''

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

__global__ void pool_lse64_kernel(const float* __restrict__ x,
                                  float* __restrict__ y,
                                  int N, int D, int H, int W) {
    const int Do = D >> 1;
    const int Ho = H >> 1;
    const int Wo = W >> 1;
    const int total = N * Do * Ho * Wo;
    const int HW = H * W;
    const int DHW = D * HW;

    for (int out = blockIdx.x * blockDim.x + threadIdx.x;
         out < total;
         out += blockDim.x * gridDim.x) {
        int t = out;
        const int ow = t % Wo; t /= Wo;
        const int oh = t % Ho; t /= Ho;
        const int od = t % Do; t /= Do;
        const int n = t;

        const int w0 = ow << 1;
        const int h0 = oh << 1;
        const int d0 = od << 1;
        const int base_sp = d0 * HW + h0 * W + w0;
        const int nbase = n * 64 * DHW;

        float m = -CUDART_INF_F;
#pragma unroll
        for (int c = 0; c < 64; ++c) {
            const float* p = x + nbase + c * DHW + base_sp;
            float mx = p[0];
            mx = fmaxf(mx, p[1]);
            mx = fmaxf(mx, p[W]);
            mx = fmaxf(mx, p[W + 1]);
            const float* q = p + HW;
            mx = fmaxf(mx, q[0]);
            mx = fmaxf(mx, q[1]);
            mx = fmaxf(mx, q[W]);
            mx = fmaxf(mx, q[W + 1]);
            m = fmaxf(m, mx);
        }

        float s = 0.0f;
#pragma unroll
        for (int c = 0; c < 64; ++c) {
            const float* p = x + nbase + c * DHW + base_sp;
            float mx = p[0];
            mx = fmaxf(mx, p[1]);
            mx = fmaxf(mx, p[W]);
            mx = fmaxf(mx, p[W + 1]);
            const float* q = p + HW;
            mx = fmaxf(mx, q[0]);
            mx = fmaxf(mx, q[1]);
            mx = fmaxf(mx, q[W]);
            mx = fmaxf(mx, q[W + 1]);
            s += expf(mx - m);
        }
        float v = logf(s) + m;
        y[out] = v > 0.0f ? v : 0.0f;
    }
}

__global__ void pool_lse_generic_kernel(const float* __restrict__ x,
                                        float* __restrict__ y,
                                        int N, int C, int D, int H, int W) {
    const int Do = D >> 1;
    const int Ho = H >> 1;
    const int Wo = W >> 1;
    const int total = N * Do * Ho * Wo;
    const int HW = H * W;
    const int DHW = D * HW;

    for (int out = blockIdx.x * blockDim.x + threadIdx.x;
         out < total;
         out += blockDim.x * gridDim.x) {
        int t = out;
        const int ow = t % Wo; t /= Wo;
        const int oh = t % Ho; t /= Ho;
        const int od = t % Do; t /= Do;
        const int n = t;
        const int w0 = ow << 1;
        const int h0 = oh << 1;
        const int d0 = od << 1;
        const int base_sp = d0 * HW + h0 * W + w0;
        const int nbase = n * C * DHW;

        float m = -CUDART_INF_F;
        for (int c = 0; c < C; ++c) {
            const float* p = x + nbase + c * DHW + base_sp;
            float mx = p[0];
            mx = fmaxf(mx, p[1]);
            mx = fmaxf(mx, p[W]);
            mx = fmaxf(mx, p[W + 1]);
            const float* q = p + HW;
            mx = fmaxf(mx, q[0]);
            mx = fmaxf(mx, q[1]);
            mx = fmaxf(mx, q[W]);
            mx = fmaxf(mx, q[W + 1]);
            m = fmaxf(m, mx);
        }
        float s = 0.0f;
        for (int c = 0; c < C; ++c) {
            const float* p = x + nbase + c * DHW + base_sp;
            float mx = p[0];
            mx = fmaxf(mx, p[1]);
            mx = fmaxf(mx, p[W]);
            mx = fmaxf(mx, p[W + 1]);
            const float* q = p + HW;
            mx = fmaxf(mx, q[0]);
            mx = fmaxf(mx, q[1]);
            mx = fmaxf(mx, q[W]);
            mx = fmaxf(mx, q[W + 1]);
            s += expf(mx - m);
        }
        float v = logf(s) + m;
        y[out] = v > 0.0f ? v : 0.0f;
    }
}

torch::Tensor fused_pool_lse_forward(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    if (!x.is_contiguous()) x = x.contiguous();

    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int D = (int)x.size(2);
    const int H = (int)x.size(3);
    const int W = (int)x.size(4);
    const int Do = D / 2, Ho = H / 2, Wo = W / 2;
    auto y = torch::empty({N, 1, Do, Ho, Wo}, x.options());

    const int total = N * Do * Ho * Wo;
    const int threads = 256;
    int blocks = (total + threads - 1) / threads;
    if (blocks > 4096) blocks = 4096;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (C == 64) {
        pool_lse64_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), N, D, H, W);
    } else {
        pool_lse_generic_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), N, C, D, H, W);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)

    def forward(self, x):
        x = F.conv3d(x, self.conv.weight, self.conv.bias, self.conv.stride, self.conv.padding, self.conv.dilation, self.conv.groups)
        return _ext.fused_pool_lse_forward(x)