import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D convolution, applies LeakyReLU, sums with a tensor, clamps, and applies GELU activation.
    """
    def __init__(self, in_channels, out_channels, kernel_size, sum_tensor_shape):
        super(Model, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.sum_tensor = nn.Parameter(torch.randn(sum_tensor_shape))

    def forward(self, x):
        x = self.conv(x)
        x = torch.nn.functional.leaky_relu(x, negative_slope=0.2)
        x = x + self.sum_tensor
        x = torch.clamp(x, min=-1.0, max=1.0)
        x = torch.nn.functional.gelu(x)
        return x

batch_size = 128
in_channels = 8
out_channels = 64
depth, height, width = 16, 64, 64
kernel_size = 3
sum_tensor_shape = (out_channels, 1, 1, 1)

def get_inputs():
    return [torch.rand(batch_size, in_channels, depth, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, sum_tensor_shape]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = r'''
torch::Tensor fused_post(torch::Tensor y, torch::Tensor s);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>
#include <algorithm>

static __device__ __forceinline__ float gelu_mid(float x) {
    return 0.5f * x * (1.0f + erff(x * 0.70710678118654752440f));
}

static __device__ __forceinline__ float post_one(float v, float addv) {
    float x = v > 0.0f ? v : 0.2f * v;
    x += addv;
    if (x >= 1.0f) return 0.8413447141647339f;
    if (x <= -1.0f) return -0.1586552560329437f;
    return gelu_mid(x);
}

__global__ void fused_post_plane_vec4_kernel(float* __restrict__ y, const float* __restrict__ s, int s4) {
    int nc = blockIdx.x;
    int c = nc & 63;
    float addv = s[c];
    float4* __restrict__ p = reinterpret_cast<float4*>(y + ((long long)nc * s4 * 4LL));
    for (int j = threadIdx.x; j < s4; j += blockDim.x) {
        float4 v = p[j];
        v.x = post_one(v.x, addv);
        v.y = post_one(v.y, addv);
        v.z = post_one(v.z, addv);
        v.w = post_one(v.w, addv);
        p[j] = v;
    }
}

__global__ void fused_post_scalar_kernel(float* __restrict__ y, const float* __restrict__ s, long long total, int channels, int spatial) {
    long long tid = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    for (long long i = tid; i < total; i += stride) {
        int c = (int)((i / spatial) % channels);
        y[i] = post_one(y[i], s[c]);
    }
}

torch::Tensor fused_post(torch::Tensor y, torch::Tensor s) {
    long long total = (long long)y.numel();
    const int n = (int)y.size(0);
    const int channels = (int)y.size(1);
    const int spatial = (int)(y.size(2) * y.size(3) * y.size(4));
    const int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if ((spatial % 4 == 0) && channels == 64) {
        int s4 = spatial / 4;
        fused_post_plane_vec4_kernel<<<n * channels, threads, 0, stream>>>(y.data_ptr<float>(), s.data_ptr<float>(), s4);
    } else {
        int blocks = (int)std::min(65535LL, (total + threads - 1LL) / threads);
        fused_post_scalar_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), s.data_ptr<float>(), total, channels, spatial);
    }
    return y;
}
'''

_ext = load_inline(
    name='fused_conv3d_post_ext_final',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['fused_post'],
    extra_cuda_cflags=['-O3'],
    extra_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    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))

    def forward(self, x):
        y = self.conv(x)
        return _ext.fused_post(y, self.sum_tensor)