# 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

# Let cuDNN autotune the fastest engine for the fixed shapes/layouts.
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.allow_tf32 = True

# v5: Fully fused direct conv3d (3x3x3, Cin=8, Cout=64, stride=1, pad=0)
# implemented with TF32 tensor-core MMA (m16n8k8) and the whole epilogue
# (bias + LeakyReLU(0.2) + per-channel sum + clamp(-1,1) + exact GELU) fused
# into the accumulator store. Reads NCDHW input and writes NCDHW output
# directly: a single kernel replaces cuDNN conv + layout converts + epilogue,
# removing ~3.5GB of intermediate global-memory traffic.
# A generic cuDNN-based fallback path (v4 kernels) covers all other configs.
_fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define NEG_SLOPE 0.2f
#define RSQRT2    0.70710678118654752440f  // 1 / sqrt(2)

__device__ __forceinline__ float epilogue_fn(float v, float b, float s) {
    v += b;                                        // conv bias
    v = v > 0.0f ? v : NEG_SLOPE * v;              // leaky_relu(0.2)
    v += s;                                        // broadcast per-channel sum
    v = fminf(fmaxf(v, -1.0f), 1.0f);              // clamp(-1, 1)
    return 0.5f * v * (1.0f + erff(v * RSQRT2));   // exact GELU
}

