# ModelNew with Custom CUDA Kernels

Key insight: `BatchNorm(x) - mean_spatial(BatchNorm(x))` simplifies algebraically to `scale_c * (x - mean_spatial(x))` where `scale_c = gamma / sqrt(var_c + eps)` (β cancels exactly). I exploit this with:
1. **Fused ConvTranspose3d + statistics kernel**: register-blocked (2×2×4 output tile/thread) direct conv-transpose (specialized for k=3, s=2, p=1) that simultaneously accumulates per-slice sums and per-channel sum/sumsq via block reductions + atomics.
2. **Tiny finalize kernel**: computes per-channel scale (batch stats in train mode with running-stat update, running stats in eval mode) and per-slice means.
3. **Fused apply kernel**: single pass computing `scale*(x - mean)`.

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

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

// ---------------------------------------------------------------------------
// Kernel A: fused ConvTranspose3d (k=3, s=2, p=1) + statistics accumulation.
// Each thread computes a 2(d) x 2(h) x 4(w) tile of outputs (16 outputs).
// Block-level partial sums -> atomics for:
//   - per (n,c) spatial sum (for the mean subtraction)
//   - per channel sum / sumsq (for batch statistics in training mode)
// ---------------------------------------------------------------------------
__global__ void convt3d_k3s2p1_stats_kernel(
    const float* __restrict__ x,        // [N, IC, ID, IH, IW]
    const float* __restrict__ wpad,     // [IC, OC, 9, 4] (kw padded to 4)
    const float* __restrict__ bias,     // [OC]
    float* __restrict__ y,              // [N, OC, OD, OH, OW]
    double* __restrict__ slice_sum,     // [N*OC]
    double* __restrict__ chan_sum,      // [OC]
    double* __restrict__ chan_sumsq,    // [OC]
    int IC, int OC,
    int ID, int IH, int IW,
    int OD, int OH, int OW,
    int tiles_d, int tiles_h, int tiles_w)
{
    const int slice = blockIdx.y;           // n*OC + oc
    const int n  = slice / OC;
    const int oc = slice - n * OC;

    const int tiles_per_slice = tiles_d * tiles_h * tiles_w;
    const int tile_idx = blockIdx.x * blockDim.x + threadIdx.x;

    float lsum = 0.f;
    float lsq  = 0.f;

    if (tile_idx < tiles_per_slice) {
        const int tw_i = tile_idx % tiles_w;
        const int th_i = (tile_idx / tiles_w) % tiles_h;
        const int td_i = tile_idx / (tiles_w * tiles_h);

        const int id_b = td_i;
        const int ih_b = th_i;
        const int iw_b = 2 * tw_i;

        const int od0 = 2 * td_i;
        const int oh0 = 2 * th_i;
        const int ow0 = 4 * tw_i;

        const float* xn = x + (long)n * IC * ID * IH * IW;
        const long xICs = (long)ID * IH * IW;

        float acc[2][2][4];
        #pragma unroll
        for (int a = 0; a < 2; ++a)
            #pragma unroll
            for (int b = 0; b < 2; ++b)
                #pragma unroll
                for (int c = 0; c < 4; ++c) acc[a][b][c] = 0.f;

        for (int ic = 0; ic < IC; ++ic) {
            const float* xic = xn + (long)ic * xICs;
            const float* wic = wpad + ((long)ic * OC + oc) * 36;

            // load 2x2x3 input tile into registers (zero-filled out of range)
            float xv[2][2][3];
            #pragma unroll
            for (int ad = 0; ad < 2; ++ad) {
                const int id = id_b + ad;
                const int idc = min(id, ID - 1);
                const bool okd = (id < ID);
                #pragma unroll
                for (int ah = 0; ah < 2; ++ah) {
                    const int ih = ih_b + ah;
                    const int ihc = min(ih, IH - 1);
                    const bool okh = okd && (ih < IH);
                    const float* xr = xic + ((long)idc * IH + ihc) * IW;
                    #pragma unroll
                    for (int aw = 0; aw < 3; ++aw) {
                        const int iw = iw_b + aw;
                        const bool okw = okh && (iw < IW);
                        xv[ad][ah][aw] = okw ? xr[min(iw, IW - 1)] : 0.f;
                    }
                }
            }

            // 27 weights (as 9 float4) applied to the tile
            #pragma unroll
            for (int kd = 0; kd < 3; ++kd) {
                const int dm = (kd == 1) ? 0 : 1;
                const int ai = (kd == 0) ? 1 : 0;
                #pragma unroll
                for (int kh = 0; kh < 3; ++kh) {
                    const int hm = (kh == 1) ? 0 : 1;
                    const int hi = (kh == 0) ? 1 : 0;
                    const float4 wv = *reinterpret_cast<const float4*>(wic + (kd * 3 + kh) * 4);
                    acc[dm][hm][0] += xv[ai][hi][0] * wv.y;   // kw = 1
                    acc[dm][hm][2] += xv[ai][hi][1] * wv.y;   // kw = 1
                    acc[dm][hm][1] += xv[ai][hi][0] * wv.z;   // kw = 2
                    acc[dm][hm][3] += xv[ai][hi][1] * wv.z;   // kw = 2
                    acc[dm][hm][1] += xv[ai][hi][1] * wv.x;   // kw = 0
                    acc[dm][hm][3] += xv[ai][hi][2] * wv.x;   // kw = 0
                }
            }
        }

        const float bv = bias[oc];
        const long ybase = (long)slice * OD * OH * OW;
        #pragma unroll
        for (int a = 0; a < 2; ++a) {
            const int od = od0 + a;
            if (od >= OD) continue;
            #pragma unroll
            for (int b = 0; b < 2; ++b) {
                const int oh = oh0 + b;
                if (oh >= OH) continue;
                #pragma unroll
                for (int c = 0; c < 4; ++c) {
                    const int ow = ow0 + c;
                    if (ow >= OW) continue;
                    const float v = acc[a][b][c] + bv;
                    y[ybase + ((long)od * OH + oh) * OW + ow] = v;
                    lsum += v;
                    lsq  += v * v;
                }
            }
        }
    }

    // block reduction of the two partial statistics
    __shared__ double s_sum[256];
    __shared__ double s_sq[256];
    const int tid = threadIdx.x;
    s_sum[tid] = (double)lsum;
    s_sq[tid]  = (double)lsq;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) {
            s_sum[tid] += s_sum[tid + s];
            s_sq[tid]  += s_sq[tid + s];
        }
        __syncthreads();
    }
    if (tid == 0) {
        atomicAdd(&slice_sum[slice], s_sum[0]);
        atomicAdd(&chan_sum[oc],     s_sum[0]);
        atomicAdd(&chan_sumsq[oc],   s_sq[0]);
    }
}

