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

# Custom fused CUDA kernel for direct 3D convolution.
# Fast path: kernel_size == 3, stride == 1, dilation == 1, groups == 1,
# float32 CUDA tensors, in_channels == 3.
#
# Strategy:
#   * Thread block of 256 threads computes a spatial tile of 32(W) x 8(H) x 4(D)
#     voxels for 4 output channels of one batch element.
#   * The (haloed) input tile and the 4x81 weights are staged in shared memory.
#   * Each thread accumulates 4 output channels x 4 adjacent voxels in registers
#     (16 FMAs per tap), reusing input registers across kernel positions.
#   * Results are re-staged through shared memory so global stores are fully
#     coalesced, and the bias add is fused into the epilogue.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define THREAD_W 8
#define TILE_H 8
#define TILE_D 4
#define VOX_W 4
#define CO_TILE 4
#define KSIZE 3

#define TILE_W (THREAD_W * VOX_W)   // 32
#define IN_W (TILE_W + KSIZE - 1)   // 34
#define IN_H (TILE_H + KSIZE - 1)   // 10
#define IN_D (TILE_D + KSIZE - 1)   // 6
#define NTHREADS (THREAD_W * TILE_H * TILE_D)  // 256
#define OUT_TILE (CO_TILE * TILE_D * TILE_H * TILE_W)  // 4096

template <int C_IN>
__global__ void __launch_bounds__(NTHREADS)
conv3d_k3_s1_kernel(const float* __restrict__ x,
                    const float* __restrict__ w,
                    const float* __restrict__ b,
                    float* __restrict__ y,
                    int N, int C_OUT,
                    int D, int H, int W,
                    int OD, int OH, int OW,
                    int pad, int numDTiles) {
    constexpr int TAPS = C_IN * KSIZE * KSIZE * KSIZE;
    constexpr int IN_TILE = C_IN * IN_D * IN_H * IN_W;

    extern __shared__ float smem[];
    float* s_in = smem;               // IN_TILE floats
    float* s_w  = smem + IN_TILE;     // CO_TILE * TAPS floats

    const int tx = threadIdx.x;  // 0..7
    const int ty = threadIdx.y;  // 0..7
    const int tz = threadIdx.z;  // 0..3
    const int tid = (tz * TILE_H + ty) * THREAD_W + tx;

    const int ow0 = blockIdx.x * TILE_W;
    const int oh0 = blockIdx.y * TILE_H;
    const int dz  = blockIdx.z % numDTiles;
    const int tmp = blockIdx.z / numDTiles;
    const int n   = tmp % N;
    const int cg  = tmp / N;
    const int od0 = dz * TILE_D;
    const int co0 = cg * CO_TILE;

    // ---- stage weights for the 4 output channels into shared memory ----
    for (int i = tid; i < CO_TILE * TAPS; i += NTHREADS) {
        const int c = i / TAPS;
        const int co = co0 + c;
        s_w[i] = (co < C_OUT) ? __ldg(w + (long)co * TAPS + (i - c * TAPS)) : 0.0f;
    }

    // ---- stage the haloed input tile into shared memory (zero padded) ----
    const long x_base = (long)n * C_IN * D * H * W;
    for (int i = tid; i < IN_TILE; i += NTHREADS) {
        const int iw = i % IN_W;
        const int ih = (i / IN_W) % IN_H;
        const int id = (i / (IN_W * IN_H)) % IN_D;
        const int ci = i / (IN_W * IN_H * IN_D);
        const int gd = od0 + id - pad;
        const int gh = oh0 + ih - pad;
        const int gw = ow0 + iw - pad;
        float v = 0.0f;
        if ((unsigned)gd < (unsigned)D && (unsigned)gh < (unsigned)H &&
            (unsigned)gw < (unsigned)W) {
            v = __ldg(x + x_base + ((long)ci * D + gd) * H * W + (long)gh * W + gw);
        }
        s_in[i] = v;
    }
    __syncthreads();

    // ---- compute: 4 output channels x 4 adjacent voxels per thread ----
    float acc[CO_TILE][VOX_W];
    #pragma unroll
    for (int c = 0; c < CO_TILE; ++c)
        #pragma unroll
        for (int v = 0; v < VOX_W; ++v)
            acc[c][v] = 0.0f;

    #pragma unroll
    for (int ci = 0; ci < C_IN; ++ci) {
        const float* s_in_c = s_in + ci * (IN_D * IN_H * IN_W);
        #pragma unroll
        for (int kd = 0; kd < KSIZE; ++kd) {
            #pragma unroll
            for (int kh = 0; kh < KSIZE; ++kh) {
                const float* row = s_in_c + (tz + kd) * (IN_H * IN_W) +
                                   (ty + kh) * IN_W + tx * VOX_W;
                float r[VOX_W + KSIZE - 1];
                #pragma unroll
                for (int j = 0; j < VOX_W + KSIZE - 1; ++j) r[j] = row[j];
                #pragma unroll
                for (int kw = 0; kw < KSIZE; ++kw) {
                    const int t = ((ci * KSIZE + kd) * KSIZE + kh) * KSIZE + kw;
                    #pragma unroll
                    for (int c = 0; c < CO_TILE; ++c) {
                        const float wv = s_w[c * TAPS + t];
                        #pragma unroll
                        for (int v = 0; v < VOX_W; ++v)
                            acc[c][v] = fmaf(r[v + kw], wv, acc[c][v]);
                    }
                }
            }
        }
    }

    // ---- re-stage through shared memory for coalesced global stores ----
    __syncthreads();
    float* s_out = s_in;  // reuse (IN_TILE >= OUT_TILE)
    #pragma unroll
    for (int c = 0; c < CO_TILE; ++c)
        #pragma unroll
        for (int v = 0; v < VOX_W; ++v)
            s_out[((c * TILE_D + tz) * TILE_H + ty) * TILE_W + tx * VOX_W + v] =
                acc[c][v];
    __syncthreads();

    #pragma unroll
    for (int i = 0; i < OUT_TILE / NTHREADS; ++i) {
        const int idx = tid + i * NTHREADS;
        const int ww = idx & (TILE_W - 1);
        const int yy = (idx >> 5) & (TILE_H - 1);
        const int zz = (idx >> 8) & (TILE_D - 1);
        const int cc = idx >> 10;
        const int co = co0 + cc;
        const int od = od0 + zz;
        const int oh = oh0 + yy;
        const int ow = ow0 + ww;
        if (co < C_OUT && od < OD && oh < OH && ow < OW) {
            const float bv = (b != nullptr) ? __ldg(b + co) : 0.0f;
            y[((((long)n * C_OUT + co) * OD + od) * OH + oh) * OW + ow] =
                s_out[idx] + bv;
        }
    }
}

