# Optimized Model with Custom CUDA Kernel

Key insight: `ConvTranspose3d` can be computed directly with a gather-based kernel where threads are linearized **by output phase** (coordinate mod stride). This makes all warps phase-uniform: threads with identical kernel-offset lists execute together, phases with empty contribution lists exit immediately (huge win for stride>1), and memory access stays coalesced. Combined with register tiling (8 output channels × 4 spatial positions per thread) and broadcast weight reads, this approaches the fp32 compute roofline.

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


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

#define CO_BLOCK 8
#define W_TILE 4
#define THREADS 256
#define MAX_K 31

// Dense-phase direct transposed 3D convolution.
// Threads are linearized so that all threads in a warp share the same output
// phase ((d,h,w) mod stride), i.e. the same set of contributing kernel taps.
// Phases with no contributing taps exit immediately (output was pre-zeroed).
__global__ void conv_transpose3d_phase_kernel(
    const float* __restrict__ x,   // (N, Ci, Di, Hi, Wi)
    const float* __restrict__ wgt, // (Ci, Co, Kd, Kh, Kw)
    float* __restrict__ y,         // (N, Co, Do, Ho, Wo)
    int N, int Ci, int Co,
    int Di, int Hi, int Wi,
    int Do, int Ho, int Wo,
    int Kd, int Kh, int Kw,
    int sd, int sh, int sw,
    int pd_, int ph_, int pw_,
    int dd, int dh, int dw,
    int Td, int Th, int Tw,
    int Tw2, int cg)
{
    // ---- decode phase (uniform per block) ----
    int phase = blockIdx.y;
    int pw = phase % sw; phase /= sw;
    int ph = phase % sh; phase /= sh;
    int pd = phase % sd;

    // ---- build per-dimension kernel-tap lists ----
    int kdl[MAX_K], khl[MAX_K], kwl[MAX_K];
    int od0[MAX_K], oh0[MAX_K], ow0[MAX_K];
    int cd = 0, ch = 0, cw = 0;
    for (int kd = 0; kd < Kd; ++kd) {
        int v = pd + pd_ - kd * dd;
        if (v % sd == 0) { kdl[cd] = kd; od0[cd] = v / sd; ++cd; }
    }
    for (int kh = 0; kh < Kh; ++kh) {
        int v = ph + ph_ - kh * dh;
        if (v % sh == 0) { khl[ch] = kh; oh0[ch] = v / sh; ++ch; }
    }
    for (int kw = 0; kw < Kw; ++kw) {
        int v = pw + pw_ - kw * dw;
        if (v % sw == 0) { kwl[cw] = kw; ow0[cw] = v / sw; ++cw; }
    }
    if (cd == 0 || ch == 0 || cw == 0) return;  // dead phase: stays zero

    // ---- decode thread -> (co group, n, td, th, tw2) ----
    long long tid = (long long)blockIdx.x * THREADS + threadIdx.x;
    int tw2 = (int)(tid % Tw2); tid /= Tw2;
    int th  = (int)(tid % Th);  tid /= Th;
    int td  = (int)(tid % Td);  tid /= Td;
    int n   = (int)(tid % N);   tid /= N;
    int cgg = (int)tid;
    if (cgg >= cg) return;

    int co0 = cgg * CO_BLOCK;
    if (co0 >= Co) return;
    int nb = Co - co0; if (nb > CO_BLOCK) nb = CO_BLOCK;

    float acc[CO_BLOCK][W_TILE];
    #pragma unroll
    for (int c = 0; c < CO_BLOCK; ++c)
        #pragma unroll
        for (int j = 0; j < W_TILE; ++j) acc[c][j] = 0.f;

    const int KHW = Kh * Kw;
    const int HiWi = Hi * Wi;

    // ---- main accumulation ----
    for (int ai = 0; ai < cd; ++ai) {
        int od = td + od0[ai];
        if (od < 0 || od >= Di) continue;
        int kd = kdl[ai];
        for (int bi = 0; bi < ch; ++bi) {
            int oh = th + oh0[bi];
            if (oh < 0 || oh >= Hi) continue;
            int kh = khl[bi];
            long long in_nrow = ((long long)n * Ci * Di + (long long)od * Hi + oh) * Wi;
            int wk_base = kd * KHW + kh * Kw;
            for (int ci = 0; ci < Ci; ++ci) {
                const float* in_ci = x + in_nrow + (long long)ci * Di * HiWi;
                const float* wp = wgt + ((long long)ci * Co + co0) * Kd * KHW + wk_base;
                for (int ci2 = 0; ci2 < cw; ++ci2) {
                    int base = ow0[ci2];
                    int kw = kwl[ci2];
                    float iv[W_TILE];
                    #pragma unroll
                    for (int j = 0; j < W_TILE; ++j) {
                        int ow = (tw2 * W_TILE + j) + base;
                        iv[j] = (ow >= 0 && ow < Wi) ? in_ci[ow] : 0.f;
                    }
                    #pragma unroll
                    for (int c = 0; c < CO_BLOCK; ++c) {
                        float wv = wp[(long long)c * KHW + kw];
                        #pragma unroll
                        for (int j = 0; j < W_TILE; ++j)
                            acc[c][j] = fmaf(iv[j], wv, acc[c][j]);
                    }
                }
            }
        }
    }

    // ---- write back ----
    const int HoWo = Ho * Wo;
    long long yrow = ((long long)n * Co * Do + (long long)(td * sd + pd) * Ho
                      + (long long)(th * sh + ph)) * Wo + pw;
    #pragma unroll
    for (int c = 0; c < CO_BLOCK; ++c) {
        if (c < nb) {
            float* yp = y + yrow + (long long)c * Do * HoWo;
            #pragma unroll
            for (int j = 0; j < W_TILE; ++j) {
                int t = tw2 * W_TILE + j;
                if (t < Tw) yp[(long long)t * sw] = acc[c][j];
            }
        }
    }
}

