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>
torch::Tensor forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor bias);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be CUDA")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

__global__ void conv_kernel(const float* __restrict__ x,const float* __restrict__ w,float* __restrict__ y,size_t total){
    size_t idx=blockIdx.x*(size_t)blockDim.x+threadIdx.x;
    if(idx>=total) return;
    int ow=idx%62; size_t t=idx/62;
    int oh=t%62; t/=62;
    int od=t%14; t/=14;
    int co=t%16; int n=t/16;
    const float* x_n=x+(size_t)n*524288;
    const float* w_c=w+(size_t)co*216;
    float acc=0.0f;
    #pragma unroll
    for(int ci=0;ci<8;++ci){
        const float* xp=x_n+(size_t)ci*65536+(size_t)od*4096+(size_t)oh*64+ow;
        const float* wp=w_c+ci*27;
        #pragma unroll
        for(int kd=0;kd<3;++kd){
            #pragma unroll
            for(int kh=0;kh<3;++kh){
                #pragma unroll
                for(int kw=0;kw<3;++kw){
                    acc=fmaf(__ldg(xp+(size_t)kd*4096+kh*64+kw), __ldg(wp+kd*9+kh*3+kw), acc);
                }
            }
        }
    }
    y[idx]=acc;
}

__global__ void pool_avg_kernel(const float* __restrict__ y,const float* __restrict__ conv_b,const float* __restrict__ bias,float* __restrict__ partial,int N){
    int nc=blockIdx.x; int n=nc>>4; int co=nc&15; int tid=threadIdx.x; const int num_pool=6727;
    const float* yc=y+((size_t)n*16+co)*14*62*62;
    float sum=0.0f;
    for(int idx=tid;idx<num_pool;idx+=256){
        int t=idx; int pw=t%31; t/=31; int ph=t%31; int pd=t/31;
        int od=pd*2, oh=ph*2, ow=pw*2;
        size_t base=(size_t)od*3844+(size_t)oh*62+ow;
        float m=yc[base];
        m=fmaxf(m,yc[base+1]);
        m=fmaxf(m,yc[base+62]);
        m=fmaxf(m,yc[base+63]);
        m=fmaxf(m,yc[base+3844]);
        m=fmaxf(m,yc[base+3845]);
        m=fmaxf(m,yc[base+3906]);
        m=fmaxf(m,yc[base+3907]);
        sum += m;
    }
    __shared__ float s[256]; s[tid]=sum; __syncthreads();
    for(int stride=128;stride>0;stride>>=1){ if(tid<stride) s[tid]+=s[tid+stride]; __syncthreads(); }
    if(tid==0) partial[n*16+co]=s[0]*(1.0f/(2.0f*6727.0f))+conv_b[co]*0.5f+bias[co];
}

__global__ void sum_channels_kernel(const float* __restrict__ partial,float* __restrict__ out,int N){
    int n=blockIdx.x*blockDim.x+threadIdx.x;
    if(n<N){
        float s=0;
        #pragma unroll
        for(int c=0;c<16;++c) s+=partial[n*16+c];
        out[n]=s;
    }
}

torch::Tensor forward_cuda(torch::Tensor x,torch::Tensor weight,torch::Tensor conv_bias,torch::Tensor bias){
    CHECK_INPUT(x);CHECK_INPUT(weight);CHECK_INPUT(conv_bias);CHECK_INPUT(bias);
    int N=x.size(0);
    auto y=torch::empty({N,16,14,62,62},x.options());
    auto partial=torch::empty({N,16},x.options());
    auto out=torch::empty({N,1,1,1},x.options());
    cudaStream_t stream=at::cuda::getCurrentCUDAStream();
    size_t total=(size_t)N*16*14*62*62;
    conv_kernel<<<(total+255)/256,256,0,stream>>>(x.data_ptr<float>(),weight.data_ptr<float>(),y.data_ptr<float>(),total);
    pool_avg_kernel<<<N*16,256,0,stream>>>(y.data_ptr<float>(),conv_bias.data_ptr<float>(),bias.data_ptr<float>(),partial.data_ptr<float>(),N);
    sum_channels_kernel<<<(N+255)/256,256,0,stream>>>(partial.data_ptr<float>(),out.data_ptr<float>(),N);
    return out;
}
'''

_ext = load_inline(name='conv_store_pool_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['forward_cuda'], 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 = 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):
        return _ext.forward_cuda(x.contiguous(), self.conv.weight.contiguous(), self.conv.bias.contiguous(), self.bias.contiguous())