import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D convolution, applies ReLU, LeakyReLU, GELU, Sigmoid activations, and bias in sequence.
    """
    def __init__(self, in_channels, out_channels, kernel_size, bias_shape):
        super(Model, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.bias = nn.Parameter(torch.randn(bias_shape)) 

    def forward(self, x):
        x = self.conv(x)
        x = torch.relu(x)
        x = torch.nn.functional.leaky_relu(x, negative_slope=0.01)
        x = torch.nn.functional.gelu(x)
        x = torch.sigmoid(x)
        x = x + self.bias
        return x

batch_size = 64
in_channels = 8
out_channels = 32
depth, height, width = 32, 64, 64
kernel_size = 3
bias_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, bias_shape]

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

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor fused_post_act_bias(torch::Tensor x, torch::Tensor bias);
"""

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

__global__ void post_kernel(float* __restrict__ x,
                            const float* __restrict__ bias,
                            long long n_elem,
                            int channels,
                            int spatial) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    const float inv_sqrt2 = 0.70710678118654752440f;

    for (; idx < n_elem; idx += stride) {
        float v = x[idx];
        // ReLU followed by leaky_relu(negative_slope=0.01) is just ReLU.
        v = v > 0.0f ? v : 0.0f;
        // Exact GELU, matching torch.nn.functional.gelu default (approximate='none').
        float g = 0.5f * v * (1.0f + erff(v * inv_sqrt2));
        float s = 1.0f / (1.0f + expf(-g));
        int c = (int)((idx / spatial) % channels);
        x[idx] = s + bias[c];
    }
}

torch::Tensor fused_post_act_bias(torch::Tensor x, torch::Tensor bias) {
    const long long n_elem = x.numel();
    if (n_elem == 0) return x;
    const int channels = (int)x.size(1);
    const int spatial = (int)(x.size(2) * x.size(3) * x.size(4));

    const int threads = 256;
    int blocks = (int)((n_elem + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    post_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(),
                                                bias.data_ptr<float>(),
                                                n_elem,
                                                channels,
                                                spatial);
    return x;
}
"""

_ext = load_inline(
    name="fused_conv3d_post_act_bias_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_post_act_bias"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Conv3d followed by fused ReLU/LeakyReLU/GELU/Sigmoid/channel-bias.
    """
    def __init__(self, in_channels, out_channels, kernel_size, bias_shape):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.bias = nn.Parameter(torch.randn(bias_shape))

    def forward(self, x):
        y = F.conv3d(x, self.conv.weight, self.conv.bias, self.conv.stride,
                     self.conv.padding, self.conv.dilation, self.conv.groups)
        return _ext.fused_post_act_bias(y, self.bias)