```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

# 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

# Key idea: cuDNN's fastest algo for this conv is a native NHWC (channels-last
# 3D) tensor-core kernel. In the previous version the NCDHW tensors forced
# cuDNN to pay for NCHW->NHWC and NHWC->NCHW transposes, and our post-process
# kernel then ran on the transposed layout (~5.7ms of pure layout traffic).
# Here we keep weights/activations channels-last end to end, and run the fused
# LeakyReLU -> +sum_tensor -> clamp -> GELU epilogue directly on the NHWC data
# (float4 spans 4 consecutive channels), so no transpose kernels remain except
# a single one-time conversion of the input.
fused_postprocess_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float fused_op(float v, float s) {
    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 (NDHWC) path: float4 spans 4 consecutive channels ----
__global__ void fused_cl_vec4_kernel(const float4* __restrict__ x,
                                     const float4* __restrict__ s4,
                                     float4* __restrict__ out,
                                     long total_vec,
                                     int c_vec,
                                     long c_mask) {
    const long stride = (long)gridDim.x * blockDim.x;
    for (long i = (long)blockIdx.x * blockDim.x + threadIdx.x; i < total_vec; i += stride) {
        const int cv = (c_mask >= 0) ? (int)(i & c_mask) : (int)(i % c_vec);
        const float4 sv = s4[cv];
        const float4 v = x[i];
        float4 r;
        r.x = fused_op(v.x, sv.x);
        r.y = fused_op(v.y, sv.y);
        r.z = fused_op(v.z, sv.z);
        r.w = fused_op(v.w, sv.w);
        out[i] = r;
    }
}

__global__ void fused_cl_scalar_kernel(const float* __restrict__ x,
                                       const float* __restrict__ s,
                                       float* __restrict__ out,
                                       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) {
        out[i] = fused_op(x[i], s[i % channels]);
    }
}

// ---- NCDHW fallback path: float4 spans 4 consecutive spatial locations ----
__global__ void fused_vec4_kernel(const float4* __restrict__ x,
                                  const float* __restrict__ sum_tensor,
                                  float4* __restrict__ out,
                                  int spatial_vec,
                                  int channels) {
    const int nc = blockIdx.y;
    const float sv = sum_tensor[nc % channels];
    const long base = (long)nc * spatial_vec;
    const float4* xs = x + base;
    float4* os = out + 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 = xs[i];
        float4 r;
        r.x = fused_op(v.x, sv);
        r.y = fused_op(v.y, sv);
        r.z = fused_op(v.z, sv);
        r.w = fused_op(v.w, sv);
        os[i] = r;
    }
}

__global__ void fused_scalar_kernel(const float* __restrict__ x,
                                    const float* __restrict__ sum_tensor,
                                    float* __restrict__ out,
                                    long spatial_size,
                                    int channels) {
    const int nc = blockIdx.y;
    const float sv = sum_tensor[nc % channels];
    const long base = (long)nc * spatial_size;
    const float* xs = x + base;
    float* os = out + base;
    const long stride = (long)gridDim.x * blockDim.x;
    for (long i = (long)blockIdx.x * blockDim.x + threadIdx.x; i < spatial_size; i += stride) {
        os[i] = fused_op(xs[i], sv);
    }
}

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

    auto sc = 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();
        if ((channels & 3) == 0) {
            const long total_vec = total >> 2;
            const int c_vec = channels >> 2;
            const long c_mask = ((c_vec & (c_vec - 1)) == 0) ? (long)(c_vec - 1) : -1L;
            long want = (total_vec + block - 1) / block;
            int gx = (int)(want > 65535 ? 65535 : want);
            fused_cl_vec4_kernel<<<gx, block, 0, stream>>>(
                reinterpret_cast<const float4*>(x.data_ptr<float>()),
                reinterpret_cast<const float4*>(sc.data_ptr<float>()),
                reinterpret_cast<float4*>(x.data_ptr<float>()),
                total_vec, 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>>>(
                x.data_ptr<float>(), sc.data_ptr<float>(), x.data_ptr<float>(),
                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;

    if (spatial % 4 == 0) {
        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);
        fused_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            sc.data_ptr<float>(),
            reinterpret_cast<float4*>(xc.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>(), xc.data_ptr<float>(),
            spatial, channels);
    }
    return xc;
}
"""

fused_postprocess_cpp_source = (
    "torch::Tensor fused_postprocess_cuda(torch::Tensor x, torch::Tensor sum_tensor);"
)

fused_postprocess = load_inline(
    name="fused_postprocess_cl",
    cpp_sources=fused_postprocess_cpp_source,
    cuda_sources=fused_postprocess_source,
    functions=["fused_postprocess_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: runs the 3D convolution natively in channels-last (NHWC)
    layout so cuDNN uses its tensor-core NHWC kernels without any layout
    transpose kernels, then applies LeakyReLU + per-channel sum + clamp + GELU
    in a single fused, vectorized CUDA kernel over the NHWC data.
    """
    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))
        # Keep conv weights in channels-last layout (one-time cost) so cuDNN
        # never inserts NCHW<->NHWC transpose kernels.
        with torch.no_grad():
            self.conv.weight.data = self.conv.weight.data.contiguous(
                memory_format=torch.channels_last_3d)

    def forward(self, x):
        x = x.contiguous(memory_format=torch.channels_last_3d)
        x = self.conv(x)
        x = fused_postprocess.fused_postprocess_cuda(
            x, self.sum_tensor.view(-1))
        return x
```