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

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

// ---------------------------------------------------------------------------
// Tensor-core (TF32) fast path for Conv3d with:
//   in_channels = 3, out_channels = 64, kernel = 3, stride = 1, pad = 0,
//   dilation = 1, groups = 1, float32 CUDA tensors.
//
// Implicit GEMM:  C[oc, p] = sum_k W[oc, k] * X[k, p]
//   M = 64 (all out channels), K = 81 real taps padded to 96,
//   N = output pixels of one (batch, oz) plane (oy*OW+ox flattened).
//
// Each block computes a 64 x 128 tile. Weights (64 x 96, tf32) are staged
// once. The im2col B panel is streamed in K-slices of 16 through a 3-stage
// cp.async pipeline. Block = 256 threads (8 warps, 2m x 4n); each warp owns
// a 32 x 32 tile = 2x4 mma.m16n8k8 ops. tf32 rounding (rna) is applied to A
// at stage time and to B fragments at load time, matching cuDNN TF32
// numerics closely.
// ---------------------------------------------------------------------------

#define KREAL 81
#define KPAD 96
#define ASTRIDE 100
#define NTILE 128
#define BSTRIDE 136
#define KSLICE 16
#define STAGES 3
#define NSLICE (KPAD / KSLICE)
#define SMEM_BYTES (64 * ASTRIDE * 4 + STAGES * KSLICE * BSTRIDE * 4)

