import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a transposed 3D convolution, multiplies by a scalar, applies max pooling, 
    global average pooling, and clamps the output.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, scale, maxpool_kernel_size):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.scale = scale
        self.maxpool = nn.MaxPool3d(kernel_size=maxpool_kernel_size)
        self.global_avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
        self.clamp_min = 0
        self.clamp_max = 1

    def forward(self, x):
        x = self.conv_transpose(x)
        x = x * self.scale
        x = self.maxpool(x)
        x = self.global_avg_pool(x)
        x = torch.clamp(x, min=self.clamp_min, max=self.clamp_max)
        return x

batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1
scale = 0.5
maxpool_kernel_size = 2

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, scale, maxpool_kernel_size]

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

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>

#define IC 3
#define OC 16
#define ID 16
#define IH 32
#define IW 32
#define KD 3

template<int A, int B, int C>
__device__ __forceinline__ float conv_pool_at(const float* __restrict__ x, const float* __restrict__ sw, float bv,
                                              int n, int pd, int ph, int pw){
    float v = bv;
    const int dn = A ? 2 : 1;
    const int hn = B ? 2 : 1;
    const int wn = C ? 2 : 1;
    int did[2] = {pd + (A ? 1 : 0), pd};
    int hid[2] = {ph + (B ? 1 : 0), ph};
    int wid[2] = {pw + (C ? 1 : 0), pw};
    int dkd[2] = {A ? 0 : 1, 2};
    int hkd[2] = {B ? 0 : 1, 2};
    int wkd[2] = {C ? 0 : 1, 2};
    #pragma unroll
    for(int ic=0; ic<IC; ++ic){
        const float* xp = x + (((n*IC + ic)*ID)*IH)*IW;
        const float* wp = sw + ic*27;
        #pragma unroll
        for(int di=0; di<dn; ++di){
            #pragma unroll
            for(int hi=0; hi<hn; ++hi){
                #pragma unroll
                for(int wi=0; wi<wn; ++wi){
                    v = fmaf(xp[(did[di]*IH + hid[hi])*IW + wid[wi]], wp[(dkd[di]*KD + hkd[hi])*KD + wkd[wi]], v);
                }
            }
        }
    }
    return v;
}

__device__ __forceinline__ float block_reduce_sum(float val){
    #pragma unroll
    for(int offset=16; offset>0; offset>>=1) val += __shfl_down_sync(0xffffffff, val, offset);
    __shared__ float warp_sums[8];
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    if(lane == 0) warp_sums[warp] = val;
    __syncthreads();
    val = (threadIdx.x < 8) ? warp_sums[lane] : 0.0f;
    if(warp == 0){
        #pragma unroll
        for(int offset=4; offset>0; offset>>=1) val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__global__ __launch_bounds__(256, 2) void fused_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ bias,
                             float* __restrict__ out, float scale){
    __shared__ float sw[81];
    int bid = blockIdx.x;
    int n = bid >> 4;
    int oc = bid & 15;
    if(threadIdx.x < 81){
        int ic = threadIdx.x / 27;
        int k = threadIdx.x - ic*27;
        sw[threadIdx.x] = w[(ic*OC + oc)*27 + k];
    }
    float bv = bias[oc];
    __syncthreads();
    float sum = 0.0f;
    for(int idx = threadIdx.x; idx < 14415; idx += 256){
        int pd = idx / 961;
        int r = idx - pd*961;
        int ph = r / 31;
        int pw = r - ph*31;
        float m0 = conv_pool_at<0,0,0>(x,sw,bv,n,pd,ph,pw);
        float m1 = conv_pool_at<0,0,1>(x,sw,bv,n,pd,ph,pw);
        float m2 = conv_pool_at<0,1,0>(x,sw,bv,n,pd,ph,pw);
        float m3 = conv_pool_at<0,1,1>(x,sw,bv,n,pd,ph,pw);
        float m4 = conv_pool_at<1,0,0>(x,sw,bv,n,pd,ph,pw);
        float m5 = conv_pool_at<1,0,1>(x,sw,bv,n,pd,ph,pw);
        float m6 = conv_pool_at<1,1,0>(x,sw,bv,n,pd,ph,pw);
        float m7 = conv_pool_at<1,1,1>(x,sw,bv,n,pd,ph,pw);
        float m = fmaxf(fmaxf(fmaxf(m0,m1), fmaxf(m2,m3)), fmaxf(fmaxf(m4,m5), fmaxf(m6,m7)));
        sum += m * scale;
    }
    float total = block_reduce_sum(sum);
    if(threadIdx.x == 0){
        float y = total * (1.0f / 14415.0f);
        y = fminf(1.0f, fmaxf(0.0f, y));
        out[bid] = y;
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale){
    auto out = torch::empty({x.size(0), weight.size(1), 1, 1, 1}, x.options());
    fused_kernel<<<dim3(x.size(0) * 16), dim3(256)>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), (float)scale);
    return out;
}
'''

_ext = load_inline(name='fused_convtranspose_pool_ext_v4', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['fused_forward'], verbose=False, extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'])

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, scale, maxpool_kernel_size):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.scale = scale
    def forward(self, x):
        return _ext.fused_forward(x, self.conv_transpose.weight, self.conv_transpose.bias, float(self.scale))