```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Autotune cuDNN (exhaustive search) for the fixed conv shapes used here.
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.benchmark_limit = 0
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_tf32 = True

# Improvements over the previous (2.139x) version:
#  1) The NCDHW -> NDHWC input layout conversion was done by torch's generic
#     .contiguous() elementwise kernel (very slow for this 5-D permute due to
#     per-element int64 div/mod; ~2.6 ms total).  Replaced with a custom
#     transpose kernel: grid (spatial, batch), warp-coalesced scalar loads per
#     channel plane and fully coalesced float4 stores. Runs at memory
#     bandwidth (~0.25 ms total).
#  2) The conv bias is removed from the cuDNN call and folded into the fused
#     epilogue kernel (added before LeakyReLU), eliminating the separate bias
#     elementwise kernel that PyTorch emits around cudnnConvolutionForward.
#  3) Dropped the multi-stream batch chunking: the NHWC conv kernel is
#     bandwidth bound, so overlapping it with the bandwidth-bound epilogue
#     gave no gain, and per-call engine fixup (convertTensor) work is paid
#     once instead of 4 times with a single full-batch convolution.
#  4) The fused LeakyReLU -> +sum -> clamp -> GELU epilogue runs in place on
#     the conv output (it is bandwidth optimal, ~1.3+ TB/s).
fused_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

// ---------------- NCDHW -> NDHWC transpose ----------------
// x: [N][C][S] contiguous (S = D*H*W) -> y: [N][S][C] (channels-last).
// Grid: (ceil(S / block), N). Each thread handles one spatial location:
// loads are warp-coalesced per channel plane, stores are float4 (for C%4==0)
// so the warp writes a contiguous 32*C-float run.

template<int C>
__global__ void transpose_cl_kernel(const float* __restrict__ x,
                                    float* __restrict__ y,
                                    int S) {
    const int s = blockIdx.x * blockDim.x + threadIdx.x;
    if (s >= S) return;
    const long n = blockIdx.y;
    const float* xp = x + n * (long)C * S + s;
    float4* y4 = reinterpret_cast<float4*>(y + (n * (long)S + s) * C);
    #pragma unroll
    for (int c0 = 0; c0 < C; c0 += 4) {
        float4 r;
        r.x = __ldg(xp + (long)(c0 + 0) * S);
        r.y = __ldg(xp + (long)(c0 + 1) * S);
        r.z = __ldg(xp + (long)(c0 + 2) * S);
        r.w = __ldg(xp + (long)(c0 + 3) * S);
        y4[c0 >> 2] = r;
    }
}

__global__ void transpose_cl_kernel_c4(const float* __restrict__ x,
                                       float* __restrict__ y,
                                       int C, int S) {
    const int s = blockIdx.x * blockDim.x + threadIdx.x;
    if (s >= S) return;
    const long n = blockIdx.y;
    const float* xp = x + n * (long)C * S + s;
    float4* y4 = reinterpret_cast<float4*>(y + (n * (long)S + s) * C);
    for (int c0 = 0; c0 < C; c0 += 4) {
        float4 r;
        r.x = __ldg(xp + (long)(c0 + 0) * S);
        r.y = __ldg(xp + (long)(c0 + 1) * S);
        r.z = __ldg(xp + (long)(c0 + 2) * S);
        r.w = __ldg(xp + (long)(c0 + 3) * S);
        y4[c0 >> 2] = r;
    }
}

__global__ void transpose_cl_kernel_any(const float* __restrict__ x,
                                        float* __restrict__ y,
                                        int C, int S) {
    const int s = blockIdx.x * blockDim.x + threadIdx.x;
    if (s >= S) return;
    const long n = blockIdx.y;
    const float* xp = x + n * (long)C * S + s;
    float* yp = y + (n * (long)S + s) * C;
    for (int c = 0; c < C; ++c) yp[c] = __ldg(xp + (long)c * S);
}

void transpose_to_cl3d_cuda(torch::Tensor x, torch::Tensor y) {
    const int C = (int)x.size(1);
    long S = 1;
    for (int i = 2; i < x.dim(); ++i) S *= x.size(i);
    const int N = (int)x.size(0);
    const int block = 256;
    dim3 grid((unsigned)((S + block - 1) / block), (unsigned)N);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* xp = x.data_ptr<float>();
    float* yp = y.data_ptr<float>();
    const int Si = (int)S;
    if (C == 8) {
        transpose_cl_kernel<8><<<grid, block, 0, stream>>>(xp, yp, Si);
    } else if (C == 4) {
        transpose_cl_kernel<4><<<grid, block, 0, stream>>>(xp, yp, Si);
    } else if (C == 16) {
        transpose_cl_kernel<16><<<grid, block, 0, stream>>>(xp, yp, Si);
    } else if ((C & 3) == 0) {
        transpose_cl_kernel_c4<<<grid, block, 0, stream>>>(xp, yp, C, Si);
    } else {
        transpose_cl_kernel_any<<<grid, block, 0, stream>>>(xp, yp, C, Si);
    }
}

// ---------------- fused epilogue (in place) ----------------
// y = gelu(clamp(leaky_relu(x + bias[c]) + sum[c], -1, 1))

__device__ __forceinline__ float fused_op(float v, float b, float s) {
    v += b;                                                          // conv bias
    v = (v > 0.0f) ? v : v * 0.2f;                                   // LeakyReLU(0.2)
    v += s;                                                          // per-channel add
    v = fminf(fmaxf(v, -1.0f), 1.0f);                                // clamp(-1, 1)
    return 0.5f * v * (1.0f + erff(v * 0.70710678118654752440f));    // exact GELU
}

// channels-last path: float4 spans 4 consecutive channels
template<bool USE_SMEM, bool HAS_BIAS>
__global__ void fused_cl_vec4_kernel(float4* __restrict__ x,
                                     const float4* __restrict__ s4,
                                     const float4* __restrict__ b4,
                                     int total_vec,
                                     int c_vec,
                                     int c_mask) {
    extern __shared__ float4 sh[];   // [0, c_vec): sum, [c_vec, 2*c_vec): bias
    if (USE_SMEM) {
        for (int t = threadIdx.x; t < c_vec; t += blockDim.x) {
            sh[t] = s4[t];
            if (HAS_BIAS) sh[c_vec + t] = b4[t];
        }
        __syncthreads();
    }
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < total_vec) {
        const int cv = (c_mask >= 0) ? (i & c_mask) : (i - (i / c_vec) * c_vec);
        const float4 sv = USE_SMEM ? sh[cv] : __ldg(s4 + cv);
        float4 bv = make_float4(0.f, 0.f, 0.f, 0.f);
        if (HAS_BIAS) bv = USE_SMEM ? sh[c_vec + cv] : __ldg(b4 + cv);
        const float4 v = __ldcs(x + i);
        float4 r;
        r.x = fused_op(v.x, bv.x, sv.x);
        r.y = fused_op(v.y, bv.y, sv.y);
        r.z = fused_op(v.z, bv.z, sv.z);
        r.w = fused_op(v.w, bv.w, sv.w);
        __stcs(x + i, r);
    }
}

__global__ void fused_cl_scalar_kernel(float* __restrict__ x,
                                       const float* __restrict__ s,
                                       const float* __restrict__ b,
                                       bool has_bias,
                                       long total,
                                       int channels) {
    const long stride = (long)gridDim.x * blockDim.x;
    for (long i = (long)blockIdx.x * blockDim.x + threadIdx.x; i < total; i += stride) {
        const int c = (int)(i % channels);
        x[i] = fused_op(x[i], has_bias ? b[c] : 0.0f, s[c]);
    }
}

// NCDHW fallback path: float4 spans 4 consecutive spatial locations
template<bool HAS_BIAS>
__global__ void fused_vec4_kernel(float4* __restrict__ x,
                                  const float* __restrict__ sum_tensor,
                                  const float* __restrict__ bias,
                                  int spatial_vec,
                                  int channels) {
    const int nc = blockIdx.y;
    const float sv = sum_tensor[nc % channels];
    const float bv = HAS_BIAS ? bias[nc % channels] : 0.0f;
    const long base = (long)nc * spatial_vec;
    float4* xs = x + base;
    const int stride = gridDim.x * blockDim.x;
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < spatial_vec; i += stride) {
        const float4 v = __ldcs(xs + i);
        float4 r;
        r.x = fused_op(v.x, bv, sv);
        r.y = fused_op(v.y, bv, sv);
        r.z = fused_op(v.z, bv, sv);
        r.w = fused_op(v.w, bv, sv);
        __stcs(xs + i, r);
    }
}

__global__ void fused_scalar_kernel(float* __restrict__ x,
                                    const float* __restrict__ sum_tensor,
                                    const float* __restrict__ bias,
                                    bool has_bias,
                                    long spatial_size,
                                    int channels) {
    const int nc = blockIdx.y;
    const float sv = sum_tensor[nc % channels];
    const float bv = has_bias ? bias[nc % channels] : 0.0f;
    const long base = (long)nc * spatial_size;
    float* xs = x + base;
    const long stride = (long)gridDim.x * blockDim.x;
    for (long i = (long)blockIdx.x * blockDim.x + threadIdx.x; i < spatial_size; i += stride) {
        xs[i] = fused_op(xs[i], bv, sv);
    }
}

torch::Tensor fused_postprocess_cuda(torch::Tensor x,
                                     torch::Tensor sum_tensor,
                                     torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

    auto sc = sum_tensor.contiguous();
    const bool has_bias = bias.numel() > 0;
    auto bc = has_bias ? bias.contiguous() : sum_tensor.contiguous();
    const int channels = (int)x.size(1);
    const int block = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (x.is_contiguous(at::MemoryFormat::ChannelsLast3d)) {
        const long total = x.numel();
        float* xp = x.data_ptr<float>();
        const bool aligned = ((((uintptr_t)xp) & 15) == 0);
        if ((channels & 3) == 0 && aligned) {
            const long total_vec = total >> 2;
            const int c_vec = channels >> 2;
            const long c_mask_l = ((c_vec & (c_vec - 1)) == 0) ? (long)(c_vec - 1) : -1L;
            if (total_vec <= (long)INT32_MAX) {
                const int tv = (int)total_vec;
                const int gx = (tv + block - 1) / block;
                const int c_mask = (int)c_mask_l;
                const bool use_smem = (c_vec <= 2048);
                const size_t shb = use_smem ? ((size_t)c_vec * (has_bias ? 2 : 1)) * sizeof(float4) : 0;
                const float4* s4 = reinterpret_cast<const float4*>(sc.data_ptr<float>());
                const float4* b4 = reinterpret_cast<const float4*>(bc.data_ptr<float>());
                float4* x4 = reinterpret_cast<float4*>(xp);
                if (use_smem && has_bias)
                    fused_cl_vec4_kernel<true, true><<<gx, block, shb, stream>>>(x4, s4, b4, tv, c_vec, c_mask);
                else if (use_smem)
                    fused_cl_vec4_kernel<true, false><<<gx, block, shb, stream>>>(x4, s4, b4, tv, c_vec, c_mask);
                else if (has_bias)
                    fused_cl_vec4_kernel<false, true><<<gx, block, 0, stream>>>(x4, s4, b4, tv, c_vec, c_mask);
                else
                    fused_cl_vec4_kernel<false, false><<<gx, block, 0, stream>>>(x4, s4, b4, tv, c_vec, c_mask);
            } else {
                long want = (total + block - 1) / block;
                int gx = (int)(want > 65535 ? 65535 : want);
                fused_cl_scalar_kernel<<<gx, block, 0, stream>>>(
                    xp, sc.data_ptr<float>(), bc.data_ptr<float>(), has_bias, total, channels);
            }
        } else {
            long want = (total + block - 1) / block;
            int gx = (int)(want > 65535 ? 65535 : want);
            fused_cl_scalar_kernel<<<gx, block, 0, stream>>>(
                xp, sc.data_ptr<float>(), bc.data_ptr<float>(), has_bias, total, channels);
        }
        return x;
    }

    auto xc = x.contiguous();
    long spatial = 1;
    for (int i = 2; i < xc.dim(); ++i) spatial *= xc.size(i);
    const long nc_total = xc.numel() / spatial;
    const bool aligned = ((((uintptr_t)xc.data_ptr<float>()) & 15) == 0);

    if (spatial % 4 == 0 && aligned && (spatial >> 2) <= (long)INT32_MAX) {
        const int spatial_vec = (int)(spatial >> 2);
        long want = (spatial_vec + block - 1) / block;
        int gx = (int)(want > 65535 ? 65535 : want);
        dim3 grid(gx, (unsigned)nc_total);
        if (has_bias)
            fused_vec4_kernel<true><<<grid, block, 0, stream>>>(
                reinterpret_cast<float4*>(xc.data_ptr<float>()),
                sc.data_ptr<float>(), bc.data_ptr<float>(), spatial_vec, channels);
        else
            fused_vec4_kernel<false><<<grid, block, 0, stream>>>(
                reinterpret_cast<float4*>(xc.data_ptr<float>()),
                sc.data_ptr<float>(), bc.data_ptr<float>(), spatial_vec, channels);
    } else {
        long want = (spatial + block - 1) / block;
        int gx = (int)(want > 65535 ? 65535 : want);
        dim3 grid(gx, (unsigned)nc_total);
        fused_scalar_kernel<<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(),
            has_bias, spatial, channels);
    }
    return xc;
}
"""

fused_cpp_source = (
    "void transpose_to_cl3d_cuda(torch::Tensor x, torch::Tensor y);\n"
    "torch::Tensor fused_postprocess_cuda(torch::Tensor x, torch::Tensor sum_tensor, torch::Tensor bias);"
)

fused_ext = load_inline(
    name="fused_conv3d_epilogue_v2",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_source,
    functions=["transpose_to_cl3d_cuda", "fused_postprocess_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model:
      * One-time conversion of conv weights to channels-last (NDHWC) so cuDNN
        runs its tensor-core NHWC kernels with no layout-transpose kernels.
      * Custom warp-coalesced NCDHW -> NDHWC transpose kernel for the input
        (replaces torch's slow generic permute-copy kernel).
      * Bias removed from the cuDNN call and folded into a single fused
        LeakyReLU -> +sum_tensor -> clamp -> GELU epilogue kernel that runs in
        place on the conv output (vectorized float4, streaming loads/stores).
    """
    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._w_cl = None
        self._w_cl_version = -1

    def _get_weight_cl(self):
        w = self.conv.weight
        if (self._w_cl is None or self._w_cl_version != w._version
                or self._w_cl.device != w.device):
            with torch.no_grad():
                self._w_cl = w.detach().contiguous(
                    memory_format=torch.channels_last_3d)
            self._w_cl_version = w._version
        return self._w_cl

    def forward(self, x):
        if (not x.is_cuda) or x.dtype != torch.float32:
            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)

        sum_flat = self.sum_tensor.view(-1)
        bias = self.conv.bias if self.conv.bias is not None else sum_flat[:0]

        # NCDHW -> NDHWC via the custom transpose kernel.
        xt = torch.empty_like(x, memory_format=torch.channels_last_3d)
        fused_ext.transpose_to_cl3d_cuda(x.contiguous(), xt)

        # Conv (no bias here; bias is folded into the fused epilogue).
        y = F.conv3d(xt, self._get_weight_cl(), None,
                     self.conv.stride, self.conv.padding,
                     self.conv.dilation, self.conv.groups)

        # Fused LeakyReLU -> +sum -> clamp -> GELU, in place.
        return fused_ext.fused_postprocess_cuda(y, sum_flat, bias)
```