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_act_bias_(torch::Tensor y, torch::Tensor bias);
"""

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

__global__ void fused_act_bias_kernel(float* __restrict__ y, const float* __restrict__ bias,
                                      long long total, int C, int spatial) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    const float k0 = 0.7978845608028654f;
    const float k1 = 0.044715f;
    for (; idx < total; idx += stride) {
        float x = y[idx];
        int c = (int)((idx / spatial) % C);
        float b = bias[c];
        if (x <= 0.0f) {
            y[idx] = 0.5f + b;
        } else {
            float x2 = x * x;
            float g = 0.5f * x * (1.0f + tanhf(k0 * (x + k1 * x * x2)));
            float s = 1.0f / (1.0f + expf(-g));
            y[idx] = s + b;
        }
    }
}

torch::Tensor fused_act_bias_(torch::Tensor y, torch::Tensor bias) {
    long long total = y.numel();
    if (total == 0) return y;
    int C = (int)y.size(1);
    int spatial = (int)(y.size(2) * y.size(3) * y.size(4));
    int threads = 256;
    int blocks = (int)((total + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;
    fused_act_bias_kernel<<<blocks, threads>>>(y.data_ptr<float>(), bias.data_ptr<float>(), total, C, spatial);
    return y;
}
"""

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

class ModelNew(nn.Module):
    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,
                     stride=self.conv.stride, padding=self.conv.padding,
                     dilation=self.conv.dilation, groups=self.conv.groups)
        return _ext.fused_act_bias_(y, self.bias)