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

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

// ---------------------------------------------------------------------------
// Fast path: Conv3d with in_channels=3, kernel=3, stride=1, pad=0, dilation=1,
// groups=1, float32, CUDA tensor.
//
// Each thread block computes a (16 z) x (8 y) x (16 x) tile of the output for
// one (batch, out_channel) pair. The corresponding (18 x 10 x 18) input tile
// for all 3 input channels plus the 81 weights are staged in shared memory.
// Each thread computes 8 outputs (2 along z, 4 along x) with full register
// reuse of the sliding convolution window.
// ---------------------------------------------------------------------------

#define FT_Z 16
#define FT_Y 8
#define FT_X 16
#define FIN_Z (FT_Z + 2)
#define FIN_Y (FT_Y + 2)
#define FIN_X (FT_X + 2)
#define TILE_ELEMS (3 * FIN_Z * FIN_Y * FIN_X)

__global__ void __launch_bounds__(256)
conv3d_k3_ic3_kernel(const float* __restrict__ in,
                     const float* __restrict__ wt,
                     const float* __restrict__ bias,
                     float* __restrict__ out,
                     int B, int OC,
                     int D, int H, int W,
                     int OD, int OH, int OW,
                     int tiles_y, int tiles_z) {
    __shared__ float s_in[3][FIN_Z][FIN_Y][FIN_X];
    __shared__ float s_w[81];

    const int tid = (threadIdx.z * blockDim.y + threadIdx.y) * blockDim.x + threadIdx.x;

    const int tx_tile = blockIdx.x;
    const int ty_tile = blockIdx.y % tiles_y;
    const int tz_tile = blockIdx.y / tiles_y;
    const int oc = blockIdx.z % OC;
    const int b  = blockIdx.z / OC;

    const int z0 = tz_tile * FT_Z;
    const int y0 = ty_tile * FT_Y;
    const int x0 = tx_tile * FT_X;

    // Stage weights for this output channel.
    if (tid < 81) s_w[tid] = wt[(long)oc * 81 + tid];

    // Stage input tile (zero-filled outside the input volume).
    const float* ib = in + (long)b * 3 * D * H * W;
    #pragma unroll 4
    for (int i = tid; i < TILE_ELEMS; i += 256) {
        int ic  = i / (FIN_Z * FIN_Y * FIN_X);
        int r   = i % (FIN_Z * FIN_Y * FIN_X);
        int zz  = r / (FIN_Y * FIN_X);
        int r2  = r % (FIN_Y * FIN_X);
        int yy  = r2 / FIN_X;
        int xx  = r2 % FIN_X;
        int gz = z0 + zz, gy = y0 + yy, gx = x0 + xx;
        float v = 0.f;
        if (gz < D && gy < H && gx < W) {
            v = __ldg(&ib[((long)ic * D + gz) * H * W + (long)gy * W + gx]);
        }
        s_in[ic][zz][yy][xx] = v;
    }
    __syncthreads();

    const int tz = threadIdx.z;  // 0..7 -> outputs z0+2*tz, z0+2*tz+1
    const int ty = threadIdx.y;  // 0..7
    const int tx = threadIdx.x;  // 0..3 -> outputs x0+4*tx .. x0+4*tx+3

    const int oz = z0 + tz * 2;
    const int oy = y0 + ty;
    const int ox = x0 + tx * 4;

    if (oz >= OD || oy >= OH || ox >= OW) return;

    float acc[2][4];
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) acc[i][j] = 0.f;
    }

    #pragma unroll
    for (int ic = 0; ic < 3; ++ic) {
        float w[27];
        #pragma unroll
        for (int i = 0; i < 27; ++i) w[i] = s_w[ic * 27 + i];

        #pragma unroll
        for (int ky = 0; ky < 3; ++ky) {
            #pragma unroll
            for (int zs = 0; zs < 4; ++zs) {
                const float* row = &s_in[ic][tz * 2 + zs][ty + ky][tx * 4];
                float v0 = row[0], v1 = row[1], v2 = row[2];
                float v3 = row[3], v4 = row[4], v5 = row[5];
                if (zs < 3) {  // contributes to output z (dz = 0), kz = zs
                    const float* wk = &w[zs * 9 + ky * 3];
                    float w0 = wk[0], w1 = wk[1], w2 = wk[2];
                    acc[0][0] = fmaf(v0, w0, acc[0][0]);
                    acc[0][0] = fmaf(v1, w1, acc[0][0]);
                    acc[0][0] = fmaf(v2, w2, acc[0][0]);
                    acc[0][1] = fmaf(v1, w0, acc[0][1]);
                    acc[0][1] = fmaf(v2, w1, acc[0][1]);
                    acc[0][1] = fmaf(v3, w2, acc[0][1]);
                    acc[0][2] = fmaf(v2, w0, acc[0][2]);
                    acc[0][2] = fmaf(v3, w1, acc[0][2]);
                    acc[0][2] = fmaf(v4, w2, acc[0][2]);
                    acc[0][3] = fmaf(v3, w0, acc[0][3]);
                    acc[0][3] = fmaf(v4, w1, acc[0][3]);
                    acc[0][3] = fmaf(v5, w2, acc[0][3]);
                }
                if (zs > 0) {  // contributes to output z+1 (dz = 1), kz = zs-1
                    const float* wk = &w[(zs - 1) * 9 + ky * 3];
                    float w0 = wk[0], w1 = wk[1], w2 = wk[2];
                    acc[1][0] = fmaf(v0, w0, acc[1][0]);
                    acc[1][0] = fmaf(v1, w1, acc[1][0]);
                    acc[1][0] = fmaf(v2, w2, acc[1][0]);
                    acc[1][1] = fmaf(v1, w0, acc[1][1]);
                    acc[1][1] = fmaf(v2, w1, acc[1][1]);
                    acc[1][1] = fmaf(v3, w2, acc[1][1]);
                    acc[1][2] = fmaf(v2, w0, acc[1][2]);
                    acc[1][2] = fmaf(v3, w1, acc[1][2]);
                    acc[1][2] = fmaf(v4, w2, acc[1][2]);
                    acc[1][3] = fmaf(v3, w0, acc[1][3]);
                    acc[1][3] = fmaf(v4, w1, acc[1][3]);
                    acc[1][3] = fmaf(v5, w2, acc[1][3]);
                }
            }
        }
    }

    if (bias != nullptr) {
        float bv = bias[oc];
        #pragma unroll
        for (int i = 0; i < 2; ++i) {
            #pragma unroll
            for (int j = 0; j < 4; ++j) acc[i][j] += bv;
        }
    }

    #pragma unroll
    for (int dz = 0; dz < 2; ++dz) {
        int z = oz + dz;
        if (z < OD) {
            float* op = out + ((((long)b * OC + oc) * OD + z) * OH + oy) * OW + ox;
            #pragma unroll
            for (int dx = 0; dx < 4; ++dx) {
                if (ox + dx < OW) op[dx] = acc[dz][dx];
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Generic fallback: handles arbitrary channels / kernel / stride / padding /
// dilation / groups. One thread per output element.
// ---------------------------------------------------------------------------
__global__ void conv3d_generic_kernel(const float* __restrict__ in,
                                      const float* __restrict__ wt,
                                      const float* __restrict__ bias,
                                      float* __restrict__ out,
                                      int B, int IC, int OC,
                                      int D, int H, int W,
                                      int OD, int OH, int OW, int K,
                                      int stride, int pad, int dil,
                                      int groups) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long total = (long)B * OC * OD * OH * OW;
    if (idx >= total) return;

    int ox = idx % OW;
    int oy = (idx / OW) % OH;
    int oz = (idx / ((long)OW * OH)) % OD;
    int oc = (idx / ((long)OW * OH * OD)) % OC;
    int b  = idx / ((long)OW * OH * OD * OC);

    int ocpg = OC / groups;
    int icpg = IC / groups;
    int g = oc / ocpg;

    float acc = (bias != nullptr) ? bias[oc] : 0.f;
    const float* ib = in + ((long)b * IC + g * icpg) * D * H * W;
    const float* wb = wt + (long)oc * icpg * K * K * K;

    for (int ic = 0; ic < icpg; ++ic) {
        for (int kz = 0; kz < K; ++kz) {
            long iz = (long)oz * stride - pad + (long)kz * dil;
            if (iz < 0 || iz >= D) continue;
            for (int ky = 0; ky < K; ++ky) {
                long iy = (long)oy * stride - pad + (long)ky * dil;
                if (iy < 0 || iy >= H) continue;
                const float* irow = &ib[(((long)ic * D + iz) * H + iy) * W];
                const float* wrow = &wb[(((long)ic * K + kz) * K + ky) * K];
                long ix0 = (long)ox * stride - pad;
                for (int kx = 0; kx < K; ++kx) {
                    long ix = ix0 + (long)kx * dil;
                    if (ix < 0 || ix >= W) continue;
                    acc = fmaf(__ldg(&irow[ix]), __ldg(&wrow[kx]), acc);
                }
            }
        }
    }
    out[idx] = acc;
}

torch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor weight,
                             torch::Tensor bias, int64_t stride,
                             int64_t pad, int64_t dilation, int64_t groups,
                             bool has_bias) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(x.dim() == 5, "expected 5D input");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    torch::Tensor bc;
    const float* bias_ptr = nullptr;
    if (has_bias) {
        bc = bias.contiguous();
        bias_ptr = bc.data_ptr<float>();
    }

    const int B  = xc.size(0);
    const int IC = xc.size(1);
    const int D  = xc.size(2);
    const int H  = xc.size(3);
    const int W  = xc.size(4);
    const int OC = wc.size(0);
    const int K  = wc.size(2);

    const int OD = (int)((D + 2 * pad - dilation * (K - 1) - 1) / stride + 1);
    const int OH = (int)((H + 2 * pad - dilation * (K - 1) - 1) / stride + 1);
    const int OW = (int)((W + 2 * pad - dilation * (K - 1) - 1) / stride + 1);

    auto out = torch::empty({B, OC, OD, OH, OW}, xc.options());
    if (OD <= 0 || OH <= 0 || OW <= 0) return out;

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

    const bool fast = (K == 3 && IC == 3 && groups == 1 && stride == 1 &&
                       pad == 0 && dilation == 1 && (long)B * OC <= 65535);

    if (fast) {
        static bool configured = []() {
            cudaFuncSetAttribute(conv3d_k3_ic3_kernel,
                                 cudaFuncAttributePreferredSharedMemoryCarveout,
                                 100);
            return true;
        }();
        (void)configured;

        const int tiles_x = (OW + FT_X - 1) / FT_X;
        const int tiles_y = (OH + FT_Y - 1) / FT_Y;
        const int tiles_z = (OD + FT_Z - 1) / FT_Z;
        dim3 grid(tiles_x, tiles_y * tiles_z, B * OC);
        dim3 block(4, 8, 8);
        conv3d_k3_ic3_kernel<<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), B, OC, D, H, W, OD, OH, OW,
            tiles_y, tiles_z);
    } else {
        const long total = (long)B * OC * OD * OH * OW;
        const int threads = 256;
        const long blocks = (total + threads - 1) / threads;
        conv3d_generic_kernel<<<blocks, threads, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), B, IC, OC, D, H, W, OD, OH, OW, K,
            (int)stride, (int)pad, (int)dilation, (int)groups);
    }
    return out;
}
"""

conv3d_cpp_source = (
    "torch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, int64_t stride, int64_t pad, int64_t dilation, "
    "int64_t groups, bool has_bias);"
)

conv3d_ext = load_inline(
    name="conv3d_custom",
    cpp_sources=conv3d_cpp_source,
    cuda_sources=conv3d_cuda_source,
    functions=["conv3d_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_ldflags=[""],
)


class ModelNew(nn.Module):
    """
    Drop-in replacement for Model that runs the 3D convolution with a custom
    CUDA kernel. Falls back to PyTorch's implementation for configurations
    outside the supported set (e.g. CPU tensors or non-float32 dtypes).
    """

    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 forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.is_cuda and x.dtype == torch.float32:
            has_bias = self.conv3d.bias is not None
            bias_arg = self.conv3d.bias if has_bias else x  # dummy when unused
            return conv3d_ext.conv3d_forward(
                x, self.conv3d.weight, bias_arg,
                self.conv3d.stride[0], self.conv3d.padding[0],
                self.conv3d.dilation[0], self.conv3d.groups, has_bias)
        return self.conv3d(x)
```