```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

# ---------------------------------------------------------------------------
# Math behind the fusion (kernel_size=3, stride=2, padding=1 transposed conv,
# followed by two stride-2 2x2x2 average pools):
#
#   convT output p receives x[i] with weight W[k] iff p = 2*i + k - 1.
#   Two avg pools with odd sizes (63 -> 31 -> 15) only use convT outputs
#   p in [0,59] per dim, i.e. final[q] = (1/64) * sum_{d in [0,3]^3} y[4q+d].
#   Gathering the 64 sub-terms gives 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]
#   where E folds W over the 64 pooled positions (and the 1/64 factor).
#   BatchNorm is per-channel affine, so it commutes with the average pool and
#   collapses into a per-channel scale/shift applied at the end.
# ---------------------------------------------------------------------------

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

#define ROWS 8
#define TOTAL_ROWS 225
#define X_TILE_FLOATS (ROWS * 3 * 3 * 3 * 31)

// Build effective kernel E[oc][ic][27] from conv-transpose W[ic][oc][27].
// 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_eff_kernel(const float* __restrict__ W,
                                 float* __restrict__ E,
                                 int IC, int OC) {
    int oc = blockIdx.x / IC;
    int ic = blockIdx.x % IC;
    int j = threadIdx.x;              // 0..26
    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]];
    E[((long)oc * IC + ic) * 27 + j] = acc * (1.0f / 64.0f);
}

// Fully fused: effective stride-2 conv + BN affine + both avg pools.
// x: [N][IC][32][32][32], E: [OC][IC][27], out: [N][OC][15][15][15]
__global__ void fused_conv_bn_pool_kernel(
        const float* __restrict__ x,
        const float* __restrict__ E,
        const float* __restrict__ bn_w,
        const float* __restrict__ bn_b,
        const float* __restrict__ bn_mean,   // running_mean (eval) or batch mean (train)
        const float* __restrict__ bn_var,    // running_var  (eval) or batch biased var (train)
        const float* __restrict__ conv_bias,
        float* __restrict__ out,
        int N, int IC, int OC, float eps) {
    __shared__ float xs[X_TILE_FLOATS];
    __shared__ float Es[16 * 3 * 27];

    const int n   = blockIdx.z;
    const int row = blockIdx.x * ROWS + (threadIdx.x >> 4);  // 0..224
    const int qx  = threadIdx.x & 15;
    const int tid = threadIdx.x;

    // load effective kernel
    for (int i = tid; i < OC * IC * 27; i += blockDim.x)
        Es[i] = E[i];

    // load input tile: rows of this block -> x[n][ic][2qz+dz][2qy+dy][0..30]
    const int row0 = blockIdx.x * ROWS;
    const int nrows = min(ROWS, TOTAL_ROWS - row0);
    const long xbase = (long)n * IC * 32768;
    for (int idx = tid; idx < nrows * 3 * 3 * 3 * 31; idx += blockDim.x) {
        int xx  = idx % 31;
        int t   = idx / 31;
        int dy  = t % 3;
        int t2  = t / 3;
        int dz  = t2 % 3;
        int ic  = t2 / 3;
        int r   = ic / 3;  ic = ic % 3;
        int qz  = (row0 + r) / 15;
        int qy  = (row0 + r) % 15;
        int gz  = 2 * qz + dz;
        int gy  = 2 * qy + dy;
        xs[idx] = x[xbase + ((long)ic * 32 + gz) * 1024 + gy * 32 + xx];
    }
    __syncthreads();

    if (qx >= 15 || row >= TOTAL_ROWS) return;
    const int qz = row / 15;
    const int qy = row % 15;
    const int r  = threadIdx.x >> 4;

    for (int oc = 0; oc < OC; ++oc) {
        const float* ep = Es + oc * IC * 27;
        float acc = 0.f;
        for (int ic = 0; ic < IC; ++ic) {
            const float* xp = xs + ((r * 3 + ic) * 9) * 31 + 2 * qx;
            #pragma unroll
            for (int j = 0; j < 27; ++j) {
                int jz = j / 9, jy = (j / 3) % 3, jx = j % 3;
                acc += ep[j] * xp[(jz * 3 + jy) * 31 + jx];
            }
            ep += 27;
        }
        float scale = bn_w[oc] * rsqrtf(bn_var[oc] + eps);
        float shift = bn_b[oc] + scale * (conv_bias[oc] - bn_mean[oc]);
        out[(((long)n * OC + oc) * 225 + row) * 15 + qx] = acc * scale + shift;
    }
}

// per-channel sum / sum-of-squares over the full convT output (train mode)
__global__ void bn_stats_kernel(const float* __restrict__ y,
                                double* __restrict__ sum,
                                double* __restrict__ sumsq,
                                long per_channel) {
    const int oc = blockIdx.y;
    const float* base = y + (long)blockIdx.x * gridDim.y * per_channel + (long)oc * per_channel;
    double s = 0.0, s2 = 0.0;
    for (long i = threadIdx.x; i < per_channel; i += blockDim.x) {
        float v = base[i];
        s  += (double)v;
        s2 += (double)v * (double)v;
    }
    __shared__ double bs[256], bs2[256];
    bs[threadIdx.x] = s; bs2[threadIdx.x] = s2;
    __syncthreads();
    for (int off = 128; off > 0; off >>= 1) {
        if (threadIdx.x < off) {
            bs[threadIdx.x]  += bs[threadIdx.x + off];
            bs2[threadIdx.x] += bs2[threadIdx.x + off];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        atomicAdd(&sum[oc], bs[0]);
        atomicAdd(&sumsq[oc], bs2[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, float eps, 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_eff(torch::Tensor W, int64_t IC, int64_t OC) {
    auto E = torch::empty({OC, IC, 27}, W.options());
    dim3 grid(OC * IC);
    build_eff_kernel<<<grid, 32>>>(W.data_ptr<float>(), E.data_ptr<float>(), (int)IC, (int)OC);
    return E;
}

torch::Tensor fused_eval(torch::Tensor x, torch::Tensor E,
                         torch::Tensor bn_w, torch::Tensor bn_b,
                         torch::Tensor bn_mean, torch::Tensor bn_var,
                         torch::Tensor conv_bias,
                         int64_t N, int64_t IC, int64_t OC, double eps) {
    auto out = torch::empty({N, OC, 15, 15, 15}, x.options());
    dim3 grid((TOTAL_ROWS + ROWS - 1) / ROWS, 1, N);
    fused_conv_bn_pool_kernel<<<grid, 128>>>(
        x.data_ptr<float>(), E.data_ptr<float>(),
        bn_w.data_ptr<float>(), bn_b.data_ptr<float>(),
        bn_mean.data_ptr<float>(), bn_var.data_ptr<float>(),
        conv_bias.data_ptr<float>(), out.data_ptr<float>(),
        (int)N, (int)IC, (int)OC, (float)eps);
    return out;
}

// train-mode helper: stats over convT output, then fused pool with batch stats
torch::Tensor fused_train(torch::Tensor x, torch::Tensor E,
                          torch::Tensor y,
                          torch::Tensor bn_w, torch::Tensor bn_b,
                          torch::Tensor running_mean, torch::Tensor running_var,
                          torch::Tensor conv_bias,
                          int64_t N, int64_t IC, int64_t OC,
                          double eps, double momentum) {
    auto opts_d = torch::TensorOptions().dtype(torch::kFloat64).device(x.device());
    auto sum    = torch::zeros({OC}, opts_d);
    auto sumsq  = torch::zeros({OC}, opts_d);
    long per_channel = 250047;  // 63*63*63
    dim3 sgrid(N, OC);
    bn_stats_kernel<<<sgrid, 256>>>(y.data_ptr<float>(),
                                    sum.data_ptr<double>(), sumsq.data_ptr<double>(),
                                    per_channel);
    auto bmean = torch::empty({OC}, x.options());
    auto bvar  = torch::empty({OC}, x.options());
    long count = (long)N * per_channel;
    bn_finalize_kernel<<<(OC + 63) / 64, 64>>>(
        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, (float)eps, (int)OC);

    auto out = torch::empty({N, OC, 15, 15, 15}, x.options());
    dim3 grid((TOTAL_ROWS + ROWS - 1) / ROWS, 1, N);
    fused_conv_bn_pool_kernel<<<grid, 128>>>(
        x.data_ptr<float>(), E.data_ptr<float>(),
        bn_w.data_ptr<float>(), bn_b.data_ptr<float>(),
        bmean.data_ptr<float>(), bvar.data_ptr<float>(),
        conv_bias.data_ptr<float>(), out.data_ptr<float>(),
        (int)N, (int)IC, (int)OC, (float)eps);
    return out;
}
"""

cpp_source = r"""
torch::Tensor build_eff(torch::Tensor W, int64_t IC, int64_t OC);
torch::Tensor fused_eval(torch::Tensor x, torch::Tensor E,
                         torch::Tensor bn_w, torch::Tensor bn_b,
                         torch::Tensor bn_mean, torch::Tensor bn_var,
                         torch::Tensor conv_bias,
                         int64_t N, int64_t IC, int64_t OC, double eps);
torch::Tensor fused_train(torch::Tensor x, torch::Tensor E,
                          torch::Tensor y,
                          torch::Tensor bn_w, torch::Tensor bn_b,
                          torch::Tensor running_mean, torch::Tensor running_var,
                          torch::Tensor conv_bias,
                          int64_t N, int64_t IC, int64_t OC,
                          double eps, double momentum);
"""

_fused = load_inline(
    name="fused_convT_bn_avgpool",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["build_eff", "fused_eval", "fused_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 fused into a single custom CUDA kernel by
    algebraically collapsing the whole chain into one valid stride-2 kernel-3
    convolution with an effective kernel, plus a per-channel affine transform.
    """

    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 forward(self, x):
        x = x.contiguous()
        w = self.conv_transpose.weight.contiguous()
        bn = self.batch_norm
        ic, oc = self._in_channels, self._out_channels
        N = x.shape[0]
        E = _fused.build_eff(w, ic, oc)
        conv_bias = self.conv_transpose.bias.contiguous()

        if not self.training:
            return _fused.fused_eval(
                x, E,
                bn.weight.contiguous(), bn.bias.contiguous(),
                bn.running_mean.contiguous(), bn.running_var.contiguous(),
                conv_bias, N, ic, oc, float(bn.eps),
            )
        else:
            # need batch statistics over the full convT output
            y = F.conv_transpose3d(
                x, w, conv_bias, stride=2, padding=1
            ).contiguous()
            return _fused.fused_train(
                x, E, y,
                bn.weight.contiguous(), bn.bias.contiguous(),
                bn.running_mean, bn.running_var,
                conv_bias, N, ic, oc, float(bn.eps), float(bn.momentum),
            )


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