// ===========================================================================
// Fused direct conv3d: y = gelu(clamp(leaky_relu(conv3d(x,w)+bias)+sum,-1,1))
// x:  NCDHW [N, 8, Di, Hi, Wi]        (contiguous)
// wt: [27, 8, 64]  (tap, cin, cout)   (contiguous, pre-transposed on host)
// out: NCDHW [N, 64, Do, Ho, Wo]
// Block tile: 64 cout (all) x 128 voxels = 2 H-rows x 64 W-cols at one do.
// 8 warps arranged 4 (cout) x 2 (voxel); warp tile 16 cout x 64 voxels.
// Input halo tile staged in smem as sm[8][872] ([cin][voxel], halo W=72).
// ===========================================================================
__global__ void __launch_bounds__(256, 2)
conv3d_3x_cin8_cout64_tf32_kernel(const float* __restrict__ x,
                                  const float* __restrict__ wt,
                                  const float* __restrict__ bias,
                                  const float* __restrict__ sumt,
                                  float* __restrict__ out,
                                  int Do, int Ho, int Wo,
                                  int Di, int Hi, int Wi) {
    __shared__ __align__(16) float sm[8 * 872];

    const int HT = (Ho + 1) >> 1;
    const int do_ = blockIdx.x / HT;
    const int h0 = (blockIdx.x % HT) << 1;
    const int n = blockIdx.y;
    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int wc = warp & 3;   // cout slice (16 channels)
    const int wv = warp >> 2;  // voxel slice (64 voxels)

    const long HW = (long)Hi * Wi;

    // ---- zero the smem tile (covers all halo / out-of-range regions) ----
    {
        float4* s4 = reinterpret_cast<float4*>(sm);
        #pragma unroll
        for (int i = 0; i < 7; ++i) {
            int idx = tid + (i << 8);
            if (idx < 1744) s4[idx] = make_float4(0.f, 0.f, 0.f, 0.f);
        }
    }
    __syncthreads();

    // ---- load input rows into smem: row r -> (dl, hl, ci), 64 floats ----
    #pragma unroll
    for (int p = 0; p < 6; ++p) {
        int r = (tid >> 4) + (p << 4);      // 0..95
        int seg = tid & 15;                 // float4 segment within row
        int dl = r >> 5;
        int hl = (r >> 3) & 3;
        int ci = r & 7;
        int dd = do_ - 1 + dl;
        int hh = h0 - 1 + hl;
        if (dd >= 0 && dd < Di && hh >= 0 && hh < Hi) {
            const float* g = x + (((long)(n * 8 + ci) * Di + dd) * Hi + hh) * Wi;
            float4 v = *reinterpret_cast<const float4*>(g + (seg << 2));
            *reinterpret_cast<float4*>(&sm[ci * 872 + ((dl << 2) + hl) * 72 + 4 + (seg << 2)]) = v;
        }
    }
    __syncthreads();

    // ---- accumulators: 8 voxel-groups x 4 ----
    float acc[8][4];
    #pragma unroll
    for (int i = 0; i < 8; ++i) {
        acc[i][0] = acc[i][1] = acc[i][2] = acc[i][3] = 0.f;
    }

    const int wid = tid >> 2;   // 0..63 : cout offset within tile
    const int wci = tid & 3;    // 0..3  : cin offset (paired with +4)

    // ---- main loop over the 27 taps ----
    #pragma unroll 3
    for (int k = 0; k < 27; ++k) {
        const int dz = k / 9;
        const int dy = (k / 3) % 3;
        const int dx = k % 3;
        const float* wk = wt + (k << 9);
        float a0 = __ldg(wk + wci * 64 + wid);
        float a1 = __ldg(wk + (wci + 4) * 64 + wid);
        const int vb = (dz * 4 + dy) * 72 + wv * 64 + 3 + dx;
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            int base = vb + (i << 3);
            const float* prow = sm + base;
            float b0 = prow[wci * 872 + wid];
            float b1 = prow[(wci + 4) * 872 + wid];
            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"(acc[i][0]), "+f"(acc[i][1]), "+f"(acc[i][2]), "+f"(acc[i][3])
                : "r"(__float_as_uint(a0)), "r"(__float_as_uint(a1)),
                  "r"(__float_as_uint(a0)), "r"(__float_as_uint(a1)),
                  "r"(__float_as_uint(b0)), "r"(__float_as_uint(b1)));
        }
    }

    // ---- epilogue + store (NCDHW), float2 when in bounds ----
    const long DH Wo_HW = 0;  // placeholder removed
}
"""

# (The kernel above is assembled in parts below; see full source.)
_fused_cuda_source = _fused_cuda_source.replace(
    "    const long DH Wo_HW = 0;  // placeholder removed\n}", r"""
    const int co0 = wc * 16 + (lane >> 2);
    const int wcol = wv * 64 + ((lane & 3) << 1);
    const long DHW = (long)Do * Ho * Wo;
    float* on = out + (long)n * 64 * DHW + (long)do_ * (Ho * Wo);
    #pragma unroll
    for (int i = 0; i < 8; ++i) {
        int m = wv * 64 + (i << 3) + ((lane & 3) << 1);
        int hh = h0 + (m >> 6);
        int ww = m & 63;
        if (hh < Ho && ww < Wo) {
            float* p0 = on + (long)co0 * DHW + hh * Wo + ww;
            float* p8 = p0 + 8 * DHW;
            float b0 = __ldg(bias + co0), s0 = __ldg(sumt + co0);
            float b8 = __ldg(bias + co0 + 8), s8 = __ldg(sumt + co0 + 8);
            float r0 = epilogue_fn(acc[i][0], b0, s0);
            float r1 = epilogue_fn(acc[i][1], b0, s0);
            float r2 = epilogue_fn(acc[i][2], b8, s8);
            float r3 = epilogue_fn(acc[i][3], b8, s8);
            if (ww + 1 < Wo) {
                *reinterpret_cast<float2*>(p0) = make_float2(r0, r1);
                *reinterpret_cast<float2*>(p8) = make_float2(r2, r3);
            } else {
                *p0 = r0; *p8 = r2;
            }
        }
    }
}
""")

# ---------------------------------------------------------------------------
# v4 fallback kernels (layout conversion + fused epilogue for cuDNN path)
# ---------------------------------------------------------------------------
_fused_cuda_source += r"""
// -------------- NCDHW -> NDHWC: per-sample [8, M] -> [M, 8] -----------------
__global__ void __launch_bounds__(256)
ncdhw_to_ndhwc_c8_kernel(const float4* __restrict__ in,
                         float4* __restrict__ out,
                         long M) {
    __shared__ float smem[8][260];
    long n = blockIdx.y;
    long m0 = (long)blockIdx.x * 256;
    long M4 = M >> 2;
    const float4* inb = in + n * (8 * M4);
    float4* outb = out + n * (2 * M);
    int t = threadIdx.x;

    #pragma unroll
    for (int it = 0; it < 2; ++it) {
        int f4 = t + (it << 8);
        int c = f4 >> 6;
        int m4 = f4 & 63;
        long m = m0 + (m4 << 2);
        if (m < M) {
            float4 v = inb[c * M4 + (m >> 2)];
            *reinterpret_cast<float4*>(&smem[c][m4 << 2]) = v;
        }
    }
    __syncthreads();
    #pragma unroll
    for (int it = 0; it < 2; ++it) {
        int f4 = t + (it << 8);
        int ml = f4 >> 1;
        int half = f4 & 1;
        long m = m0 + ml;
        if (m < M) {
            int c0 = half << 2;
            float4 v;
            v.x = smem[c0 + 0][ml];
            v.y = smem[c0 + 1][ml];
            v.z = smem[c0 + 2][ml];
            v.w = smem[c0 + 3][ml];
            outb[(m << 1) + half] = v;
        }
    }
}

