```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 ideas on top of the previous (2.139x) version:
#  1) Keep the conv natively in channels-last (NDHWC) so cuDNN runs its
#     tensor-core NHWC kernels with no layout-transpose kernels.
#  2) Single fused LeakyReLU -> +sum_tensor -> clamp -> GELU epilogue kernel,
#     now with streaming (evict-first) loads/stores, shared-memory caching of
#     the per-channel sum tensor, and one float4 per thread.
#  3) Chunk the batch across several CUDA streams in a software pipeline:
#     while chunk i's bandwidth-bound epilogue runs, chunk i+1's compute-bound
#     convolution executes concurrently.  The epilogue writes directly into
#     batch-slices of a preallocated channels-last output, so no torch.cat
#     copy is needed.
fused_postprocess_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

__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 ----
template<bool USE_SMEM>
__global__ void fused_cl_vec4_kernel(const float4* __restrict__ x,
                                     const float4* __restrict__ s4,
                                     float4* __restrict__ out,
                                     int total_vec,
                                     int c_vec,
                                     int c_mask) {
    extern __shared__ float4 sh[];
    if (USE_SMEM) {
        for (int t = threadIdx.x; t < c_vec; t += blockDim.x) sh[t] = s4[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);
        const float4 v = __ldcs(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);
        __stcs(out + i, r);
    }
}

__global__ void fused_cl_vec4_long_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 = __ldcs(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);
        __stcs(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 = __ldcs(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);
        __stcs(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::Tensor out) {
    TORCH_CHECK(x.is_cuda() && out.is_cuda(), "x/out must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && out.scalar_type() == at::kFloat, "float32 only");
    TORCH_CHECK(x.numel() == out.numel(), "x/out shape mismatch");

    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) &&
        out.is_contiguous(at::MemoryFormat::ChannelsLast3d)) {
        const long total = x.numel();
        const float* xp = x.data_ptr<float>();
        float* op = out.data_ptr<float>();
        const bool aligned = ((((uintptr_t)xp) & 15) == 0) && ((((uintptr_t)op) & 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;
                if (c_vec <= 2048) {
                    fused_cl_vec4_kernel<true><<<gx, block, (size_t)c_vec * sizeof(float4), stream>>>(
                        reinterpret_cast<const float4*>(xp),
                        reinterpret_cast<const float4*>(sc.data_ptr<float>()),
                        reinterpret_cast<float4*>(op), tv, c_vec, c_mask);
                } else {
                    fused_cl_vec4_kernel<false><<<gx, block, 0, stream>>>(
                        reinterpret_cast<const float4*>(xp),
                        reinterpret_cast<const float4*>(sc.data_ptr<float>()),
                        reinterpret_cast<float4*>(op), tv, c_vec, c_mask);
                }
            } else {
                long want = (total_vec + block - 1) / block;
                int gx = (int)(want > 65535 ? 65535 : want);
                fused_cl_vec4_long_kernel<<<gx, block, 0, stream>>>(
                    reinterpret_cast<const float4*>(xp),
                    reinterpret_cast<const float4*>(sc.data_ptr<float>()),
                    reinterpret_cast<float4*>(op), total_vec, c_vec, c_mask_l);
            }
        } 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>(), op, total, channels);
        }
        return out;
    }

    auto xc = x.contiguous();
    auto oc = out.is_contiguous() ? out : out.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) &&
                         ((((uintptr_t)oc.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);
        fused_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            sc.data_ptr<float>(),
            reinterpret_cast<float4*>(oc.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>(), oc.data_ptr<float>(),
            spatial, channels);
    }
    if (oc.data_ptr() != out.data_ptr()) out.copy_(oc);
    return out;
}
"""

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

fused_postprocess = load_inline(
    name="fused_postprocess_pipe",
    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 (NDHWC)
    layout so cuDNN uses its tensor-core NHWC kernels without layout transpose
    kernels, pipelines batch chunks across CUDA streams to overlap the
    compute-bound convolutions with the bandwidth-bound epilogues, and applies
    LeakyReLU + per-channel sum + clamp + GELU in one fused, vectorized CUDA
    kernel written directly into a preallocated channels-last output.
    """
    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))
        # One-time: keep conv weights channels-last so cuDNN never inserts
        # NCDHW<->NDHWC transpose kernels.
        with torch.no_grad():
            self.conv.weight.data = self.conv.weight.data.contiguous(
                memory_format=torch.channels_last_3d)
        self._streams = None
        self._stream_dev = None
        self._max_chunks = 4

    def _out_spatial(self, x):
        k = self.conv.kernel_size
        st = self.conv.stride
        p = self.conv.padding
        dl = self.conv.dilation
        D, H, W = x.size(2), x.size(3), x.size(4)
        od = (D + 2 * p[0] - dl[0] * (k[0] - 1) - 1) // st[0] + 1
        oh = (H + 2 * p[1] - dl[1] * (k[1] - 1) - 1) // st[1] + 1
        ow = (W + 2 * p[2] - dl[2] * (k[2] - 1) - 1) // st[2] + 1
        return od, oh, ow

    def forward(self, x):
        sum_flat = self.sum_tensor.view(-1)
        n = x.size(0)

        n_chunks = self._max_chunks
        while n_chunks > 1 and (n % n_chunks) != 0:
            n_chunks //= 2

        # Single-chunk path: transpose once, conv, in-place fused epilogue.
        if n_chunks <= 1 or not x.is_cuda:
            xc = x.contiguous(memory_format=torch.channels_last_3d)
            y = self.conv(xc)
            return fused_postprocess.fused_postprocess_cuda(y, sum_flat, y)

        dev = x.device
        if (self._streams is None or len(self._streams) != n_chunks
                or self._stream_dev != dev):
            self._streams = [torch.cuda.Stream(device=dev) for _ in range(n_chunks)]
            self._stream_dev = dev
        cur = torch.cuda.current_stream(dev)

        od, oh, ow = self._out_spatial(x)
        out = torch.empty((n, self.conv.out_channels, od, oh, ow),
                          device=dev, dtype=x.dtype,
                          memory_format=torch.channels_last_3d)

        chunk = n // n_chunks
        for s in self._streams:
            s.wait_stream(cur)
        for i, s in enumerate(self._streams):
            with torch.cuda.stream(s):
                xi = x.narrow(0, i * chunk, chunk).contiguous(
                    memory_format=torch.channels_last_3d)
                yi = self.conv(xi)
                oi = out.narrow(0, i * chunk, chunk)
                fused_postprocess.fused_postprocess_cuda(yi, sum_flat, oi)
        for s in self._streams:
            cur.wait_stream(s)
        return out
```