```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, W = 64, float32 CUDA tensors.
//
// Implicit GEMM:  C[oc, p] = sum_k W[oc, k] * X[k, p]
//   M = 64 (all out channels), K = 81 taps padded to 88,
//   N = output pixels of one (batch, oz) plane (oy*OW+ox flattened), tiled
//   by 128 columns per block.
//
// Key optimization vs. a classic im2col panel: the kx expansion of the im2col
// B matrix makes every input element be fetched from L2 three times. Instead
// we stage *raw* input rows (3ic x 3kz x 6iy rows of the full 64-wide input
// row) plus one zero row, and recover B[k][n] = raw[rowbase(k)+oy_rel(n)]
// [ox(n)+kx(k)] via a tiny per-tap offset table T[k] kept in shared memory.
// This cuts L2->SMEM traffic ~4.5x, shrinks the B panel ~3x, and removes the
// multi-stage cp.async pipeline entirely (one bulk 16B-chunk cp.async stage).
//
// Block = 256 threads (8 warps, 2m x 4n); each warp owns a 32 x 32 tile =
// 2x4 mma.m16n8k8.tf32 ops per k8 step, 11 k-steps (K padded 81 -> 88 with
// zero taps). tf32 rounding (rna) matches cuDNN TF32 numerics closely.
// ---------------------------------------------------------------------------

#define KREAL 81
#define KPADA 88            // multiple of 8
#define ASTRIDE 92          // multiple of 4, conflict-free for ldmatrix
#define NTILE 128
#define BROWS 54            // 3ic * 3kz * 6iy raw rows
#define ZROW 54             // extra all-zero row (for padding taps / OOB)
#define BSTRIDE 68          // multiple of 4, spreads banks across rows
#define CPR 16              // 16B chunks per 64-float input row (W=64)
#define SMEM_FLOATS (64 * ASTRIDE + (BROWS + 1) * BSTRIDE + KPADA)
#define SMEM_BYTES (SMEM_FLOATS * 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_async16(void* sdst, const float* gsrc,
                                           int src_bytes) {
    asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;\n" ::"r"(
                     smem_u32(sdst)),
                 "l"(gsrc), "r"(src_bytes));
}

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

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

extern __shared__ float smem[];

__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;             // [BROWS+1][BSTRIDE]
    int*   Ts = reinterpret_cast<int*>(Bs + (BROWS + 1) * BSTRIDE);  // [KPADA]

    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 oy0    = n_base / OW;              // first output row in tile
    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

    const float* in_b = in + (long)b * 3 * D * H * W;

    // ---- stage per-tap offset table T[k] = rowbase(k)*BSTRIDE + kx(k) ----
    if (tid < KPADA) {
        int t;
        if (tid < KREAL) {
            const int ic = tid / 27;
            const int r  = tid - ic * 27;
            const int kz = r / 9;
            const int rr = r - kz * 9;
            const int ky = rr / 3;
            const int kx = rr - ky * 3;
            t = ((ic * 3 + kz) * 6 + ky) * BSTRIDE + kx;
        } else {
            t = ZROW * BSTRIDE;
        }
        Ts[tid] = t;
    }

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

    // ---- stage raw B rows (55 rows x 64 floats) as 16B chunks ----
    #pragma unroll
    for (int c = tid; c < (BROWS + 1) * CPR; c += 256) {
        const int row  = c >> 4;
        const int part = c & 15;
        const float* src = in_b;
        int bytes = 0;
        if (row < BROWS) {
            const int ic  = row / 18;
            const int rr  = row - ic * 18;
            const int kz  = rr / 6;
            const int iyr = rr - kz * 6;         // 0..5
            const int iy  = oy0 + iyr;
            if (iy < H) {
                src = in_b + (((long)(ic * D + oz + kz)) * H + iy) * W +
                      (part << 2);
                bytes = 16;
            }
        }
        cp_async16(&Bs[row * BSTRIDE + (part << 2)], src, bytes);
    }
    cp_commit();

    // ---- per-thread mma coordinates (fixed for whole kernel) ----
    const int g   = lane >> 2;                   // 0..7
    const int tig = lane & 3;                    // 0..3

    int c0[4];                                   // per-nt smem column offsets
    #pragma unroll
    for (int nt = 0; nt < 4; ++nt) {
        const int n  = n_base + warp_n * 32 + nt * 8 + g;
        const int oy = n / OW;
        const int ox = n - oy * OW;
        c0[nt] = (oy - oy0) * BSTRIDE + ox;
    }

    cp_wait_all();
    __syncthreads();

    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 kl = 0; kl < KPADA / 8; ++kl) {
        const int k0 = kl * 8;
        uint32_t a[2][4];
        #pragma unroll
        for (int mt = 0; mt < 2; ++mt) {
            const 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 T0 = Ts[k0 + tig];
        const int T1 = Ts[k0 + tig + 4];
        #pragma unroll
        for (int nt = 0; nt < 4; ++nt) {
            const uint32_t b0 =
                __float_as_uint(to_tf32(Bs[T0 + c0[nt]]));
            const uint32_t b1 =
                __float_as_uint(to_tf32(Bs[T1 + c0[nt]]));
            #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 c0v = acc[mt][nt][0] + bv0;
                const float c1v = acc[mt][nt][1] + bv0;
                const float c2v = acc[mt][nt][2] + bv8;
                const float c3v = acc[mt][nt][3] + bv8;
                if (col + 1 < rem) {
                    *reinterpret_cast<float2*>(op + col)  = make_float2(c0v, c1v);
                    *reinterpret_cast<float2*>(op8 + col) = make_float2(c2v, c3v);
                } else {
                    op[col]  = c0v;
                    op8[col] = c2v;
                }
            }
        }
    }
}

// ---------------------------------------------------------------------------
// 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 && W == 64 &&
                       OW == W - 2 && (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_rawrow_v4",
    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, W=64),
    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)
```