import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self, num_classes=1000):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=96, kernel_size=11, stride=4, padding=2)
    
    def forward(self, x):
        x = self.conv1(x)
        return x

# Test code
batch_size = 256
num_classes = 1000

def get_inputs():
    return [torch.rand(batch_size, 3, 224, 224)]

def get_init_inputs():
    return [num_classes]

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

_cpp_src = r"""
torch::Tensor conv1_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
"""

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

#define H_IN 224
#define W_IN 224
#define C_IN 3
#define C_OUT 96
#define K 11
#define STRIDE 4
#define PAD 2
#define H_OUT 55
#define W_OUT 55
#define HW_OUT 3025
#define OC_TILE 16
#define OC_GROUPS 6
#define W_OC_STRIDE 363

__device__ __forceinline__ void accum16(float xv, const float* __restrict__ wp,
                                        float &a0, float &a1, float &a2, float &a3,
                                        float &a4, float &a5, float &a6, float &a7,
                                        float &a8, float &a9, float &a10, float &a11,
                                        float &a12, float &a13, float &a14, float &a15) {
    a0  = fmaf(xv, wp[0  * W_OC_STRIDE], a0);
    a1  = fmaf(xv, wp[1  * W_OC_STRIDE], a1);
    a2  = fmaf(xv, wp[2  * W_OC_STRIDE], a2);
    a3  = fmaf(xv, wp[3  * W_OC_STRIDE], a3);
    a4  = fmaf(xv, wp[4  * W_OC_STRIDE], a4);
    a5  = fmaf(xv, wp[5  * W_OC_STRIDE], a5);
    a6  = fmaf(xv, wp[6  * W_OC_STRIDE], a6);
    a7  = fmaf(xv, wp[7  * W_OC_STRIDE], a7);
    a8  = fmaf(xv, wp[8  * W_OC_STRIDE], a8);
    a9  = fmaf(xv, wp[9  * W_OC_STRIDE], a9);
    a10 = fmaf(xv, wp[10 * W_OC_STRIDE], a10);
    a11 = fmaf(xv, wp[11 * W_OC_STRIDE], a11);
    a12 = fmaf(xv, wp[12 * W_OC_STRIDE], a12);
    a13 = fmaf(xv, wp[13 * W_OC_STRIDE], a13);
    a14 = fmaf(xv, wp[14 * W_OC_STRIDE], a14);
    a15 = fmaf(xv, wp[15 * W_OC_STRIDE], a15);
}

__global__ __launch_bounds__(256, 2)
void conv1_kernel(const float* __restrict__ x,
                  const float* __restrict__ w,
                  const float* __restrict__ b,
                  float* __restrict__ y,
                  int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * OC_GROUPS * HW_OUT;
    if (idx >= total) return;

    int spatial = idx % HW_OUT;
    int t = idx / HW_OUT;
    int g = t % OC_GROUPS;
    int n = t / OC_GROUPS;

    int oh = spatial / W_OUT;
    int ow = spatial - oh * W_OUT;
    int ih0 = oh * STRIDE - PAD;
    int iw0 = ow * STRIDE - PAD;
    int oc = g * OC_TILE;

    float a0  = b[oc + 0];
    float a1  = b[oc + 1];
    float a2  = b[oc + 2];
    float a3  = b[oc + 3];
    float a4  = b[oc + 4];
    float a5  = b[oc + 5];
    float a6  = b[oc + 6];
    float a7  = b[oc + 7];
    float a8  = b[oc + 8];
    float a9  = b[oc + 9];
    float a10 = b[oc + 10];
    float a11 = b[oc + 11];
    float a12 = b[oc + 12];
    float a13 = b[oc + 13];
    float a14 = b[oc + 14];
    float a15 = b[oc + 15];

    const float* xbase = x + n * C_IN * H_IN * W_IN;
    const float* wbase0 = w + oc * W_OC_STRIDE;

    if (oh > 0 && oh < 54 && ow > 0 && ow < 54) {
        #pragma unroll
        for (int c = 0; c < C_IN; ++c) {
            const float* xc = xbase + c * H_IN * W_IN + ih0 * W_IN + iw0;
            const float* wc = wbase0 + c * K * K;
            #pragma unroll
            for (int kh = 0; kh < K; ++kh) {
                const float* xr = xc + kh * W_IN;
                const float* wr = wc + kh * K;
                #pragma unroll
                for (int kw = 0; kw < K; ++kw) {
                    float xv = xr[kw];
                    const float* wp = wr + kw;
                    accum16(xv, wp, a0,a1,a2,a3,a4,a5,a6,a7,a8,a9,a10,a11,a12,a13,a14,a15);
                }
            }
        }
    } else {
        int khs = ih0 < 0 ? -ih0 : 0;
        int khe = (H_IN - ih0) < K ? (H_IN - ih0) : K;
        int kws = iw0 < 0 ? -iw0 : 0;
        int kwe = (W_IN - iw0) < K ? (W_IN - iw0) : K;
        #pragma unroll
        for (int c = 0; c < C_IN; ++c) {
            const float* xc = xbase + c * H_IN * W_IN + ih0 * W_IN + iw0;
            const float* wc = wbase0 + c * K * K;
            for (int kh = khs; kh < khe; ++kh) {
                const float* xr = xc + kh * W_IN;
                const float* wr = wc + kh * K;
                for (int kw = kws; kw < kwe; ++kw) {
                    float xv = xr[kw];
                    const float* wp = wr + kw;
                    accum16(xv, wp, a0,a1,a2,a3,a4,a5,a6,a7,a8,a9,a10,a11,a12,a13,a14,a15);
                }
            }
        }
    }

    float* out = y + (n * C_OUT + oc) * HW_OUT + spatial;
    out[0  * HW_OUT] = a0;
    out[1  * HW_OUT] = a1;
    out[2  * HW_OUT] = a2;
    out[3  * HW_OUT] = a3;
    out[4  * HW_OUT] = a4;
    out[5  * HW_OUT] = a5;
    out[6  * HW_OUT] = a6;
    out[7  * HW_OUT] = a7;
    out[8  * HW_OUT] = a8;
    out[9  * HW_OUT] = a9;
    out[10 * HW_OUT] = a10;
    out[11 * HW_OUT] = a11;
    out[12 * HW_OUT] = a12;
    out[13 * HW_OUT] = a13;
    out[14 * HW_OUT] = a14;
    out[15 * HW_OUT] = a15;
}

torch::Tensor conv1_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    const int64_t N64 = x.size(0);
    auto y = torch::empty({N64, C_OUT, H_OUT, W_OUT}, x.options());
    int N = (int)N64;
    int total = N * OC_GROUPS * HW_OUT;
    const int threads = 256;
    const int blocks = (total + threads - 1) / threads;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    conv1_kernel<<<blocks, threads, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>(), N);
    return y;
}
"""

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

class ModelNew(nn.Module):
    def __init__(self, num_classes=1000):
        super(ModelNew, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=96, kernel_size=11, stride=4, padding=2)

    def forward(self, x):
        return _ext.conv1_forward(x, self.conv1.weight, self.conv1.bias)