import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D convolution, scales the output, applies tanh, multiplies by a scaling factor, and applies sigmoid.
    """
    def __init__(self, in_channels, out_channels, kernel_size, scaling_factor, bias_shape):
        super(Model, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.scaling_factor = nn.Parameter(torch.randn(bias_shape))
        self.bias = nn.Parameter(torch.randn(bias_shape)) 

    def forward(self, x):
        x = self.conv(x)
        x = x * self.scaling_factor 
        x = torch.tanh(x)
        x = x * self.bias
        x = torch.sigmoid(x)
        return x

batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 64, 64
kernel_size = 3
scaling_factor = 2
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, scaling_factor, bias_shape]

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

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

static inline void check_inputs(const torch::Tensor& x, const torch::Tensor& w) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(w.is_cuda(), "w must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only fp32 supported");
}

__device__ __forceinline__ float sigmoidf_fast(float x) {
    return 1.0f / (1.0f + expf(-x));
}

// Specialized for the benchmark shape: C=3, K=3, O arbitrary (16 here), contiguous NCDHW.
__global__ __launch_bounds__(256, 2)
void conv3d_fused_k3_c3_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ conv_b,
    const float* __restrict__ scale,
    const float* __restrict__ post_b,
    float* __restrict__ y,
    long total,
    int N, int O, int D, int H, int W,
    int Do, int Ho, int Wo)
{
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long stride = (long)blockDim.x * gridDim.x;
    const int HW = H * W;
    const int DHW = D * HW;
    const int out_spatial = Do * Ho * Wo;
    const int out_hw = Ho * Wo;

    for (; idx < total; idx += stride) {
        long t = idx;
        int ow = (int)(t % Wo); t /= Wo;
        int oh = (int)(t % Ho); t /= Ho;
        int od = (int)(t % Do); t /= Do;
        int oc = (int)(t % O);  t /= O;
        int n  = (int)t;

        const float* xp = x + ((long)n * 3 * D + od) * HW + oh * W + ow;
        const float* wp = w + (long)oc * 3 * 27;
        float acc = conv_b ? __ldg(conv_b + oc) : 0.0f;

        #pragma unroll
        for (int c = 0; c < 3; ++c) {
            const float* xpc = xp + (long)c * DHW;
            const float* wpc = wp + c * 27;
            // kd = 0
            acc = fmaf(__ldg(xpc + 0*HW + 0*W + 0), __ldg(wpc + 0), acc);
            acc = fmaf(__ldg(xpc + 0*HW + 0*W + 1), __ldg(wpc + 1), acc);
            acc = fmaf(__ldg(xpc + 0*HW + 0*W + 2), __ldg(wpc + 2), acc);
            acc = fmaf(__ldg(xpc + 0*HW + 1*W + 0), __ldg(wpc + 3), acc);
            acc = fmaf(__ldg(xpc + 0*HW + 1*W + 1), __ldg(wpc + 4), acc);
            acc = fmaf(__ldg(xpc + 0*HW + 1*W + 2), __ldg(wpc + 5), acc);
            acc = fmaf(__ldg(xpc + 0*HW + 2*W + 0), __ldg(wpc + 6), acc);
            acc = fmaf(__ldg(xpc + 0*HW + 2*W + 1), __ldg(wpc + 7), acc);
            acc = fmaf(__ldg(xpc + 0*HW + 2*W + 2), __ldg(wpc + 8), acc);
            // kd = 1
            acc = fmaf(__ldg(xpc + 1*HW + 0*W + 0), __ldg(wpc + 9), acc);
            acc = fmaf(__ldg(xpc + 1*HW + 0*W + 1), __ldg(wpc + 10), acc);
            acc = fmaf(__ldg(xpc + 1*HW + 0*W + 2), __ldg(wpc + 11), acc);
            acc = fmaf(__ldg(xpc + 1*HW + 1*W + 0), __ldg(wpc + 12), acc);
            acc = fmaf(__ldg(xpc + 1*HW + 1*W + 1), __ldg(wpc + 13), acc);
            acc = fmaf(__ldg(xpc + 1*HW + 1*W + 2), __ldg(wpc + 14), acc);
            acc = fmaf(__ldg(xpc + 1*HW + 2*W + 0), __ldg(wpc + 15), acc);
            acc = fmaf(__ldg(xpc + 1*HW + 2*W + 1), __ldg(wpc + 16), acc);
            acc = fmaf(__ldg(xpc + 1*HW + 2*W + 2), __ldg(wpc + 17), acc);
            // kd = 2
            acc = fmaf(__ldg(xpc + 2*HW + 0*W + 0), __ldg(wpc + 18), acc);
            acc = fmaf(__ldg(xpc + 2*HW + 0*W + 1), __ldg(wpc + 19), acc);
            acc = fmaf(__ldg(xpc + 2*HW + 0*W + 2), __ldg(wpc + 20), acc);
            acc = fmaf(__ldg(xpc + 2*HW + 1*W + 0), __ldg(wpc + 21), acc);
            acc = fmaf(__ldg(xpc + 2*HW + 1*W + 1), __ldg(wpc + 22), acc);
            acc = fmaf(__ldg(xpc + 2*HW + 1*W + 2), __ldg(wpc + 23), acc);
            acc = fmaf(__ldg(xpc + 2*HW + 2*W + 0), __ldg(wpc + 24), acc);
            acc = fmaf(__ldg(xpc + 2*HW + 2*W + 1), __ldg(wpc + 25), acc);
            acc = fmaf(__ldg(xpc + 2*HW + 2*W + 2), __ldg(wpc + 26), acc);
        }
        float v = acc * __ldg(scale + oc);
        v = tanhf(v);
        v *= __ldg(post_b + oc);
        y[idx] = sigmoidf_fast(v);
    }
}

// Generic direct fused kernel for less common constructor variants.
__global__ void conv3d_fused_generic_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ conv_b,
    const float* __restrict__ scale,
    const float* __restrict__ post_b,
    float* __restrict__ y,
    long total,
    int N, int C, int O, int D, int H, int W, int Kd, int Kh, int Kw,
    int Do, int Ho, int Wo)
{
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long stride = (long)blockDim.x * gridDim.x;
    const int HW = H * W;
    const int DHW = D * HW;
    const int out_hw = Ho * Wo;
    const int KHW = Kh * Kw;
    const int Ktot = Kd * KHW;
    for (; idx < total; idx += stride) {
        long t = idx;
        int ow = (int)(t % Wo); t /= Wo;
        int oh = (int)(t % Ho); t /= Ho;
        int od = (int)(t % Do); t /= Do;
        int oc = (int)(t % O);  t /= O;
        int n = (int)t;
        float acc = conv_b ? __ldg(conv_b + oc) : 0.0f;
        for (int c=0; c<C; ++c) {
            for (int kd=0; kd<Kd; ++kd) {
                for (int kh=0; kh<Kh; ++kh) {
                    const float* xp = x + ((long)n*C + c)*DHW + (od+kd)*HW + (oh+kh)*W + ow;
                    const float* wp = w + ((long)oc*C + c)*Ktot + kd*KHW + kh*Kw;
                    for (int kw=0; kw<Kw; ++kw) acc = fmaf(__ldg(xp+kw), __ldg(wp+kw), acc);
                }
            }
        }
        float v = acc * __ldg(scale + oc);
        v = tanhf(v);
        v *= __ldg(post_b + oc);
        y[idx] = 1.0f / (1.0f + expf(-v));
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor scale, torch::Tensor post_bias) {
    check_inputs(x, weight);
    TORCH_CHECK(x.dim() == 5 && weight.dim() == 5, "expected 5D tensors");
    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto sc = scale.contiguous();
    auto pc = post_bias.contiguous();
    torch::Tensor bc = conv_bias.defined() ? conv_bias.contiguous() : conv_bias;

    int N = (int)xc.size(0), C = (int)xc.size(1), D = (int)xc.size(2), H = (int)xc.size(3), W = (int)xc.size(4);
    int O = (int)wc.size(0), Kd = (int)wc.size(2), Kh = (int)wc.size(3), Kw = (int)wc.size(4);
    int Do = D - Kd + 1, Ho = H - Kh + 1, Wo = W - Kw + 1;
    auto y = torch::empty({N, O, Do, Ho, Wo}, xc.options());
    long total = y.numel();
    if (total == 0) return y;
    int threads = 256;
    int blocks = (int)((total + threads - 1) / threads);
    int maxBlocks = 65535;
    if (blocks > maxBlocks) blocks = maxBlocks;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* bp = bc.defined() ? bc.data_ptr<float>() : nullptr;
    if (C == 3 && Kd == 3 && Kh == 3 && Kw == 3) {
        conv3d_fused_k3_c3_kernel<<<blocks, threads, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bp, sc.data_ptr<float>(), pc.data_ptr<float>(), y.data_ptr<float>(),
            total, N, O, D, H, W, Do, Ho, Wo);
    } else {
        conv3d_fused_generic_kernel<<<blocks, threads, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bp, sc.data_ptr<float>(), pc.data_ptr<float>(), y.data_ptr<float>(),
            total, N, C, O, D, H, W, Kd, Kh, Kw, Do, Ho, Wo);
    }
    return y;
}
'''

_ext = load_inline(
    name='fused_conv3d_tanh_sigmoid_ext_v1',
    cpp_sources='',
    cuda_sources=_cuda_src,
    functions=['fused_forward'],
    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, scaling_factor, bias_shape):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.scaling_factor = nn.Parameter(torch.randn(bias_shape))
        self.bias = nn.Parameter(torch.randn(bias_shape))

    def forward(self, x):
        return _ext.fused_forward(x, self.conv.weight, self.conv.bias, self.scaling_factor, self.bias)