import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self, num_input_features: int, num_output_features: int):
        """
        :param num_input_features: The number of input feature maps
        :param num_output_features: The number of output feature maps
        """
        super(Model, self).__init__()
        self.transition = nn.Sequential(
            nn.BatchNorm2d(num_input_features),
            nn.ReLU(inplace=True),
            nn.Conv2d(num_input_features, num_output_features, kernel_size=1, bias=False),
            nn.AvgPool2d(kernel_size=2, stride=2)
        )

    def forward(self, x):
        """
        :param x: Input tensor of shape (batch_size, num_input_features, height, width)
        :return: Downsampled tensor with reduced number of feature maps
        """
        return self.transition(x)

batch_size = 128
num_input_features = 32
num_output_features = 64
height, width = 256, 256

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

def get_init_inputs():
    return [num_input_features, num_output_features]

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

cpp_src = r'''
#include <torch/extension.h>
torch::Tensor trans_forward_train_cuda(torch::Tensor x, torch::Tensor bn_w, torch::Tensor bn_b, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor conv_w);
torch::Tensor trans_forward_eval_cuda(torch::Tensor x, torch::Tensor bn_w, torch::Tensor bn_b, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor conv_w);
torch::Tensor trans_forward_train(torch::Tensor x, torch::Tensor bn_w, torch::Tensor bn_b, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor conv_w){return trans_forward_train_cuda(x,bn_w,bn_b,running_mean,running_var,conv_w);} 
torch::Tensor trans_forward_eval(torch::Tensor x, torch::Tensor bn_w, torch::Tensor bn_b, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor conv_w){return trans_forward_eval_cuda(x,bn_w,bn_b,running_mean,running_var,conv_w);} 
'''

cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#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 reduce_sum_sumsq_kernel(const float* __restrict__ x, float* __restrict__ psum, float* __restrict__ psq, int N, int IC, int HW, int M, int NB){int c=blockIdx.x,bid=blockIdx.y,tid=threadIdx.x;float s=0.f,q=0.f;for(int m=bid*blockDim.x+tid;m<M;m+=NB*blockDim.x){int n=m/HW;int p=m-n*HW;float v=x[(n*IC+c)*HW+p];s+=v;q=fmaf(v,v,q);}__shared__ float ss[256],qq[256];ss[tid]=s;qq[tid]=q;__syncthreads();for(int off=128;off>0;off>>=1){if(tid<off){ss[tid]+=ss[tid+off];qq[tid]+=qq[tid+off];}__syncthreads();}if(tid==0){psum[c*NB+bid]=ss[0];psq[c*NB+bid]=qq[0];}}
__global__ void finalize_stats_kernel(const float* __restrict__ psum,const float* __restrict__ psq,float* __restrict__ mean,float* __restrict__ invstd,float* __restrict__ running_mean,float* __restrict__ running_var,int M,int NB){int c=blockIdx.x,tid=threadIdx.x;float s=0.f,q=0.f;for(int i=tid;i<NB;i+=blockDim.x){s+=psum[c*NB+i];q+=psq[c*NB+i];}__shared__ float ss[256],qq[256];ss[tid]=s;qq[tid]=q;__syncthreads();for(int off=128;off>0;off>>=1){if(tid<off){ss[tid]+=ss[tid+off];qq[tid]+=qq[tid+off];}__syncthreads();}if(tid==0){float m=ss[0]/(float)M;float var=qq[0]/(float)M-m*m;var=var>0.f?var:0.f;mean[c]=m;invstd[c]=rsqrtf(var+1.0e-5f);running_mean[c]=running_mean[c]*0.9f+0.1f*m;running_var[c]=running_var[c]*0.9f+0.1f*(var*((float)M/(float)(M-1)));}}
__global__ void fused16_kernel(const float* __restrict__ x,const float* __restrict__ bn_w,const float* __restrict__ bn_b,const float* __restrict__ mean,const float* __restrict__ invstd,const float* __restrict__ cw,float* __restrict__ out,int N,int H,int W){const int IC=32,OC=64;const int OH=H>>1,OW=W>>1,SP=OH*OW;int pix=blockIdx.x*blockDim.x+threadIdx.x;int ob=blockIdx.y,n=blockIdx.z;if(pix>=SP)return;int oh=pix/OW,ow=pix-oh*OW;int base_x=((n*IC)*H+(oh<<1))*W+(ow<<1);float a0=0.f,a1=0.f,a2=0.f,a3=0.f,a4=0.f,a5=0.f,a6=0.f,a7=0.f,a8=0.f,a9=0.f,a10=0.f,a11=0.f,a12=0.f,a13=0.f,a14=0.f,a15=0.f;int oc0=ob<<4;
#pragma unroll
for(int ic=0;ic<32;++ic){float scale=bn_w[ic]*invstd[ic];float bias=bn_b[ic]-mean[ic]*scale;int idx=base_x+ic*H*W;float v0=fmaf(x[idx],scale,bias),v1=fmaf(x[idx+1],scale,bias),v2=fmaf(x[idx+W],scale,bias),v3=fmaf(x[idx+W+1],scale,bias);v0=v0>0.f?v0:0.f;v1=v1>0.f?v1:0.f;v2=v2>0.f?v2:0.f;v3=v3>0.f?v3:0.f;float s=(v0+v1+v2+v3)*0.25f;const float* wp=cw+oc0*32+ic;a0=fmaf(wp[0*32],s,a0);a1=fmaf(wp[1*32],s,a1);a2=fmaf(wp[2*32],s,a2);a3=fmaf(wp[3*32],s,a3);a4=fmaf(wp[4*32],s,a4);a5=fmaf(wp[5*32],s,a5);a6=fmaf(wp[6*32],s,a6);a7=fmaf(wp[7*32],s,a7);a8=fmaf(wp[8*32],s,a8);a9=fmaf(wp[9*32],s,a9);a10=fmaf(wp[10*32],s,a10);a11=fmaf(wp[11*32],s,a11);a12=fmaf(wp[12*32],s,a12);a13=fmaf(wp[13*32],s,a13);a14=fmaf(wp[14*32],s,a14);a15=fmaf(wp[15*32],s,a15);}int o=n*OC*SP+pix;out[o+(oc0+0)*SP]=a0;out[o+(oc0+1)*SP]=a1;out[o+(oc0+2)*SP]=a2;out[o+(oc0+3)*SP]=a3;out[o+(oc0+4)*SP]=a4;out[o+(oc0+5)*SP]=a5;out[o+(oc0+6)*SP]=a6;out[o+(oc0+7)*SP]=a7;out[o+(oc0+8)*SP]=a8;out[o+(oc0+9)*SP]=a9;out[o+(oc0+10)*SP]=a10;out[o+(oc0+11)*SP]=a11;out[o+(oc0+12)*SP]=a12;out[o+(oc0+13)*SP]=a13;out[o+(oc0+14)*SP]=a14;out[o+(oc0+15)*SP]=a15;}
__global__ void fused16_eval_kernel(const float* __restrict__ x,const float* __restrict__ bn_w,const float* __restrict__ bn_b,const float* __restrict__ rm,const float* __restrict__ rv,const float* __restrict__ cw,float* __restrict__ out,int N,int H,int W){const int IC=32,OC=64;const int OH=H>>1,OW=W>>1,SP=OH*OW;int pix=blockIdx.x*blockDim.x+threadIdx.x;int ob=blockIdx.y,n=blockIdx.z;if(pix>=SP)return;int oh=pix/OW,ow=pix-oh*OW;int base_x=((n*IC)*H+(oh<<1))*W+(ow<<1);float a0=0.f,a1=0.f,a2=0.f,a3=0.f,a4=0.f,a5=0.f,a6=0.f,a7=0.f,a8=0.f,a9=0.f,a10=0.f,a11=0.f,a12=0.f,a13=0.f,a14=0.f,a15=0.f;int oc0=ob<<4;
#pragma unroll
for(int ic=0;ic<32;++ic){float scale=bn_w[ic]*rsqrtf(rv[ic]+1.0e-5f);float bias=bn_b[ic]-rm[ic]*scale;int idx=base_x+ic*H*W;float v0=fmaf(x[idx],scale,bias),v1=fmaf(x[idx+1],scale,bias),v2=fmaf(x[idx+W],scale,bias),v3=fmaf(x[idx+W+1],scale,bias);v0=v0>0.f?v0:0.f;v1=v1>0.f?v1:0.f;v2=v2>0.f?v2:0.f;v3=v3>0.f?v3:0.f;float s=(v0+v1+v2+v3)*0.25f;const float* wp=cw+oc0*32+ic;a0=fmaf(wp[0*32],s,a0);a1=fmaf(wp[1*32],s,a1);a2=fmaf(wp[2*32],s,a2);a3=fmaf(wp[3*32],s,a3);a4=fmaf(wp[4*32],s,a4);a5=fmaf(wp[5*32],s,a5);a6=fmaf(wp[6*32],s,a6);a7=fmaf(wp[7*32],s,a7);a8=fmaf(wp[8*32],s,a8);a9=fmaf(wp[9*32],s,a9);a10=fmaf(wp[10*32],s,a10);a11=fmaf(wp[11*32],s,a11);a12=fmaf(wp[12*32],s,a12);a13=fmaf(wp[13*32],s,a13);a14=fmaf(wp[14*32],s,a14);a15=fmaf(wp[15*32],s,a15);}int o=n*OC*SP+pix;out[o+(oc0+0)*SP]=a0;out[o+(oc0+1)*SP]=a1;out[o+(oc0+2)*SP]=a2;out[o+(oc0+3)*SP]=a3;out[o+(oc0+4)*SP]=a4;out[o+(oc0+5)*SP]=a5;out[o+(oc0+6)*SP]=a6;out[o+(oc0+7)*SP]=a7;out[o+(oc0+8)*SP]=a8;out[o+(oc0+9)*SP]=a9;out[o+(oc0+10)*SP]=a10;out[o+(oc0+11)*SP]=a11;out[o+(oc0+12)*SP]=a12;out[o+(oc0+13)*SP]=a13;out[o+(oc0+14)*SP]=a14;out[o+(oc0+15)*SP]=a15;}
torch::Tensor trans_forward_train_cuda(torch::Tensor x,torch::Tensor bn_w,torch::Tensor bn_b,torch::Tensor running_mean,torch::Tensor running_var,torch::Tensor conv_w){CHECK_INPUT(x);CHECK_INPUT(bn_w);CHECK_INPUT(bn_b);CHECK_INPUT(running_mean);CHECK_INPUT(running_var);CHECK_INPUT(conv_w);int N=x.size(0),IC=x.size(1),H=x.size(2),W=x.size(3),HW=H*W,M=N*HW,OC=conv_w.size(0);auto out=torch::empty({N,OC,H/2,W/2},x.options());const int NB=256;auto psum=torch::empty({IC,NB},x.options()),psq=torch::empty({IC,NB},x.options()),mean=torch::empty({IC},x.options()),invstd=torch::empty({IC},x.options());reduce_sum_sumsq_kernel<<<dim3(IC,NB),256>>>(x.data_ptr<float>(),psum.data_ptr<float>(),psq.data_ptr<float>(),N,IC,HW,M,NB);finalize_stats_kernel<<<IC,256>>>(psum.data_ptr<float>(),psq.data_ptr<float>(),mean.data_ptr<float>(),invstd.data_ptr<float>(),running_mean.data_ptr<float>(),running_var.data_ptr<float>(),M,NB);int SP=(H/2)*(W/2);fused16_kernel<<<dim3((SP+127)/128,(OC+15)/16,N),128>>>(x.data_ptr<float>(),bn_w.data_ptr<float>(),bn_b.data_ptr<float>(),mean.data_ptr<float>(),invstd.data_ptr<float>(),conv_w.data_ptr<float>(),out.data_ptr<float>(),N,H,W);return out;}
torch::Tensor trans_forward_eval_cuda(torch::Tensor x,torch::Tensor bn_w,torch::Tensor bn_b,torch::Tensor running_mean,torch::Tensor running_var,torch::Tensor conv_w){CHECK_INPUT(x);CHECK_INPUT(bn_w);CHECK_INPUT(bn_b);CHECK_INPUT(running_mean);CHECK_INPUT(running_var);CHECK_INPUT(conv_w);int N=x.size(0),H=x.size(2),W=x.size(3),OC=conv_w.size(0);auto out=torch::empty({N,OC,H/2,W/2},x.options());int SP=(H/2)*(W/2);fused16_eval_kernel<<<dim3((SP+127)/128,(OC+15)/16,N),128>>>(x.data_ptr<float>(),bn_w.data_ptr<float>(),bn_b.data_ptr<float>(),running_mean.data_ptr<float>(),running_var.data_ptr<float>(),conv_w.data_ptr<float>(),out.data_ptr<float>(),N,H,W);return out;}
'''
_ext=load_inline(name='transition_fused_train_final', cpp_sources=[cpp_src], cuda_sources=[cuda_src], functions=['trans_forward_train','trans_forward_eval'], extra_cuda_cflags=['-O3','--use_fast_math'], extra_cflags=['-O3'], verbose=False)
class ModelNew(nn.Module):
    def __init__(self,num_input_features:int,num_output_features:int):
        super(ModelNew,self).__init__(); self.transition=nn.Sequential(nn.BatchNorm2d(num_input_features),nn.ReLU(inplace=True),nn.Conv2d(num_input_features,num_output_features,kernel_size=1,bias=False),nn.AvgPool2d(kernel_size=2,stride=2))
    def forward(self,x):
        bn=self.transition[0]; conv=self.transition[2]
        if self.training: return _ext.trans_forward_train(x.contiguous(),bn.weight,bn.bias,bn.running_mean,bn.running_var,conv.weight)
        return _ext.trans_forward_eval(x.contiguous(),bn.weight,bn.bias,bn.running_mean,bn.running_var,conv.weight)