import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a convolution, applies tanh, scaling, adds a bias term, and then max-pools.
    """
    def __init__(self, in_channels, out_channels, kernel_size, scaling_factor, bias_shape, pool_kernel_size):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.scaling_factor = scaling_factor
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self.max_pool = nn.MaxPool2d(pool_kernel_size)

    def forward(self, x):
        # Convolution
        x = self.conv(x)
        # Tanh activation
        x = torch.tanh(x)
        # Scaling
        x = x * self.scaling_factor
        # Bias addition
        x = x + self.bias
        # Max-pooling
        x = self.max_pool(x)
        return x

batch_size = 128
in_channels = 8
out_channels = 64
height, width = 256, 256
kernel_size = 3
scaling_factor = 2.0
bias_shape = (out_channels, 1, 1)
pool_kernel_size = 4

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

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

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

cpp_source = r"""
#include <torch/extension.h>
torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor conv_b, torch::Tensor bias, double scaling);
"""

cuda_source = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

__global__ void fused_kernel(const float* __restrict__ x,
                             const float* __restrict__ w,
                             const float* __restrict__ cb,
                             const float* __restrict__ bias,
                             float* __restrict__ out,
                             float scaling,
                             int total) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    const int W = 256, H = 63, PW = 63, OC = 64, IC = 8;

    // Thread order makes oc contiguous within a warp.  This gives coalesced
    // weight/bias loads while reusing identical input addresses through L1/L2.
    int t = idx;
    int oc = t & 63;  t >>= 6;
    int pw = t % PW;  t /= PW;
    int ph = t % H;   t /= H;
    int n = t;

    int base_y = ph * 4, base_x = pw * 4;
    const float* xbase = x + ((n * IC) * W + base_y) * W + base_x;
    const float* wbase = w + oc * (IC * 9);

    float a0 = __ldg(cb + oc), a1 = a0, a2 = a0, a3 = a0;
    float a4 = a0, a5 = a0, a6 = a0, a7 = a0;
    float a8 = a0, a9 = a0, a10 = a0, a11 = a0;
    float a12 = a0, a13 = a0, a14 = a0, a15 = a0;

    #pragma unroll
    for (int c = 0; c < 8; ++c) {
        const float* xp = xbase + c * W * W;
        const float* wp = wbase + c * 9;
        #pragma unroll
        for (int ky = 0; ky < 3; ++ky) {
            #pragma unroll
            for (int kx = 0; kx < 3; ++kx) {
                float ww = __ldg(wp + ky * 3 + kx);
                const float* p = xp + ky * W + kx;
                a0  = fmaf(__ldg(p + 0*W + 0), ww, a0);
                a1  = fmaf(__ldg(p + 0*W + 1), ww, a1);
                a2  = fmaf(__ldg(p + 0*W + 2), ww, a2);
                a3  = fmaf(__ldg(p + 0*W + 3), ww, a3);
                a4  = fmaf(__ldg(p + 1*W + 0), ww, a4);
                a5  = fmaf(__ldg(p + 1*W + 1), ww, a5);
                a6  = fmaf(__ldg(p + 1*W + 2), ww, a6);
                a7  = fmaf(__ldg(p + 1*W + 3), ww, a7);
                a8  = fmaf(__ldg(p + 2*W + 0), ww, a8);
                a9  = fmaf(__ldg(p + 2*W + 1), ww, a9);
                a10 = fmaf(__ldg(p + 2*W + 2), ww, a10);
                a11 = fmaf(__ldg(p + 2*W + 3), ww, a11);
                a12 = fmaf(__ldg(p + 3*W + 0), ww, a12);
                a13 = fmaf(__ldg(p + 3*W + 1), ww, a13);
                a14 = fmaf(__ldg(p + 3*W + 2), ww, a14);
                a15 = fmaf(__ldg(p + 3*W + 3), ww, a15);
            }
        }
    }

    float m;
    if (scaling >= 0.0f) {
        m = fmaxf(fmaxf(fmaxf(a0,a1), fmaxf(a2,a3)), fmaxf(fmaxf(a4,a5), fmaxf(a6,a7)));
        float m2 = fmaxf(fmaxf(fmaxf(a8,a9), fmaxf(a10,a11)), fmaxf(fmaxf(a12,a13), fmaxf(a14,a15)));
        m = fmaxf(m, m2);
    } else {
        m = fminf(fminf(fminf(a0,a1), fminf(a2,a3)), fminf(fminf(a4,a5), fminf(a6,a7)));
        float m2 = fminf(fminf(fminf(a8,a9), fminf(a10,a11)), fminf(fminf(a12,a13), fminf(a14,a15)));
        m = fminf(m, m2);
    }

    int out_idx = ((n * OC + oc) * H + ph) * PW + pw;
    out[out_idx] = tanhf(m) * scaling + __ldg(bias + oc);
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor conv_b, torch::Tensor bias, double scaling) {
    auto out = torch::empty({x.size(0), w.size(0), 63, 63}, x.options());
    int total = x.size(0) * w.size(0) * 63 * 63;
    const int threads = 256;
    fused_kernel<<<(total + threads - 1) / threads, threads>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), conv_b.data_ptr<float>(),
        bias.data_ptr<float>(), out.data_ptr<float>(), (float)scaling, total);
    return out;
}
"""

_ext = load_inline(name="fused_conv_tanh_pool_ext_final", cpp_sources=cpp_source, cuda_sources=cuda_source, 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, pool_kernel_size):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.scaling_factor = scaling_factor
        self.bias = nn.Parameter(torch.randn(bias_shape))

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