import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a transposed convolution, multiplies by a scalar, applies global average pooling, 
    another global average pooling
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, multiplier):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)
        self.multiplier = multiplier

    def forward(self, x):
        x = self.conv_transpose(x)
        x = x * self.multiplier
        x = torch.mean(x, dim=[2, 3], keepdim=True)  # First global average pooling
        x = torch.mean(x, dim=[2, 3], keepdim=True)  # Second global average pooling
        return x

batch_size = 16
in_channels = 64
out_channels = 128
height, width = 128, 128
kernel_size = 3
stride = 2
padding = 1
output_padding = 1
multiplier = 0.5

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

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

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

_cpp_src = r'''
#include <torch/extension.h>
void deconv_gap_launch(const float* x, const float* weight, const float* bias, float* agg, float* out,
                       int N, int Cin, int H, int W, int Cout, float multiplier);
torch::Tensor deconv_gap_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double multiplier) {
    int N=(int)x.size(0), Cin=(int)x.size(1), H=(int)x.size(2), W=(int)x.size(3), Cout=(int)weight.size(1);
    auto out=torch::empty({N,Cout,1,1}, x.options());
    auto agg=torch::empty({N,Cin,4}, x.options());
    deconv_gap_launch(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), agg.data_ptr<float>(), out.data_ptr<float>(), N,Cin,H,W,Cout,(float)multiplier);
    return out;
}
'''
_cuda_src = r'''
#include <cuda.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>

__device__ __forceinline__ void warp_reduce3(float &a, float &b, float &c) {
    unsigned mask = 0xffffffffu;
    #pragma unroll
    for (int off=16; off>0; off>>=1) {
        a += __shfl_down_sync(mask, a, off);
        b += __shfl_down_sync(mask, b, off);
        c += __shfl_down_sync(mask, c, off);
    }
}
__device__ __forceinline__ float warp_reduce1(float a) {
    unsigned mask = 0xffffffffu;
    #pragma unroll
    for (int off=16; off>0; off>>=1) a += __shfl_down_sync(mask, a, off);
    return a;
}

__global__ void aggregate_kernel(const float* __restrict__ x, float* __restrict__ agg, int Cin, int H, int W) {
    int nci=blockIdx.x, tid=threadIdx.x;
    int HW=H*W;
    const float* base=x+(long long)nci*HW;
    float sum=0.f;
    for(int idx=tid; idx<HW; idx+=blockDim.x) sum += base[idx];
    float row=0.f;
    for(int j=tid; j<W; j+=blockDim.x) row += base[j];
    float col=0.f;
    for(int i=tid; i<H; i+=blockDim.x) col += base[(long long)i*W];
    warp_reduce3(sum,row,col);
    __shared__ float ss[8], sr[8], sc[8];
    int lane=tid&31, warp=tid>>5;
    if(lane==0){ ss[warp]=sum; sr[warp]=row; sc[warp]=col; }
    __syncthreads();
    if(warp==0){
        sum = (lane < (blockDim.x>>5)) ? ss[lane] : 0.f;
        row = (lane < (blockDim.x>>5)) ? sr[lane] : 0.f;
        col = (lane < (blockDim.x>>5)) ? sc[lane] : 0.f;
        warp_reduce3(sum,row,col);
        if(lane==0){
            long long o=(long long)nci*4;
            agg[o+0]=sum; agg[o+1]=row; agg[o+2]=col; agg[o+3]=base[0];
        }
    }
}

__global__ void output_kernel(const float* __restrict__ agg, const float* __restrict__ weight, const float* __restrict__ bias,
                              float* __restrict__ out, int Cin, int H, int W, int Cout, float multiplier) {
    int nco=blockIdx.x, co=nco % Cout, n=nco / Cout, tid=threadIdx.x;
    float val=0.f;
    for(int ci=tid; ci<Cin; ci+=blockDim.x){
        const float* a=agg+(((long long)n*Cin+ci)*4);
        float all=a[0], row=a[1], col=a[2], corner=a[3];
        const float* wb=weight+(((long long)ci*Cout+co)*9);
        float w00=wb[0], w01=wb[1], w02=wb[2];
        float w10=wb[3], w11=wb[4], w12=wb[5];
        float w20=wb[6], w21=wb[7], w22=wb[8];
        float sall=w00+w01+w02+w10+w11+w12+w20+w21+w22;
        float srow=w10+w11+w12+w20+w21+w22;
        float scol=w01+w02+w11+w12+w21+w22;
        float scorner=w11+w12+w21+w22;
        val += all*sall + row*(srow-sall) + col*(scol-sall) + corner*(scorner-srow-scol+sall);
    }
    val = warp_reduce1(val);
    __shared__ float sh[2];
    int lane=tid&31, warp=tid>>5;
    if(lane==0) sh[warp]=val;
    __syncthreads();
    if(tid==0){
        float total=sh[0]+sh[1];
        float area=(float)(4*H*W);
        out[(long long)n*Cout+co]=multiplier*(bias[co]+total/area);
    }
}

void deconv_gap_launch(const float* x, const float* weight, const float* bias, float* agg, float* out,
                       int N, int Cin, int H, int W, int Cout, float multiplier) {
    cudaStream_t stream=at::cuda::getCurrentCUDAStream();
    aggregate_kernel<<<N*Cin,256,0,stream>>>(x,agg,Cin,H,W);
    output_kernel<<<N*Cout,64,0,stream>>>(agg,weight,bias,out,Cin,H,W,Cout,multiplier);
}
'''
_ext=load_inline(name='deconv_gap_ext_v2', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['deconv_gap_forward'], verbose=False, extra_cuda_cflags=['-O3'], extra_cflags=['-O3'])
class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, multiplier):
        super(ModelNew,self).__init__()
        self.conv_transpose=nn.ConvTranspose2d(in_channels,out_channels,kernel_size,stride=stride,padding=padding,output_padding=output_padding)
        self.multiplier=multiplier
    def forward(self,x):
        return _ext.deconv_gap_forward(x,self.conv_transpose.weight,self.conv_transpose.bias,float(self.multiplier))