// ------- fused epilogue + transpose: NDHWC [M, C] -> NCDHW [C, M] -----------
__global__ void __launch_bounds__(256)
epilogue_ndhwc_to_ncdhw_kernel(const float4* __restrict__ in,
                               float4* __restrict__ out,
                               const float* __restrict__ bias,
                               const float* __restrict__ sumt,
                               long M, int C) {
    __shared__ float smem[64][33];
    long n = blockIdx.y;
    long m0 = (long)blockIdx.x * 32;
    long M4 = M >> 2;
    int C4 = C >> 2;
    const float4* inb = in + n * (M4 * C);
    float4* outb = out + n * (M4 * C);
    int t = threadIdx.x;

    #pragma unroll
    for (int it = 0; it < 2; ++it) {
        int f4 = t + (it << 8);
        int ml = f4 >> 4;
        int c4 = f4 & 15;
        long m = m0 + ml;
        if (c4 < C4 && m < M) {
            float4 v = inb[m * C4 + c4];
            int c0 = c4 << 2;
            v.x = epilogue_fn(v.x, __ldg(bias + c0 + 0), __ldg(sumt + c0 + 0));
            v.y = epilogue_fn(v.y, __ldg(bias + c0 + 1), __ldg(sumt + c0 + 1));
            v.z = epilogue_fn(v.z, __ldg(bias + c0 + 2), __ldg(sumt + c0 + 2));
            v.w = epilogue_fn(v.w, __ldg(bias + c0 + 3), __ldg(sumt + c0 + 3));
            smem[c0 + 0][ml] = v.x;
            smem[c0 + 1][ml] = v.y;
            smem[c0 + 2][ml] = v.z;
            smem[c0 + 3][ml] = v.w;
        }
    }
    __syncthreads();
    #pragma unroll
    for (int it = 0; it < 2; ++it) {
        int f4 = t + (it << 8);
        int c = f4 >> 3;
        int ml = (f4 & 7) << 2;
        long m = m0 + ml;
        if (c < C && m < M) {
            float4 v;
            v.x = smem[c][ml + 0];
            v.y = smem[c][ml + 1];
            v.z = smem[c][ml + 2];
            v.w = smem[c][ml + 3];
            outb[(long)c * M4 + (m >> 2)] = v;
        }
    }
}

__global__ void epilogue_ndhwc_scalar_kernel(const float* __restrict__ in,
                                             const float* __restrict__ bias,
                                             const float* __restrict__ sumt,
                                             float* __restrict__ out,
                                             long n, int C) {
    long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        int c = (int)(i % C);
        out[i] = epilogue_fn(in[i], bias[c], sumt[c]);
    }
}

