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>
#include <ATen/ops/conv_transpose2d.h>
#include <vector>

torch::Tensor bias_tanh_inplace_cuda(torch::Tensor y, torch::Tensor bias);

torch::Tensor forward(torch::Tensor x,
                      torch::Tensor weight,
                      torch::Tensor conv_bias,
                      torch::Tensor bias,
                      int64_t stride,
                      int64_t padding,
                      int64_t output_padding) {
    std::vector<int64_t> st{stride, stride};
    std::vector<int64_t> pad{padding, padding};
    std::vector<int64_t> outpad{output_padding, output_padding};
    std::vector<int64_t> dil{1, 1};

    auto y = at::conv_transpose2d(x, weight, conv_bias, st, pad, outpad, 1, dil);
    return bias_tanh_inplace_cuda(y, bias);
}
'''

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

#ifndef CHECK_CUDA
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#endif
#ifndef CHECK_CONTIGUOUS
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#endif

static constexpr int BLOCK = 256;
static constexpr int ITEMS = 4;

__global__ void bias_tanh_kernel(float* __restrict__ y,
                                 const float* __restrict__ b,
                                 long long hw,
                                 int channels) {
    const int plane = blockIdx.y;              // n*C + c
    const int c = plane % channels;
    const float bv = __ldg(b + c);
    const long long base = ((long long)plane) * hw +
                           ((long long)blockIdx.x) * (BLOCK * ITEMS) + threadIdx.x;

    #pragma unroll
    for (int i = 0; i < ITEMS; ++i) {
        long long off = base + (long long)i * BLOCK;
        if (off < ((long long)plane + 1LL) * hw) {
            float v = y[off] - bv;
            y[off] = tanhf(v);
        }
    }
}

torch::Tensor bias_tanh_inplace_cuda(torch::Tensor y, torch::Tensor bias) {
    CHECK_CUDA(y);
    CHECK_CUDA(bias);
    CHECK_CONTIGUOUS(y);
    CHECK_CONTIGUOUS(bias);
    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.dim() == 4, "expected NCHW output");

    const int64_t n = y.size(0);
    const int64_t c = y.size(1);
    const int64_t h = y.size(2);
    const int64_t w = y.size(3);
    const long long hw = (long long)h * (long long)w;
    const dim3 block(BLOCK);
    const dim3 grid((unsigned int)((hw + BLOCK * ITEMS - 1) / (BLOCK * ITEMS)),
                    (unsigned int)(n * c));

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    bias_tanh_kernel<<<grid, block, 0, stream>>>(y.data_ptr<float>(), bias.data_ptr<float>(), hw, (int)c);
    return y;
}
'''

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


def _as_int(v):
    if isinstance(v, tuple) or isinstance(v, list):
        return int(v[0])
    return int(v)


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_i = _as_int(stride)
        self._padding_i = _as_int(padding)
        self._output_padding_i = _as_int(output_padding)

    def forward(self, x):
        return _ext.forward(x, self.conv_transpose.weight, self.conv_transpose.bias,
                            self.bias, self._stride_i, self._padding_i,
                            self._output_padding_i)