import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a convolution, applies ReLU, and applies HardSwish activation.
    """
    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.relu(x)
        x = x * torch.clamp((x + 3) / 6, 0, 1)
        return x

batch_size = 128
in_channels = 8
out_channels = 64
height, width = 128, 128
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 fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b);
'''

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

#define ROWS 32
#define COUT_TILE 8
#define KSIZE_FLAT 72

__global__ __launch_bounds__(256, 2)
void conv3x3_relu_hswish_kernel(const float* __restrict__ x,
                                const float* __restrict__ w,
                                const float* __restrict__ b,
                                float* __restrict__ out,
                                int total_rows, int H, int W, int OH, int OW) {
    const int tx = threadIdx.x;   // 0..31: consecutive output spatial rows
    const int ty = threadIdx.y;   // 0..7 : channel within tile
    const int tid = ty * ROWS + tx;

    const int row = blockIdx.x * ROWS + tx;

    __shared__ float sh_in[KSIZE_FLAT][ROWS];
    __shared__ float sh_w[KSIZE_FLAT][COUT_TILE];

    int n = 0, rem = 0, oh = 0, ow0 = 0;
    if (row < total_rows) {
        const int area = OH * OW;
        n = row / area;
        rem = row - n * area;
        oh = rem / OW;
        ow0 = rem - oh * OW;
    }

    // Load the 3x3x8 input patch for 32 consecutive output positions once;
    // it is reused for all 64 output channels computed by this block.
    #pragma unroll
    for (int t = tid; t < KSIZE_FLAT * ROWS; t += 256) {
        const int kk = t / ROWS;
        const int rr = t - kk * ROWS;
        const int rrow = blockIdx.x * ROWS + rr;
        float v = 0.0f;
        if (rrow < total_rows) {
            const int area = OH * OW;
            const int nn = rrow / area;
            const int rrem = rrow - nn * area;
            const int ooh = rrem / OW;
            const int oow = rrem - ooh * OW;
            const int ci = kk / 9;
            const int k = kk - ci * 9;
            const int ky = k / 3;
            const int kx = k - ky * 3;
            v = __ldg(x + (((nn * 8 + ci) * H + (ooh + ky)) * W + (oow + kx)));
        }
        sh_in[kk][rr] = v;
    }
    __syncthreads();

    // Compute all 64 output channels in 8 channel tiles.
    #pragma unroll
    for (int cbase = 0; cbase < 64; cbase += COUT_TILE) {
        // Load weights for this channel tile: layout [64, 8, 3, 3].
        #pragma unroll
        for (int t = tid; t < KSIZE_FLAT * COUT_TILE; t += 256) {
            const int kk = t / COUT_TILE;
            const int cc = t - kk * COUT_TILE;
            const int ci = kk / 9;
            const int k = kk - ci * 9;
            const int ky = k / 3;
            const int kx = k - ky * 3;
            sh_w[kk][cc] = __ldg(w + ((((cbase + cc) * 8 + ci) * 3 + ky) * 3 + kx));
        }
        __syncthreads();

        if (row < total_rows) {
            const int co = cbase + ty;
            float acc = __ldg(b + co);
            #pragma unroll
            for (int kk = 0; kk < KSIZE_FLAT; ++kk) {
                acc = fmaf(sh_in[kk][tx], sh_w[kk][ty], acc);
            }

            float y = 0.0f;
            if (acc > 0.0f) {
                float m = (acc + 3.0f) * 0.166666666666666657f;
                if (m > 1.0f) m = 1.0f;
                y = acc * m;
            }
            out[((n * 64 + co) * OH * OW) + rem] = y;
        }
        __syncthreads();
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b) {
    const int N = (int)x.size(0);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int OH = H - 2;
    const int OW = W - 2;
    auto out = torch::empty({N, 64, OH, OW}, x.options());
    const int total_rows = N * OH * OW;

    dim3 block(ROWS, COUT_TILE, 1);
    dim3 grid((total_rows + ROWS - 1) / ROWS, 1, 1);
    conv3x3_relu_hswish_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(),
        total_rows, H, W, OH, OW);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''

_ext = load_inline(
    name='fused_conv3x3_relu_hswish_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['fused_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    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):
        return _ext.fused_forward(x, self.conv.weight, self.conv.bias)