```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.
//
// The convolution is evaluated as an implicit GEMM:
//   C[oc, (b,oz,oy,ox)] = sum_{ic,kz,ky,kx} W[oc, k] * X[b,ic,oz+kz,oy+ky,ox+kx]
// with M = 64 (out channels), K = 81 (3*3*3*3) padded to 88, N = output
// pixels. Each thread block computes a 64(M) x 64(N) tile where N covers one
// output row (ox0 .. ox0+63) for a single (b, oz, oy). The full weight matrix
// (64 x 88) and the im2col B tile (88 x 64) are staged in shared memory (as
// tf32), then consumed with mma.m16n8k8 tensor-core instructions.
//
// Block: 128 threads = 4 warps arranged 2(m) x 2(n); each warp owns a
// 32 x 32 tile = 2x4 mma tiles of 16x8.
// ---------------------------------------------------------------------------

#define KPAD 88          // 81 real k-values padded to a multiple of 8
#define ASTRIDE 92       // conflict-free stride for ldmatrix on As
#define BSTRIDE 72       // conflict-free stride for B-fragment lds on Bs
#define SMEM_BYTES (64 * ASTRIDE * 4 + KPAD * 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));
}

extern __shared__ float smem[];

__global__ void __launch_bounds__(128)
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;     // [KPAD][BSTRIDE]

    const int ox0 = blockIdx.x * 64;
    const int oy  = blockIdx.y;
    const int oz  = blockIdx.z % OD;
    const int b   = blockIdx.z / OD;

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

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

    // ---- stage im2col B tile (88 x 64) ----
    const float* ib = in + (long)b * 3 * D * H * W;
    const long HW = (long)H * W;
    for (int i = tid; i < KPAD * 64; i += 128) {
        int k = i >> 6;
        int n = i & 63;
        float v = 0.f;
        if (k < 81) {
            int ic = k / 27;
            int r  = k - ic * 27;
            int kz = r / 9;
            int r2 = r - kz * 9;
            int ky = r2 / 3;
            int kx = r2 - ky * 3;
            int gx = ox0 + n + kx;
            if (gx < W) {
                v = __ldg(&ib[(ic * D + (oz + kz)) * HW +
                              (long)(oy + ky) * W + gx]);
            }
        }
        Bs[k * BSTRIDE + n] = to_tf32(v);
    }
    __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;

    // ldmatrix address decomposition (fixed per thread)
    const int a_row0 = warp_m * 32 + ((lane >> 3) & 1) * 8 + (lane & 7);
    const int a_colf = (lane >> 4) * 4;

    #pragma unroll
    for (int ks = 0; ks < KPAD / 8; ++ks) {
        const int k0 = ks * 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);
        }
        #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(Bs[(k0 + tig) * BSTRIDE + nb]);
            const uint32_t b1 =
                __float_as_uint(Bs[(k0 + 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 ----
    const bool even_ow = ((OW & 1) == 0);
    const long plane = (long)OD * OH * OW;
    #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;
        #pragma unroll
        for (int nt = 0; nt < 4; ++nt) {
            const int col = warp_n * 32 + nt * 8 + tig * 2;
            if (col < OW) {
                float* op0 = out + (((long)b * 64 + mrow) * OD + oz) *
                                       ((long)OH * OW) +
                                   (long)oy * OW + ox0 + col;
                float* op8 = op0 + 8 * plane;
                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 (even_ow && col + 1 < OW) {
                    *reinterpret_cast<float2*>(op0) = make_float2(c0, c1);
                    *reinterpret_cast<float2*>(op8) = make_float2(c2, c3);
                } else {
                    op0[0] = c0;
                    if (col + 1 < OW) op0[1] = c1;
                    op8[0] = c2;
                    if (col + 1 < OW) op8[1] = c3;
                }
            }
        }
    }
}

// ---------------------------------------------------------------------------
// 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 && OH <= 65535 &&
                       D >= 3 && H >= 3);

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

        const int tiles_x = (OW + 63) / 64;
        dim3 grid(tiles_x, OH, B * OD);
        conv3d_tc_kernel<<<grid, 128, 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",
    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)
```