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

torch.backends.cudnn.benchmark = True

# Fused CUDA kernel:
#   conv-bias add -> LeakyReLU(0.2) -> add per-channel sum_tensor -> clamp(-1, 1) -> exact GELU
# The convolution itself is delegated to cuDNN, but it is launched WITHOUT its bias
# so the bias addition is fused into this single element-wise pass, eliminating
# one full read/write sweep over the (large) activation tensor.
fused_post_ops_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

#define BLOCK_SIZE 256

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

__global__ void fused_post_ops_scalar_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    const float* __restrict__ sum_t,
    float* __restrict__ out,
    long long total,
    int channels,
    long long dhw)
{
    long long start  = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)gridDim.x * blockDim.x;
    for (long long idx = start; idx < total; idx += stride) {
        int c = (int)((idx / dhw) % channels);
        out[idx] = fused_op(x[idx], bias[c], sum_t[c]);
    }
}

__global__ void fused_post_ops_vec4_kernel(
    const float4* __restrict__ x,
    const float* __restrict__ bias,
    const float* __restrict__ sum_t,
    float4* __restrict__ out,
    long long total4,
    int channels,
    long long dhw)
{
    long long start  = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)gridDim.x * blockDim.x;
    for (long long idx = start; idx < total4; idx += stride) {
        long long base = idx * 4;
        int c = (int)((base / dhw) % channels);   // valid because dhw % 4 == 0
        float b = bias[c];
        float s = sum_t[c];
        float4 v = x[idx];
        v.x = fused_op(v.x, b, s);
        v.y = fused_op(v.y, b, s);
        v.z = fused_op(v.z, b, s);
        v.w = fused_op(v.w, b, s);
        out[idx] = v;
    }
}

torch::Tensor fused_post_ops_cuda(torch::Tensor x, torch::Tensor conv_bias, torch::Tensor sum_tensor) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "x must be float32");

    auto x_c   = x.contiguous();
    auto sum_c = sum_tensor.contiguous();
    auto out   = torch::empty_like(x_c);

    long long total = x_c.numel();
    if (total == 0) return out;

    int channels = x_c.size(1);
    long long dhw = 1;
    for (int i = 2; i < x_c.dim(); ++i) dhw *= x_c.size(i);

    torch::Tensor bias;
    if (conv_bias.defined() && conv_bias.numel() == channels) {
        bias = conv_bias.contiguous();
    } else {
        bias = torch::zeros({channels}, x_c.options());
    }

    const long long max_blocks = 1048576;

    if ((total % 4 == 0) && (dhw % 4 == 0)) {
        long long total4 = total / 4;
        long long needed = (total4 + BLOCK_SIZE - 1) / BLOCK_SIZE;
        int num_blocks = (int)(needed < max_blocks ? needed : max_blocks);
        fused_post_ops_vec4_kernel<<<num_blocks, BLOCK_SIZE>>>(
            reinterpret_cast<const float4*>(x_c.data_ptr<float>()),
            bias.data_ptr<float>(),
            sum_c.data_ptr<float>(),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            total4, channels, dhw);
    } else {
        long long needed = (total + BLOCK_SIZE - 1) / BLOCK_SIZE;
        int num_blocks = (int)(needed < max_blocks ? needed : max_blocks);
        fused_post_ops_scalar_kernel<<<num_blocks, BLOCK_SIZE>>>(
            x_c.data_ptr<float>(),
            bias.data_ptr<float>(),
            sum_c.data_ptr<float>(),
            out.data_ptr<float>(),
            total, channels, dhw);
    }
    return out;
}
"""

fused_post_ops_cpp_source = (
    "torch::Tensor fused_post_ops_cuda(torch::Tensor x, torch::Tensor conv_bias, torch::Tensor sum_tensor);"
)

fused_post_ops = load_inline(
    name="fused_post_ops",
    cpp_sources=fused_post_ops_cpp_source,
    cuda_sources=fused_post_ops_source,
    functions=["fused_post_ops_cuda"],
    verbose=True,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuDNN 3D convolution (bias stripped) followed by a single
    fused CUDA kernel performing bias-add + LeakyReLU + tensor-sum + clamp + GELU.
    """
    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.fused_post_ops = fused_post_ops

    def forward(self, x):
        # Run convolution without bias; the bias is folded into the fused kernel
        # to avoid an extra memory pass over the activation tensor.
        conv_out = F.conv3d(
            x,
            self.conv.weight,
            None,
            self.conv.stride,
            self.conv.padding,
            self.conv.dilation,
            self.conv.groups,
        )
        bias = self.conv.bias
        if bias is None:
            bias = torch.empty(0, device=x.device, dtype=x.dtype)
        return self.fused_post_ops.fused_post_ops_cuda(conv_out, bias, self.sum_tensor)
```