```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# ---------------------------------------------------------------------------
# Algebra behind the fusion (convT k=3,s=2,p=1 -> BN -> avgpool2 -> avgpool2):
#
#   convT output p receives x[i]*W[k] iff p = 2*i + k - 1.
#   The two stride-2 avg pools (63 -> 31 -> 15, odd sizes) only read convT
#   outputs p in [0,59], and final[q] = (1/64) * sum_{d in [0,3]^3} y[4q+d].
#   Gathering those 64 terms yields a *valid* stride-2 kernel-3 convolution
#   on x[0..30] with an effective kernel E:
#       final[oc][q] = sum_ic sum_{j in [0,2]^3} E[oc][ic][j] * x[ic][2q+j]
#   BatchNorm is a per-channel affine map and commutes with the average pools,
#   so it folds into E (scale) and a conv bias (shift).  The whole pooled path
#   is therefore ONE cuDNN conv3d (TF32 tensor cores allowed).
#
#   Train-mode batch statistics of the *full* 63^3 convT output are computed
#   by a dedicated kernel that evaluates y on the fly (each output element
#   costs <= 8 taps x IC MACs, parity of p determines the taps) using a
#   shared-memory input tile + register-blocked x values, so the ~1 GB
#   intermediate tensor is never materialized.
# ---------------------------------------------------------------------------

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

#define TILE_OUT 16
#define TILE_IN 9
#define NREG 5

// Build effective kernel E2[oc][ic][27] (scaled by BN) and shift[oc].
// Valid (d,j)->k mapping per dim: j=0 -> k={1,2}, j=1 -> k={0,1,2}, j=2 -> k={0}.
__global__ void build_params_kernel(const float* __restrict__ W,
                                    const float* __restrict__ bn_w,
                                    const float* __restrict__ bn_b,
                                    const float* __restrict__ mean,
                                    const float* __restrict__ var,
                                    const float* __restrict__ conv_bias,
                                    float* __restrict__ E2,
                                    float* __restrict__ shift,
                                    int IC, int OC, float eps) {
    int oc = blockIdx.x / IC;
    int ic = blockIdx.x % IC;
    int j  = threadIdx.x;
    if (j > 27) return;

    float scale = bn_w[oc] * rsqrtf(var[oc] + eps);
    if (threadIdx.x == 0)
        shift[oc] = bn_b[oc] + scale * (conv_bias[oc] - mean[oc]);
    if (j == 27) return;

    int jz = j / 9, jy = (j / 3) % 3, jx = j % 3;
    int kzl[3], kyl[3], kxl[3], nz, ny, nx;
    if (jz == 0)      { kzl[0]=1; kzl[1]=2; nz=2; }
    else if (jz == 1) { kzl[0]=0; kzl[1]=1; kzl[2]=2; nz=3; }
    else              { kzl[0]=0; nz=1; }
    if (jy == 0)      { kyl[0]=1; kyl[1]=2; ny=2; }
    else if (jy == 1) { kyl[0]=0; kyl[1]=1; kyl[2]=2; ny=3; }
    else              { kyl[0]=0; ny=1; }
    if (jx == 0)      { kxl[0]=1; kxl[1]=2; nx=2; }
    else if (jx == 1) { kxl[0]=0; kxl[1]=1; kxl[2]=2; nx=3; }
    else              { kxl[0]=0; nx=1; }

    const float* wb = W + ((long)ic * OC + oc) * 27;
    float acc = 0.f;
    for (int a = 0; a < nz; ++a)
        for (int b = 0; b < ny; ++b)
            for (int c = 0; c < nx; ++c)
                acc += wb[kzl[a] * 9 + kyl[b] * 3 + kxl[c]];
    E2[((long)oc * IC + ic) * 27 + j] = acc * (scale * (1.0f / 64.0f));
}

// On-the-fly batch statistics of the full convT output y (63^3 per slice).
// y[n][oc][p] = bias[oc] + sum_ic sum_{valid k} x[n][ic][i]*W[ic][oc][k],
// p = 2i + k - 1 per dim (p even: single tap k=1; p odd: taps k=0,2).
// Each block handles a 16x16x16 output tile for one batch element n and all
// output channels (x input tile 3*9*9*9 kept in shared memory, W in shared,
// inner x values register-blocked).  Per-block double partials are atomically
// accumulated per channel.  y itself is never written to memory.
__global__ void stats_kernel(const float* __restrict__ x,
                             const float* __restrict__ W,
                             const float* __restrict__ bias,
                             double* __restrict__ sum,
                             double* __restrict__ sumsq,
                             int IC, int OC, int XSZ) {
    __shared__ float xs[3 * TILE_IN * TILE_IN * TILE_IN];
    __shared__ float Ws[16 * 3 * 27];
    __shared__ double bs[16];
    __shared__ double bs2[16];

    const int tid  = threadIdx.x;
    const int n    = blockIdx.y;
    const int pz0  = blockIdx.x * TILE_OUT;          // 0, 16, 32, 48
    const int iz0  = (pz0 + 1) >> 1;
    const int zmax = min(63, pz0 + TILE_OUT);

    // load x tile [ic][dz][dy][dx] -> x[n][ic][iz0+dz][ty0+dy][tx0+dx]
    const long xn = (long)n * IC * XSZ * XSZ * XSZ;
    for (int idx = tid; idx < IC * 729; idx += 256) {
        int dx = idx % TILE_IN;
        int t  = idx / TILE_IN;
        int dy = t % TILE_IN;
        int dz = (t / TILE_IN) % TILE_IN;
        int ic = t / (TILE_IN * TILE_IN);
        xs[idx] = x[xn + ((long)ic * XSZ + (iz0 + dz)) * (XSZ * XSZ)
                    + (long)(ty0 + dy) * XSZ + (tx0 + dx)];
    }
    for (int idx = tid; idx < OC * IC * 27; idx += 256)
        Ws[idx] = W[idx];
    __syncthreads();

    double accS[16], accQ[16];
    #pragma unroll
    for (int c = 0; c < 16; ++c) { accS[c] = 0.0; accQ[c] = 0.0; }

    // parity of pz0 / ty0 / tx0 is even (multiples of 16 and 8)
    int kzmin = (pz0 == 0) ? 1 : 0;
    int kymax = (zmax == 63) ? 1 : 2;
    int kxmax = 2;

    const int r  = tid >> 6;          // 0..3  -> pz = pz0 + 2r + oz
    const int ty = (tid >> 4) & 3;    // 0..3
    const int tx = tid & 15;          // 0..15 -> covers 32 px in pairs
    const int pz_b = pz0 + 2 * r;
    const int py_b = ty0 + 2 * ty;
    const int px_b = tx0 + 2 * tx;
    const int izb = (pz_b + 1) >> 1;  // pz_b even -> (pz_b+2)/2
    const int iyb = (py_b + 1) >> 1;
    const int ixb = (px_b + 1) >> 1;

    for (int oz = 0; oz < 2; ++oz) {
        int pz = pz_b + oz;
        if (pz >= zmax) break;
        int kz0, kz1, dzbase;
        if (oz == 0) { kz0 = kzmin; kz1 = 2;     dzbase = izb - 1 - iz0; }
        else         { kz0 = 0;     kz1 = kzmax; dzbase = izb - iz0;     }
        for (int oy = 0; oy < 2; ++oy) {
            int ky0, ky1, dybase;
            if (oy == 0) { ky0 = 0; ky1 = 2;     dybase = iyb - 1 - ty0; }
            else         { ky0 = 0; ky1 = kymax; dybase = iyb - ty0;     }

            // load register-blocked x values: [dz][dy][5 dx][ic]
            float xr[2][2][NREG][3];
            #pragma unroll
            for (int a = 0; a < 2; ++a) {
                #pragma unroll
                for (int b = 0; b < 2; ++b) {
                    const float* xrow = xs + ((a * 2 + dzbase) * TILE_IN + (b * 2 + dybase)) * TILE_IN + (ixb - 1 - tx0);
                    #pragma unroll
                    for (int c = 0; c < NREG; ++c) {
                        #pragma unroll
                        for (int ic = 0; ic < 3; ++ic)
                            xr[a][b][c][ic] = xrow[ic * 729 + c];
                    }
                }
            }

            float y0[16], y1[16];
            #pragma unroll
            for (int c = 0; c < 16; ++c) {
                float bb = bias[c];
                y0[c] = bb; y1[c] = bb;
            }

            for (int ic = 0; ic < IC; ++ic) {
                #pragma unroll
                for (int a = 0; a < 2; ++a) {
                    #pragma unroll
                    for (int kz = kz0; kz <= kz1; ++kz) {
                        #pragma unroll
                        for (int b = 0; b < 2; ++b) {
                            #pragma unroll
                            for (int ky = ky0; ky <= ky1; ++ky) {
                                int wbase = ic * 27 + kz * 9 + ky * 3;
                                #pragma unroll
                                for (int c = 0; c < NREG; ++c) {
                                    float xv = xr[a][b][c][ic];
                                    #pragma unroll
                                    for (int kx = 0; kx <= kxmax; ++kx) {
                                        float w0 = Ws[wbase + kx];
                                        #pragma unroll
                                        for (int oc = 0; oc < 16; ++oc) {
                                            float w = Ws[oc * 81 + wbase + kx];
                                            if (kx == 2 - c) y1[oc] += xv * w;
                                            else             y0[oc] += xv * w;
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
            #pragma unroll
            for (int c = 0; c < 16; ++c) {
                accS[c] += (double)(y0[c] + y1[c]);
                accQ[c] += (double)(y0[c] * y0[c] + y1[c] * y1[c]);
            }
        }
    }

    for (int oc = 0; oc < OC; ++oc) {
        if (tid == 0) { bs[oc] = 0.0; bs2[oc] = 0.0; }
        __syncthreads();
        double vs = accS[oc], vq = accQ[oc];
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            vs += __shfl_down_sync(0xffffffffu, vs, off);
            vq += __shfl_down_sync(0xffffffffu, vq, off);
        }
        if ((tid & 31) == 0) {
            atomicAdd(&bs[oc], vs);
            atomicAdd(&bs2[oc], vq);
        }
        __syncthreads();
        if (tid == 0) {
            atomicAdd(&sum[oc], bs[oc]);
            atomicAdd(&sumsq[oc], bs2[oc]);
        }
        __syncthreads();
    }
}

// zero the stats accumulators
__global__ void zero_stats_kernel(double* __restrict__ sum,
                                  double* __restrict__ sumsq, int C) {
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c < C) { sum[c] = 0.0; sumsq[c] = 0.0; }
}

// compute batch mean/var and update running stats like aten batch_norm
__global__ void bn_finalize_kernel(const double* __restrict__ sum,
                                   const double* __restrict__ sumsq,
                                   float* __restrict__ running_mean,
                                   float* __restrict__ running_var,
                                   float* __restrict__ out_mean,
                                   float* __restrict__ out_var,
                                   long count, double momentum, int C) {
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= C) return;
    double m  = sum[c] / (double)count;
    double sq = sumsq[c] / (double)count;
    double v  = sq - m * m;
    if (v < 0.0) v = 0.0;
    out_mean[c] = (float)m;
    out_var[c]  = (float)v;
    double unb = (count > 1) ? v * ((double)count / (double)(count - 1)) : 0.0;
    running_mean[c] = (float)((1.0 - momentum) * (double)running_mean[c] + momentum * m);
    running_var[c]  = (float)((1.0 - momentum) * (double)running_var[c]  + momentum * unb);
}

torch::Tensor build_params(torch::Tensor W, torch::Tensor bn_w, torch::Tensor bn_b,
                           torch::Tensor mean, torch::Tensor var, torch::Tensor conv_bias,
                           int64_t IC, int64_t OC, double eps) {
    auto E2    = torch::empty({OC, IC, 3, 3, 3}, W.options());
    auto shift = torch::empty({OC}, W.options());
    auto stream = at::cuda::getCurrentCUDAStream();
    dim3 grid((unsigned)(OC * IC));
    build_params_kernel<<<grid, 32, 0, stream>>>(
        W.data_ptr<float>(), bn_w.data_ptr<float>(), bn_b.data_ptr<float>(),
        mean.data_ptr<float>(), var.data_ptr<float>(), conv_bias.data_ptr<float>(),
        E2.data_ptr<float>(), shift.data_ptr<float>(), (int)IC, (int)OC, (float)eps);
    return torch::cat({E2.reshape({OC, -1}), shift.reshape({OC, 1})}, 1);
}

// train-mode: on-the-fly convT batch stats -> batch mean/var (+running update)
std::vector<torch::Tensor> stats_train(torch::Tensor x, torch::Tensor W, torch::Tensor conv_bias,
                                       torch::Tensor running_mean, torch::Tensor running_var,
                                       int64_t N, int64_t IC, int64_t OC,
                                       double momentum, int64_t XSZ) {
    auto stream = at::cuda::getCurrentCUDAStream();
    auto opts_d = torch::TensorOptions().dtype(torch::kFloat64).device(x.device());
    auto sum   = torch::empty({OC}, opts_d);
    auto sumsq = torch::empty({OC}, opts_d);

    zero_stats_kernel<<<1, 64, 0, stream>>>(sum.data_ptr<double>(), sumsq.data_ptr<double>(), (int)OC);

    dim3 sgrid(4, (unsigned)N, 1);
    stats_kernel<<<sgrid, 256, 0, stream>>>(
        x.data_ptr<float>(), W.data_ptr<float>(), conv_bias.data_ptr<float>(),
        sum.data_ptr<double>(), sumsq.data_ptr<double>(), (int)IC, (int)OC, (int)XSZ);

    auto bmean = torch::empty({OC}, x.options());
    auto bvar  = torch::empty({OC}, x.options());
    long count = (long)N * 250047;   // 63*63*63
    bn_finalize_kernel<<<1, 64, 0, stream>>>(
        sum.data_ptr<double>(), sumsq.data_ptr<double>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        bmean.data_ptr<float>(), bvar.data_ptr<float>(),
        count, momentum, (int)OC);
    return {bmean, bvar};
}
"""

cpp_source = r"""
torch::Tensor build_params(torch::Tensor W, torch::Tensor bn_w, torch::Tensor bn_b,
                           torch::Tensor mean, torch::Tensor var, torch::Tensor conv_bias,
                           int64_t IC, int64_t OC, double eps);
std::vector<torch::Tensor> stats_train(torch::Tensor x, torch::Tensor W, torch::Tensor conv_bias,
                                       torch::Tensor running_mean, torch::Tensor running_var,
                                       int64_t N, int64_t IC, int64_t OC,
                                       double momentum, int64_t XSZ);
"""

_fused = load_inline(
    name="fused_convT_bn_avgpool_v3",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["build_params", "stats_train"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.  The conv-transpose (k=3, s=2, p=1), batch-norm
    and both 2x2x2 average pools are collapsed algebraically into a single
    valid stride-2 kernel-3 convolution with an effective kernel (BN folded
    into its weights/bias), executed with one cuDNN conv3d.  In training mode
    the batch statistics of the transposed-conv output are computed on the fly
    by a custom CUDA kernel (the 63^3-per-channel intermediate is never
    materialized), removing the dominant cuDNN dgrad + memset + stats passes.
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias_shape):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size, stride=stride, padding=padding
        )
        self.batch_norm = nn.BatchNorm3d(out_channels)
        self._in_channels = in_channels
        self._out_channels = out_channels

    def _fallback(self, x):
        x = self.conv_transpose(x)
        x = self.batch_norm(x)
        x = F.avg_pool3d(x, kernel_size=2)
        x = F.avg_pool3d(x, kernel_size=2)
        return x

    def forward(self, x):
        ic, oc = self._in_channels, self._out_channels
        bn = self.batch_norm
        ct = self.conv_transpose

        fast_ok = (
            x.is_cuda and x.dtype == torch.float32 and x.dim() == 5
            and x.size(1) == ic and x.size(2) == 32 and x.size(3) == 32 and x.size(4) == 32
            and ic == 3 and oc == 16
            and ct.bias is not None and bn.affine and bn.track_running_stats
        )
        if not fast_ok:
            return self._fallback(x)

        x = x.contiguous()
        w = ct.weight.contiguous()
        conv_bias = ct.bias.contiguous()
        N = x.shape[0]

        if self.training:
            bmean, bvar = _fused.stats_train(
                x, w, conv_bias, bn.running_mean, bn.running_var,
                N, ic, oc, float(bn.momentum), 32,
            )
            mean_t, var_t = bmean, bvar
        else:
            mean_t = bn.running_mean
            var_t = bn.running_var

        params = _fused.build_params(
            w, bn.weight.contiguous(), bn.bias.contiguous(),
            mean_t, var_t, conv_bias, ic, oc, float(bn.eps),
        )
        E2 = params[:, : ic * 27].reshape(oc, ic, 3, 3, 3)
        shift = params[:, ic * 27]
        return F.conv3d(x, E2, shift, stride=2)


batch_size = 64
in_channels = 3
out_channels = 16
depth, height, width = 32, 32, 32
kernel_size = 3
stride = 2
padding = 1
bias_shape = (out_channels, 1, 1, 1)


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, bias_shape]
```