```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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


_fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

namespace {

constexpr int C_IN = 3;
constexpr int C_OUT = 16;
constexpr int IN_DIM = 32;
constexpr int OUT_DIM = 15;
constexpr int K = 3;
constexpr int TILE = 4;
constexpr int TILE_IN = 2 * TILE + 1;
constexpr int SPATIAL_OUT = OUT_DIM * OUT_DIM * OUT_DIM;

__global__ void fused_convtranspose_bn_pool_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ conv_b,
    const float* __restrict__ bn_w,
    const float* __restrict__ bn_b,
    const float* __restrict__ bn_m,
    const float* __restrict__ bn_v,
    float* __restrict__ out,
    const float eps) {
    __shared__ float s_x[C_IN][TILE_IN][TILE_IN][TILE_IN];
    __shared__ float s_w[C_OUT][C_IN][K * K * K];
    __shared__ float s_factor[C_OUT];
    __shared__ float s_add[C_OUT];

    const int tx = threadIdx.x;        // local output x, 0..3
    const int ty = threadIdx.y;        // local output y, 0..3
    const int tz = threadIdx.z;        // 0..63: local z = tz & 3, channel = tz >> 2
    const int lx = tx;
    const int ly = ty;
    const int lz = tz & 3;
    const int oc = tz >> 2;
    const int tid = tx + (ty << 2) + (tz << 4);
    const int nthreads = 4 * 4 * 64;

    const int tile_x = blockIdx.x;
    const int tile_y = blockIdx.z & 3;
    const int tile_z = blockIdx.z >> 2;
    const int n = blockIdx.y;

    const int in_x0 = tile_x * 8;
    const int in_y0 = tile_y * 8;
    const int in_z0 = tile_z * 8;

    for (int idx = tid; idx < C_IN * TILE_IN * TILE_IN * TILE_IN; idx += nthreads) {
        const int ic = idx / (TILE_IN * TILE_IN * TILE_IN);
        const int r = idx - ic * (TILE_IN * TILE_IN * TILE_IN);
        const int z = r / (TILE_IN * TILE_IN);
        const int y = (r / TILE_IN) % TILE_IN;
        const int xx = r % TILE_IN;

        const int gx = in_x0 + xx;
        const int gy = in_y0 + y;
        const int gz = in_z0 + z;

        float v = 0.0f;
        if (gx < IN_DIM && gy < IN_DIM && gz < IN_DIM) {
            const long src = (((long)n * C_IN + ic) * IN_DIM + gz) * IN_DIM + gy;
            v = x[src * IN_DIM + gx];
        }
        s_x[ic][z][y][xx] = v;
    }

    for (int idx = tid; idx < C_OUT * C_IN * K * K * K; idx += nthreads) {
        const int occ = idx / (C_IN * K * K * K);
        const int rem = idx - occ * (C_IN * K * K * K);
        const int ic = rem / (K * K * K);
        const int kk = rem - ic * (K * K * K);
        s_w[occ][ic][kk] = w[((long)ic * C_OUT + occ) * (K * K * K) + kk];
    }

    if (tid < C_OUT) {
        const float scale = bn_w[tid] * rsqrtf(bn_v[tid] + eps);
        s_factor[tid] = scale * (1.0f / 64.0f);
        s_add[tid] = bn_b[tid] + scale * (conv_b[tid] - bn_m[tid]);
    }

    __syncthreads();

    const int ox = tile_x * TILE + lx;
    const int oy = tile_y * TILE + ly;
    const int oz = tile_z * TILE + lz;
    if (ox >= OUT_DIM || oy >= OUT_DIM || oz >= OUT_DIM) {
        return;
    }

    const int bx = 2 * lx;
    const int by = 2 * ly;
    const int bz = 2 * lz;

    constexpr int CNT[4] = {1, 2, 1, 2};
    constexpr int CD[4][2] = {{0, 0}, {0, 1}, {1, 0}, {1, 2}};
    constexpr int CK[4][2] = {{1, 0}, {2, 0}, {1, 0}, {2, 0}};

    float acc = 0.0f;

#pragma unroll
    for (int qz = 0; qz < 4; ++qz) {
#pragma unroll
        for (int az = 0; az < CNT[qz]; ++az) {
            const int rz = bz + CD[qz][az];
            const int kz = CK[qz][az];
#pragma unroll
            for (int qy = 0; qy < 4; ++qy) {
#pragma unroll
                for (int ay = 0; ay < CNT[qy]; ++ay) {
                    const int ry = by + CD[qy][ay];
                    const int ky = CK[qy][ay];
#pragma unroll
                    for (int qx = 0; qx < 4; ++qx) {
#pragma unroll
                        for (int ax = 0; ax < CNT[qx]; ++ax) {
                            const int rx = bx + CD[qx][ax];
                            const int kx = CK[qx][ax];
                            const int kk = kz * 9 + ky * 3 + kx;

                            const float x0 = s_x[0][rz][ry][rx];
                            const float x1 = s_x[1][rz][ry][rx];
                            const float x2 = s_x[2][rz][ry][rx];

                            acc += x0 * s_w[oc][0][kk]
                                 + x1 * s_w[oc][1][kk]
                                 + x2 * s_w[oc][2][kk];
                        }
                    }
                }
            }
        }
    }

    const long sp = ((long)oz * OUT_DIM + oy) * OUT_DIM + ox;
    out[((long)n * C_OUT + oc) * SPATIAL_OUT + sp] = acc * s_factor[oc] + s_add[oc];
}

}  // namespace

torch::Tensor fused_convtranspose_bn_pool_cuda(
    torch::Tensor x,
    torch::Tensor conv_weight,
    torch::Tensor conv_bias,
    torch::Tensor bn_weight,
    torch::Tensor bn_bias,
    torch::Tensor bn_mean,
    torch::Tensor bn_var,
    double eps) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(conv_weight.is_cuda(), "conv_weight must be CUDA");
    TORCH_CHECK(conv_bias.is_cuda(), "conv_bias must be CUDA");
    TORCH_CHECK(bn_weight.is_cuda(), "bn_weight must be CUDA");
    TORCH_CHECK(bn_bias.is_cuda(), "bn_bias must be CUDA");
    TORCH_CHECK(bn_mean.is_cuda(), "bn_mean must be CUDA");
    TORCH_CHECK(bn_var.is_cuda(), "bn_var must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

    auto xc = x.contiguous();
    auto wc = conv_weight.contiguous();
    auto cb = conv_bias.contiguous();
    auto bw = bn_weight.contiguous();
    auto bb = bn_bias.contiguous();
    auto bm = bn_mean.contiguous();
    auto bv = bn_var.contiguous();

    TORCH_CHECK(xc.dim() == 5, "x must be NCDHW");
    TORCH_CHECK(xc.size(1) == C_IN, "expected 3 input channels");
    TORCH_CHECK(xc.size(2) == IN_DIM && xc.size(3) == IN_DIM && xc.size(4) == IN_DIM,
                "expected 32x32x32 input");
    TORCH_CHECK(xc.size(0) <= 65535, "batch too large for this kernel launch");
    TORCH_CHECK(wc.dim() == 5 && wc.size(0) == C_IN && wc.size(1) == C_OUT &&
                wc.size(2) == K && wc.size(3) == K && wc.size(4) == K,
                "unexpected conv_transpose weight shape");
    TORCH_CHECK(cb.numel() == C_OUT && bw.numel() == C_OUT && bb.numel() == C_OUT &&
                bm.numel() == C_OUT && bv.numel() == C_OUT,
                "unexpected parameter shape");

    auto out = torch::empty({xc.size(0), C_OUT, OUT_DIM, OUT_DIM, OUT_DIM}, xc.options());

    constexpr int tiles = (OUT_DIM + TILE - 1) / TILE;
    const dim3 block(4, 4, 64);
    const dim3 grid(tiles, static_cast<unsigned int>(xc.size(0)), tiles * tiles);

    auto stream = at::cuda::getCurrentCUDAStream();
    fused_convtranspose_bn_pool_kernel<<<grid, block, 0, stream.stream()>>>(
        xc.data_ptr<float>(),
        wc.data_ptr<float>(),
        cb.data_ptr<float>(),
        bw.data_ptr<float>(),
        bb.data_ptr<float>(),
        bm.data_ptr<float>(),
        bv.data_ptr<float>(),
        out.data_ptr<float>(),
        static_cast<float>(eps));
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out;
}
"""

_fused_cpp_source = (
    "torch::Tensor fused_convtranspose_bn_pool_cuda("
    "torch::Tensor x, torch::Tensor conv_weight, torch::Tensor conv_bias, "
    "torch::Tensor bn_weight, torch::Tensor bn_bias, torch::Tensor bn_mean, "
    "torch::Tensor bn_var, double eps);"
)

_fused_op = load_inline(
    name="fused_convtranspose3d_bn_avgpool3d_v1",
    cpp_sources=_fused_cpp_source,
    cuda_sources=_fused_cuda_source,
    functions=["fused_convtranspose_bn_pool_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    A model that performs a 3D transposed convolution, followed by batch normalization,
    two average pooling layers.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias_shape):
        super(ModelNew, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.bias_shape = bias_shape

        self.conv_transpose = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size, stride=stride, padding=padding
        )
        self.batch_norm = nn.BatchNorm3d(out_channels)
        self.avg_pool1 = nn.AvgPool3d(kernel_size=2)
        self.avg_pool2 = nn.AvgPool3d(kernel_size=2)

    def forward(self, x):
        bn = self.batch_norm
        use_fused = (
            (not bn.training)
            and bn.track_running_stats
            and x.is_cuda
            and x.dtype == torch.float32
            and x.dim() == 5
            and x.size(0) <= 65535
            and x.size(1) == 3
            and x.size(2) == 32
            and x.size(3) == 32
            and x.size(4) == 32
            and self.in_channels == 3
            and self.out_channels == 16
            and self.kernel_size == 3
            and self.stride == 2
            and self.padding == 1
            and self.conv_transpose.bias is not None
            and self.conv_transpose.groups == 1
            and self.conv_transpose.output_padding == 0
            and self.conv_transpose.dilation == 1
            and bn.weight is not None
            and bn.bias is not None
            and bn.running_mean is not None
            and bn.running_var is not None
        )

        if use_fused:
            return _fused_op.fused_convtranspose_bn_pool_cuda(
                x,
                self.conv_transpose.weight,
                self.conv_transpose.bias,
                bn.weight,
                bn.bias,
                bn.running_mean,
                bn.running_var,
                float(bn.eps),
            )

        x = self.conv_transpose(x)
        x = self.batch_norm(x)
        x = self.avg_pool1(x)
        x = self.avg_pool2(x)
        return x
```