import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, layer normalization, GELU activation, and scaling.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias=True, eps=1e-5, scaling_factor=1.0):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, bias=bias)
        self.layer_norm = nn.LayerNorm(out_channels, eps=eps)
        self.scaling_factor = scaling_factor

    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_transpose(x)
        x = self.layer_norm(x)
        x = torch.nn.functional.gelu(x)
        x = x * self.scaling_factor
        return x

batch_size = 32
in_channels = 32
out_channels = 64
D, H, W = 16, 32, 32
kernel_size = 4
stride = 2
padding = 1
bias = True
eps = 1e-5
scaling_factor = 1.0

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, bias, eps, scaling_factor]

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_ln_gelu_inplace(torch::Tensor input,
                                   torch::Tensor gamma,
                                   torch::Tensor beta,
                                   double eps,
                                   double scale);
'''

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

static __forceinline__ __device__ float warp_sum(float v) {
    unsigned mask = 0xffffffffu;
    v += __shfl_down_sync(mask, v, 16);
    v += __shfl_down_sync(mask, v, 8);
    v += __shfl_down_sync(mask, v, 4);
    v += __shfl_down_sync(mask, v, 2);
    v += __shfl_down_sync(mask, v, 1);
    return v;
}

__global__ __launch_bounds__(256, 2)
void ln_gelu64_kernel(float* __restrict__ x,
                      const float* __restrict__ gamma,
                      const float* __restrict__ beta,
                      long long rows,
                      float eps,
                      float scale) {
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warp = tid >> 5;          // 8 warps per block
    long long row = (long long)blockIdx.x * 8LL + warp;
    if (row >= rows) return;

    long long base = row * 64LL;
    float v0 = x[base + lane];
    float v1 = x[base + lane + 32];

    float s = warp_sum(v0 + v1);
    float mean = __shfl_sync(0xffffffffu, s, 0) * 0.015625f; // /64

    float d0 = v0 - mean;
    float d1 = v1 - mean;
    float ss = warp_sum(d0 * d0 + d1 * d1);
    float var = __shfl_sync(0xffffffffu, ss, 0) * 0.015625f;
    float inv_std = rsqrtf(var + eps);

    float y0 = d0 * inv_std;
    float y1 = d1 * inv_std;

    y0 = y0 * gamma[lane] + beta[lane];
    y1 = y1 * gamma[lane + 32] + beta[lane + 32];

    // exact GELU: 0.5*x*(1+erf(x/sqrt(2))) then scaling
    const float inv_sqrt2 = 0.70710678118654752440f;
    float g0 = 0.5f * y0 * (1.0f + erff(y0 * inv_sqrt2));
    float g1 = 0.5f * y1 * (1.0f + erff(y1 * inv_sqrt2));

    x[base + lane] = g0 * scale;
    x[base + lane + 32] = g1 * scale;
}

torch::Tensor fused_ln_gelu_inplace(torch::Tensor input,
                                   torch::Tensor gamma,
                                   torch::Tensor beta,
                                   double eps,
                                   double scale) {
    TORCH_CHECK(input.is_cuda(), "input must be CUDA");
    TORCH_CHECK(gamma.is_cuda() && beta.is_cuda(), "gamma/beta must be CUDA");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(gamma.scalar_type() == at::kFloat && beta.scalar_type() == at::kFloat, "only float32 gamma/beta supported");
    TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
    TORCH_CHECK(gamma.is_contiguous() && beta.is_contiguous(), "gamma/beta must be contiguous");
    TORCH_CHECK(input.dim() == 5, "expected NCDHW tensor");
    TORCH_CHECK(input.size(4) == 64, "optimized kernel requires last dimension 64");
    TORCH_CHECK(gamma.numel() == 64 && beta.numel() == 64, "gamma/beta must have 64 elements");

    long long rows = input.numel() / 64LL;
    if (rows > 0) {
        dim3 block(256);
        dim3 grid((unsigned int)((rows + 7LL) / 8LL));
        ln_gelu64_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
            input.data_ptr<float>(),
            gamma.data_ptr<float>(),
            beta.data_ptr<float>(),
            rows,
            (float)eps,
            (float)scale);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }
    return input;
}
'''

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


class ModelNew(nn.Module):
    """
    ConvTranspose3d followed by LayerNorm over the last dimension, exact GELU,
    and scaling.  The expensive convolution is delegated to PyTorch/cuDNN while
    the normalization/activation/scaling tail is fused into one CUDA pass.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding,
                 bias=True, eps=1e-5, scaling_factor=1.0):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size,
            stride=stride, padding=padding, bias=bias
        )
        self.layer_norm = nn.LayerNorm(out_channels, eps=eps)
        self.scaling_factor = scaling_factor
        self.eps = eps

    def forward(self, x):
        y = F.conv_transpose3d(
            x,
            self.conv_transpose.weight,
            self.conv_transpose.bias,
            self.conv_transpose.stride,
            self.conv_transpose.padding,
            self.conv_transpose.output_padding,
            self.conv_transpose.groups,
            self.conv_transpose.dilation,
        )
        # ConvTranspose3d returns contiguous NCDHW for this configuration.  Make
        # parameters contiguous defensively; they are already so in normal use.
        return _ext.fused_ln_gelu_inplace(
            y,
            self.layer_norm.weight.contiguous(),
            self.layer_norm.bias.contiguous(),
            float(self.layer_norm.eps),
            float(self.scaling_factor),
        )