torch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor w,
                             torch::Tensor bias, int64_t pad) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");

    const int N = x.size(0);
    const int C_IN = x.size(1);
    const int D = x.size(2);
    const int H = x.size(3);
    const int W = x.size(4);
    const int C_OUT = w.size(0);
    const int p = (int)pad;

    const int OD = D + 2 * p - (KSIZE - 1);
    const int OH = H + 2 * p - (KSIZE - 1);
    const int OW = W + 2 * p - (KSIZE - 1);
    TORCH_CHECK(OD > 0 && OH > 0 && OW > 0, "invalid output size");

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

    const int numDTiles = (OD + TILE_D - 1) / TILE_D;
    dim3 grid((OW + TILE_W - 1) / TILE_W,
              (OH + TILE_H - 1) / TILE_H,
              numDTiles * N * ((C_OUT + CO_TILE - 1) / CO_TILE));
    TORCH_CHECK(grid.z <= 65535, "grid.z too large");
    dim3 block(THREAD_W, TILE_H, TILE_D);

    const size_t shmem =
        (size_t)(C_IN * IN_D * IN_H * IN_W + CO_TILE * C_IN * 27) * sizeof(float);
    const float* bptr =
        (bias.defined() && bias.numel() > 0) ? bias.data_ptr<float>() : nullptr;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    TORCH_CHECK(C_IN == 3, "unsupported in_channels for fast kernel");
    conv3d_k3_s1_kernel<3><<<grid, block, shmem, stream>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bptr, y.data_ptr<float>(),
        N, C_OUT, D, H, W, OD, OH, OW, p, numDTiles);

    return y;
}
"""

cpp_source = "torch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias, int64_t pad);"

_conv3d_ext = load_inline(
    name="conv3d_direct_k3",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv3d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class _Conv3dFast(torch.autograd.Function):
    """Autograd wrapper around the custom CUDA kernel (backward via ATen)."""

    @staticmethod
    def forward(ctx, x, w, b, pad):
        ctx.save_for_backward(x, w, b if b is not None else torch.empty(0, device=x.device, dtype=x.dtype))
        ctx.pad = pad
        ctx.has_bias = b is not None
        b_arg = b if b is not None else torch.empty(0, device=x.device, dtype=x.dtype)
        return _conv3d_ext.conv3d_forward(x, w, b_arg, pad)

    @staticmethod
    def backward(ctx, grad_out):
        x, w, b = ctx.saved_tensors
        grad_out = grad_out.contiguous()
        gx = gw = gb = None
        if ctx.needs_input_grad[0]:
            gx = torch.nn.grad.conv3d_input(x.shape, w, grad_out, stride=1, padding=ctx.pad)
        if ctx.needs_input_grad[1]:
            gw = torch.nn.grad.conv3d_weight(x, w.shape, grad_out, stride=1, padding=ctx.pad)
        if ctx.has_bias and ctx.needs_input_grad[2]:
            gb = grad_out.sum(dim=(0, 2, 3, 4))
        return gx, gw, gb, None


class ModelNew(nn.Module):
    """
    3D convolution accelerated with a custom fused CUDA kernel.

    The custom kernel handles the common configuration (3x3x3 kernel, unit
    stride/dilation, no groups, 3 input channels, float32 CUDA tensors);
    every other configuration transparently falls back to cuDNN so the module
    remains fully general.
    """

    def __init__(self, in_channels: int, out_channels: int, kernel_size: int,
                 stride: int = 1, padding: int = 0, dilation: int = 1,
                 groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv3d = nn.Conv3d(
            in_channels, out_channels,
            (kernel_size, kernel_size, kernel_size),
            stride=stride, padding=padding, dilation=dilation,
            groups=groups, bias=bias)

    def _use_fast_path(self, x: torch.Tensor) -> bool:
        if not (x.is_cuda and x.dtype == torch.float32):
            return False
        c = self.conv3d
        if c.groups != 1 or c.in_channels != 3:
            return False
        if c.kernel_size != (3, 3, 3) or c.stride != (1, 1, 1) or c.dilation != (1, 1, 1):
            return False
        N, _, D, H, W = x.shape
        p = c.padding[0]
        OD, OH, OW = D + 2 * p - 2, H + 2 * p - 2, W + 2 * p - 2
        if OD <= 0 or OH <= 0 or OW <= 0:
            return False
        grid_z = ((OD + 3) // 4) * N * ((c.out_channels + 3) // 4)
        return grid_z <= 65535

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._use_fast_path(x):
            return _Conv3dFast.apply(
                x.contiguous(),
                self.conv3d.weight.contiguous(),
                self.conv3d.bias,
                self.conv3d.padding[0])
        return self.conv3d(x)
```