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
from torch.utils.cpp_extension import load_inline
import numpy as _np

# Chebyshev approximation coefficients are generated once at module load and
# embedded into the CUDA source before compilation.  This replaces the two
# expensive activation kernels with a single lightweight in-place kernel after
# cuDNN's convolution.
_xs = _np.linspace(-5.0, 5.0, 300001, dtype=_np.float64)
_ys = _np.tanh(_xs * _np.tanh(_np.log1p(_np.exp(_xs))))
_coef = _np.polynomial.chebyshev.chebfit(_xs / 5.0, _ys, 36).astype(_np.float32)
_lines = ['float t = x * 0.2f;', 'float b0 = 0.0f, b1 = 0.0f, b2 = 0.0f;']
for _c in _coef[:0:-1]:
    _lines.append('b2 = b1; b1 = b0; b0 = fmaf(2.0f*t, b1, %.9g - b2);' % float(_c))
_lines.append('return fmaf(t, b0, %.9g - b1);' % float(_coef[0]))
_poly_body = '\n        '.join(_lines)

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

torch::Tensor conv_mish_tanh(torch::Tensor input, torch::Tensor weight, torch::Tensor bias,
                             int64_t stride, int64_t padding);
'''

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

__device__ __forceinline__ float exact_act(float x) {
    float sp;
    if (x > 20.0f) sp = x;
    else if (x < -20.0f) sp = expf(x);
    else sp = log1pf(expf(x));
    return tanhf(x * tanhf(sp));
}

__device__ __forceinline__ float poly_act(float x) {
    if (x < -5.0f || x > 5.0f) return exact_act(x);
        ''' + _poly_body + r'''
}

__global__ void mish_tanh_kernel(float* __restrict__ data, int64_t n) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t grid_stride = (int64_t)blockDim.x * gridDim.x;
    for (int64_t i = idx; i < n; i += grid_stride) {
        data[i] = poly_act(data[i]);
    }
}

torch::Tensor conv_mish_tanh(torch::Tensor input, torch::Tensor weight, torch::Tensor bias,
                             int64_t stride, int64_t padding) {
    auto y = torch::conv3d(input, weight, bias,
                           {stride, stride, stride},
                           {padding, padding, padding},
                           {1, 1, 1}, 1);
    const int64_t n = y.numel();
    if (n > 0) {
        const int threads = 256;
        int blocks = (int)((n + threads - 1) / threads);
        if (blocks > 4096) blocks = 4096;
        mish_tanh_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(y.data_ptr<float>(), n);
    }
    return y;
}
'''

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

class ModelNew(nn.Module):
    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)
        self.stride = stride if isinstance(stride, int) else stride[0]
        self.padding = padding if isinstance(padding, int) else padding[0]

    def forward(self, x):
        return _ext.conv_mish_tanh(x, self.conv.weight, self.conv.bias, int(self.stride), int(self.padding))