```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) persistent kernel 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 per (batch, output-plane pair, 128-wide column tile):
//   C[oc, p] = sum_k W[oc, k] * X[k, p],  M = 64, K = 81 (padded to 88),
//   N = 128 flattened output pixels, two output planes (oz, oz+1) per item.
//
// Improvements over the v4 design:
//  * Persistent CTAs: the 64x88 tf32 weight panel is staged in shared memory
//    exactly once per SM-resident block instead of once per output tile.
//  * Each work item computes TWO adjacent output planes.  The raw-row B panel
//    is shared between them: plane p taps read row offset T[k] + p*6*BSTRIDE,
//    so weight A-fragments (ldmatrix) and the staged panel are reused twice,
//    doubling mma work per shared-memory byte and per staging pass.
//  * Double-buffered cp.async B panel: the next item's raw input rows are
//    prefetched while the current item is being computed.
//  * BSTRIDE = 80 (bank shift of 16 per row) to reduce B-fragment conflicts.
//
// Block = 256 threads (8 warps in 2m x 4n); per warp per k8 step:
// 2 ldmatrix.x4 (A), 16 scalar B loads + tf32 converts, 16 mma.m16n8k8.tf32
// (2 planes x 2 mt x 4 nt), accumulating into 64 registers per thread.
// ---------------------------------------------------------------------------

#define KREAL 81
#define KPADA 88            // multiple of 8
#define ASTRIDE 92          // multiple of 4, conflict-free for ldmatrix
#define NTILE 128
#define IZR 4               // iz rows spanned by 2 planes with kz=0..2
#define IYR 6               // iy rows staged (oy span <= 3, plus ky=0..2)
#define BROWS (3 * IZR * IYR)   // 72 raw input rows
#define ZROW BROWS              // extra all-zero row
#define BSTRIDE 80          // bank shift 16 per row; >= 64 data floats
#define CPR 16              // 16B chunks per 64-float input row (W=64)
#define TPLANE (6 * BSTRIDE)    // row offset between plane taps
#define SMEM_FLOATS (64 * ASTRIDE + 2 * (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" ::);
}

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

extern __shared__ float smem[];

__global__ void __launch_bounds__(256, 2)
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,
                 int has_bias) {
    float* As = smem;                              // [64][ASTRIDE]
    float* Bs0 = As + 64 * ASTRIDE;                // two buffers
    float* Bs1 = Bs0 + (BROWS + 1) * BSTRIDE;
    int*   Ts = reinterpret_cast<int*>(Bs1 + (BROWS + 1) * BSTRIDE);  // [KPADA]

    const int plane = OH * OW;
    const int xt_n  = (plane + NTILE - 1) / NTILE;
    const int pp_n  = (OD + 1) >> 1;
    const int total = B * pp_n * xt_n;

    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 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 * IZR + kz) * IYR + ky) * BSTRIDE + kx;
        } else {
            t = ZROW * BSTRIDE;
        }
        Ts[tid] = t;
    }

    // ---- stage weights (64 x 88) once, 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);
    }

    // ---- per-thread mma coordinates (fixed for whole kernel) ----
    const int g   = lane >> 2;                     // 0..7
    const int tig = lane & 3;                      // 0..3
    const int a_row0 = warp_m * 32 + ((lane >> 3) & 1) * 8 + (lane & 7);
    const int a_colf = (lane >> 4) * 4;

    const int oc_r  = warp_m * 32 + g;
    const int oc_r8 = oc_r + 8;
    const float bias0  = has_bias ? __ldg(&bias[oc_r])      : 0.f;
    const float bias8  = has_bias ? __ldg(&bias[oc_r8])     : 0.f;
    const float bias16 = has_bias ? __ldg(&bias[oc_r + 16]) : 0.f;
    const float bias24 = has_bias ? __ldg(&bias[oc_r8 + 16]): 0.f;

    const long planeL = (long)plane;

    // ---- stage B panel (raw rows) for one item into buffer ----
    auto stage_B = [&](float* Bs, int item) {
        const int xt = item % xt_n;
        int r = item / xt_n;
        const int pp = r % pp_n;
        const int b  = r / pp_n;
        const int oy0 = (xt * NTILE) / OW;
        const int oz0 = pp * 2;
        const float* in_b = in + (long)b * 3 * D * H * W + (long)oz0 * H * W;
        #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 / (IZR * IYR);
                const int rr  = row - ic * (IZR * IYR);
                const int izr = rr / IYR;
                const int iyr = rr - izr * IYR;
                const int iy  = oy0 + iyr;
                const int iz  = izr;             // relative to oz0
                if (iy < H && (oz0 + iz) < D) {
                    src = in_b + (((long)ic * D + iz) * H + iy) * W +
                          (part << 2);
                    bytes = 16;
                }
            }
            cp_async16(&Bs[row * BSTRIDE + (part << 2)], src, bytes);
        }
        cp_commit();
    };

    int item = blockIdx.x;
    if (item >= total) return;
    stage_B(Bs0, item);

    float* bufs[2] = {Bs0, Bs1};
    int buf = 0;
    bool first = true;

    for (; item < total; item += gridDim.x, buf ^= 1) {
        const int next = item + gridDim.x;
        float* Bs = bufs[buf];
        if (next < total) {
            stage_B(bufs[buf ^ 1], next);
            cp_wait_group<1>();
        } else {
            cp_wait_group<0>();
        }
        __syncthreads();

        const int xt = item % xt_n;
        int r = item / xt_n;
        const int pp = r % pp_n;
        const int b  = r / pp_n;
        const int n_base = xt * NTILE;
        const int oz0 = pp * 2;
        const int oy0 = n_base / OW;
        const int rem = plane - n_base;          // >= 1
        const bool p1v = (oz0 + 1) < OD;

        int c0[4];
        #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;
        }

        float acc[2][2][4][4];                   // [plane][mt][nt][4]
        #pragma unroll
        for (int p = 0; p < 2; ++p)
            #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[p][mt][nt][c] = 0.f;

        #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 p = 0; p < 2; ++p) {
                const int T0p = T0 + p * TPLANE;
                const int T1p = T1 + p * TPLANE;
                #pragma unroll
                for (int nt = 0; nt < 4; ++nt) {
                    const uint32_t b0 =
                        __float_as_uint(to_tf32(Bs[T0p + c0[nt]]));
                    const uint32_t b1 =
                        __float_as_uint(to_tf32(Bs[T1p + c0[nt]]));
                    #pragma unroll
                    for (int mt = 0; mt < 2; ++mt) {
                        mma_tf32(acc[p][mt][nt][0], acc[p][mt][nt][1],
                                 acc[p][mt][nt][2], acc[p][mt][nt][3],
                                 a[mt][0], a[mt][1], a[mt][2], a[mt][3],
                                 b0, b1);
                    }
                }
            }
        }

        // ---- epilogue: bias + coalesced float2 stores (plane-flattened) ----
        float* out_bz = out + ((long)b * 64 * OD + oz0) * planeL;
        #pragma unroll
        for (int p = 0; p < 2; ++p) {
            if (p == 1 && !p1v) break;
            float* outp = out_bz + (long)p * planeL;
            #pragma unroll
            for (int mt = 0; mt < 2; ++mt) {
                const float bv0 = mt ? bias16 : bias0;
                const float bv8 = mt ? bias24 : bias8;
                float* op  = outp + (long)(warp_m * 32 + mt * 16 + g) * 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[p][mt][nt][0] + bv0;
                        const float c1v = acc[p][mt][nt][1] + bv0;
                        const float c2v = acc[p][mt][nt][2] + bv8;
                        const float c3v = acc[p][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;
                        }
                    }
                }
            }
        }
        __syncthreads();   // everyone done with Bs before it is restaged
        first = false;
    }
    (void)first;
}

// ---------------------------------------------------------------------------
// 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);

    if (fast) {
        static int nblocks = []() {
            cudaFuncSetAttribute(conv3d_tc_kernel,
                                 cudaFuncAttributeMaxDynamicSharedMemorySize,
                                 SMEM_BYTES);
            cudaFuncSetAttribute(conv3d_tc_kernel,
                                 cudaFuncAttributePreferredSharedMemoryCarveout,
                                 100);
            int dev = 0;
            cudaGetDevice(&dev);
            int sms = 0;
            cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
            int per_sm = 0;
            cudaOccupancyMaxActiveBlocksPerMultiprocessor(
                &per_sm, conv3d_tc_kernel, 256, SMEM_BYTES);
            if (per_sm < 1) per_sm = 1;
            if (per_sm > 2) per_sm = 2;
            return sms * per_sm;
        }();

        const int plane = OH * OW;
        const long total = (long)B * ((OD + 1) / 2) *
                           ((plane + NTILE - 1) / NTILE);
        int grid = nblocks;
        if ((long)grid > total) grid = (int)total;
        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, has_bias ? 1 : 0);
    } 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_persistent_v5",
    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
    persistent 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)
```