import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a transposed convolution, subtracts a bias term, and applies tanh activation.
    """
    def __init__(self, in_channels, out_channels, kernel_size, bias_shape, stride=2, padding=1, output_padding=1):
        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)) 

    def forward(self, x):
        x = self.conv_transpose(x)
        x = x - self.bias
        x = torch.tanh(x)
        return x

batch_size = 32
in_channels  = 64  
out_channels = 64  
height = width = 256 
kernel_size = 4
bias_shape = (out_channels, 1, 1)

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

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

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

cpp_src = r'''
#include <torch/extension.h>
#include <ATen/ATen.h>
torch::Tensor convtrans_fused_bias_tanh(torch::Tensor x, torch::Tensor w, torch::Tensor conv_b, torch::Tensor sub_b,
                                  int64_t stride, int64_t padding, int64_t output_padding);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <vector>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void tiled4_bias_tanh_kernel(float* __restrict__ y, const float* __restrict__ cb, const float* __restrict__ sb, int C, long HW) {
    int nc = blockIdx.x;
    int c = nc % C;
    long p = ((long)blockIdx.y * blockDim.x + threadIdx.x) * 4;
    long off = (long)nc * HW + p;
    float add = cb[c] - sb[c];
    if (p + 3 < HW) {
        float v0 = y[off] + add;
        float v1 = y[off+1] + add;
        float v2 = y[off+2] + add;
        float v3 = y[off+3] + add;
        y[off] = tanhf(v0);
        y[off+1] = tanhf(v1);
        y[off+2] = tanhf(v2);
        y[off+3] = tanhf(v3);
    } else {
        for (int j=0; j<4 && p+j<HW; ++j) {
            float v = y[off+j] + add;
            y[off+j] = tanhf(v);
        }
    }
}

torch::Tensor convtrans_fused_bias_tanh(torch::Tensor x, torch::Tensor w, torch::Tensor conv_b, torch::Tensor sub_b,
                                  int64_t stride, int64_t padding, int64_t output_padding) {
    c10::cuda::CUDAGuard device_guard(x.device());
    std::vector<int64_t> strides = {stride, stride};
    std::vector<int64_t> pads = {padding, padding};
    std::vector<int64_t> outpads = {output_padding, output_padding};
    std::vector<int64_t> dilation = {1, 1};
    c10::optional<torch::Tensor> no_bias = c10::nullopt;
    auto y = at::conv_transpose2d(x, w, no_bias, strides, pads, outpads, 1, dilation);
    const int N = (int)y.size(0);
    const int C = (int)y.size(1);
    const long HW = (long)y.size(2) * y.size(3);
    dim3 grid(N*C, (unsigned int)((HW + 1024 - 1) / 1024));
    tiled4_bias_tanh_kernel<<<grid, 256, 0, at::cuda::getCurrentCUDAStream()>>>(y.data_ptr<float>(), conv_b.data_ptr<float>(), sub_b.data_ptr<float>(), C, HW);
    return y;
}
'''

_ext = load_inline(name='convtrans_bias_tanh_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=['convtrans_fused_bias_tanh'], 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, bias_shape, stride=2, padding=1, output_padding=1):
        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.stride = stride
        self.padding = padding
        self.output_padding = output_padding
    def forward(self, x):
        return _ext.convtrans_fused_bias_tanh(x, self.conv_transpose.weight, self.conv_transpose.bias, self.bias, self.stride, self.padding, self.output_padding)