import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D convolution, divides by a constant, applies max pooling,
    global average pooling, adds a bias term, and sums along a specific dimension.
    """
    def __init__(self, in_channels, out_channels, kernel_size, divisor, pool_size, bias_shape, sum_dim):
        super(Model, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.divisor = divisor
        self.max_pool = nn.MaxPool3d(pool_size)
        self.global_avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self.sum_dim = sum_dim

    def forward(self, x):
        x = self.conv(x)
        x = x / self.divisor
        x = self.max_pool(x)
        x = self.global_avg_pool(x)
        x = x + self.bias
        x = torch.sum(x, dim=self.sum_dim)
        return x

batch_size   = 128  
in_channels  = 8            
out_channels = 16  
depth = 16; height = width = 64 
kernel_size = (3, 3, 3)
divisor = 2.0
pool_size = (2, 2, 2)
bias_shape = (out_channels, 1, 1, 1)
sum_dim = 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, divisor, pool_size, bias_shape, sum_dim]

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

cpp_src = r'''
#include <torch/extension.h>
#include <ATen/ATen.h>
#include <vector>

torch::Tensor postprocess_cuda(torch::Tensor conv, torch::Tensor bias, double divisor);

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor bias, double divisor) {
    auto y = at::conv3d(x, weight, conv_bias, {1,1,1}, {0,0,0}, {1,1,1}, 1);
    return postprocess_cuda(y, bias, divisor);
}
'''

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

__global__ void post_kernel(const float* __restrict__ y, const float* __restrict__ bias,
                            float* __restrict__ out,
                            int N, int C, int D, int H, int W, float divisor) {
    int n = blockIdx.x;
    int tid = threadIdx.x;
    const int PD = D / 2;
    const int PH = H / 2;
    const int PW = W / 2;
    int total = C * PD * PH * PW;
    float local = 0.0f;
    bool pos = divisor > 0.0f;
    for (int idx = tid; idx < total; idx += blockDim.x) {
        int t = idx;
        int pw = t % PW; t /= PW;
        int ph = t % PH; t /= PH;
        int pd = t % PD; t /= PD;
        int c = t;
        int d0 = pd << 1;
        int h0 = ph << 1;
        int w0 = pw << 1;
        const float* p = y + (((n * C + c) * D + d0) * H + h0) * W + w0;
        float v0 = p[0];
        float v1 = p[1];
        float v2 = p[W];
        float v3 = p[W + 1];
        const float* q = p + H * W;
        float v4 = q[0];
        float v5 = q[1];
        float v6 = q[W];
        float v7 = q[W + 1];
        float m;
        if (pos) {
            m = fmaxf(fmaxf(fmaxf(v0, v1), fmaxf(v2, v3)), fmaxf(fmaxf(v4, v5), fmaxf(v6, v7)));
        } else {
            m = fminf(fminf(fminf(v0, v1), fminf(v2, v3)), fminf(fminf(v4, v5), fminf(v6, v7)));
        }
        local += m;
    }
    __shared__ float smem[256];
    smem[tid] = local;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }
    if (tid == 0) {
        float bsum = 0.0f;
        for (int c = 0; c < C; ++c) bsum += bias[c];
        float denom = (float)(PD * PH * PW) * divisor;
        out[n] = smem[0] / denom + bsum;
    }
}

torch::Tensor postprocess_cuda(torch::Tensor conv, torch::Tensor bias, double divisor) {
    TORCH_CHECK(conv.is_cuda(), "conv must be CUDA");
    int N = (int)conv.size(0);
    int C = (int)conv.size(1);
    int D = (int)conv.size(2);
    int H = (int)conv.size(3);
    int W = (int)conv.size(4);
    auto out = torch::empty({N, 1, 1, 1}, conv.options());
    post_kernel<<<dim3(N), dim3(256), 0, at::cuda::getCurrentCUDAStream()>>>(
        conv.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), N, C, D, H, W, (float)divisor);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''

_ext = load_inline(name='fused_conv_post_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=['fused_forward'], extra_cflags=['-O3'], extra_cuda_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, divisor, pool_size, bias_shape, sum_dim):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.divisor = float(divisor)
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self.sum_dim = sum_dim

    def forward(self, x):
        return _ext.fused_forward(x, self.conv.weight, self.conv.bias, self.bias, self.divisor)