import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a convolution, applies minimum operation, Tanh, and another Tanh.
    """
    def __init__(self, in_channels, out_channels, kernel_size):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)

    def forward(self, x):
        x = self.conv(x)
        x = torch.min(x, dim=1, keepdim=True)[0] # Apply minimum operation along the channel dimension
        x = torch.tanh(x)
        x = torch.tanh(x)
        return x

batch_size = 128
in_channels = 16
out_channels = 64
height = width = 256
kernel_size = 3

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

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

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

_cpp_src = r'''
#include <torch/extension.h>

torch::Tensor conv_min_forward_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b);

torch::Tensor conv_min_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b) {
    return conv_min_forward_cuda(x, w, b);
}
'''

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

// Transpose the small convolution weight tensor from [OC, IC, KH, KW]
// to [IC*KH*KW, OC].  For the fixed benchmark this is only 144*64 floats,
// and it makes the main kernel's per-warp weight loads contiguous.
__global__ void transpose_w_kernel(const float* __restrict__ w,
                                   float* __restrict__ wt,
                                   int OC, int Ktot) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = OC * Ktot;
    if (idx < total) {
        int oc = idx / Ktot;
        int k  = idx - oc * Ktot;
        wt[k * OC + oc] = w[idx];
    }
}

// One group of 64 threads computes all output channels for one output pixel.
// A block computes 8 pixels (512 threads).  The convolution, channel-minimum,
// and the two tanh operations are fused, avoiding materializing the 64-channel
// convolution output (~2 GB for the benchmark shape).
template<int PIXELS_PER_BLOCK>
__global__ void conv3x3_min_tanh_kernel(const float* __restrict__ x,
                                        const float* __restrict__ wt,
                                        const float* __restrict__ bias,
                                        float* __restrict__ y,
                                        int N, int H, int W) {
    constexpr int IC = 16;
    constexpr int OC = 64;
    constexpr int KH = 3;
    constexpr int KW = 3;
    constexpr int Ktot = IC * KH * KW;

    const int OH = H - 2;
    const int OW = W - 2;
    const int HWout = OH * OW;
    const int64_t total_pix = (int64_t)N * HWout;

    int tid = threadIdx.x;
    int pix_in_block = tid >> 6;       // /64
    int oc = tid & 63;                 // %64
    int64_t p = (int64_t)blockIdx.x * PIXELS_PER_BLOCK + pix_in_block;

    __shared__ float vals[PIXELS_PER_BLOCK][64];

    float acc = 0.0f;
    if (p < total_pix) {
        int n = (int)(p / HWout);
        int r = (int)(p - (int64_t)n * HWout);
        int oh = r / OW;
        int ow = r - oh * OW;

        acc = bias[oc];
        const float* xbase = x + ((int64_t)n * IC * H * W) + (int64_t)oh * W + ow;

        #pragma unroll
        for (int ic = 0; ic < IC; ++ic) {
            const float* xp = xbase + (int64_t)ic * H * W;
            int kk = ic * 9;
            float v0 = xp[0];
            float v1 = xp[1];
            float v2 = xp[2];
            float v3 = xp[W];
            float v4 = xp[W + 1];
            float v5 = xp[W + 2];
            float v6 = xp[2 * W];
            float v7 = xp[2 * W + 1];
            float v8 = xp[2 * W + 2];
            acc = fmaf(v0, wt[(kk + 0) * OC + oc], acc);
            acc = fmaf(v1, wt[(kk + 1) * OC + oc], acc);
            acc = fmaf(v2, wt[(kk + 2) * OC + oc], acc);
            acc = fmaf(v3, wt[(kk + 3) * OC + oc], acc);
            acc = fmaf(v4, wt[(kk + 4) * OC + oc], acc);
            acc = fmaf(v5, wt[(kk + 5) * OC + oc], acc);
            acc = fmaf(v6, wt[(kk + 6) * OC + oc], acc);
            acc = fmaf(v7, wt[(kk + 7) * OC + oc], acc);
            acc = fmaf(v8, wt[(kk + 8) * OC + oc], acc);
        }
        vals[pix_in_block][oc] = acc;
    } else {
        vals[pix_in_block][oc] = INFINITY;
    }
    __syncthreads();

    // Reduce the 64 channel values for each pixel.  Only one thread per pixel
    // writes the final scalar, so output is contiguous [N,1,OH,OW].
    if (oc == 0 && p < total_pix) {
        float m = vals[pix_in_block][0];
        #pragma unroll
        for (int c = 1; c < 64; ++c) {
            float v = vals[pix_in_block][c];
            m = v < m ? v : m;
        }
        float t = tanhf(m);
        t = tanhf(t);
        y[p] = t;
    }
}

torch::Tensor conv_min_forward_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(w.is_cuda(), "w must be CUDA");
    TORCH_CHECK(b.is_cuda(), "b must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");

    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int OC = (int)w.size(0);
    const int KH = (int)w.size(2);
    const int KW = (int)w.size(3);
    TORCH_CHECK(C == 16 && OC == 64 && KH == 3 && KW == 3, "this optimized kernel expects C=16, OC=64, K=3");

    const int OH = H - 2;
    const int OW = W - 2;
    auto y = torch::empty({N, 1, OH, OW}, x.options());

    auto wt = torch::empty({C * KH * KW, OC}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    int Ktot = C * KH * KW;
    int total_w = OC * Ktot;
    transpose_w_kernel<<<(total_w + 255) / 256, 256, 0, stream>>>(
        w.data_ptr<float>(), wt.data_ptr<float>(), OC, Ktot);

    constexpr int PPB = 8;
    int64_t total_pix = (int64_t)N * OH * OW;
    int blocks = (int)((total_pix + PPB - 1) / PPB);
    conv3x3_min_tanh_kernel<PPB><<<blocks, PPB * 64, 0, stream>>>(
        x.data_ptr<float>(), wt.data_ptr<float>(), b.data_ptr<float>(), y.data_ptr<float>(), N, H, W);

    return y;
}
'''

_ext = load_inline(
    name='fused_conv_min_tanh_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['conv_min_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    with_cuda=True,
    verbose=False,
)

class ModelNew(nn.Module):
    """
    Fused replacement for Conv2d -> min(channel) -> tanh -> tanh.
    """
    def __init__(self, in_channels, out_channels, kernel_size):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)

    def forward(self, x):
        if not x.is_contiguous():
            x = x.contiguous()
        return _ext.conv_min_forward(x, self.conv.weight, self.conv.bias)