import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D convolution, applies Mish activation, and then applies Tanh activation.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
        super(Model, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, D, H, W).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, D', H', W').
        """
        x = self.conv(x)
        x = torch.nn.functional.mish(x)
        x = torch.tanh(x)
        return x

batch_size = 16
in_channels = 32
out_channels = 64
D, H, W = 32, 64, 64
kernel_size = 3

def get_inputs():
    return [torch.rand(batch_size, in_channels, D, H, W)]

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

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 mish_tanh_inplace(torch::Tensor x);
"""

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

__device__ __forceinline__ float fast_tanh_f(float x) {
    if (x > 10.0f) return 1.0f;
    if (x < -10.0f) return -1.0f;
    float ax = fabsf(x);
    float e = __expf(-2.0f * ax);
    float y = (1.0f - e) / (1.0f + e);
    return copysignf(y, x);
}

__device__ __forceinline__ float softplus_f(float x) {
    // Stable softplus.  The tails also avoid unnecessary log work.
    if (x > 20.0f) return x;
    if (x < -20.0f) return __expf(x);
    return log1pf(__expf(x));
}

__global__ void mish_tanh_kernel(float* __restrict__ p, int64_t n) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (; idx < n; idx += stride) {
        float x = p[idx];
        float sp = softplus_f(x);
        float mish = x * fast_tanh_f(sp);
        p[idx] = fast_tanh_f(mish);
    }
}

torch::Tensor mish_tanh_inplace(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(x.is_contiguous(), "input must be contiguous");
    const int64_t n = x.numel();
    if (n == 0) return x;
    const int threads = 256;
    int blocks = (int)((n + threads - 1) / threads);
    // Enough blocks to cover the device while keeping launch overhead small.
    if (blocks > 65535) blocks = 65535;
    mish_tanh_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), n);
    return x;
}
"""

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


class ModelNew(nn.Module):
    """
    3D convolution followed by fused Mish and Tanh activation.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)

    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)
        # Conv3d returns a contiguous fresh tensor; apply both pointwise activations in one CUDA pass.
        return _ext.mish_tanh_inplace(y)