import torch
import torch.nn as nn

class Model(nn.Module):
    """
    This model performs a convolution, average pooling, applies sigmoid, and sums the result.
    """
    def __init__(self, in_channels, out_channels, kernel_size, pool_kernel_size):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.avg_pool = nn.AvgPool2d(pool_kernel_size)

    def forward(self, x):
        x = self.conv(x)
        x = self.avg_pool(x)
        x = torch.sigmoid(x)
        x = torch.sum(x, dim=[1,2,3]) # Sum over all spatial dimensions
        return x

batch_size = 128
in_channels = 8
out_channels = 64
height, width = 384, 384
kernel_size = 3
pool_kernel_size = 4

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

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

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cpp_source = r'''
#include <torch/extension.h>
void fused_convpool_sigmoid_sum_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor ew, torch::Tensor out);
void fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor ew, torch::Tensor out) { fused_convpool_sigmoid_sum_cuda(x,w,b,ew,out); }
'''
cuda_source = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define H 384
#define W 384
#define IC 8
#define OC 64
#define PH 95
#define PW 95
#define P (PH*PW)
#define TPB 256
#define ITEMS 1
#define OCG 4
__global__ void make_eff_kernel(const float* __restrict__ w, float* __restrict__ ew){
 int idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx>=OC*IC*36) return; int t=idx; int dx=t%6; t/=6; int dy=t%6; t/=6; int ic=t%IC; t/=IC; int oc=t; float s=0.f;
 #pragma unroll
 for(int ky=0;ky<3;++ky){
  int qy=dy-ky; if((unsigned)qy<4u){
   #pragma unroll
   for(int kx=0;kx<3;++kx){int qx=dx-kx; if((unsigned)qx<4u) s += w[((oc*IC+ic)*3+ky)*3+kx];}
  }
 }
 ew[idx]=s*0.0625f;
}
__forceinline__ __device__ float warp_reduce_sum(float v){
 #pragma unroll
 for(int off=16;off>0;off>>=1) v += __shfl_down_sync(0xffffffff,v,off); return v;
}
__forceinline__ __device__ float sigm(float x){return 1.0f/(1.0f+__expf(-x));}
__forceinline__ __device__ float compute_group4(const float* __restrict__ x,const float* __restrict__ ew,const float* __restrict__ b,int n,int oc0,int p){
 int py=p/PW; int px=p-py*PW; int by=py<<2; int bx=px<<2;
 float a0=b[oc0], a1=b[oc0+1], a2=b[oc0+2], a3=b[oc0+3];
 const float* xp0=x+((n*IC)*H+by)*W+bx;
 const float* w0=ew+(oc0*IC)*36; const float* w1=w0+IC*36; const float* w2=w1+IC*36; const float* w3=w2+IC*36;
 #pragma unroll
 for(int ic=0;ic<IC;++ic){ const float* xp=xp0+ic*H*W; const float* q0=w0+ic*36; const float* q1=w1+ic*36; const float* q2=w2+ic*36; const float* q3=w3+ic*36;
  #pragma unroll
  for(int dy=0;dy<6;++dy){ const float* xr=xp+dy*W; int off=dy*6;
   #pragma unroll
   for(int dx=0;dx<6;++dx){ float xv=xr[dx]; int k=off+dx; a0=fmaf(xv,q0[k],a0); a1=fmaf(xv,q1[k],a1); a2=fmaf(xv,q2[k],a2); a3=fmaf(xv,q3[k],a3); }
  }
 }
 return sigm(a0)+sigm(a1)+sigm(a2)+sigm(a3);
}
__global__ void main_kernel(const float* __restrict__ x,const float* __restrict__ ew,const float* __restrict__ b,float* __restrict__ out){
 int oc0=blockIdx.x*OCG; int n=blockIdx.y; int tile=blockIdx.z; int tid=threadIdx.x; int p0=tile*(TPB*ITEMS)+tid; float val=0.f;
 #pragma unroll
 for(int it=0;it<ITEMS;++it){int p=p0+it*TPB; if(p<P) val += compute_group4(x,ew,b,n,oc0,p);}
 val=warp_reduce_sum(val); __shared__ float smem[TPB/32]; if((tid&31)==0) smem[tid>>5]=val; __syncthreads();
 if(tid<TPB/32){float v=smem[tid]; v=warp_reduce_sum(v); if(tid==0) atomicAdd(out+n,v);} }
void fused_convpool_sigmoid_sum_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor ew, torch::Tensor out){
 auto stream=at::cuda::getCurrentCUDAStream(); cudaMemsetAsync(out.data_ptr<float>(),0,out.numel()*sizeof(float),stream);
 make_eff_kernel<<<(OC*IC*36+255)/256,256,0,stream>>>(w.data_ptr<float>(),ew.data_ptr<float>());
 dim3 grid(OC/OCG,x.size(0),(P+TPB*ITEMS-1)/(TPB*ITEMS)); main_kernel<<<grid,TPB,0,stream>>>(x.data_ptr<float>(),ew.data_ptr<float>(),b.data_ptr<float>(),out.data_ptr<float>());
}
'''
_ext=load_inline(name='fused_convpool_sigmoid_sum_ext_g4', cpp_sources=cpp_source, cuda_sources=cuda_source, functions=['fused_forward'], extra_cuda_cflags=['-O3','--use_fast_math'], extra_cflags=['-O3'], verbose=False)
class ModelNew(nn.Module):
 def __init__(self,in_channels,out_channels,kernel_size,pool_kernel_size):
  super(ModelNew,self).__init__(); self.conv=nn.Conv2d(in_channels,out_channels,kernel_size); self.avg_pool=nn.AvgPool2d(pool_kernel_size); self.register_buffer('_eff_w',torch.empty((out_channels,in_channels,6,6),dtype=torch.float32))
 def forward(self,x):
  out=torch.empty((x.shape[0],),device=x.device,dtype=torch.float32); _ext.fused_forward(x,self.conv.weight,self.conv.bias,self._eff_w,out); return out