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>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>

torch::Tensor fused_deconv_pool_avg_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale);

torch::Tensor fused_deconv_pool_avg(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale) {
    return fused_deconv_pool_avg_cuda(x, weight, bias, scale);
}
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <float.h>

__global__ void fused_kernel_ks3_s2_p1_pool2(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ out,
    int N, int D, int H, int W, int Cout, float scale)
{
    // One block computes one (batch, output-channel) result.
    int bc = blockIdx.x;
    int oc = bc % Cout;
    int n  = bc / Cout;
    int tid = threadIdx.x;

    __shared__ float smem[81 + 32];  // weights for 3 input channels * 3*3*3, then warp sums
    float* sw = smem;
    float* warp_sums = smem + 81;

    // Cache the small weight tile for this output channel.  PyTorch ConvTranspose3d
    // weight layout is [Cin, Cout, kD, kH, kW].
    if (tid < 81) {
        int t = tid;
        int kw = t % 3; t /= 3;
        int kh = t % 3; t /= 3;
        int kd = t % 3; t /= 3;
        int ic = t;
        sw[tid] = weight[((((ic * Cout + oc) * 3 + kd) * 3 + kh) * 3 + kw)];
    }
    __syncthreads();

    int PD = D - 1;
    int PH = H - 1;
    int PW = W - 1;
    int P = PD * PH * PW;

    float local = 0.0f;
    int plane = H * W;
    int nbase = n * 3 * D * H * W;

    for (int idx = tid; idx < P; idx += blockDim.x) {
        int tmp = idx;
        int pw = tmp % PW; tmp /= PW;
        int ph = tmp % PH;
        int pd = tmp / PH;

        float v000 = 0.0f, v001 = 0.0f, v010 = 0.0f, v011 = 0.0f;
        float v100 = 0.0f, v101 = 0.0f, v110 = 0.0f, v111 = 0.0f;

        #pragma unroll
        for (int ic = 0; ic < 3; ++ic) {
            const float* xp = x + nbase + ic * D * H * W + pd * plane + ph * W + pw;
            float x000 = __ldg(xp);
            float x001 = __ldg(xp + 1);
            float x010 = __ldg(xp + W);
            float x011 = __ldg(xp + W + 1);
            float x100 = __ldg(xp + plane);
            float x101 = __ldg(xp + plane + 1);
            float x110 = __ldg(xp + plane + W);
            float x111 = __ldg(xp + plane + W + 1);

            const float* w = sw + ic * 27;
            // Names are v[output parity d,h,w].  For odd transposed-conv output
            // positions, either neighboring input element contributes.
            v000 = fmaf(x000, w[13], v000); // 111

            v001 = fmaf(x000, w[14], v001); // 112
            v001 = fmaf(x001, w[12], v001); // 110

            v010 = fmaf(x000, w[16], v010); // 121
            v010 = fmaf(x010, w[10], v010); // 101

            v011 = fmaf(x000, w[17], v011); // 122
            v011 = fmaf(x001, w[15], v011); // 120
            v011 = fmaf(x010, w[11], v011); // 102
            v011 = fmaf(x011, w[9],  v011); // 100

            v100 = fmaf(x000, w[22], v100); // 211
            v100 = fmaf(x100, w[4],  v100); // 011

            v101 = fmaf(x000, w[23], v101); // 212
            v101 = fmaf(x001, w[21], v101); // 210
            v101 = fmaf(x100, w[5],  v101); // 012
            v101 = fmaf(x101, w[3],  v101); // 010

            v110 = fmaf(x000, w[25], v110); // 221
            v110 = fmaf(x010, w[19], v110); // 201
            v110 = fmaf(x100, w[7],  v110); // 021
            v110 = fmaf(x110, w[1],  v110); // 001

            v111 = fmaf(x000, w[26], v111); // 222
            v111 = fmaf(x001, w[24], v111); // 220
            v111 = fmaf(x010, w[20], v111); // 202
            v111 = fmaf(x011, w[18], v111); // 200
            v111 = fmaf(x100, w[8],  v111); // 022
            v111 = fmaf(x101, w[6],  v111); // 020
            v111 = fmaf(x110, w[2],  v111); // 002
            v111 = fmaf(x111, w[0],  v111); // 000
        }

        float m01 = fmaxf(v000, v001);
        float m23 = fmaxf(v010, v011);
        float m45 = fmaxf(v100, v101);
        float m67 = fmaxf(v110, v111);
        float m = fmaxf(fmaxf(m01, m23), fmaxf(m45, m67));
        local += m;
    }

    // Block reduction.
    unsigned mask = 0xffffffffu;
    for (int off = 16; off > 0; off >>= 1) {
        local += __shfl_down_sync(mask, local, off);
    }
    int lane = tid & 31;
    int warp = tid >> 5;
    if (lane == 0) warp_sums[warp] = local;
    __syncthreads();

    if (warp == 0) {
        float sum = (tid < (blockDim.x >> 5)) ? warp_sums[lane] : 0.0f;
        for (int off = 16; off > 0; off >>= 1) {
            sum += __shfl_down_sync(mask, sum, off);
        }
        if (tid == 0) {
            float y = scale * (bias[oc] + sum / (float)P);
            y = fminf(1.0f, fmaxf(0.0f, y));
            out[((n * Cout + oc))] = y;
        }
    }
}

torch::Tensor fused_deconv_pool_avg_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double scale) {
    const c10::cuda::CUDAGuard device_guard(x.device());
    int N = (int)x.size(0);
    int D = (int)x.size(2);
    int H = (int)x.size(3);
    int W = (int)x.size(4);
    int Cout = (int)weight.size(1);
    auto out = torch::empty({N, Cout, 1, 1, 1}, x.options());

    dim3 block(256);
    dim3 grid(N * Cout);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fused_kernel_ks3_s2_p1_pool2<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(),
        N, D, H, W, Cout, (float)scale);
    return out;
}
"""

_ext = load_inline(
    name="fused_deconv3d_pool_avg_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_deconv_pool_avg"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

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
        self.maxpool_kernel_size = maxpool_kernel_size
        self.clamp_min = 0
        self.clamp_max = 1

    def forward(self, x):
        return _ext.fused_deconv_pool_avg(x, self.conv_transpose.weight, self.conv_transpose.bias, float(self.scale))