// Fill output with bias values (used when bias is present, since dead phases
// are never touched by the compute kernel).
__global__ void fill_bias_kernel(float* __restrict__ y,
                                 const float* __restrict__ bias,
                                 long long total, int Co, long long sHW)
{
    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < total) {
        int co = (int)((i / sHW) % Co);
        y[i] = bias[co];
    }
}

at::Tensor conv_transpose3d_forward(
    at::Tensor x, at::Tensor wgt, c10::optional<at::Tensor> bias_opt,
    int64_t stride, int64_t pad, int64_t dil)
{
    TORCH_CHECK(x.is_cuda() && wgt.is_cuda(), "inputs must be CUDA");
    auto xc = x.contiguous();
    auto wc = wgt.contiguous();

    const int N  = (int)xc.size(0);
    const int Ci = (int)xc.size(1);
    const int Di = (int)xc.size(2);
    const int Hi = (int)xc.size(3);
    const int Wi = (int)xc.size(4);
    const int Co = (int)wc.size(1);
    const int Kd = (int)wc.size(2);
    const int Kh = (int)wc.size(3);
    const int Kw = (int)wc.size(4);
    const int s = (int)stride, p = (int)pad, d = (int)dil;

    TORCH_CHECK(Kd <= MAX_K && Kh <= MAX_K && Kw <= MAX_K, "kernel too large");

    const int Do = (Di - 1) * s - 2 * p + d * (Kd - 1) + 1;
    const int Ho = (Hi - 1) * s - 2 * p + d * (Kh - 1) + 1;
    const int Wo = (Wi - 1) * s - 2 * p + d * (Kw - 1) + 1;

    auto y = at::empty({N, Co, Do, Ho, Wo}, xc.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    long long total = (long long)N * Co * Do * Ho * Wo;
    if (total == 0) return y;

    const bool has_bias = bias_opt.has_value() && bias_opt->defined();
    at::Tensor bc;
    if (has_bias) {
        bc = bias_opt->contiguous();
        long long sHW = (long long)Do * Ho * Wo;
        int threads = 256;
        long long blocks = (total + threads - 1) / threads;
        fill_bias_kernel<<<(unsigned)blocks, threads, 0, stream>>>(
            y.data_ptr<float>(), bc.data_ptr<float>(), total, Co, sHW);
    } else {
        cudaMemsetAsync(y.data_ptr<float>(), 0, (size_t)total * sizeof(float), stream);
    }

    const int Td = (Do + s - 1) / s;
    const int Th = (Ho + s - 1) / s;
    const int Tw = (Wo + s - 1) / s;
    const int Tw2 = (Tw + W_TILE - 1) / W_TILE;
    const int cg = (Co + CO_BLOCK - 1) / CO_BLOCK;

    long long per_phase = (long long)N * Td * Th * Tw2 * cg;
    int bx = (int)((per_phase + THREADS - 1) / THREADS);
    int phases = s * s * s;
    dim3 grid(bx, phases, 1);

    conv_transpose3d_phase_kernel<<<grid, THREADS, 0, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), y.data_ptr<float>(),
        N, Ci, Co, Di, Hi, Wi, Do, Ho, Wo, Kd, Kh, Kw,
        s, s, s, p, p, p, d, d, d,
        Td, Th, Tw, Tw2, cg);

    return y;
}
"""

_cpp_src = r"""
#include <torch/extension.h>
at::Tensor conv_transpose3d_forward(
    at::Tensor x, at::Tensor wgt, c10::optional<at::Tensor> bias_opt,
    int64_t stride, int64_t pad, int64_t dil);
"""

_mod = load_inline(
    name="conv_transpose3d_phase",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["conv_transpose3d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int,
                 stride: int = 1, padding: int = 0, dilation: int = 1,
                 bias: bool = False):
        super(ModelNew, self).__init__()
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.conv_transpose3d = nn.ConvTranspose3d(
            in_channels, out_channels,
            kernel_size=(kernel_size, kernel_size, kernel_size),
            stride=stride, padding=padding, dilation=dilation, bias=bias)
        self._fallback = (
            kernel_size > 31 or stride < 1
            or self.conv_transpose3d.groups != 1
            or self.conv_transpose3d.output_padding != (0, 0, 0)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (not self._fallback) and x.is_cuda and x.dtype == torch.float32:
            return _mod.conv_transpose3d_forward(
                x, self.conv_transpose3d.weight, self.conv_transpose3d.bias,
                self.stride, self.padding, self.dilation)
        return self.conv_transpose3d(x)
```

**Optimization summary:**
- **Phase-dense linearization**: output positions are grouped by `(d,h,w) mod stride`. Every warp shares one kernel-tap list → zero divergence from stride sparsity, and the 7/8 of phases that are identically zero (for stride 2, dilation 2, padding 1) are skipped entirely at full memory bandwidth via `cudaMemsetAsync`.
- **Register tiling**: each thread accumulates 8 output channels × 4 spatial positions (32 accumulators), so each loaded weight feeds 4 FMAs and each input value feeds 8 FMAs, keeping the kernel FMA-pipe bound rather than LSU bound.
- **Coalescing & broadcasts**: consecutive threads read consecutive input voxels; weight reads are warp-uniform broadcasts served from L1.
- **Generality**: any stride/padding/dilation/kernel ≤ 31 and optional bias are handled, with a safe fallback to the native op for exotic configurations (groups, output padding, non-fp32).