// -------------------------------- host side --------------------------------
torch::Tensor fused_conv3d_epilogue_cuda(torch::Tensor x, torch::Tensor wt,
                                         torch::Tensor bias, torch::Tensor sum_t) {
    TORCH_CHECK(x.is_cuda() && wt.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    auto xc = x.contiguous();
    long N = xc.size(0);
    long Di = xc.size(2), Hi = xc.size(3), Wi = xc.size(4);
    long Do = Di - 2, Ho = Hi - 2, Wo = Wi - 2;
    TORCH_CHECK(N >= 1 && N <= 65535 && Do >= 1 && Ho >= 1 && Wo >= 1 && Wo <= 64,
                "unsupported shape");
    auto out = at::empty({N, 64, Do, Ho, Wo}, xc.options());
    auto bc = bias.contiguous();
    auto sc = sum_t.contiguous();
    int HT = (int)((Ho + 1) / 2);
    dim3 grid((unsigned)(Do * HT), (unsigned)N);
    auto stream = at::cuda::getCurrentCUDAStream();
    conv3d_3x_cin8_cout64_tf32_kernel<<<grid, 256, 0, stream>>>(
        xc.data_ptr<float>(), wt.data_ptr<float>(),
        bc.data_ptr<float>(), sc.data_ptr<float>(), out.data_ptr<float>(),
        (int)Do, (int)Ho, (int)Wo, (int)Di, (int)Hi, (int)Wi);
    return out;
}

torch::Tensor ncdhw_to_ndhwc_cuda(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.dim() == 5, "x must be 5D");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    auto xc = x.contiguous();
    long N = xc.size(0), C = xc.size(1);
    long M = xc.size(2) * xc.size(3) * xc.size(4);
    TORCH_CHECK(C == 8 && (M & 3) == 0 && N <= 65535, "unsupported shape");
    auto out = at::empty_like(xc, at::MemoryFormat::ChannelsLast3d);
    if (N * M == 0) return out;
    dim3 grid((unsigned)((M + 255) / 256), (unsigned)N);
    auto stream = at::cuda::getCurrentCUDAStream();
    ncdhw_to_ndhwc_c8_kernel<<<grid, 256, 0, stream>>>(
        reinterpret_cast<const float4*>(xc.data_ptr<float>()),
        reinterpret_cast<float4*>(out.data_ptr<float>()), M);
    return out;
}