// ---------------------------------------------------------------------------
// Kernel B: finalize statistics -> per-slice means, per-channel scales,
// and (optionally) update running statistics.
// ---------------------------------------------------------------------------
__global__ void finalize_stats_kernel(
    const double* __restrict__ slice_sum,
    const double* __restrict__ chan_sum,
    const double* __restrict__ chan_sumsq,
    const float* __restrict__ gamma,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    float* __restrict__ slice_mean,  // [N*OC]
    float* __restrict__ scale,       // [OC]
    int NSLICE, int OC,
    double inv_spatial, double inv_total, double total,
    float eps, float momentum,
    int training, int update_running)
{
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < NSLICE) {
        slice_mean[i] = (float)((double)slice_sum[i] * inv_spatial);
    }
    if (i < OC) {
        float s;
        if (training) {
            const double mu = (double)chan_sum[i] * inv_total;
            double var = (double)chan_sumsq[i] * inv_total - mu * mu;
            if (var < 0.0) var = 0.0;
            s = gamma[i] * (float)rsqrt(var + (double)eps);
            if (update_running) {
                const float rm = running_mean[i];
                running_mean[i] = (1.f - momentum) * rm + momentum * (float)mu;
                const double unb = var * (total / (total - 1.0));
                const float rv = running_var[i];
                running_var[i] = (1.f - momentum) * rv + momentum * (float)unb;
            }
        } else {
            s = gamma[i] * rsqrtf(running_var[i] + eps);
        }
        scale[i] = s;
    }
}

// ---------------------------------------------------------------------------
// Kernel C: out = scale[c] * (y - slice_mean[n,c])  (fused BN + mean subtract)
// ---------------------------------------------------------------------------
__global__ void apply_scale_mean_kernel(
    const float* __restrict__ y,
    float* __restrict__ out,
    const float* __restrict__ scale,
    const float* __restrict__ slice_mean,
    int OC, int slice_elems)
{
    const int slice = blockIdx.y;
    const int oc = slice % OC;
    const float s = scale[oc];
    const float m = slice_mean[slice];
    const long base = (long)slice * slice_elems;
    for (int i = blockIdx.x * blockDim.x + threadIdx.x;
         i < slice_elems; i += blockDim.x * gridDim.x) {
        out[base + i] = s * (y[base + i] - m);
    }
}

