import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a transposed convolution, adds a bias term, clamps, scales, clamps, and divides.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, bias_shape, scaling_factor):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)
        self.bias = nn.Parameter(torch.randn(bias_shape)) 
        self.scaling_factor = scaling_factor

    def forward(self, x):
        x = self.conv_transpose(x)
        x = x + self.bias
        x = torch.clamp(x, min=0.0, max=1.0)
        x = x * self.scaling_factor
        x = torch.clamp(x, min=0.0, max=1.0)
        x = x / self.scaling_factor
        return x

batch_size = 128
in_channels  = 64  
out_channels = 64  
height = width = 128 
kernel_size = 3
stride = 2
padding = 1
output_padding = 1
bias_shape = (out_channels, 1, 1)
scaling_factor = 2.0

def get_inputs():
    return [torch.rand(batch_size, in_channels, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, output_padding, bias_shape, 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_deconv_postprocess(torch::Tensor y, torch::Tensor bias, double scaling_factor);
"""

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

static __device__ __forceinline__ float post_op(float v, float b, float s) {
    v += b;
    // torch.clamp(v, 0, 1), preserving NaNs (comparisons with NaN are false)
    v = (v < 0.0f) ? 0.0f : v;
    v = (v > 1.0f) ? 1.0f : v;
    v = v * s;
    v = (v < 0.0f) ? 0.0f : v;
    v = (v > 1.0f) ? 1.0f : v;
    v = v / s;
    return v;
}

__global__ void post_vec4_kernel(float* __restrict__ y,
                                 const float* __restrict__ bias,
                                 long long hw,
                                 int C,
                                 float s) {
    int nc = blockIdx.x;
    int c = nc % C;
    float b = bias[c];
    long long base = (long long)nc * hw;
    float4* __restrict__ y4 = reinterpret_cast<float4*>(y + base);
    long long n4 = hw >> 2;

    for (long long i = threadIdx.x; i < n4; i += blockDim.x) {
        float4 v = y4[i];
        v.x = post_op(v.x, b, s);
        v.y = post_op(v.y, b, s);
        v.z = post_op(v.z, b, s);
        v.w = post_op(v.w, b, s);
        y4[i] = v;
    }
}

__global__ void post_scalar_kernel(float* __restrict__ y,
                                   const float* __restrict__ bias,
                                   long long hw,
                                   int C,
                                   float s) {
    int nc = blockIdx.x;
    int c = nc % C;
    float b = bias[c];
    long long base = (long long)nc * hw;
    for (long long i = threadIdx.x; i < hw; i += blockDim.x) {
        y[base + i] = post_op(y[base + i], b, s);
    }
}

torch::Tensor fused_deconv_postprocess(torch::Tensor y, torch::Tensor bias, double scaling_factor) {
    const c10::cuda::CUDAGuard device_guard(y.device());
    TORCH_CHECK(y.is_cuda(), "y must be a CUDA tensor");
    TORCH_CHECK(bias.is_cuda(), "bias must be a CUDA tensor");
    TORCH_CHECK(y.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(bias.scalar_type() == at::kFloat, "only float32 bias is supported");
    TORCH_CHECK(y.is_contiguous(), "y must be contiguous");
    TORCH_CHECK(bias.is_contiguous(), "bias must be contiguous");
    TORCH_CHECK(y.dim() == 4, "y must be NCHW");

    int N = (int)y.size(0);
    int C = (int)y.size(1);
    long long H = (long long)y.size(2);
    long long W = (long long)y.size(3);
    long long hw = H * W;
    int nc = N * C;
    if (nc == 0 || hw == 0) return y;

    const int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    float s = (float)scaling_factor;

    if ((hw & 3LL) == 0) {
        post_vec4_kernel<<<nc, threads, 0, stream>>>(
            y.data_ptr<float>(), bias.data_ptr<float>(), hw, C, s);
    } else {
        post_scalar_kernel<<<nc, threads, 0, stream>>>(
            y.data_ptr<float>(), bias.data_ptr<float>(), hw, C, s);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
"""

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


class ModelNew(nn.Module):
    """
    Same module as the reference.  The transposed convolution is computed by
    cuDNN/PyTorch, and the following bias/add/clamp/scale/clamp/divide chain is
    fused into one custom CUDA pass over the large output tensor.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, bias_shape, scaling_factor):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(
            in_channels, out_channels, kernel_size,
            stride=stride, padding=padding, output_padding=output_padding
        )
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self.scaling_factor = scaling_factor
        self._stride = stride
        self._padding = padding
        self._output_padding = output_padding

    def forward(self, x):
        y = F.conv_transpose2d(
            x,
            self.conv_transpose.weight,
            self.conv_transpose.bias,
            stride=self.conv_transpose.stride,
            padding=self.conv_transpose.padding,
            output_padding=self.conv_transpose.output_padding,
            groups=self.conv_transpose.groups,
            dilation=self.conv_transpose.dilation,
        )
        return _ext.fused_deconv_postprocess(y, self.bias, float(self.scaling_factor))