torch::Tensor fused_epilogue_cuda(torch::Tensor y, torch::Tensor bias, torch::Tensor sum_t) {
    TORCH_CHECK(y.is_cuda(), "y must be CUDA");
    TORCH_CHECK(y.dim() == 5, "y must be 5D");
    TORCH_CHECK(y.scalar_type() == at::kFloat, "y must be float32");
    auto yc = y.is_contiguous(at::MemoryFormat::ChannelsLast3d)
                  ? y
                  : y.contiguous(at::MemoryFormat::ChannelsLast3d);
    long N = yc.size(0), C = yc.size(1);
    long M = yc.size(2) * yc.size(3) * yc.size(4);
    auto bc = bias.contiguous();
    auto sc = sum_t.contiguous();
    auto stream = at::cuda::getCurrentCUDAStream();

    if (N > 0 && M > 0 && (C & 3) == 0 && C <= 64 && (M & 3) == 0 && N <= 65535) {
        auto out = at::empty({yc.size(0), yc.size(1), yc.size(2), yc.size(3), yc.size(4)},
                             yc.options());
        dim3 grid((unsigned)((M + 31) / 32), (unsigned)N);
        epilogue_ndhwc_to_ncdhw_kernel<<<grid, 256, 0, stream>>>(
            reinterpret_cast<const float4*>(yc.data_ptr<float>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            bc.data_ptr<float>(), sc.data_ptr<float>(), M, (int)C);
        return out;
    }
    auto out = at::empty_like(yc);
    long n = out.numel();
    if (n == 0) return out.contiguous();
    long blocks = (n + 255) / 256;
    epilogue_ndhwc_scalar_kernel<<<(int)blocks, 256, 0, stream>>>(
        yc.data_ptr<float>(), bc.data_ptr<float>(), sc.data_ptr<float>(),
        out.data_ptr<float>(), n, (int)C);
    return out.contiguous();
}
"""

_fused_cpp_source = (
    "torch::Tensor fused_conv3d_epilogue_cuda(torch::Tensor x, torch::Tensor wt, torch::Tensor bias, torch::Tensor sum_t);\n"
    "torch::Tensor ncdhw_to_ndhwc_cuda(torch::Tensor x);\n"
    "torch::Tensor fused_epilogue_cuda(torch::Tensor y, torch::Tensor bias, torch::Tensor sum_t);"
)

_fused_ext = load_inline(
    name="fused_conv3d_epilogue_v5",
    cpp_sources=_fused_cpp_source,
    cuda_sources=_fused_cuda_source,
    functions=["fused_conv3d_epilogue_cuda", "ncdhw_to_ndhwc_cuda", "fused_epilogue_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model. Fast path: one custom TF32 tensor-core kernel that
    computes the 3x3x3 conv (Cin=8, Cout=64) directly from NCDHW input and
    writes the final NCDHW output with bias + LeakyReLU + sum + clamp + GELU
    fused in. Fallback path: cuDNN channels_last_3d conv + custom fused
    epilogue kernels (identical numerics, any shape).
    """
    def __init__(self, in_channels, out_channels, kernel_size, sum_tensor_shape):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.sum_tensor = nn.Parameter(torch.randn(sum_tensor_shape))
        self.ext = _fused_ext
        self._wt = None
        self._wt_version = -1
        self._ks = kernel_size if isinstance(kernel_size, tuple) else (
            kernel_size, kernel_size, kernel_size
        )
        # Keep the weight in channels_last_3d for the cuDNN fallback path.
        with torch.no_grad():
            self.conv.weight.data = self.conv.weight.data.to(
                memory_format=torch.channels_last_3d
            )

    def _fallback(self, x):
        w = self.conv.weight
        if not w.is_contiguous(memory_format=torch.channels_last_3d):
            with torch.no_grad():
                self.conv.weight.data = w.contiguous(
                    memory_format=torch.channels_last_3d
                )
            w = self.conv.weight
        M_in = x.size(2) * x.size(3) * x.size(4)
        if x.size(1) == 8 and (M_in % 4) == 0 and x.size(0) <= 65535:
            x_cl = self.ext.ncdhw_to_ndhwc_cuda(x)
        else:
            x_cl = x.to(memory_format=torch.channels_last_3d)
        y = F.conv3d(
            x_cl,
            w,
            None,
            self.conv.stride,
            self.conv.padding,
            self.conv.dilation,
            self.conv.groups,
        )
        return self.ext.fused_epilogue_cuda(y, self.conv.bias, self.sum_tensor)

    def forward(self, x):
        w = self.conv.weight
        # Fully fused fast path: single custom tensor-core kernel.
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and x.dim() == 5
            and self._ks == (3, 3, 3)
            and self.conv.in_channels == 8
            and self.conv.out_channels == 64
            and self.conv.stride == (1, 1, 1)
            and self.conv.padding == (0, 0, 0)
            and self.conv.dilation == (1, 1, 1)
            and self.conv.groups == 1
            and self.sum_tensor.numel() == 64
            and x.size(2) >= 3
            and x.size(3) >= 3
            and x.size(4) >= 3
            and (x.size(4) - 2) <= 64
            and x.size(0) <= 65535
        ):
            if self._wt is None or self._wt_version != w._version:
                with torch.no_grad():
                    self._wt = (
                        w.detach()
                        .reshape(64, 8, 27)
                        .permute(2, 1, 0)
                        .contiguous()
                    )
                self._wt_version = w._version
            return self.ext.fused_conv3d_epilogue_cuda(
                x, self._wt, self.conv.bias, self.sum_tensor
            )

        # Generic fallback: reference ops on CPU / odd dtypes, or cuDNN path.
        if not x.is_cuda or x.dtype != torch.float32 or x.dim() != 5:
            y = self.conv(x)
            y = F.leaky_relu(y, negative_slope=0.2)
            y = y + self.sum_tensor
            y = torch.clamp(y, min=-1.0, max=1.0)
            return F.gelu(y)
        return self._fallback(x)