// ---------------------------------------------------------------------------
// Host wrapper
// ---------------------------------------------------------------------------
torch::Tensor fused_convt_bn_meansub(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    torch::Tensor gamma,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    double eps, double momentum,
    bool training, bool track_running, bool has_bias, bool has_affine)
{
    TORCH_CHECK(x.is_cuda() && x.dtype() == torch::kFloat32);
    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto gc = has_affine ? gamma.contiguous()
                         : torch::ones({weight.size(1)}, x.options());

    const int N  = xc.size(0);
    const int IC = xc.size(1);
    const int ID = xc.size(2);
    const int IH = xc.size(3);
    const int IW = xc.size(4);
    const int OC = wc.size(1);

    const int OD = (ID - 1) * 2 - 2 + 3;
    const int OH = (IH - 1) * 2 - 2 + 3;
    const int OW = (IW - 1) * 2 - 2 + 3;

    auto y   = torch::empty({N, OC, OD, OH, OW}, x.options());
    auto out = torch::empty_like(y);

    auto opts_d = torch::TensorOptions().dtype(torch::kFloat64).device(x.device());
    auto slice_sum  = torch::zeros({(long)N * OC}, opts_d);
    auto chan_sum   = torch::zeros({OC}, opts_d);
    auto chan_sumsq = torch::zeros({OC}, opts_d);
    auto slice_mean = torch::empty({(long)N * OC}, x.options());
    auto scale      = torch::empty({OC}, x.options());

    auto wpad = torch::constant_pad_nd(wc.view({IC, OC, 9, 3}), {0, 1}, 0.0)
                    .contiguous();

    auto bc = has_bias ? bias.contiguous()
                       : torch::zeros({OC}, x.options());

    auto stream = at::cuda::getCurrentCUDAStream();

    // Kernel A
    const int tiles_d = (OD + 1) / 2;
    const int tiles_h = (OH + 1) / 2;
    const int tiles_w = (OW + 3) / 4;
    const int tiles_per_slice = tiles_d * tiles_h * tiles_w;
    const int threads = 256;
    const int blocks_x = (tiles_per_slice + threads - 1) / threads;
    dim3 gridA(blocks_x, N * OC);
    convt3d_k3s2p1_stats_kernel<<<gridA, threads, 0, stream>>>(
        xc.data_ptr<float>(), wpad.data_ptr<float>(), bc.data_ptr<float>(),
        y.data_ptr<float>(),
        slice_sum.data_ptr<double>(), chan_sum.data_ptr<double>(),
        chan_sumsq.data_ptr<double>(),
        IC, OC, ID, IH, IW, OD, OH, OW,
        tiles_d, tiles_h, tiles_w);

    // Kernel B
    const int NSLICE = N * OC;
    const double spatial = (double)OD * OH * OW;
    const double total = spatial * (double)N;
    const int work = NSLICE > OC ? NSLICE : OC;
    const int blocksB = (work + 255) / 256;
    const int update_running = (training && track_running) ? 1 : 0;
    finalize_stats_kernel<<<blocksB, 256, 0, stream>>>(
        slice_sum.data_ptr<double>(), chan_sum.data_ptr<double>(),
        chan_sumsq.data_ptr<double>(),
        gc.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        slice_mean.data_ptr<float>(), scale.data_ptr<float>(),
        NSLICE, OC, 1.0 / spatial, 1.0 / total, total,
        (float)eps, (float)momentum, training ? 1 : 0, update_running);

    // Kernel C
    const int slice_elems = OD * OH * OW;
    int cx = (slice_elems + threads - 1) / threads;
    if (cx > 65535) cx = 65535;
    dim3 gridC(cx, NSLICE);
    apply_scale_mean_kernel<<<gridC, threads, 0, stream>>>(
        y.data_ptr<float>(), out.data_ptr<float>(),
        scale.data_ptr<float>(), slice_mean.data_ptr<float>(),
        OC, slice_elems);

    return out;
}
"""

cpp_source = """
torch::Tensor fused_convt_bn_meansub(
    torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
    torch::Tensor gamma, torch::Tensor running_mean, torch::Tensor running_var,
    double eps, double momentum,
    bool training, bool track_running, bool has_bias, bool has_affine);
"""

fused_op = load_inline(
    name="fused_convt_bn_meansub",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_convt_bn_meansub"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    ConvTranspose3d + BatchNorm3d + spatial-mean subtraction, fused into
    custom CUDA kernels. Falls back to native PyTorch ops for unsupported
    configurations.
    """

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

        ct = self.conv_transpose
        self._use_fused = (
            ct.kernel_size == (3, 3, 3)
            and ct.stride == (2, 2, 2)
            and ct.padding == (1, 1, 1)
            and ct.dilation == (1, 1, 1)
            and ct.groups == 1
            and ct.output_padding == (0, 0, 0)
            and ct.padding_mode == "zeros"
        )

    def forward(self, x):
        bn = self.batch_norm
        if (self._use_fused and x.is_cuda and x.dtype == torch.float32):
            ct = self.conv_transpose
            track = bn.track_running_stats
            if not track and not bn.training:
                # no running stats available in eval mode -> native fallback
                pass
            else:
                rm = bn.running_mean if track else torch.zeros(
                    bn.num_features, device=x.device)
                rv = bn.running_var if track else torch.ones(
                    bn.num_features, device=x.device)
                return fused_op.fused_convt_bn_meansub(
                    x, ct.weight,
                    ct.bias if ct.bias is not None else torch.empty(0, device=x.device),
                    bn.weight if bn.affine else torch.empty(0, device=x.device),
                    rm, rv,
                    float(bn.eps), float(bn.momentum),
                    bool(bn.training), bool(track),
                    ct.bias is not None, bool(bn.affine))

        # fallback: original PyTorch path
        x = self.conv_transpose(x)
        x = self.batch_norm(x)
        x = x - torch.mean(x, dim=(2, 3, 4), keepdim=True)
        return x
```