__device__ __forceinline__ float to_tf32(float v) {
    uint32_t u;
    asm("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(u) : "f"(v));
    return __uint_as_float(u);
}

__device__ __forceinline__ uint32_t smem_u32(const void* p) {
    return (uint32_t)__cvta_generic_to_shared(p);
}

__device__ __forceinline__ void ldmatrix_x4(uint32_t& r0, uint32_t& r1,
                                            uint32_t& r2, uint32_t& r3,
                                            uint32_t addr) {
    asm volatile(
        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
        : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
        : "r"(addr));
}

__device__ __forceinline__ void mma_tf32(float& c0, float& c1, float& c2,
                                         float& c3, uint32_t a0, uint32_t a1,
                                         uint32_t a2, uint32_t a3, uint32_t b0,
                                         uint32_t b1) {
    asm volatile(
        "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
        : "+f"(c0), "+f"(c1), "+f"(c2), "+f"(c3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}

__device__ __forceinline__ void cp_async4(void* sdst, const float* gsrc,
                                          int src_bytes) {
    asm volatile("cp.async.ca.shared.global [%0], [%1], 4, %2;\n" ::"r"(
                     smem_u32(sdst)),
                 "l"(gsrc), "r"(src_bytes));
}

__device__ __forceinline__ void cp_commit() {
    asm volatile("cp.async.commit_group;\n" ::);
}

template <int N>
__device__ __forceinline__ void cp_wait() {
    asm volatile("cp.async.wait_group %0;\n" ::"n"(N));
}

extern __shared__ float smem[];

__device__ __forceinline__ void stage_slice(float* Bs, const float* in_b,
                                            int s, int oz, int oy, int ox,
                                            bool n_ok, int kh, int n_col,
                                            int D, int H, int W) {
    float* buf = Bs + (s % STAGES) * (KSLICE * BSTRIDE);
    #pragma unroll
    for (int j = 0; j < 8; ++j) {
        const int kk = 2 * j + kh;        // slice-local k, 0..15
        const int k = s * KSLICE + kk;    // global k, 0..95
        const int ic = k / 27;
        const int r = k - ic * 27;
        const int kz = r / 9;
        const int rr = r - kz * 9;
        const int ky = rr / 3;
        const int kx = rr - ky * 3;
        const bool ok = n_ok && (k < KREAL);
        const float* src =
            in_b + (((long)(ic * D + oz + kz) * H) + (oy + ky)) * W + ox + kx;
        cp_async4(&buf[kk * BSTRIDE + n_col], src, ok ? 4 : 0);
    }
}

__global__ void __launch_bounds__(256)
conv3d_tc_kernel(const float* __restrict__ in, const float* __restrict__ wt,
                 const float* __restrict__ bias, float* __restrict__ out,
                 int B, int D, int H, int W, int OD, int OH, int OW) {
    float* As = smem;                    // [64][ASTRIDE]
    float* Bs = smem + 64 * ASTRIDE;     // [STAGES][KSLICE][BSTRIDE]

    const int plane  = OH * OW;
    const int n_base = blockIdx.x * NTILE;
    const int oz     = blockIdx.y % OD;
    const int b      = blockIdx.y / OD;
    const int rem    = plane - n_base;   // >= 1

    const int tid    = threadIdx.x;
    const int lane   = tid & 31;
    const int warp   = tid >> 5;
    const int warp_m = warp >> 2;        // 0..1
    const int warp_n = warp & 3;         // 0..3

    // ---- stage weights (64 x 96), converting to tf32 ----
    for (int i = tid; i < 64 * KPAD; i += 256) {
        int oc = i / KPAD;
        int k  = i - oc * KPAD;
        float v = (k < KREAL) ? __ldg(&wt[oc * KREAL + k]) : 0.f;
        As[oc * ASTRIDE + k] = to_tf32(v);
    }

    // ---- per-thread staging coordinates (fixed across slices) ----
    const int n_col = tid & (NTILE - 1);
    const int kh    = tid >> 7;          // 0..1
    const int p0    = n_base + n_col;
    const bool n_ok = p0 < plane;
    const int oy    = p0 / OW;
    const int ox    = p0 - oy * OW;
    const float* in_b = in + (long)b * 3 * D * H * W;

    // ---- pipeline prologue ----
    stage_slice(Bs, in_b, 0, oz, oy, ox, n_ok, kh, n_col, D, H, W);
    cp_commit();
    stage_slice(Bs, in_b, 1, oz, oy, ox, n_ok, kh, n_col, D, H, W);
    cp_commit();
    __syncthreads();

    const int g   = lane >> 2;           // 0..7
    const int tig = lane & 3;            // 0..3

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

    const int a_row0 = warp_m * 32 + ((lane >> 3) & 1) * 8 + (lane & 7);
    const int a_colf = (lane >> 4) * 4;

    #pragma unroll
    for (int s = 0; s < NSLICE; ++s) {
        if (s + 1 < NSLICE) cp_wait<1>(); else cp_wait<0>();
        __syncthreads();
        if (s + 2 < NSLICE) {
            stage_slice(Bs, in_b, s + 2, oz, oy, ox, n_ok, kh, n_col, D, H, W);
            cp_commit();
        }
        const float* buf = Bs + (s % STAGES) * (KSLICE * BSTRIDE);
        #pragma unroll
        for (int kl = 0; kl < KSLICE / 8; ++kl) {
            const int k0 = s * KSLICE + kl * 8;
            uint32_t a[2][4];
            #pragma unroll
            for (int mt = 0; mt < 2; ++mt) {
                uint32_t addr = smem_u32(
                    &As[(a_row0 + mt * 16) * ASTRIDE + k0 + a_colf]);
                ldmatrix_x4(a[mt][0], a[mt][1], a[mt][2], a[mt][3], addr);
            }
            const int kr = kl * 8;
            #pragma unroll
            for (int nt = 0; nt < 4; ++nt) {
                const int nb = warp_n * 32 + nt * 8 + g;
                const uint32_t b0 = __float_as_uint(
                    to_tf32(buf[(kr + tig) * BSTRIDE + nb]));
                const uint32_t b1 = __float_as_uint(
                    to_tf32(buf[(kr + tig + 4) * BSTRIDE + nb]));
                #pragma unroll
                for (int mt = 0; mt < 2; ++mt) {
                    mma_tf32(acc[mt][nt][0], acc[mt][nt][1], acc[mt][nt][2],
                             acc[mt][nt][3], a[mt][0], a[mt][1], a[mt][2],
                             a[mt][3], b0, b1);
                }
            }
        }
    }

    // ---- epilogue: bias + coalesced float2 stores (plane-flattened) ----
    const long planeL = (long)plane;
    float* out_bz = out + ((long)b * 64 * OD + oz) * planeL;
    #pragma unroll
    for (int mt = 0; mt < 2; ++mt) {
        const int mrow = warp_m * 32 + mt * 16 + g;
        const float bv0 = (bias != nullptr) ? __ldg(&bias[mrow]) : 0.f;
        const float bv8 = (bias != nullptr) ? __ldg(&bias[mrow + 8]) : 0.f;
        float* op = out_bz + (long)mrow * OD * planeL + n_base;
        float* op8 = op + 8 * OD * planeL;
        #pragma unroll
        for (int nt = 0; nt < 4; ++nt) {
            const int col = warp_n * 32 + nt * 8 + tig * 2;
            if (col < rem) {
                const float c0 = acc[mt][nt][0] + bv0;
                const float c1 = acc[mt][nt][1] + bv0;
                const float c2 = acc[mt][nt][2] + bv8;
                const float c3 = acc[mt][nt][3] + bv8;
                if (col + 1 < rem) {
                    *reinterpret_cast<float2*>(op + col)  = make_float2(c0, c1);
                    *reinterpret_cast<float2*>(op8 + col) = make_float2(c2, c3);
                } else {
                    op[col]  = c0;
                    op8[col] = c2;
                }
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Generic fallback: 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 && OC == 64 && groups == 1 &&
                       stride == 1 && pad == 0 && dilation == 1 &&
                       (long)B * OD <= 65535);

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

        const int plane = OH * OW;
        dim3 grid((plane + NTILE - 1) / NTILE, B * OD);
        conv3d_tc_kernel<<<grid, 256, SMEM_BYTES, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), B, D, H, W, OD, OH, OW);
    } 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_tc_v2",
    cpp_sources=conv3d_cpp_source,
    cuda_sources=conv3d_cuda_source,
    functions=["conv3d_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=[""],
)


class ModelNew(nn.Module):
    """
    Drop-in replacement for Model that runs the 3D convolution with a custom
    tensor-core (TF32) CUDA kernel for the target configuration
    (in_channels=3, out_channels=64, kernel=3, stride=1, padding=0), a
    generic custom CUDA kernel for other CUDA float32 configurations, and the
    original PyTorch op otherwise (e.g. CPU tensors).
    """

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